addressed comments

Author: elizjo

ads/aqua/extension/deployment_handler.py

@@ -58,12 +58,11 @@ def get(self, id: Union[str, List[str]] = None):
                 model_id=id.split(",") if "," in id else id
             )
         elif paths.startswith("aqua/deployments/recommend_shapes"):
-            id = id or self.get_argument("model_id", default=None)
             if not id or not isinstance(id, str):
                 raise HTTPError(
                     400,
                     f"Invalid request format for {self.request.path}. "
-                    "Expected a single model OCID",
+                    "Expected a single model OCID specified as --model_id",
                 )
             id = id.replace(" ", "")
             return self.get_recommend_shape(model_id=id)
@@ -189,14 +188,10 @@ def get_recommend_shape(self, model_id: str):
 
         compartment_id = self.get_argument("compartment_id", default=COMPARTMENT_OCID)
 
-        generate_table = (
-            self.get_argument("generate_table", default="True").lower() == "true"
-        )
-
         recommend_report = app.recommend_shape(
             model_id=model_id,
             compartment_id=compartment_id,
-            generate_table=generate_table,
+            generate_table=False,
         )
 
         return self.finish(recommend_report)

ads/aqua/modeldeployment/constants.py

@@ -12,4 +12,3 @@
 DEFAULT_WAIT_TIME = 12000
 DEFAULT_POLL_INTERVAL = 10
 
-SHAPE_MAP = {"NVIDIA_GPU": "GPU"}

ads/aqua/modeldeployment/deployment.py

@@ -67,7 +67,6 @@
 from ads.aqua.modeldeployment.constants import (
     DEFAULT_POLL_INTERVAL,
     DEFAULT_WAIT_TIME,
-    SHAPE_MAP,
 )
 from ads.aqua.modeldeployment.entities import (
     AquaDeployment,
@@ -77,7 +76,10 @@
 )
 from ads.aqua.modeldeployment.model_group_config import ModelGroupConfig
 from ads.aqua.shaperecommend.recommend import AquaShapeRecommend
-from ads.aqua.shaperecommend.shape_report import ShapeRecommendationReport
+from ads.aqua.shaperecommend.shape_report import (
+    RequestRecommend,
+    ShapeRecommendationReport,
+)
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.common.utils import UNKNOWN, get_log_links
 from ads.common.work_request import DataScienceWorkRequest
@@ -1250,60 +1252,6 @@ def validate_deployment_params(
             )
         return {"valid": True}
 
-    def valid_compute_shapes(self, **kwargs) -> List["ComputeShapeSummary"]:
-        """
-        Returns a filtered list of GPU-only ComputeShapeSummary objects by reading and parsing a JSON file.
-
-        Parameters
-        ----------
-        file : str
-            Path to the JSON file containing shape data.
-
-        Returns
-        -------
-        List[ComputeShapeSummary]
-            List of ComputeShapeSummary objects passing the checks.
-
-        Raises
-        ------
-        ValueError
-            If the file cannot be opened, parsed, or the 'shapes' key is missing.
-        """
-        compartment_id = kwargs.pop("compartment_id", COMPARTMENT_OCID)
-        oci_shapes: list[ModelDeploymentShapeSummary] = self.list_resource(
-            self.ds_client.list_model_deployment_shapes,
-            compartment_id=compartment_id,
-            **kwargs,
-        )
-        set_user_shapes = {shape.name: shape for shape in oci_shapes}
-
-        gpu_shapes_metadata = load_gpu_shapes_index().shapes
-
-        valid_shapes = []
-        # only loops through GPU shapes, update later to include CPU shapes
-        for name, spec in gpu_shapes_metadata.items():
-            if name in set_user_shapes:
-                oci_shape = set_user_shapes.get(name)
-
-                compute_shape = ComputeShapeSummary(
-                    available=True,
-                    core_count=oci_shape.core_count,
-                    memory_in_gbs=oci_shape.memory_in_gbs,
-                    shape_series=SHAPE_MAP.get(oci_shape.shape_series, "GPU"),
-                    name=oci_shape.name,
-                    gpu_specs=spec,
-                )
-            else:
-                compute_shape = ComputeShapeSummary(
-                    available=False, name=name, shape_series="GPU", gpu_specs=spec
-                )
-            valid_shapes.append(compute_shape)
-
-        valid_shapes.sort(
-            key=lambda shape: shape.gpu_specs.gpu_memory_in_gbs, reverse=True
-        )
-        return valid_shapes
-
     def recommend_shape(self, **kwargs) -> Union[Table, ShapeRecommendationReport]:
         """
         For the CLI (set generate_table = True), generates the table (in rich diff) with valid
@@ -1335,13 +1283,16 @@ def recommend_shape(self, **kwargs) -> Union[Table, ShapeRecommendationReport]:
         AquaValueError
             If model type is unsupported by tool (no recommendation report generated)
         """
