fixed unit tests

Author: elizjo

ads/aqua/common/utils.py

@@ -1253,24 +1253,24 @@ def load_gpu_shapes_index(
     file_name = "gpu_shapes_index.json"
 
     # Try remote load
-    # remote_data: Dict[str, Any] = {}
-    # if CONDA_BUCKET_NS:
-    #     try:
-    #         auth = auth or authutil.default_signer()
-    #         storage_path = (
-    #             f"oci://{CONDA_BUCKET_NAME}@{CONDA_BUCKET_NS}/service_pack/{file_name}"
-    #         )
-    #         logger.debug(
-    #             "Loading GPU shapes index from Object Storage: %s", storage_path
-    #         )
-    #         with fsspec.open(storage_path, mode="r", **auth) as f:
-    #             remote_data = json.load(f)
-    #         logger.debug(
-    #             "Loaded %d shapes from Object Storage",
-    #             len(remote_data.get("shapes", {})),
-    #         )
-    #     except Exception as ex:
-    #         logger.debug("Remote load failed (%s); falling back to local", ex)
+    remote_data: Dict[str, Any] = {}
+    if CONDA_BUCKET_NS:
+        try:
+            auth = auth or authutil.default_signer()
+            storage_path = (
+                f"oci://{CONDA_BUCKET_NAME}@{CONDA_BUCKET_NS}/service_pack/{file_name}"
+            )
+            logger.debug(
+                "Loading GPU shapes index from Object Storage: %s", storage_path
+            )
+            with fsspec.open(storage_path, mode="r", **auth) as f:
+                remote_data = json.load(f)
+            logger.debug(
+                "Loaded %d shapes from Object Storage",
+                len(remote_data.get("shapes", {})),
+            )
+        except Exception as ex:
+            logger.debug("Remote load failed (%s); falling back to local", ex)
 
     # Load local copy
     local_data: Dict[str, Any] = {}

ads/aqua/shaperecommend/recommend.py

@@ -61,7 +61,7 @@ def which_gpu(self, **kwargs) -> ShapeRecommendationReport:
 
         Parameters
         ----------
-        model_ocid : str
+        ocid : str
            OCID of the model to recommend feasible compute shapes.
 
         Returns
@@ -77,19 +77,23 @@ def which_gpu(self, **kwargs) -> ShapeRecommendationReport:
         """
         try:
             request = RequestRecommend(**kwargs)
-            data, model_name = self.get_model_config(request.model_ocid)
+            ds_model = self.validate_model_ocid(request.model_ocid)
+            data = self.get_model_config(ds_model)
 
             llm_config = LLMConfig.from_raw_config(data)
 
             available_shapes = self.valid_compute_shapes()
+
+            model_name = ds_model.display_name if ds_model.display_name else ""
+
             recommendations = self.summarize_shapes_for_seq_lens(
                 llm_config, available_shapes, model_name
             )
 
         # custom error to catch model incompatibility issues
         except AquaRecommendationError as error:
             return ShapeRecommendationReport(
-            recommendations=[], troubleshoot=str(error)
+                recommendations=[], troubleshoot=str(error)
             )
 
         except ValidationError as ex:
@@ -115,10 +119,16 @@ def rich_diff_table(shape_report: ShapeRecommendationReport) -> Table:
         Returns:
             Table: A rich Table displaying model deployment recommendations.
         """
-        logger.debug("Starting to generate rich diff table from ShapeRecommendationReport.")
+        logger.debug(
+            "Starting to generate rich diff table from ShapeRecommendationReport."
+        )
 
-        name = shape_report.model_name
-        header = f"Model Deployment Recommendations: {name}" if name else "Model Deployment Recommendations"
+        name = shape_report.display_name
+        header = (
+            f"Model Deployment Recommendations: {name}"
+            if name
+            else "Model Deployment Recommendations"
+        )
         logger.debug(f"Table header set to: {header!r}")
 
         if shape_report.troubleshoot:
@@ -167,13 +177,12 @@ def rich_diff_table(shape_report: ShapeRecommendationReport) -> Table:
                 str(model.total_model_gb),
                 deploy.quantization,
                 str(deploy.max_model_len),
-                full_recommendation
+                full_recommendation,
             )
 
         logger.debug("Completed populating table with recommendation rows.")
         return table
 
