Merge branch 'main' into ODSC-76209/GPU-Shape-Recommendation

Author: elizjo

ads/aqua/constants.py

@@ -56,6 +56,9 @@
 MODEL_BY_REFERENCE_OSS_PATH_KEY = "artifact_location"
 
 AQUA_CHAT_TEMPLATE_METADATA_KEY = "chat_template"
+UNKNOWN_ENUM_VALUE = "UNKNOWN_ENUM_VALUE"
+MODEL_GROUP = "MODEL_GROUP"
+SINGLE_MODEL_FLEX = "SINGLE_MODEL_FLEX"
 
 CONSOLE_LINK_RESOURCE_TYPE_MAPPING = {
     "datasciencemodel": "models",

ads/aqua/modeldeployment/deployment.py

@@ -45,8 +45,11 @@
     AQUA_MODEL_TYPE_SERVICE,
     AQUA_MULTI_MODEL_CONFIG,
     MODEL_BY_REFERENCE_OSS_PATH_KEY,
+    MODEL_GROUP,
     MODEL_NAME_DELIMITER,
+    SINGLE_MODEL_FLEX,
     UNKNOWN_DICT,
+    UNKNOWN_ENUM_VALUE,
 )
 from ads.aqua.data import AquaResourceIdentifier
 from ads.aqua.model import AquaModelApp
@@ -873,21 +876,26 @@ def list(self, **kwargs) -> List["AquaDeployment"]:
 
             if oci_aqua:
                 # skipping the AQUA model deployments that are created from model group
-                # TODO: remove this checker after AQUA deployment is integrated with model group
-                aqua_model_id = model_deployment.freeform_tags.get(
-                    Tags.AQUA_MODEL_ID_TAG, UNKNOWN
-                )
                 if (
-                    "datasciencemodelgroup" in aqua_model_id
-                    or model_deployment.model_deployment_configuration_details.deployment_type
-                    == "UNKNOWN_ENUM_VALUE"
+                    model_deployment.model_deployment_configuration_details.deployment_type
+                    in [UNKNOWN_ENUM_VALUE, MODEL_GROUP, SINGLE_MODEL_FLEX]
                 ):
                     continue
-                results.append(
-                    AquaDeployment.from_oci_model_deployment(
-                        model_deployment, self.region
+                try:
+                    results.append(
+                        AquaDeployment.from_oci_model_deployment(
+                            model_deployment, self.region
+                        )
                     )
-                )
+                except Exception as e:
+                    logger.error(
+                        f"There was an issue processing the list of model deployments . Error: {str(e)}",
+                        exc_info=True,
+                    )
+                    raise AquaRuntimeError(
+                        f"There was an issue processing the list of model deployments . Error: {str(e)}"
+                    ) from e
+
                 # log telemetry if MD is in active or failed state
                 deployment_id = model_deployment.id
                 state = model_deployment.lifecycle_state.upper()

ads/aqua/shaperecommend/constants.py

@@ -94,6 +94,9 @@
 DEFAULT_WEIGHT_SIZE = "bfloat16"
 DEFAULT_MAX_SEQ_LEN = 4096
 
+DEFAULT_WEIGHT_SIZE = "float32"
+
+
 BITS_AND_BYTES_8BIT = "8bit"
 BITS_AND_BYTES_4BIT = "4bit"
 

ads/aqua/shaperecommend/recommend.py

@@ -111,7 +111,17 @@ def which_shapes(
                 shape_recommendation_report = self._summarize_shapes_for_seq_lens(
                     llm_config, shapes, model_name
                 )
+                
+            data = self._get_model_config(ds_model)
 
+            llm_config = LLMConfig.from_raw_config(data)
+
+            model_name = ds_model.display_name if ds_model.display_name else ""
+
+            shape_recommendation_report = self._summarize_shapes_for_seq_lens(
+                llm_config, shapes, model_name
+            )
+            
             if request.generate_table and shape_recommendation_report.recommendations:
                 shape_recommendation_report = self._rich_diff_table(
                     shape_recommendation_report
@@ -257,6 +267,7 @@ def _rich_diff_table(shape_report: ShapeRecommendationReport) -> Table:
             else:
                 total_memory = f"CPU: {str(shape.memory_in_gbs)}"
 
+
             if model:
                 model_size = str(model.total_model_gb)
             else:

ads/aqua/shaperecommend/shape_report.py

@@ -3,6 +3,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import json
+
 from typing import List, Optional
 
 from pydantic import BaseModel, Field
@@ -17,6 +18,7 @@
     VLLM_ENV_KEY,
     VLLM_PARAMS_KEY,
 )
+from ads.aqua.shaperecommend.constants import QUANT_MAPPING
 from ads.aqua.shaperecommend.estimator import MemoryEstimator
 from ads.config import COMPARTMENT_OCID
 
@@ -56,6 +58,8 @@ class DeploymentParams(BaseModel):  # noqa: N801
         None, description="Type of quantization (e.g. 4bit)."
     )
     max_model_len: Optional[int] = Field(None, description="Maximum length of input sequence.")
+    max_model_len: int = Field(..., description="Maximum length of input sequence.")
+
     params: str = Field(
         ..., description="Runtime parameters for deployment with vLLM, etc."
     )
@@ -88,6 +92,12 @@ class ModelConfig(BaseModel):
 
     recommendation: Optional[str] = Field("", description="GPU recommendation for the model.")
 
+    model_details: ModelDetail = Field(..., description="Details about the model.")
+    deployment_params: DeploymentParams = Field(
+        ..., description="Parameters for deployment."
+    )
+    recommendation: str = Field(..., description="GPU recommendation for the model.")
+
     class Config:
         protected_namespaces = ()
 
@@ -246,7 +256,6 @@ class ShapeRecommendationReport(BaseModel):
         description="Details for troubleshooting if no shapes fit the current model.",
     )
 
