From a1c0e1267aabe6cc654ac92320c4119d9765f490 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Mon, 8 Sep 2025 17:02:40 -0700 Subject: [PATCH 01/24] Add Hugging Face model support to Shape Recommender --- ads/aqua/shaperecommend/recommend.py | 86 +++++++++++++++++-- ads/aqua/shaperecommend/shape_report.py | 2 +- .../with_extras/aqua/test_recommend.py | 59 ++++++++++++- 3 files changed, 135 insertions(+), 12 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 1b93598e0..83dee909b 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -3,7 +3,11 @@ # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ import shutil -from typing import List, Union +import os +import re +import json +import requests +from typing import List, Union, Optional, Dict, Any from pydantic import ValidationError from rich.table import Table @@ -42,6 +46,45 @@ OCIDataScienceModelDeployment, ) +class HuggingFaceModelFetcher: + """ + Utility class to fetch model configurations from HuggingFace. + """ + HUGGINGFACE_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json" + + @classmethod + def is_huggingface_model_id(cls, model_id: str) -> bool: + if model_id.startswith("ocid1."): + return False + hf_pattern = r'^[a-zA-Z0-9_-]+(/[a-zA-Z0-9_.-]+)?$' + return bool(re.match(hf_pattern, model_id)) + + @classmethod + def get_hf_token(cls) -> Optional[str]: + return os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_TOKEN") + + @classmethod + def fetch_config_only(cls, model_id: str) -> Dict[str, Any]: + try: + config_url = cls.HUGGINGFACE_CONFIG_URL.format(model_id=model_id) + headers = {} + token = cls.get_hf_token() + if token: + headers["Authorization"] = f"Bearer {token}" + response = requests.get(config_url, headers=headers, timeout=10) + if response.status_code == 401: + raise AquaValueError( + f"Model '{model_id}' requires authentication. Please set your HuggingFace token." + ) + 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}") + return response.json() + except requests.RequestException as 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 class AquaShapeRecommend: """ @@ -91,14 +134,8 @@ def which_shapes( """ try: 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) - + data, model_name = self._get_model_config_and_name(request.model_id, request.compartment_id) 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 ) @@ -127,6 +164,39 @@ def which_shapes( return shape_recommendation_report + 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.") + try: + 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.") + 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]: + """ + 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) + 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: str) -> List["ComputeShapeSummary"]: """ Returns a filtered list of GPU-only ComputeShapeSummary objects by reading and parsing a JSON file. diff --git a/ads/aqua/shaperecommend/shape_report.py b/ads/aqua/shaperecommend/shape_report.py index e3d9854f0..c845d7c68 100644 --- a/ads/aqua/shaperecommend/shape_report.py +++ b/ads/aqua/shaperecommend/shape_report.py @@ -18,7 +18,7 @@ class RequestRecommend(BaseModel): """ model_id: str = Field( - ..., description="The OCID 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( diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index cb61dae86..fb481e4ac 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -8,11 +8,11 @@ import os import re from unittest.mock import MagicMock - +from unittest.mock import patch import pytest from ads.aqua.common.entities import ComputeShapeSummary -from ads.aqua.common.errors import AquaRecommendationError +from ads.aqua.common.errors import AquaRecommendationError, AquaValueError from ads.aqua.shaperecommend.estimator import ( LlamaMemoryEstimator, MemoryEstimator, @@ -20,7 +20,7 @@ get_estimator, ) from ads.aqua.shaperecommend.llm_config import LLMConfig -from ads.aqua.shaperecommend.recommend import AquaShapeRecommend +from ads.aqua.shaperecommend.recommend import AquaShapeRecommend, HuggingFaceModelFetcher from ads.aqua.shaperecommend.shape_report import ( DeploymentParams, ModelConfig, @@ -454,3 +454,56 @@ def test_shape_report_pareto_front(self): assert c and d in pf assert a and b not in pf assert len(pf) == 2 +# --- Tests for HuggingFaceModelFetcher --- +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), + ]) + def test_is_huggingface_model_id(self, model_id, expected): + assert HuggingFaceModelFetcher.is_huggingface_model_id(model_id) == expected + + @patch('requests.get') + def test_fetch_config_only_success(self, mock_get): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"model_type": "llama"} + mock_get.return_value = mock_response + + config = HuggingFaceModelFetcher.fetch_config_only("some/model") + assert config == {"model_type": "llama"} + mock_get.assert_called_once() + + @patch('requests.get') + def test_fetch_config_only_not_found(self, mock_get): + mock_response = MagicMock() + mock_response.status_code = 404 + mock_get.return_value = mock_response + + with pytest.raises(AquaValueError, match="not found on HuggingFace"): + HuggingFaceModelFetcher.fetch_config_only("non/existent") + + @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" + + # Add this method inside the TestHuggingFaceModelFetcher class + + @pytest.mark.network + def test_fetch_config_only_real_call_success(self): + """ + Tests a real network call to fetch a public model's configuration. + This test requires an internet connection. + """ + # Use a well-known, small, public model to minimize test flakiness + model_id = "distilbert-base-uncased" + + try: + config = HuggingFaceModelFetcher.fetch_config_only(model_id) + # Assert that we got a dictionary with expected keys + 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}") \ No newline at end of file From 2126fd035705e8c29696cb8ae755f437a723f980 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Mon, 8 Sep 2025 17:05:45 -0700 Subject: [PATCH 02/24] Add Hugging Face model support to Shape Recommender --- tests/unitary/with_extras/aqua/test_recommend.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index fb481e4ac..7e0231d8d 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -454,7 +454,7 @@ def test_shape_report_pareto_front(self): assert c and d in pf assert a and b not in pf assert len(pf) == 2 -# --- Tests for HuggingFaceModelFetcher --- + class TestHuggingFaceModelFetcher: @pytest.mark.parametrize("model_id, expected", [ ("meta-llama/Llama-2-7b-hf", True), @@ -488,7 +488,6 @@ def test_fetch_config_only_not_found(self, mock_get): def test_get_hf_token(self): assert HuggingFaceModelFetcher.get_hf_token() == "test_token_123" - # Add this method inside the TestHuggingFaceModelFetcher class @pytest.mark.network def test_fetch_config_only_real_call_success(self): @@ -496,12 +495,10 @@ def test_fetch_config_only_real_call_success(self): Tests a real network call to fetch a public model's configuration. This test requires an internet connection. """ - # Use a well-known, small, public model to minimize test flakiness model_id = "distilbert-base-uncased" try: config = HuggingFaceModelFetcher.fetch_config_only(model_id) - # Assert that we got a dictionary with expected keys assert isinstance(config, dict) assert "model_type" in config assert "dim" in config From b24938e79f8d5eeadc8e8035e5b71d481dffa58f Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Tue, 9 Sep 2025 14:20:38 -0700 Subject: [PATCH 03/24] added HUGGINGFACE_CONFIG_URL and is_valid_ocid, and fixed error message --- ads/aqua/shaperecommend/constants.py | 2 ++ ads/aqua/shaperecommend/recommend.py | 9 +++++---- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/ads/aqua/shaperecommend/constants.py b/ads/aqua/shaperecommend/constants.py index 08f2f2133..ee8cf52c2 100644 --- a/ads/aqua/shaperecommend/constants.py +++ b/ads/aqua/shaperecommend/constants.py @@ -114,3 +114,5 @@ "ARM": "CPU", "UNKNOWN_ENUM_VALUE": "N/A", } + +HUGGINGFACE_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json" diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 83dee909b..c3e0bf7ac 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -24,6 +24,7 @@ get_resource_type, load_config, load_gpu_shapes_index, + is_valid_ocid, ) from ads.aqua.shaperecommend.constants import ( BITS_AND_BYTES_4BIT, @@ -32,6 +33,7 @@ SHAPE_MAP, TEXT_GENERATION, TROUBLESHOOT_MSG, + HUGGINGFACE_CONFIG_URL, ) from ads.aqua.shaperecommend.estimator import get_estimator from ads.aqua.shaperecommend.llm_config import LLMConfig @@ -50,11 +52,10 @@ class HuggingFaceModelFetcher: """ Utility class to fetch model configurations from HuggingFace. """ - HUGGINGFACE_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json" @classmethod def is_huggingface_model_id(cls, model_id: str) -> bool: - if model_id.startswith("ocid1."): + if is_valid_ocid(model_id): return False hf_pattern = r'^[a-zA-Z0-9_-]+(/[a-zA-Z0-9_.-]+)?$' return bool(re.match(hf_pattern, model_id)) @@ -66,7 +67,7 @@ def get_hf_token(cls) -> Optional[str]: @classmethod def fetch_config_only(cls, model_id: str) -> Dict[str, Any]: try: - config_url = cls.HUGGINGFACE_CONFIG_URL.format(model_id=model_id) + config_url = HUGGINGFACE_CONFIG_URL.format(model_id=model_id) headers = {} token = cls.get_hf_token() if token: @@ -74,7 +75,7 @@ def fetch_config_only(cls, model_id: str) -> Dict[str, Any]: response = requests.get(config_url, headers=headers, timeout=10) if response.status_code == 401: raise AquaValueError( - f"Model '{model_id}' requires authentication. Please set your HuggingFace token." + f"Model '{model_id}' requires authentication. Please set your HuggingFace access token as an environment variable." ) elif response.status_code == 404: raise AquaValueError(f"Model '{model_id}' not found on HuggingFace.") From ff37df964f57114e3a2bc297b89548e728c3f3bb Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 11:58:39 -0700 Subject: [PATCH 04/24] added support for no compartment id provided and unit test for the same --- ads/aqua/shaperecommend/recommend.py | 19 ++++++++++++++++--- .../with_extras/aqua/test_recommend.py | 18 +++++++++++++++++- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index c3e0bf7ac..bd7c55157 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -198,7 +198,7 @@ def _search_model_in_catalog(self, model_id: str, compartment_id: str) -> Option logger.warning(f"Could not search for model '{model_id}' in catalog: {e}") return None - def valid_compute_shapes(self, compartment_id: str) -> 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. @@ -214,9 +214,22 @@ def valid_compute_shapes(self, compartment_id: str) -> List["ComputeShapeSummary Raises ------ - ValueError - If the file cannot be opened, parsed, or the 'shapes' key is missing. + AquaValueError + If a compartment_id is not provided and cannot be found in the + environment variables. """ + if not compartment_id: + 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}") + + if not compartment_id: + raise AquaValueError( + "A compartment OCID is required to list available shapes. " + "Please provide it as a parameter or set the 'NB_SESSION_COMPARTMENT_OCID' " + "or 'PROJECT_COMPARTMENT_OCID' environment variable." + ) + oci_shapes = OCIDataScienceModelDeployment.shapes(compartment_id=compartment_id) set_user_shapes = {shape.name: shape for shape in oci_shapes} diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index 7e0231d8d..a32122571 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -503,4 +503,20 @@ def test_fetch_config_only_real_call_success(self): assert "model_type" in config assert "dim" in config except AquaValueError as e: - pytest.fail(f"Real network call to Hugging Face failed: {e}") \ No newline at end of file + pytest.fail(f"Real network call to Hugging Face failed: {e}") + + + @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() \ No newline at end of file From 4c2c6a2db5986c3d206152aa3f8af48e06121817 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Mon, 8 Sep 2025 17:02:40 -0700 Subject: [PATCH 05/24] Add Hugging Face model support to Shape Recommender --- ads/aqua/shaperecommend/recommend.py | 86 +++++++++++++++++-- ads/aqua/shaperecommend/shape_report.py | 2 +- .../with_extras/aqua/test_recommend.py | 59 ++++++++++++- 3 files changed, 135 insertions(+), 12 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 1b93598e0..83dee909b 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -3,7 +3,11 @@ # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ import shutil -from typing import List, Union +import os +import re +import json +import requests +from typing import List, Union, Optional, Dict, Any from pydantic import ValidationError from rich.table import Table @@ -42,6 +46,45 @@ OCIDataScienceModelDeployment, ) +class HuggingFaceModelFetcher: + """ + Utility class to fetch model configurations from HuggingFace. + """ + HUGGINGFACE_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json" + + @classmethod + def is_huggingface_model_id(cls, model_id: str) -> bool: + if model_id.startswith("ocid1."): + return False + hf_pattern = r'^[a-zA-Z0-9_-]+(/[a-zA-Z0-9_.-]+)?$' + return bool(re.match(hf_pattern, model_id)) + + @classmethod + def get_hf_token(cls) -> Optional[str]: + return os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_TOKEN") + + @classmethod + def fetch_config_only(cls, model_id: str) -> Dict[str, Any]: + try: + config_url = cls.HUGGINGFACE_CONFIG_URL.format(model_id=model_id) + headers = {} + token = cls.get_hf_token() + if token: + headers["Authorization"] = f"Bearer {token}" + response = requests.get(config_url, headers=headers, timeout=10) + if response.status_code == 401: + raise AquaValueError( + f"Model '{model_id}' requires authentication. Please set your HuggingFace token." + ) + 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}") + return response.json() + except requests.RequestException as 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 class AquaShapeRecommend: """ @@ -91,14 +134,8 @@ def which_shapes( """ try: 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) - + data, model_name = self._get_model_config_and_name(request.model_id, request.compartment_id) 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 ) @@ -127,6 +164,39 @@ def which_shapes( return shape_recommendation_report + 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.") + try: + 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.") + 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]: + """ + 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) + 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: str) -> List["ComputeShapeSummary"]: """ Returns a filtered list of GPU-only ComputeShapeSummary objects by reading and parsing a JSON file. diff --git a/ads/aqua/shaperecommend/shape_report.py b/ads/aqua/shaperecommend/shape_report.py index e3d9854f0..c845d7c68 100644 --- a/ads/aqua/shaperecommend/shape_report.py +++ b/ads/aqua/shaperecommend/shape_report.py @@ -18,7 +18,7 @@ class RequestRecommend(BaseModel): """ model_id: str = Field( - ..., description="The OCID 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( diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index cb61dae86..fb481e4ac 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -8,11 +8,11 @@ import os import re from unittest.mock import MagicMock - +from unittest.mock import patch import pytest from ads.aqua.common.entities import ComputeShapeSummary -from ads.aqua.common.errors import AquaRecommendationError +from ads.aqua.common.errors import AquaRecommendationError, AquaValueError from ads.aqua.shaperecommend.estimator import ( LlamaMemoryEstimator, MemoryEstimator, @@ -20,7 +20,7 @@ get_estimator, ) from ads.aqua.shaperecommend.llm_config import LLMConfig -from ads.aqua.shaperecommend.recommend import AquaShapeRecommend +from ads.aqua.shaperecommend.recommend import AquaShapeRecommend, HuggingFaceModelFetcher from ads.aqua.shaperecommend.shape_report import ( DeploymentParams, ModelConfig, @@ -454,3 +454,56 @@ def test_shape_report_pareto_front(self): assert c and d in pf assert a and b not in pf assert len(pf) == 2 +# --- Tests for HuggingFaceModelFetcher --- +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), + ]) + def test_is_huggingface_model_id(self, model_id, expected): + assert HuggingFaceModelFetcher.is_huggingface_model_id(model_id) == expected + + @patch('requests.get') + def test_fetch_config_only_success(self, mock_get): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"model_type": "llama"} + mock_get.return_value = mock_response + + config = HuggingFaceModelFetcher.fetch_config_only("some/model") + assert config == {"model_type": "llama"} + mock_get.assert_called_once() + + @patch('requests.get') + def test_fetch_config_only_not_found(self, mock_get): + mock_response = MagicMock() + mock_response.status_code = 404 + mock_get.return_value = mock_response + + with pytest.raises(AquaValueError, match="not found on HuggingFace"): + HuggingFaceModelFetcher.fetch_config_only("non/existent") + + @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" + + # Add this method inside the TestHuggingFaceModelFetcher class + + @pytest.mark.network + def test_fetch_config_only_real_call_success(self): + """ + Tests a real network call to fetch a public model's configuration. + This test requires an internet connection. + """ + # Use a well-known, small, public model to minimize test flakiness + model_id = "distilbert-base-uncased" + + try: + config = HuggingFaceModelFetcher.fetch_config_only(model_id) + # Assert that we got a dictionary with expected keys + 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}") \ No newline at end of file From d2bf70967ea9032ad02631504ed5944807151472 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Mon, 8 Sep 2025 17:05:45 -0700 Subject: [PATCH 06/24] Add Hugging Face model support to Shape Recommender --- tests/unitary/with_extras/aqua/test_recommend.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index fb481e4ac..7e0231d8d 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -454,7 +454,7 @@ def test_shape_report_pareto_front(self): assert c and d in pf assert a and b not in pf assert len(pf) == 2 -# --- Tests for HuggingFaceModelFetcher --- + class TestHuggingFaceModelFetcher: @pytest.mark.parametrize("model_id, expected", [ ("meta-llama/Llama-2-7b-hf", True), @@ -488,7 +488,6 @@ def test_fetch_config_only_not_found(self, mock_get): def test_get_hf_token(self): assert HuggingFaceModelFetcher.get_hf_token() == "test_token_123" - # Add this method inside the TestHuggingFaceModelFetcher class @pytest.mark.network def test_fetch_config_only_real_call_success(self): @@ -496,12 +495,10 @@ def test_fetch_config_only_real_call_success(self): Tests a real network call to fetch a public model's configuration. This test requires an internet connection. """ - # Use a well-known, small, public model to minimize test flakiness model_id = "distilbert-base-uncased" try: config = HuggingFaceModelFetcher.fetch_config_only(model_id) - # Assert that we got a dictionary with expected keys assert isinstance(config, dict) assert "model_type" in config assert "dim" in config From e78ca08c3516d0ea6efba24963824456647ca0ae Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Tue, 9 Sep 2025 14:20:38 -0700 Subject: [PATCH 07/24] added HUGGINGFACE_CONFIG_URL and is_valid_ocid, and fixed error message --- ads/aqua/shaperecommend/constants.py | 2 ++ ads/aqua/shaperecommend/recommend.py | 9 +++++---- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/ads/aqua/shaperecommend/constants.py b/ads/aqua/shaperecommend/constants.py index 08f2f2133..ee8cf52c2 100644 --- a/ads/aqua/shaperecommend/constants.py +++ b/ads/aqua/shaperecommend/constants.py @@ -114,3 +114,5 @@ "ARM": "CPU", "UNKNOWN_ENUM_VALUE": "N/A", } + +HUGGINGFACE_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json" diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 83dee909b..c3e0bf7ac 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -24,6 +24,7 @@ get_resource_type, load_config, load_gpu_shapes_index, + is_valid_ocid, ) from ads.aqua.shaperecommend.constants import ( BITS_AND_BYTES_4BIT, @@ -32,6 +33,7 @@ SHAPE_MAP, TEXT_GENERATION, TROUBLESHOOT_MSG, + HUGGINGFACE_CONFIG_URL, ) from ads.aqua.shaperecommend.estimator import get_estimator from ads.aqua.shaperecommend.llm_config import LLMConfig @@ -50,11 +52,10 @@ class HuggingFaceModelFetcher: """ Utility class to fetch model configurations from HuggingFace. """ - HUGGINGFACE_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json" @classmethod def is_huggingface_model_id(cls, model_id: str) -> bool: - if model_id.startswith("ocid1."): + if is_valid_ocid(model_id): return False hf_pattern = r'^[a-zA-Z0-9_-]+(/[a-zA-Z0-9_.-]+)?$' return bool(re.match(hf_pattern, model_id)) @@ -66,7 +67,7 @@ def get_hf_token(cls) -> Optional[str]: @classmethod def fetch_config_only(cls, model_id: str) -> Dict[str, Any]: try: - config_url = cls.HUGGINGFACE_CONFIG_URL.format(model_id=model_id) + config_url = HUGGINGFACE_CONFIG_URL.format(model_id=model_id) headers = {} token = cls.get_hf_token() if token: @@ -74,7 +75,7 @@ def fetch_config_only(cls, model_id: str) -> Dict[str, Any]: response = requests.get(config_url, headers=headers, timeout=10) if response.status_code == 401: raise AquaValueError( - f"Model '{model_id}' requires authentication. Please set your HuggingFace token." + f"Model '{model_id}' requires authentication. Please set your HuggingFace access token as an environment variable." ) elif response.status_code == 404: raise AquaValueError(f"Model '{model_id}' not found on HuggingFace.") From 8554da2a925f1f115caef2bcd12360c8325f89e2 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 11:58:39 -0700 Subject: [PATCH 08/24] added support for no compartment id provided and unit test for the same --- ads/aqua/shaperecommend/recommend.py | 19 ++++++++++++++++--- .../with_extras/aqua/test_recommend.py | 18 +++++++++++++++++- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index c3e0bf7ac..bd7c55157 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -198,7 +198,7 @@ def _search_model_in_catalog(self, model_id: str, compartment_id: str) -> Option logger.warning(f"Could not search for model '{model_id}' in catalog: {e}") return None - def valid_compute_shapes(self, compartment_id: str) -> 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. @@ -214,9 +214,22 @@ def valid_compute_shapes(self, compartment_id: str) -> List["ComputeShapeSummary Raises ------ - ValueError - If the file cannot be opened, parsed, or the 'shapes' key is missing. + AquaValueError + If a compartment_id is not provided and cannot be found in the + environment variables. """ + if not compartment_id: + 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}") + + if not compartment_id: + raise AquaValueError( + "A compartment OCID is required to list available shapes. " + "Please provide it as a parameter or set the 'NB_SESSION_COMPARTMENT_OCID' " + "or 'PROJECT_COMPARTMENT_OCID' environment variable." + ) + oci_shapes = OCIDataScienceModelDeployment.shapes(compartment_id=compartment_id) set_user_shapes = {shape.name: shape for shape in oci_shapes} diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index 7e0231d8d..a32122571 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -503,4 +503,20 @@ def test_fetch_config_only_real_call_success(self): assert "model_type" in config assert "dim" in config except AquaValueError as e: - pytest.fail(f"Real network call to Hugging Face failed: {e}") \ No newline at end of file + pytest.fail(f"Real network call to Hugging Face failed: {e}") + + + @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() \ No newline at end of file From 367ce703da5f12f6c8b6f0f9e58cc54efac08698 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 15:23:33 -0700 Subject: [PATCH 09/24] removing real network call --- .../with_extras/aqua/test_recommend.py | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index a32122571..8d410fda0 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -489,21 +489,21 @@ 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): - """ - Tests a real network call to fetch a public model's configuration. - This test requires an internet connection. - """ - model_id = "distilbert-base-uncased" + # @pytest.mark.network + # def test_fetch_config_only_real_call_success(self): + # """ + # Tests a real network call to fetch a public model's configuration. + # 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}") + # 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') From 29bd22cf8d9cac85a54a02197e97348624e46a50 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 15:44:14 -0700 Subject: [PATCH 10/24] added black formatting --- ads/aqua/shaperecommend/recommend.py | 53 ++++++++++++++----- ads/aqua/shaperecommend/shape_report.py | 3 +- .../with_extras/aqua/test_recommend.py | 37 +++++++------ 3 files changed, 63 insertions(+), 30 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index bd7c55157..78508c90a 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/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,7 +175,9 @@ 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. """ @@ -173,24 +185,35 @@ def _get_model_config_and_name(self, model_id: str, compartment_id: str) -> (dic 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] @@ -198,7 +221,9 @@ def _search_model_in_catalog(self, model_id: str, compartment_id: str) -> Option 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}") diff --git a/ads/aqua/shaperecommend/shape_report.py b/ads/aqua/shaperecommend/shape_report.py index c845d7c68..92b733330 100644 --- a/ads/aqua/shaperecommend/shape_report.py +++ b/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( diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index 8d410fda0..2d7058c05 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/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,7 +502,7 @@ 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) @@ -504,9 +510,8 @@ def test_get_hf_token(self): # 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): """ @@ -514,9 +519,9 @@ def test_valid_compute_shapes_raises_error_no_compartment(self, mock_oci_shapes) 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() \ No newline at end of file + mock_oci_shapes.assert_not_called() From efe09535a8f0dd5b1a102b3022353521cda39773 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 17:16:56 -0700 Subject: [PATCH 11/24] fixed comments and using _get_model_config to do the same checks as when model ocid is provided instead of directly load_config --- ads/aqua/shaperecommend/recommend.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 78508c90a..f0380c047 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -79,7 +79,9 @@ def fetch_config_only(cls, model_id: str) -> Dict[str, Any]: f"Model '{model_id}' requires authentication. Please set your HuggingFace access token as an environment variable." ) elif response.status_code == 404: - raise AquaValueError(f"Model '{model_id}' not found on HuggingFace.") + raise AquaValueError( + f"Model '{model_id}' not found on HuggingFace. Please check the name for typos." + ) elif response.status_code != 200: raise AquaValueError( f"Failed to fetch config for '{model_id}'. Status: {response.status_code}" @@ -184,13 +186,13 @@ def _get_model_config_and_name( 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: + if ds_model: logger.info( "Loading configuration from existing model catalog artifact." ) try: return ( - load_config(ds_model.artifact, "config.json"), + self._get_model_config(ds_model), ds_model.display_name, ) except AquaFileNotFoundError: @@ -207,7 +209,7 @@ 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. + Searches for a 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. From 19ddce117422b4ab9d9730c3caee772e101a7533 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 17:47:19 -0700 Subject: [PATCH 12/24] using huggingface_hub.hf_hub_download and design changes --- ads/aqua/shaperecommend/recommend.py | 108 +++++++++++---------------- 1 file changed, 45 insertions(+), 63 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index f0380c047..9dceed068 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -2,15 +2,18 @@ # Copyright (c) 2025 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ +#!/usr/bin/env python +# Copyright (c) 2025 Oracle and/or its affiliates. +# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ + import shutil import os -import re import json -import requests -from typing import List, Union, Optional, Dict, Any +from typing import List, Union, Optional, Dict, Any, Tuple from pydantic import ValidationError from rich.table import Table +from huggingface_hub import hf_hub_download, HfHubHTTPError from ads.aqua.app import logger from ads.aqua.common.entities import ComputeShapeSummary @@ -21,7 +24,6 @@ ) from ads.aqua.common.utils import ( build_pydantic_error_message, - get_resource_type, load_config, load_gpu_shapes_index, is_valid_ocid, @@ -33,7 +35,6 @@ SHAPE_MAP, TEXT_GENERATION, TROUBLESHOOT_MSG, - HUGGINGFACE_CONFIG_URL, ) from ads.aqua.shaperecommend.estimator import get_estimator from ads.aqua.shaperecommend.llm_config import LLMConfig @@ -49,54 +50,6 @@ ) -class HuggingFaceModelFetcher: - """ - Utility class to fetch model configurations from HuggingFace. - """ - - @classmethod - 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_.-]+)?$" - return bool(re.match(hf_pattern, model_id)) - - @classmethod - def get_hf_token(cls) -> Optional[str]: - return os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_TOKEN") - - @classmethod - def fetch_config_only(cls, model_id: str) -> Dict[str, Any]: - try: - config_url = HUGGINGFACE_CONFIG_URL.format(model_id=model_id) - headers = {} - token = cls.get_hf_token() - if token: - headers["Authorization"] = f"Bearer {token}" - response = requests.get(config_url, headers=headers, timeout=10) - if response.status_code == 401: - raise AquaValueError( - f"Model '{model_id}' requires authentication. Please set your HuggingFace access token as an environment variable." - ) - elif response.status_code == 404: - raise AquaValueError( - f"Model '{model_id}' not found on HuggingFace. Please check the name for typos." - ) - elif response.status_code != 200: - 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 - except json.JSONDecodeError as e: - raise AquaValueError( - f"Invalid config format for model '{model_id}'." - ) from e - - class AquaShapeRecommend: """ Interface for recommending GPU shapes for machine learning model deployments @@ -146,7 +99,7 @@ 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 + model_id=request.model_id, compartment_id=request.compartment_id ) llm_config = LLMConfig.from_raw_config(data) shape_recommendation_report = self._summarize_shapes_for_seq_lens( @@ -183,8 +136,15 @@ def _get_model_config_and_name( """ 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.") + if is_valid_ocid(model_id): + 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 + + logger.info( + f"'{model_id}' is not an OCID, treating as a Hugging Face model ID." + ) + if compartment_id: ds_model = self._search_model_in_catalog(model_id, compartment_id) if ds_model: logger.info( @@ -199,23 +159,45 @@ def _get_model_config_and_name( 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 + + return self._fetch_hf_config(model_id), model_id + + def _fetch_hf_config(self, model_id: str) -> Dict: + """ + Downloads a model's config.json from Hugging Face Hub using the + huggingface_hub library. + """ + try: + config_path = hf_hub_download(repo_id=model_id, filename="config.json") + with open(config_path, "r", encoding="utf-8") as f: + return json.load(f) + except HfHubHTTPError as e: + if "401" in str(e): + raise AquaValueError( + f"Model '{model_id}' requires authentication. Please set your HuggingFace access token as an environment variable (HF_TOKEN). cli command: export HF_TOKEN=" + ) + elif "404" in str(e) or "not found" in str(e).lower(): + raise AquaValueError( + f"Model '{model_id}' not found on HuggingFace. Please check the name for typos." + ) + raise AquaValueError( + f"Failed to download config for '{model_id}': {e}" + ) from e def _search_model_in_catalog( self, model_id: str, compartment_id: str ) -> Optional[DataScienceModel]: """ - Searches for a model in the Data Science model catalog by display name. + Searches for a model in the Data Science catalog by its 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 ) + if len(models) > 1: + logger.warning( + f"Found multiple models with the name '{model_id}'. Using the first one found." + ) if models: logger.info(f"Found model '{model_id}' in the Data Science catalog.") return models[0] From 3d935bb075a4add465ee60f3d55d1dec6992e65d Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 17:48:36 -0700 Subject: [PATCH 13/24] an example also how to provide the compartment in the parameters. --- ads/aqua/shaperecommend/recommend.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 9dceed068..3591a7a1f 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -239,6 +239,7 @@ def valid_compute_shapes( "A compartment OCID is required to list available shapes. " "Please provide it as a parameter or set the 'NB_SESSION_COMPARTMENT_OCID' " "or 'PROJECT_COMPARTMENT_OCID' environment variable." + "cli command: export NB_SESSION_COMPARTMENT_OCID=" ) oci_shapes = OCIDataScienceModelDeployment.shapes(compartment_id=compartment_id) From e3241430ffaeb831bba4f3600e4f5fdff0694ba9 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 17:57:51 -0700 Subject: [PATCH 14/24] added compartment id logic _get_model_config as a sanity check --- ads/aqua/shaperecommend/recommend.py | 41 ++++++++++++++++++---------- 1 file changed, 26 insertions(+), 15 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 3591a7a1f..6eb73851e 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -131,7 +131,7 @@ def which_shapes( return shape_recommendation_report def _get_model_config_and_name( - self, model_id: str, compartment_id: str + self, model_id: str, compartment_id: Optional[str] = None ) -> (dict, str): """ Loads model configuration, handling OCID and Hugging Face model IDs. @@ -144,21 +144,32 @@ def _get_model_config_and_name( logger.info( f"'{model_id}' is not an OCID, treating as a Hugging Face model ID." ) - if compartment_id: - ds_model = self._search_model_in_catalog(model_id, compartment_id) - if ds_model: - logger.info( - "Loading configuration from existing model catalog artifact." + if not compartment_id: + 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}") + if not compartment_id: + raise AquaValueError( + "A compartment OCID is required to list available shapes. " + "Please provide it as a parameter or set the 'NB_SESSION_COMPARTMENT_OCID' " + "or 'PROJECT_COMPARTMENT_OCID' environment variable." + "cli command: export NB_SESSION_COMPARTMENT_OCID=" + ) + + ds_model = self._search_model_in_catalog(model_id, compartment_id) + if ds_model: + logger.info("Loading configuration from existing model catalog artifact.") + try: + return ( + self._get_model_config(ds_model), + ds_model.display_name, + ) + except AquaFileNotFoundError: + logger.warning( + "config.json not found in artifact, fetching from Hugging Face Hub." ) - try: - return ( - self._get_model_config(ds_model), - ds_model.display_name, - ) - except AquaFileNotFoundError: - logger.warning( - "config.json not found in artifact, fetching from Hugging Face Hub." - ) return self._fetch_hf_config(model_id), model_id From 2251e3cdcb2e03989bc7a817fb0a76168f3c1961 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Wed, 10 Sep 2025 18:13:26 -0700 Subject: [PATCH 15/24] added docstrings --- ads/aqua/shaperecommend/recommend.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 6eb73851e..f07ac7138 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -131,10 +131,25 @@ def which_shapes( return shape_recommendation_report def _get_model_config_and_name( - self, model_id: str, compartment_id: Optional[str] = None - ) -> (dict, str): + self, model_id: str, compartment_id: Optional[str] + ) -> Tuple[Dict, str]: """ - Loads model configuration, handling OCID and Hugging Face model IDs. + Loads model configuration by trying OCID logic first, then falling back + to treating the model_id as a Hugging Face Hub ID. + + Parameters + ---------- + model_id : str + The model OCID or Hugging Face model ID. + compartment_id : Optional[str] + The compartment OCID, used for searching the model catalog. + + Returns + ------- + Tuple[Dict, str] + A tuple containing: + - The model configuration dictionary. + - The display name for the model. """ if is_valid_ocid(model_id): logger.info(f"'{model_id}' identified as a model OCID.") From c3b5afa70c0fd8c25b42759d223cf5917f1fe715 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Thu, 11 Sep 2025 12:05:44 -0700 Subject: [PATCH 16/24] commented out search_model_in_catalog logic --- ads/aqua/shaperecommend/recommend.py | 104 ++++++++++++++------------- 1 file changed, 54 insertions(+), 50 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index f07ac7138..014263542 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -98,8 +98,11 @@ def which_shapes( """ try: shapes = self.valid_compute_shapes(compartment_id=request.compartment_id) + # data, model_name = self._get_model_config_and_name( + # model_id=request.model_id, compartment_id=request.compartment_id + # ) data, model_name = self._get_model_config_and_name( - model_id=request.model_id, compartment_id=request.compartment_id + model_id=request.model_id, ) llm_config = LLMConfig.from_raw_config(data) shape_recommendation_report = self._summarize_shapes_for_seq_lens( @@ -131,7 +134,8 @@ def which_shapes( return shape_recommendation_report def _get_model_config_and_name( - self, model_id: str, compartment_id: Optional[str] + self, + model_id: str, ) -> Tuple[Dict, str]: """ Loads model configuration by trying OCID logic first, then falling back @@ -141,8 +145,8 @@ def _get_model_config_and_name( ---------- model_id : str The model OCID or Hugging Face model ID. - compartment_id : Optional[str] - The compartment OCID, used for searching the model catalog. + # compartment_id : Optional[str] + # The compartment OCID, used for searching the model catalog. Returns ------- @@ -159,32 +163,32 @@ def _get_model_config_and_name( logger.info( f"'{model_id}' is not an OCID, treating as a Hugging Face model ID." ) - if not compartment_id: - 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}") - if not compartment_id: - raise AquaValueError( - "A compartment OCID is required to list available shapes. " - "Please provide it as a parameter or set the 'NB_SESSION_COMPARTMENT_OCID' " - "or 'PROJECT_COMPARTMENT_OCID' environment variable." - "cli command: export NB_SESSION_COMPARTMENT_OCID=" - ) - - ds_model = self._search_model_in_catalog(model_id, compartment_id) - if ds_model: - logger.info("Loading configuration from existing model catalog artifact.") - try: - return ( - self._get_model_config(ds_model), - ds_model.display_name, - ) - except AquaFileNotFoundError: - logger.warning( - "config.json not found in artifact, fetching from Hugging Face Hub." - ) + # if not compartment_id: + # 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}") + # if not compartment_id: + # raise AquaValueError( + # "A compartment OCID is required to list available shapes. " + # "Please provide it as a parameter or set the 'NB_SESSION_COMPARTMENT_OCID' " + # "or 'PROJECT_COMPARTMENT_OCID' environment variable." + # "cli command: export NB_SESSION_COMPARTMENT_OCID=" + # ) + + # ds_model = self._search_model_in_catalog(model_id, compartment_id) + # if ds_model: + # logger.info("Loading configuration from existing model catalog artifact.") + # try: + # return ( + # self._get_model_config(ds_model), + # ds_model.display_name, + # ) + # except AquaFileNotFoundError: + # logger.warning( + # "config.json not found in artifact, fetching from Hugging Face Hub." + # ) return self._fetch_hf_config(model_id), model_id @@ -210,26 +214,26 @@ def _fetch_hf_config(self, model_id: str) -> Dict: f"Failed to download config for '{model_id}': {e}" ) from e - def _search_model_in_catalog( - self, model_id: str, compartment_id: str - ) -> Optional[DataScienceModel]: - """ - Searches for a model in the Data Science catalog by its display name. - """ - try: - models = DataScienceModel.list( - compartment_id=compartment_id, display_name=model_id - ) - if len(models) > 1: - logger.warning( - f"Found multiple models with the name '{model_id}'. Using the first one found." - ) - 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 _search_model_in_catalog( + # self, model_id: str, compartment_id: str + # ) -> Optional[DataScienceModel]: + # """ + # Searches for a model in the Data Science catalog by its display name. + # """ + # try: + # models = DataScienceModel.list( + # compartment_id=compartment_id, display_name=model_id + # ) + # if len(models) > 1: + # logger.warning( + # f"Found multiple models with the name '{model_id}'. Using the first one found." + # ) + # 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 From d68e501d3c35d341cd387b67ed7681a93b633626 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Thu, 11 Sep 2025 12:29:08 -0700 Subject: [PATCH 17/24] added get_resource_type as an import --- ads/aqua/shaperecommend/recommend.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 014263542..17517750a 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -27,6 +27,7 @@ load_config, load_gpu_shapes_index, is_valid_ocid, + get_resource_type, ) from ads.aqua.shaperecommend.constants import ( BITS_AND_BYTES_4BIT, From 79d6f5fda4f39d0451e631938e1dcf1860755f1c Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Thu, 11 Sep 2025 12:39:28 -0700 Subject: [PATCH 18/24] fixed imports --- ads/aqua/shaperecommend/recommend.py | 4 +- .../with_extras/aqua/test_recommend.py | 69 ------------------- 2 files changed, 3 insertions(+), 70 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 17517750a..d428e2f44 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -13,8 +13,10 @@ from pydantic import ValidationError from rich.table import Table -from huggingface_hub import hf_hub_download, HfHubHTTPError +# In ads/aqua/shaperecommend/recommend.py +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import HfHubHTTPError from ads.aqua.app import logger from ads.aqua.common.entities import ComputeShapeSummary from ads.aqua.common.errors import ( diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index 2d7058c05..37a141748 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -22,7 +22,6 @@ from ads.aqua.shaperecommend.llm_config import LLMConfig from ads.aqua.shaperecommend.recommend import ( AquaShapeRecommend, - HuggingFaceModelFetcher, ) from ads.aqua.shaperecommend.shape_report import ( DeploymentParams, @@ -457,71 +456,3 @@ def test_shape_report_pareto_front(self): assert c and d in pf 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), - ], - ) - def test_is_huggingface_model_id(self, model_id, expected): - assert HuggingFaceModelFetcher.is_huggingface_model_id(model_id) == expected - - @patch("requests.get") - def test_fetch_config_only_success(self, mock_get): - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = {"model_type": "llama"} - mock_get.return_value = mock_response - - config = HuggingFaceModelFetcher.fetch_config_only("some/model") - assert config == {"model_type": "llama"} - mock_get.assert_called_once() - - @patch("requests.get") - def test_fetch_config_only_not_found(self, mock_get): - mock_response = MagicMock() - mock_response.status_code = 404 - mock_get.return_value = mock_response - - with pytest.raises(AquaValueError, match="not found on HuggingFace"): - HuggingFaceModelFetcher.fetch_config_only("non/existent") - - @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): - # """ - # Tests a real network call to fetch a public model's configuration. - # 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.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() From a362351dd8e133c7d65a1b5579998a2c2788b4e1 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Thu, 11 Sep 2025 14:40:13 -0700 Subject: [PATCH 19/24] fixed imports --- ads/aqua/shaperecommend/recommend.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index d428e2f44..a712b317a 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -14,7 +14,6 @@ from pydantic import ValidationError from rich.table import Table -# In ads/aqua/shaperecommend/recommend.py from huggingface_hub import hf_hub_download from huggingface_hub.utils import HfHubHTTPError from ads.aqua.app import logger From 49962ff2d8d2c0a49332802e408f3f8cdf1c2f66 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Fri, 12 Sep 2025 15:13:27 -0700 Subject: [PATCH 20/24] resolving comments --- ads/aqua/shaperecommend/constants.py | 2 - ads/aqua/shaperecommend/recommend.py | 108 ++++++--------------------- 2 files changed, 24 insertions(+), 86 deletions(-) diff --git a/ads/aqua/shaperecommend/constants.py b/ads/aqua/shaperecommend/constants.py index ee8cf52c2..08f2f2133 100644 --- a/ads/aqua/shaperecommend/constants.py +++ b/ads/aqua/shaperecommend/constants.py @@ -114,5 +114,3 @@ "ARM": "CPU", "UNKNOWN_ENUM_VALUE": "N/A", } - -HUGGINGFACE_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json" diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index a712b317a..4c2625d16 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -2,20 +2,17 @@ # Copyright (c) 2025 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ -#!/usr/bin/env python -# Copyright (c) 2025 Oracle and/or its affiliates. -# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ - -import shutil -import os import json -from typing import List, Union, Optional, Dict, Any, Tuple +import os +import shutil +from typing import Dict, List, Optional, Tuple, Union from pydantic import ValidationError from rich.table import Table from huggingface_hub import hf_hub_download from huggingface_hub.utils import HfHubHTTPError + from ads.aqua.app import logger from ads.aqua.common.entities import ComputeShapeSummary from ads.aqua.common.errors import ( @@ -25,14 +22,15 @@ ) from ads.aqua.common.utils import ( build_pydantic_error_message, + get_resource_type, + is_valid_ocid, load_config, load_gpu_shapes_index, - is_valid_ocid, - get_resource_type, + format_hf_custom_error_message, ) from ads.aqua.shaperecommend.constants import ( - BITS_AND_BYTES_4BIT, BITSANDBYTES, + BITS_AND_BYTES_4BIT, SAFETENSORS, SHAPE_MAP, TEXT_GENERATION, @@ -46,6 +44,7 @@ ShapeRecommendationReport, ShapeReport, ) +from ads.config import COMPARTMENT_OCID from ads.model.datascience_model import DataScienceModel from ads.model.service.oci_datascience_model_deployment import ( OCIDataScienceModelDeployment, @@ -100,9 +99,6 @@ def which_shapes( """ try: shapes = self.valid_compute_shapes(compartment_id=request.compartment_id) - # data, model_name = self._get_model_config_and_name( - # model_id=request.model_id, compartment_id=request.compartment_id - # ) data, model_name = self._get_model_config_and_name( model_id=request.model_id, ) @@ -158,41 +154,18 @@ def _get_model_config_and_name( - The display name for the model. """ if is_valid_ocid(model_id): - logger.info(f"'{model_id}' identified as a model OCID.") + logger.info(f"Detected OCID: Fetching OCI model config for '{model_id}'.") ds_model = self._validate_model_ocid(model_id) - return self._get_model_config(ds_model), ds_model.display_name + config = self._fetch_hf_config(model_id) + model_name = ds_model.display_name + else: + logger.info( + f"Assuming Hugging Face model ID: Fetching config for '{model_id}'." + ) + config = self._fetch_hf_config(model_id) + model_name = model_id - logger.info( - f"'{model_id}' is not an OCID, treating as a Hugging Face model ID." - ) - # if not compartment_id: - # 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}") - # if not compartment_id: - # raise AquaValueError( - # "A compartment OCID is required to list available shapes. " - # "Please provide it as a parameter or set the 'NB_SESSION_COMPARTMENT_OCID' " - # "or 'PROJECT_COMPARTMENT_OCID' environment variable." - # "cli command: export NB_SESSION_COMPARTMENT_OCID=" - # ) - - # ds_model = self._search_model_in_catalog(model_id, compartment_id) - # if ds_model: - # logger.info("Loading configuration from existing model catalog artifact.") - # try: - # return ( - # self._get_model_config(ds_model), - # ds_model.display_name, - # ) - # except AquaFileNotFoundError: - # logger.warning( - # "config.json not found in artifact, fetching from Hugging Face Hub." - # ) - - return self._fetch_hf_config(model_id), model_id + return config, model_name def _fetch_hf_config(self, model_id: str) -> Dict: """ @@ -204,38 +177,7 @@ def _fetch_hf_config(self, model_id: str) -> Dict: with open(config_path, "r", encoding="utf-8") as f: return json.load(f) except HfHubHTTPError as e: - if "401" in str(e): - raise AquaValueError( - f"Model '{model_id}' requires authentication. Please set your HuggingFace access token as an environment variable (HF_TOKEN). cli command: export HF_TOKEN=" - ) - elif "404" in str(e) or "not found" in str(e).lower(): - raise AquaValueError( - f"Model '{model_id}' not found on HuggingFace. Please check the name for typos." - ) - raise AquaValueError( - f"Failed to download config for '{model_id}': {e}" - ) from e - - # def _search_model_in_catalog( - # self, model_id: str, compartment_id: str - # ) -> Optional[DataScienceModel]: - # """ - # Searches for a model in the Data Science catalog by its display name. - # """ - # try: - # models = DataScienceModel.list( - # compartment_id=compartment_id, display_name=model_id - # ) - # if len(models) > 1: - # logger.warning( - # f"Found multiple models with the name '{model_id}'. Using the first one found." - # ) - # 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 + format_hf_custom_error_message(e) def valid_compute_shapes( self, compartment_id: Optional[str] = None @@ -260,18 +202,16 @@ def valid_compute_shapes( environment variables. """ if not compartment_id: - compartment_id = os.environ.get( - "NB_SESSION_COMPARTMENT_OCID" - ) or os.environ.get("PROJECT_COMPARTMENT_OCID") + compartment_id = COMPARTMENT_OCID if compartment_id: logger.info(f"Using compartment_id from environment: {compartment_id}") if not compartment_id: raise AquaValueError( "A compartment OCID is required to list available shapes. " - "Please provide it as a parameter or set the 'NB_SESSION_COMPARTMENT_OCID' " - "or 'PROJECT_COMPARTMENT_OCID' environment variable." - "cli command: export NB_SESSION_COMPARTMENT_OCID=" + "Please specify it using the --compartment_id parameter.\n\n" + "Example:\n" + 'ads aqua deployment recommend_shape --model_id "" --compartment_id ""' ) oci_shapes = OCIDataScienceModelDeployment.shapes(compartment_id=compartment_id) From a1bde9f8e6fa2b8265f6f1dd482e8d186a4e5d52 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Tue, 16 Sep 2025 15:27:40 -0700 Subject: [PATCH 21/24] added changes to tests --- .../with_extras/aqua/test_recommend.py | 36 ++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index 37a141748..311c682b9 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -7,7 +7,7 @@ import json import os import re -from unittest.mock import MagicMock +from unittest.mock import MagicMock, mock_open from unittest.mock import patch import pytest @@ -277,6 +277,40 @@ def create(config_file=""): class TestAquaShapeRecommend: + + @patch("ads.aqua.shaperecommend.recommend.hf_hub_download") + @patch("builtins.open", new_callable=mock_open) + def test_fetch_hf_config_success(self, mock_file, mock_download): + """Test successful config fetch from Hugging Face""" + app = AquaShapeRecommend() + model_id = "test/model" + config_path = "/fake/path/config.json" + expected_config = {"model_type": "llama", "hidden_size": 4096} + + mock_download.return_value = config_path + mock_file.return_value.read.return_value = json.dumps(expected_config) + + result = app._fetch_hf_config(model_id) + + assert result == expected_config + mock_download.assert_called_once_with(repo_id=model_id, filename="config.json") + + @patch("ads.aqua.shaperecommend.recommend.hf_hub_download") + @patch("ads.aqua.shaperecommend.recommend.format_hf_custom_error_message") + def test_fetch_hf_config_http_error(self, mock_format_error, mock_download): + """Test error handling when Hugging Face request fails""" + from huggingface_hub.utils import HfHubHTTPError + + app = AquaShapeRecommend() + model_id = "nonexistent/model" + http_error = HfHubHTTPError("Model not found") + mock_download.side_effect = http_error + + with pytest.raises(HfHubHTTPError): + app._fetch_hf_config(model_id) + + mock_format_error.assert_called_once_with(http_error) + @pytest.mark.parametrize( "config, expected_recs, expected_troubleshoot", [ From 287b406ff656b5bb22516fa9499ce4302ba67023 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Tue, 16 Sep 2025 15:32:47 -0700 Subject: [PATCH 22/24] unittests for huggingface fetching --- tests/unitary/with_extras/aqua/test_recommend.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index 6a8dc9258..cec261b64 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -307,9 +307,10 @@ def test_fetch_hf_config_http_error(self, mock_format_error, mock_download): http_error = HfHubHTTPError("Model not found") mock_download.side_effect = http_error - with pytest.raises(HfHubHTTPError): - app._fetch_hf_config(model_id) + # The method doesn't re-raise, so it returns None + result = app._fetch_hf_config(model_id) + assert result is None mock_format_error.assert_called_once_with(http_error) @pytest.mark.parametrize( From 5088b8509a220cace9d7dba8eb23b0e7fa994981 Mon Sep 17 00:00:00 2001 From: Liz Johnson Date: Wed, 17 Sep 2025 13:49:12 -0700 Subject: [PATCH 23/24] modified unit tests --- ads/aqua/shaperecommend/recommend.py | 26 +++++++------------ .../with_extras/aqua/test_recommend.py | 7 ++--- 2 files changed, 13 insertions(+), 20 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index 89e153824..fa5128a83 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -4,14 +4,14 @@ import json import os +import re import shutil from typing import Dict, List, Optional, Tuple, Union -from pydantic import ValidationError -from rich.table import Table - from huggingface_hub import hf_hub_download from huggingface_hub.utils import HfHubHTTPError +from pydantic import ValidationError +from rich.table import Table from ads.aqua.app import logger from ads.aqua.common.entities import ComputeShapeSummary @@ -22,15 +22,15 @@ ) from ads.aqua.common.utils import ( build_pydantic_error_message, + format_hf_custom_error_message, get_resource_type, is_valid_ocid, load_config, load_gpu_shapes_index, - format_hf_custom_error_message, ) from ads.aqua.shaperecommend.constants import ( - BITSANDBYTES, BITS_AND_BYTES_4BIT, + BITSANDBYTES, SAFETENSORS, SHAPE_MAP, TEXT_GENERATION, @@ -98,14 +98,10 @@ def which_shapes( """ try: shapes = self.valid_compute_shapes(compartment_id=request.compartment_id) - + data, model_name = self._get_model_config_and_name( model_id=request.model_id, ) - llm_config = LLMConfig.from_raw_config(data) - shape_recommendation_report = self._summarize_shapes_for_seq_lens( - llm_config, shapes, model_name - ) if request.deployment_config: shape_recommendation_report = ( @@ -115,16 +111,11 @@ def which_shapes( ) else: - ds_model = self._get_data_science_model(request.model_id) - - data = self._get_model_config(ds_model) - llm_config = LLMConfig.from_raw_config(data) 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( @@ -174,8 +165,8 @@ def _get_model_config_and_name( """ if is_valid_ocid(model_id): logger.info(f"Detected OCID: Fetching OCI model config for '{model_id}'.") - ds_model = self._validate_model_ocid(model_id) - config = self._fetch_hf_config(model_id) + ds_model = self._get_data_science_model(model_id) + config = self._get_model_config(ds_model) model_name = ds_model.display_name else: logger.info( @@ -403,6 +394,7 @@ def _get_model_config(model: DataScienceModel): """ model_task = model.freeform_tags.get("task", "").lower() + model_task = re.sub(r"-", "_", model_task) model_format = model.freeform_tags.get("model_format", "").lower() logger.info(f"Current model task type: {model_task}") diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index cec261b64..1aa3712cc 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -436,10 +436,11 @@ def test_which_shapes_valid_from_file( )[1], ) - raw = load_config(config_file) + mock_raw_config = load_config(config_file) + mock_ds_model_name = mock_model.display_name if service_managed_model: - config = AquaDeploymentConfig(**raw) + config = AquaDeploymentConfig(**mock_raw_config) request = RequestRecommend( model_id="ocid1.datasciencemodel.oc1.TEST", @@ -447,7 +448,7 @@ def test_which_shapes_valid_from_file( deployment_config=config, ) else: - monkeypatch.setattr(app, "_get_model_config", lambda _: raw) + monkeypatch.setattr(app, "_get_model_config_and_name", lambda _: (mock_ds_model_name, mock_raw_config)) request = RequestRecommend( model_id="ocid1.datasciencemodel.oc1.TEST", generate_table=False From 2a9d84334e84f242fb987ac57a09927ae51ec718 Mon Sep 17 00:00:00 2001 From: Aryan Gosaliya Date: Thu, 18 Sep 2025 14:19:46 -0700 Subject: [PATCH 24/24] fixed unit tests --- ads/aqua/shaperecommend/recommend.py | 15 ++++++++++----- tests/unitary/with_extras/aqua/test_recommend.py | 6 +++++- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py index fa5128a83..e83a0aa50 100644 --- a/ads/aqua/shaperecommend/recommend.py +++ b/ads/aqua/shaperecommend/recommend.py @@ -99,11 +99,13 @@ def which_shapes( try: shapes = self.valid_compute_shapes(compartment_id=request.compartment_id) - data, model_name = self._get_model_config_and_name( - model_id=request.model_id, - ) - if request.deployment_config: + if is_valid_ocid(request.model_id): + ds_model = self._get_data_science_model(request.model_id) + model_name = ds_model.display_name + else: + model_name = request.model_id + shape_recommendation_report = ( ShapeRecommendationReport.from_deployment_config( request.deployment_config, model_name, shapes @@ -111,6 +113,9 @@ def which_shapes( ) else: + data, model_name = self._get_model_config_and_name( + model_id=request.model_id, + ) llm_config = LLMConfig.from_raw_config(data) shape_recommendation_report = self._summarize_shapes_for_seq_lens( @@ -394,7 +399,7 @@ def _get_model_config(model: DataScienceModel): """ model_task = model.freeform_tags.get("task", "").lower() - model_task = re.sub(r"-", "_", model_task) + model_task = re.sub(r"-", "_", model_task) model_format = model.freeform_tags.get("model_format", "").lower() logger.info(f"Current model task type: {model_task}") diff --git a/tests/unitary/with_extras/aqua/test_recommend.py b/tests/unitary/with_extras/aqua/test_recommend.py index 1aa3712cc..33b2503a8 100644 --- a/tests/unitary/with_extras/aqua/test_recommend.py +++ b/tests/unitary/with_extras/aqua/test_recommend.py @@ -448,7 +448,11 @@ def test_which_shapes_valid_from_file( deployment_config=config, ) else: - monkeypatch.setattr(app, "_get_model_config_and_name", lambda _: (mock_ds_model_name, mock_raw_config)) + monkeypatch.setattr( + app, + "_get_model_config_and_name", + lambda model_id: (mock_raw_config, mock_ds_model_name), + ) request = RequestRecommend( model_id="ocid1.datasciencemodel.oc1.TEST", generate_table=False