-
     def shapes(self, **kwargs) -> Table:
         """
         For the CLI, generates the table (in rich diff) with valid GPU deployment shapes
@@ -203,12 +212,31 @@ def shapes(self, **kwargs) -> Table:
             if shape_recommend_report.troubleshoot:
                 raise AquaValueError(shape_recommend_report.troubleshoot)
             else:
-                raise AquaValueError("Unable to generate recommendations from model. Please ensure model is registered and is a decoder-only text-generation model.")
+                raise AquaValueError(
+                    "Unable to generate recommendations from model. Please ensure model is registered and is a decoder-only text-generation model."
+                )
 
         return self.rich_diff_table(shape_recommend_report)
 
     @staticmethod
-    def get_model_config(ocid: str):
+    def validate_model_ocid(ocid: str) -> DataScienceModel:
+        """
+        Ensures the OCID passed is valid for referencing a DataScienceModel resource.
+        """
+        resource_type = get_resource_type(ocid)
+
+        if resource_type != "datasciencemodel":
+            raise AquaValueError(
+                f"The provided OCID '{ocid}' is not a valid Oracle Cloud Data Science Model OCID. "
+                "Please provide an OCID corresponding to a Data Science model resource. "
+                "Tip: Data Science model OCIDs typically start with 'ocid1.datasciencemodel...'."
+            )
+
+        model = DataScienceModel.from_id(ocid)
+        return model
+
+    @staticmethod
+    def get_model_config(model: DataScienceModel):
         """
         Loads the configuration for a given Oracle Cloud Data Science model.
 
@@ -218,8 +246,8 @@ def get_model_config(ocid: str):
 
         Parameters
         ----------
-        ocid : str
-            The OCID of the Data Science model.
+        model : DataScienceModel
+            The DataScienceModel representation of the model used in recommendations
 
         Returns
         -------
@@ -235,18 +263,6 @@ def get_model_config(ocid: str):
         AquaRecommendationError
             If the model OCID provided is not supported (only text-generation decoder models in safetensor format supported).
         """
-        resource_type = get_resource_type(ocid)
-
-        if resource_type != "datasciencemodel":
-            raise AquaValueError(
-                f"The provided OCID '{ocid}' is not a valid Oracle Cloud Data Science Model OCID. "
-                "Please provide an OCID corresponding to a Data Science model resource. "
-                "Tip: Data Science model OCIDs typically start with 'ocid1.datasciencemodel...'."
-            )
-
-        model = DataScienceModel.from_id(ocid)
-
-        model_name = model.display_name
 
         model_task = model.freeform_tags.get("task", "").lower()
         model_format = model.freeform_tags.get("model_format", "").lower()
@@ -283,7 +299,7 @@ def get_model_config(ocid: str):
                 "Please ensure your model follows the Hugging Face format and includes a 'config.json' with the necessary architecture parameters."
             ) from e
 
-        return data, model_name
+        return data
 
     @staticmethod
     def valid_compute_shapes() -> List["ComputeShapeSummary"]:
@@ -444,5 +460,7 @@ def summarize_shapes_for_seq_lens(
                 )
 
         return ShapeRecommendationReport(
-            model_name=name, recommendations=recommendations, troubleshoot=troubleshoot_msg
+            display_name=name,
+            recommendations=recommendations,
+            troubleshoot=troubleshoot_msg,
         )

ads/aqua/shaperecommend/shape_report.py

@@ -16,7 +16,12 @@ class RequestRecommend(BaseModel):
     A request to recommend compute shapes and parameters for a given model.
     """
 
-    model_ocid: str = Field(..., description="The OCID of the model to recommend feasible compute shapes.")
+    model_ocid: str = Field(
+        ..., description="The OCID of the model to recommend feasible compute shapes."
+    )
+
+    class Config:
+        protected_namespaces = ()
 
 
 class DeploymentParams(BaseModel):  # noqa: N801
@@ -42,6 +47,9 @@ class ModelDetail(BaseModel):
         ..., description="Total size of model and cache in GB."
     )
 
+    class Config:
+        protected_namespaces = ()
+
 
 class ModelConfig(BaseModel):
     """
@@ -54,8 +62,13 @@ class ModelConfig(BaseModel):
     )
     recommendation: str = Field(..., description="GPU recommendation for the model.")
 
+    class Config:
+        protected_namespaces = ()
+
     @classmethod
-    def constuct_model_config(cls, estimator: MemoryEstimator, allowed_gpu_memory: float) -> "ModelConfig":
+    def constuct_model_config(
+        cls, estimator: MemoryEstimator, allowed_gpu_memory: float
+    ) -> "ModelConfig":
         """
         Assembles a complete ModelConfig, including model details, deployment parameters (vLLM), and recommendations.
 