-
     @classmethod
     def from_deployment_config(cls, deployment_config: AquaDeploymentConfig, model_name: str, valid_shapes: List[ComputeShapeSummary]) -> "ShapeRecommendationReport":
         """

ads/jobs/builders/infrastructure/dsc_job.py

@@ -1751,6 +1751,7 @@ def is_multi_node_job(runtime):
         return (
             MULTI_NODE_JOB_SUPPORT
             and isinstance(runtime, MultiNodeRuntime)
+            and runtime.replica
             and runtime.replica > 1
         )
 

ads/jobs/builders/infrastructure/dsc_job_runtime.py

@@ -365,6 +365,11 @@ def _get_node_group(self, dsc_job):
             dsc_job,
             "job_node_configuration_details.job_node_group_configuration_details_list",
         )
+        if node_groups is None:
+            node_groups = get_value(
+                dsc_job,
+                "job_node_configuration_details.jobNodeGroupConfigurationDetailsList",
+            )
         if node_groups and len(node_groups) == 1:
             return node_groups[0]
         return None
@@ -373,6 +378,7 @@ def _get_replica(self, dsc_job, envs):
         node_group = self._get_node_group(dsc_job)
         if node_group:
             replica = get_value(node_group, "replicas")
+            envs.pop(self.CONST_NODE_COUNT, None)
         elif not envs:
             replica = None
         elif self.CONST_WORKER_COUNT in envs:
@@ -399,7 +405,9 @@ def _extract_envs(self, dsc_job):
         env_attr = "job_configuration_details.environment_variables"
         node_group = self._get_node_group(dsc_job)
         if node_group:
-            envs = get_value(node_group, env_attr)
+            envs = get_value(node_group, env_attr) or get_value(
+                node_group, "jobConfigurationDetails.environment_variables"
+            )
         else:
             envs = get_value(dsc_job, env_attr)
         if envs:

ads/pipeline/ads_pipeline.py

@@ -1728,15 +1728,19 @@ def __step_details(self, pipeline_details: Dict) -> list:
 
     def __step_infrastructure_configuration_details(self, step) -> dict:
         step_infrastructure_configuration_details = {}
-        step_infrastructure_configuration_details[
-            "blockStorageSizeInGBs"
-        ] = step.infrastructure.block_storage_size
-        step_infrastructure_configuration_details[
-            "shapeName"
-        ] = step.infrastructure.shape_name
-        step_infrastructure_configuration_details[
-            "shapeConfigDetails"
-        ] = step.infrastructure.shape_config_details
+        step_infrastructure_configuration_details["blockStorageSizeInGBs"] = (
+            step.infrastructure.block_storage_size
+        )
+        step_infrastructure_configuration_details["shapeName"] = (
+            step.infrastructure.shape_name
+        )
+        step_infrastructure_configuration_details["shapeConfigDetails"] = (
+            step.infrastructure.shape_config_details
+        )
+        if getattr(step.infrastructure, "subnet_id", ""):
+            step_infrastructure_configuration_details["subnetId"] = (
+                step.infrastructure.subnet_id
+            )
         return step_infrastructure_configuration_details
 
     def __step_configuration_details(self, pipeline_details: Dict, step) -> dict:

tests/unitary/default_setup/jobs/test_jobs_pytorch_ddp.py

@@ -10,6 +10,9 @@
 from unittest import mock
 
 from ads.jobs import DataScienceJob, DataScienceJobRun, PyTorchDistributedRuntime
+from ads.jobs.builders.infrastructure.dsc_job_runtime import (
+    MULTI_NODE_JOB_SUPPORT,
+)
 from ads.jobs.builders.infrastructure.dsc_job_runtime import (
     PyTorchDistributedRuntimeHandler as Handler,
 )
@@ -133,23 +136,26 @@ def test_create_job_runs(self, patched_run, *args):
         self.assertIsInstance(main_run, DataScienceJobRun)
         self.assertEqual(main_run.id, test_ocid)
         kwarg_list = [call_args.kwargs for call_args in patched_run.call_args_list]
-        self.assertEqual(
-            kwarg_list,
-            [
-                {
-                    "display_name": "None-0",
-                    "environment_variables": {"NODE_RANK": "0", "NODE_COUNT": "2"},
-                },
-                {
-                    "display_name": "None-1",
-                    "environment_variables": {
-                        "NODE_RANK": "1",
-                        "NODE_COUNT": "2",
-                        "MAIN_JOB_RUN_OCID": test_ocid,
+        if MULTI_NODE_JOB_SUPPORT:
+            self.assertEqual(kwarg_list, [{}])
+        else:
+            self.assertEqual(
+                kwarg_list,
+                [
+                    {
+                        "display_name": "None-0",
+                        "environment_variables": {"NODE_RANK": "0", "NODE_COUNT": "2"},
                     },
-                },
-            ],
-        )
+                    {
+                        "display_name": "None-1",
+                        "environment_variables": {
+                            "NODE_RANK": "1",
+                            "NODE_COUNT": "2",
+                            "MAIN_JOB_RUN_OCID": test_ocid,
+                        },
+                    },
+                ],
+            )
 
     @mock.patch.dict(
         os.environ, {utils.CONST_ENV_INPUT_MAPPINGS: json.dumps({INPUT_SRC: INPUT_DST})}

tests/unitary/with_extras/aqua/test_data/recommend/Devstral-Small-2507-GQA.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "MistralForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 11,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 32768,
+  "max_position_embeddings": 131072,
+  "model_type": "mistral",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 40,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 1000000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.53.1",
+  "use_cache": true,
+  "vocab_size": 131072
+}