-        compartment_id = kwargs.get("compartment_id", COMPARTMENT_OCID)
-
-        kwargs["shapes"] = self.valid_compute_shapes(compartment_id=compartment_id)
+        try:
+            request = RequestRecommend(**kwargs)
+        except ValidationError as e:
+            custom_error = build_pydantic_error_message(e)
+            raise AquaValueError(  # noqa: B904
+                f"Failed to request shape recommendation due to invalid input parameters: {custom_error}"
+            )
 
         shape_recommend = AquaShapeRecommend()
-
-        shape_recommend_report = shape_recommend.which_shapes(**kwargs)
+        shape_recommend_report = shape_recommend.which_shapes(request)
 
         return shape_recommend_report
 

ads/aqua/shaperecommend/constants.py

@@ -67,3 +67,12 @@
     "4bit": 0.5,
     "int4": 0.5,
 }
+
+SHAPE_MAP = {
+    "NVIDIA_GPU": "GPU",
+    "AMD_ROME": "CPU",
+    "GENERIC": "CPU",
+    "LEGACY": "CPU",
+    "ARM": "CPU",
+    "UNKNOWN_ENUM_VALUE": "N/A",
+}

ads/aqua/shaperecommend/llm_config.py

@@ -212,7 +212,7 @@ def validate_model_support(cls, raw: dict) -> ValueError:
         ):
             raise AquaRecommendationError(
                 "Please provide a decoder-only text-generation model (ex. Llama, Falcon, etc). "
-                "Encoder-decoder models (ex. T5, Gemma) and encoder-only (BERT) are not supported in this tool at this time."
+                "Encoder-decoder models (ex. T5, Gemma) and encoder-only (BERT) are not supported at this time."
             )
 
     @classmethod

ads/aqua/shaperecommend/recommend.py

@@ -15,9 +15,11 @@
     build_pydantic_error_message,
     get_resource_type,
     load_config,
+    load_gpu_shapes_index,
 )
 from ads.aqua.shaperecommend.constants import (
     SAFETENSORS,
+    SHAPE_MAP,
     TEXT_GENERATION,
     TROUBLESHOOT_MSG,
 )
@@ -30,9 +32,12 @@
     ShapeReport,
 )
 from ads.model.datascience_model import DataScienceModel
+from ads.model.service.oci_datascience_model_deployment import (
+    OCIDataScienceModelDeployment,
+)
 
 
-class AquaShapeRecommend(BaseModel):
+class AquaShapeRecommend:
     """
     Interface for recommending GPU shapes for machine learning model deployments
     on Oracle Cloud Infrastructure Data Science service.
@@ -42,7 +47,7 @@ class AquaShapeRecommend(BaseModel):
     Must be used within a properly configured and authenticated OCI environment.
     """
 
-    def which_shapes(self, **kwargs) -> Union[ShapeRecommendationReport, Table]:
+    def which_shapes(self, request: RequestRecommend) -> Union[ShapeRecommendationReport, Table]:
         """
         Lists valid GPU deployment shapes for the provided model and configuration.
 
@@ -77,7 +82,8 @@ def which_shapes(self, **kwargs) -> Union[ShapeRecommendationReport, Table]:
             If parameters are missing or invalid, or if no valid sequence length is requested.
         """
         try:
-            request = RequestRecommend(**kwargs)
+            shapes = self.valid_compute_shapes(compartment_id=request.compartment_id)
+
             ds_model = self._validate_model_ocid(request.model_id)
             data = self._get_model_config(ds_model)
 
@@ -86,7 +92,7 @@ def which_shapes(self, **kwargs) -> Union[ShapeRecommendationReport, Table]:
             model_name = ds_model.display_name if ds_model.display_name else ""
 
             shape_recommendation_report = self._summarize_shapes_for_seq_lens(
-                llm_config, request.shapes, model_name
+                llm_config, shapes, model_name
             )
 
             if request.generate_table and shape_recommendation_report.recommendations:
@@ -107,10 +113,61 @@ def which_shapes(self, **kwargs) -> Union[ShapeRecommendationReport, Table]:
             ) from ex
         except AquaValueError as ex:
             logger.error(f"Error with LLM config: {ex}")
-            raise
+            raise AquaValueError(  # noqa: B904
+                f"An error occured while producing recommendations: {ex}"
+            )
 
         return shape_recommendation_report
 