@@ -78,32 +91,33 @@ def constuct_model_config(cls, estimator: MemoryEstimator, allowed_gpu_memory: f
         """
         deployment_params = DeploymentParams(
             quantization=getattr(estimator.llm_config, "quantization", None),
-            max_model_len=getattr(estimator, "seq_len", None)
+            max_model_len=getattr(estimator, "seq_len", None),
         )
         model_detail = ModelDetail(
             model_size_gb=round(getattr(estimator, "model_memory", 0.0), 2),
             kv_cache_size_gb=round(getattr(estimator, "kv_cache_memory", 0.0), 2),
-            total_model_gb=round(getattr(estimator, "total_memory", 0.0), 2)
+            total_model_gb=round(getattr(estimator, "total_memory", 0.0), 2),
         )
         return ModelConfig(
             model_details=model_detail,
             deployment_params=deployment_params,
-            recommendation= estimator.limiting_factor(allowed_gpu_memory)
+            recommendation=estimator.limiting_factor(allowed_gpu_memory),
         )
 
 
 class ShapeReport(BaseModel):
     """
     The feasible deployment configurations for the model per shape.
     """
-    shape_details: 'ComputeShapeSummary' = Field(
+
+    shape_details: "ComputeShapeSummary" = Field(
         ..., description="Details about the compute shape (ex. VM.GPU.A10.2)."
     )
-    configurations: List['ModelConfig'] = Field(
+    configurations: List["ModelConfig"] = Field(
         default_factory=list, description="List of model configurations."
     )
 
-    def is_dominated(self, others: List['ShapeReport']) -> bool:
+    def is_dominated(self, others: List["ShapeReport"]) -> bool:
         """
         Determines whether this shape is dominated by any other shape in a Pareto sense.
 
@@ -128,31 +142,35 @@ def is_dominated(self, others: List['ShapeReport']) -> bool:
 
         cand_cost = self.shape_details.gpu_specs.ranking.cost
         cand_perf = self.shape_details.gpu_specs.ranking.performance
-        cand_quant = QUANT_MAPPING.get(self.configurations[0].deployment_params.quantization, 0)
+        cand_quant = QUANT_MAPPING.get(
+            self.configurations[0].deployment_params.quantization, 0
+        )
         cand_maxlen = self.configurations[0].deployment_params.max_model_len
 
         for other in others:
             other_cost = other.shape_details.gpu_specs.ranking.cost
             other_perf = other.shape_details.gpu_specs.ranking.performance
-            other_quant = QUANT_MAPPING.get(other.configurations[0].deployment_params.quantization, 0)
+            other_quant = QUANT_MAPPING.get(
+                other.configurations[0].deployment_params.quantization, 0
+            )
             other_maxlen = other.configurations[0].deployment_params.max_model_len
             if (
-                other_cost <= cand_cost and
-                other_perf >= cand_perf and
-                other_quant >= cand_quant and
-                other_maxlen >= cand_maxlen and
-                (
-                    other_cost < cand_cost or
-                    other_perf > cand_perf or
-                    other_quant > cand_quant or
-                    other_maxlen > cand_maxlen
+                other_cost <= cand_cost
+                and other_perf >= cand_perf
+                and other_quant >= cand_quant
+                and other_maxlen >= cand_maxlen
+                and (
+                    other_cost < cand_cost
+                    or other_perf > cand_perf
+                    or other_quant > cand_quant
+                    or other_maxlen > cand_maxlen
                 )
             ):
                 return True
         return False
 
     @classmethod
-    def pareto_front(cls, shapes: List['ShapeReport']) -> List['ShapeReport']:
+    def pareto_front(cls, shapes: List["ShapeReport"]) -> List["ShapeReport"]:
         """
         Filters a list of shapes/configurations to those on the Pareto frontier.
 
@@ -171,7 +189,11 @@ def pareto_front(cls, shapes: List['ShapeReport']) -> List['ShapeReport']:
         The returned set contains non-dominated deployments for maximizing
         performance, quantization, and model length, while minimizing cost.
         """
-        return [shape for shape in shapes if not shape.is_dominated([s for s in shapes if s != shape])]
+        return [
+            shape
+            for shape in shapes
+            if not shape.is_dominated([s for s in shapes if s != shape])
+        ]
 
 
 class ShapeRecommendationReport(BaseModel):
@@ -184,7 +206,8 @@ class ShapeRecommendationReport(BaseModel):
         troubleshoot (Optional[TroubleshootShapeSummary]): Troubleshooting information
             if no valid deployment shapes are available.
     """
-    model_name: Optional[str] = Field(
+
+    display_name: Optional[str] = Field(
         "", description="Name of the model used for recommendations."
     )
     recommendations: List[ShapeReport] = Field(