added black formatting

Author: Aryanag2

ads/aqua/shaperecommend/recommend.py

@@ -48,6 +48,7 @@
     OCIDataScienceModelDeployment,
 )
 
+
 class HuggingFaceModelFetcher:
     """
     Utility class to fetch model configurations from HuggingFace.
@@ -57,7 +58,7 @@ class HuggingFaceModelFetcher:
     def is_huggingface_model_id(cls, model_id: str) -> bool:
         if is_valid_ocid(model_id):
             return False
-        hf_pattern = r'^[a-zA-Z0-9_-]+(/[a-zA-Z0-9_.-]+)?$'
+        hf_pattern = r"^[a-zA-Z0-9_-]+(/[a-zA-Z0-9_.-]+)?$"
         return bool(re.match(hf_pattern, model_id))
 
     @classmethod
@@ -80,12 +81,19 @@ def fetch_config_only(cls, model_id: str) -> Dict[str, Any]:
             elif response.status_code == 404:
                 raise AquaValueError(f"Model '{model_id}' not found on HuggingFace.")
             elif response.status_code != 200:
-                raise AquaValueError(f"Failed to fetch config for '{model_id}'. Status: {response.status_code}")
+                raise AquaValueError(
+                    f"Failed to fetch config for '{model_id}'. Status: {response.status_code}"
+                )
             return response.json()
         except requests.RequestException as e:
-            raise AquaValueError(f"Network error fetching config for {model_id}: {e}") from e
+            raise AquaValueError(
+                f"Network error fetching config for {model_id}: {e}"
+            ) from e
         except json.JSONDecodeError as e:
-            raise AquaValueError(f"Invalid config format for model '{model_id}'.") from e
+            raise AquaValueError(
+                f"Invalid config format for model '{model_id}'."
+            ) from e
+
 
 class AquaShapeRecommend:
     """