+    def valid_compute_shapes(self, compartment_id: str) -> List["ComputeShapeSummary"]:
+        """
+        Returns a filtered list of GPU-only ComputeShapeSummary objects by reading and parsing a JSON file.
+
+        Parameters
+        ----------
+        file : str
+            Path to the JSON file containing shape data.
+
+        Returns
+        -------
+        List[ComputeShapeSummary]
+            List of ComputeShapeSummary objects passing the checks.
+
+        Raises
+        ------
+        ValueError
+            If the file cannot be opened, parsed, or the 'shapes' key is missing.
+        """
+        oci_shapes = OCIDataScienceModelDeployment.shapes(compartment_id=compartment_id)
+        set_user_shapes = {shape.name: shape for shape in oci_shapes}
+
+        gpu_shapes_metadata = load_gpu_shapes_index().shapes
+
+        valid_shapes = []
+        # only loops through GPU shapes, update later to include CPU shapes
+        for name, spec in gpu_shapes_metadata.items():
+            if name in set_user_shapes:
+                oci_shape = set_user_shapes.get(name)
+
+                compute_shape = ComputeShapeSummary(
+                    available=True,
+                    core_count=oci_shape.core_count,
+                    memory_in_gbs=oci_shape.memory_in_gbs,
+                    shape_series=SHAPE_MAP.get(oci_shape.shape_series, "GPU"),
+                    name=oci_shape.name,
+                    gpu_specs=spec,
+                )
+            else:
+                compute_shape = ComputeShapeSummary(
+                    available=False, name=name, shape_series="GPU", gpu_specs=spec
+                )
+            valid_shapes.append(compute_shape)
+
+        valid_shapes.sort(
+            key=lambda shape: shape.gpu_specs.gpu_memory_in_gbs, reverse=True
+        )
+        return valid_shapes
+
     @staticmethod
     def _rich_diff_table(shape_report: ShapeRecommendationReport) -> Table:
         """
@@ -321,7 +378,7 @@ def _summarize_shapes_for_seq_lens(
         recommendations = []
 
         if not shapes:
-            raise ValueError(
+            raise AquaValueError(
                 "No GPU shapes were passed for recommendation. Ensure shape parsing succeeded."
             )
 

ads/aqua/shaperecommend/shape_report.py

@@ -20,9 +20,6 @@ class RequestRecommend(BaseModel):
     model_id: str = Field(
         ..., description="The OCID of the model to recommend feasible compute shapes."
     )
-    shapes: List[ComputeShapeSummary] = Field(
-        ..., description="The list of shapes on OCI."
-    )
     generate_table: Optional[bool] = (
         Field(
             True,

tests/unitary/with_extras/aqua/test_deployment_handler.py

@@ -284,3 +284,74 @@ def test_get_model_list(self, mock_get, mock_finish):
         mock_finish.side_effect = lambda x: x
         result = self.aqua_model_list_handler.get(model_id="test_model_id")
         mock_get.assert_called()
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from tornado.web import HTTPError
+
+from ads.aqua.extension.base_handler import AquaAPIhandler
+from ads.aqua.extension.errors import Errors
+from ads.aqua.extension.recommend_handler import AquaRecommendHandler
+
+
+@pytest.fixture
+def handler():
+    # Patch AquaAPIhandler.__init__ for unit test stubbing
+    AquaAPIhandler.__init__ = lambda self, *args, **kwargs: None
+    h = AquaRecommendHandler(MagicMock(), MagicMock())
+    h.finish = MagicMock()
+    h.request = MagicMock()
+    # Set required Tornado internal fields
+    h._headers = {}
+    h._write_buffer = []
+    return h
+
+
+def test_post_valid_input(monkeypatch, handler):
+    input_data = {"model_ocid": "ocid1.datasciencemodel.oc1.XYZ"}
+    expected = {"recommendations": ["VM.GPU.A10.1"], "troubleshoot": ""}
+
+    # Patch class on correct import path, so handler sees our fake implementation
+    class FakeAquaRecommendApp:
+        def which_gpu(self, **kwargs):
+            return expected
+
+    monkeypatch.setattr(
+        "ads.aqua.extension.recommend_handler.AquaRecommendApp", FakeAquaRecommendApp
+    )
+
+    handler.get_json_body = MagicMock(return_value=input_data)
+    handler.post()
+    handler.finish.assert_called_once_with(expected)
+
+
+def test_post_no_input(handler):
+    handler.get_json_body = MagicMock(return_value=None)
+    handler._headers = {}
+    handler._write_buffer = []
+    handler.write_error = MagicMock()
+    handler.post()
+    handler.write_error.assert_called_once()
+    exc_info = handler.write_error.call_args.kwargs.get("exc_info")
+    assert exc_info is not None
+    exc_type, exc_value, _ = exc_info
+    assert exc_type is HTTPError
+    assert exc_value.status_code == 400
+    assert exc_value.log_message == Errors.NO_INPUT_DATA
+
+
+def test_post_invalid_input(handler):
+    handler.get_json_body = MagicMock(side_effect=Exception("bad input"))
+    handler._headers = {}
+    handler._write_buffer = []
+    handler.write_error = MagicMock()
+    handler.post()
+    handler.write_error.assert_called_once()
+    exc_info = handler.write_error.call_args.kwargs.get("exc_info")
+    assert exc_info is not None
+    exc_type, exc_value, _ = exc_info
+    assert exc_type is HTTPError
+    assert exc_value.status_code == 400
+    assert exc_value.log_message == Errors.INVALID_INPUT_DATA_FORMAT
+