@@ -135,7 +143,9 @@ def which_shapes(
         """
         try:
             shapes = self.valid_compute_shapes(compartment_id=request.compartment_id)
-            data, model_name = self._get_model_config_and_name(request.model_id, request.compartment_id)
+            data, model_name = self._get_model_config_and_name(
+                request.model_id, request.compartment_id
+            )
             llm_config = LLMConfig.from_raw_config(data)
             shape_recommendation_report = self._summarize_shapes_for_seq_lens(
                 llm_config, shapes, model_name
@@ -165,40 +175,55 @@ def which_shapes(
 
         return shape_recommendation_report
 
-    def _get_model_config_and_name(self, model_id: str, compartment_id: str) -> (dict, str):
+    def _get_model_config_and_name(
+        self, model_id: str, compartment_id: str
+    ) -> (dict, str):
         """
         Loads model configuration, handling OCID and Hugging Face model IDs.
         """
         if HuggingFaceModelFetcher.is_huggingface_model_id(model_id):
             logger.info(f"'{model_id}' identified as a Hugging Face model ID.")
             ds_model = self._search_model_in_catalog(model_id, compartment_id)
             if ds_model and ds_model.artifact:
-                logger.info("Loading configuration from existing model catalog artifact.")
+                logger.info(
+                    "Loading configuration from existing model catalog artifact."
+                )
                 try:
-                    return load_config(ds_model.artifact, "config.json"), ds_model.display_name
+                    return (
+                        load_config(ds_model.artifact, "config.json"),
+                        ds_model.display_name,
+                    )
                 except AquaFileNotFoundError:
-                    logger.warning("config.json not found in artifact, fetching from Hugging Face Hub.")
+                    logger.warning(
+                        "config.json not found in artifact, fetching from Hugging Face Hub."
+                    )
             return HuggingFaceModelFetcher.fetch_config_only(model_id), model_id
         else:
             logger.info(f"'{model_id}' identified as a model OCID.")
             ds_model = self._validate_model_ocid(model_id)
             return self._get_model_config(ds_model), ds_model.display_name
 
-    def _search_model_in_catalog(self, model_id: str, compartment_id: str) -> Optional[DataScienceModel]:
+    def _search_model_in_catalog(
+        self, model_id: str, compartment_id: str
+    ) -> Optional[DataScienceModel]:
         """
         Searches for a Hugging Face model in the Data Science model catalog by display name.
         """
         try:
             # This should work since the SDK's list method can filter by display_name.
-            models = DataScienceModel.list(compartment_id=compartment_id, display_name=model_id)
+            models = DataScienceModel.list(
+                compartment_id=compartment_id, display_name=model_id
+            )
             if models:
                 logger.info(f"Found model '{model_id}' in the Data Science catalog.")
                 return models[0]
         except Exception as e:
             logger.warning(f"Could not search for model '{model_id}' in catalog: {e}")
         return None
 
-    def valid_compute_shapes(self, compartment_id: Optional[str] = None) -> List["ComputeShapeSummary"]:
+    def valid_compute_shapes(
+        self, compartment_id: Optional[str] = None
+    ) -> List["ComputeShapeSummary"]:
         """
         Returns a filtered list of GPU-only ComputeShapeSummary objects by reading and parsing a JSON file.
 
@@ -219,7 +244,9 @@ def valid_compute_shapes(self, compartment_id: Optional[str] = None) -> List["Co
             environment variables.
         """
         if not compartment_id:
-            compartment_id = os.environ.get("NB_SESSION_COMPARTMENT_OCID") or os.environ.get("PROJECT_COMPARTMENT_OCID")
+            compartment_id = os.environ.get(
+                "NB_SESSION_COMPARTMENT_OCID"
+            ) or os.environ.get("PROJECT_COMPARTMENT_OCID")
             if compartment_id:
                 logger.info(f"Using compartment_id from environment: {compartment_id}")
 

ads/aqua/shaperecommend/shape_report.py

@@ -18,7 +18,8 @@ class RequestRecommend(BaseModel):
     """
 
     model_id: str = Field(
-        ..., description="The OCID or Hugging Face ID of the model to recommend feasible compute shapes."
+        ...,
+        description="The OCID or Hugging Face ID of the model to recommend feasible compute shapes.",
     )
     generate_table: Optional[bool] = (
         Field(

tests/unitary/with_extras/aqua/test_recommend.py

@@ -20,7 +20,10 @@
     get_estimator,
 )
 from ads.aqua.shaperecommend.llm_config import LLMConfig
-from ads.aqua.shaperecommend.recommend import AquaShapeRecommend, HuggingFaceModelFetcher
+from ads.aqua.shaperecommend.recommend import (
+    AquaShapeRecommend,
+    HuggingFaceModelFetcher,
+)
 from ads.aqua.shaperecommend.shape_report import (
     DeploymentParams,
     ModelConfig,
@@ -455,16 +458,20 @@ def test_shape_report_pareto_front(self):
         assert a and b not in pf
         assert len(pf) == 2
 
+
 class TestHuggingFaceModelFetcher:
-    @pytest.mark.parametrize("model_id, expected", [
-        ("meta-llama/Llama-2-7b-hf", True),
-        ("mistralai/Mistral-7B-v0.1", True),
-        ("ocid1.datasciencemodel.oc1.iad.xxxxxxxx", False),
-    ])
+    @pytest.mark.parametrize(
+        "model_id, expected",
+        [
+            ("meta-llama/Llama-2-7b-hf", True),
+            ("mistralai/Mistral-7B-v0.1", True),
+            ("ocid1.datasciencemodel.oc1.iad.xxxxxxxx", False),
+        ],
+    )
     def test_is_huggingface_model_id(self, model_id, expected):
         assert HuggingFaceModelFetcher.is_huggingface_model_id(model_id) == expected
 
-    @patch('requests.get')
+    @patch("requests.get")
     def test_fetch_config_only_success(self, mock_get):
         mock_response = MagicMock()
         mock_response.status_code = 200
@@ -475,7 +482,7 @@ def test_fetch_config_only_success(self, mock_get):
         assert config == {"model_type": "llama"}
         mock_get.assert_called_once()
 
-    @patch('requests.get')
+    @patch("requests.get")
     def test_fetch_config_only_not_found(self, mock_get):
         mock_response = MagicMock()
         mock_response.status_code = 404
@@ -487,7 +494,6 @@ def test_fetch_config_only_not_found(self, mock_get):
     @patch.dict(os.environ, {"HF_TOKEN": "test_token_123"}, clear=True)
     def test_get_hf_token(self):
         assert HuggingFaceModelFetcher.get_hf_token() == "test_token_123"
-        
 
     # @pytest.mark.network
     # def test_fetch_config_only_real_call_success(self):
@@ -496,27 +502,26 @@ def test_get_hf_token(self):
     #     This test requires an internet connection.
     #     """
     #     model_id = "distilbert-base-uncased"
-        
+
     #     try:
     #         config = HuggingFaceModelFetcher.fetch_config_only(model_id)
     #         assert isinstance(config, dict)
     #         assert "model_type" in config
     #         assert "dim" in config
     #     except AquaValueError as e:
     #         pytest.fail(f"Real network call to Hugging Face failed: {e}")
-            
-            
-    @patch('ads.aqua.shaperecommend.recommend.OCIDataScienceModelDeployment.shapes')
+
+    @patch("ads.aqua.shaperecommend.recommend.OCIDataScienceModelDeployment.shapes")
     @patch.dict(os.environ, {}, clear=True)
     def test_valid_compute_shapes_raises_error_no_compartment(self, mock_oci_shapes):
         """
         Tests that valid_compute_shapes raises a ValueError when no compartment ID is
         provided and none can be found in the environment.
         """
         app = AquaShapeRecommend()
-        
+
         with pytest.raises(AquaValueError, match="A compartment OCID is required"):
             app.valid_compute_shapes(compartment_id=None)
-        
+
         # Verify that the OCI SDK was not called because the check failed early
-        mock_oci_shapes.assert_not_called()
+        mock_oci_shapes.assert_not_called()