diff --git a/agentplatform/_genai/model_garden.py b/agentplatform/_genai/model_garden.py index 662284507f..a0538087e9 100644 --- a/agentplatform/_genai/model_garden.py +++ b/agentplatform/_genai/model_garden.py @@ -119,6 +119,20 @@ def _ListPublisherModelsRequestParameters_to_vertex( return to_object +def _RecommendSpecRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["parent"]) is not None: + setv(to_object, ["_url", "parent"], getv(from_object, ["parent"])) + + if getv(from_object, ["gcs_uri"]) is not None: + setv(to_object, ["gcsUri"], getv(from_object, ["gcs_uri"])) + + return to_object + + class ModelGarden(_api_module.BaseModule): """Model Garden module.""" @@ -267,6 +281,80 @@ def _get_publisher_model( self._api_client._verify_response(return_value) return return_value + def _recommend_spec( + self, + *, + parent: str, + gcs_uri: str, + config: Optional[types.RecommendSpecConfigOrDict] = None, + ) -> types.RecommendSpecResponse: + """ + Recommends spec for a custom model (internal). + """ + + parameter_model = types._RecommendSpecRequestParameters( + parent=parent, + gcs_uri=gcs_uri, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RecommendSpecRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}:recommendSpec".format_map(request_url_dict) + else: + path = "{parent}:recommendSpec" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.RecommendSpecResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + @staticmethod def _build_filter_str( model_filter: Optional[str], @@ -771,6 +859,139 @@ def list_publisher_model_deploy_options( return options + @staticmethod + def _extract_recommend_spec(spec) -> dict[str, Any]: + """Extracts machine spec fields from a single recommend-spec entry. + + Args: + spec: A ``RecommendSpecResponseMachineAndModelContainerSpec`` describing a + recommended machine and container configuration. + + Returns: + A dict with the ``machine_type``, ``accelerator_type`` and + ``accelerator_count`` of the spec (values may be ``None``). + """ + machine_spec = spec.machine_spec + return { + "machine_type": getattr(machine_spec, "machine_type", None), + "accelerator_type": getattr( + getattr(machine_spec, "accelerator_type", None), "name", None + ), + "accelerator_count": getattr(machine_spec, "accelerator_count", None), + } + + @staticmethod + def _extract_recommendation(recommendation) -> dict[str, Any]: + """Extracts the spec, region and user quota state from a recommendation. + + Args: + recommendation: A ``RecommendSpecResponseRecommendation`` returned when + machine availability is requested. + + Returns: + A dict with the machine spec fields plus ``region`` and, when known, the + ``user_quota_state``. + """ + extracted_spec = ModelGarden._extract_recommend_spec(recommendation.spec) + extracted_spec["region"] = getattr(recommendation, "region", None) + if ( + recommendation.user_quota_state + and recommendation.user_quota_state + != types.QuotaState.QUOTA_STATE_UNSPECIFIED + ): + extracted_spec["user_quota_state"] = getattr( + getattr(recommendation, "user_quota_state", None), "name", None + ) + return extracted_spec + + @staticmethod + def _format_custom_deploy_options(options: list[dict[str, Any]]) -> str: + """Formats custom model deploy options into a human-readable string. + + Mirrors the legacy ``vertexai.model_garden`` ``CustomModel`` SDK output: + each option is rendered as an ``[Option N]`` block followed by its non-null + fields; ``accelerator_count`` is rendered unquoted. + + Args: + options: The extracted deploy option dicts to format. + + Returns: + A human-readable, multi-line string describing the deploy options. + """ + return "\n\n".join( + f"[Option {i + 1}]\n" + + ",\n".join( + f' {k}="{v}"' if k != "accelerator_count" else f" {k}={v}" + for k, v in option.items() + if v is not None + ) + for i, option in enumerate(options) + ) + + def list_custom_model_deploy_options( + self, + src: str, + config: Optional[types.ListCustomModelDeployOptionsConfigOrDict] = None, + ) -> str: + """Lists the recommended deploy options for a Model Garden custom model. + + Args: + src: The Google Cloud Storage URI of the custom model, storing the model + weights and config files (e.g. ``'gs://my-bucket/weights/'``). + config: Optional configuration. Accepts a + ``ListCustomModelDeployOptionsConfig`` instance or an equivalent dict. + + Returns: + A human-readable string describing the recommended deploy options + (machine type, accelerator type/count, region and, when available, the + user quota state). + + Raises: + ValueError: If ``src`` is not specified. + """ + if not src: + raise ValueError("src must be specified.") + if config is None: + config = types.ListCustomModelDeployOptionsConfig() + if isinstance(config, dict): + config = types.ListCustomModelDeployOptionsConfig.model_validate(config) + + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + + api_config = types.RecommendSpecConfig( + check_machine_availability=config.available_machines, + check_user_quota=config.filter_by_user_quota, + ) + + response = self._recommend_spec( + parent=parent, + gcs_uri=src, + config=api_config, + ) + + options = [] + if response.recommendations: + options = [ + self._extract_recommendation(recommendation) + for recommendation in response.recommendations + if recommendation.spec + ] + if config.filter_by_user_quota: + options = [ + option + for option in options + if option.get("user_quota_state") == "QUOTA_STATE_USER_HAS_QUOTA" + ] + elif response.specs: + options = [ + self._extract_recommend_spec(spec) for spec in response.specs if spec + ] + + return self._format_custom_deploy_options(options) + class AsyncModelGarden(_api_module.BaseModule): """Model Garden module.""" @@ -924,6 +1145,82 @@ async def _get_publisher_model( self._api_client._verify_response(return_value) return return_value + async def _recommend_spec( + self, + *, + parent: str, + gcs_uri: str, + config: Optional[types.RecommendSpecConfigOrDict] = None, + ) -> types.RecommendSpecResponse: + """ + Recommends spec for a custom model (internal). + """ + + parameter_model = types._RecommendSpecRequestParameters( + parent=parent, + gcs_uri=gcs_uri, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RecommendSpecRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}:recommendSpec".format_map(request_url_dict) + else: + path = "{parent}:recommendSpec" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.RecommendSpecResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + async def _list_all_publisher_models( self, api_config: types.ListPublisherModelsConfig, @@ -1108,3 +1405,69 @@ async def list_publisher_model_deploy_options( return ModelGarden._format_concise_deploy_options(options) return options + + async def list_custom_model_deploy_options( + self, + src: str, + config: Optional[types.ListCustomModelDeployOptionsConfigOrDict] = None, + ) -> str: + """Lists the recommended deploy options for a Model Garden custom model. + + Args: + src: The Google Cloud Storage URI of the custom model, storing the model + weights and config files (e.g. ``'gs://my-bucket/weights/'``). + config: Optional configuration. Accepts a + ``ListCustomModelDeployOptionsConfig`` instance or an equivalent dict. + + Returns: + A human-readable string describing the recommended deploy options + (machine type, accelerator type/count, region and, when available, the + user quota state). + + Raises: + ValueError: If ``src`` is not specified. + """ + if not src: + raise ValueError("src must be specified.") + if config is None: + config = types.ListCustomModelDeployOptionsConfig() + if isinstance(config, dict): + config = types.ListCustomModelDeployOptionsConfig.model_validate(config) + + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + + api_config = types.RecommendSpecConfig( + check_machine_availability=config.available_machines, + check_user_quota=config.filter_by_user_quota, + ) + + response = await self._recommend_spec( + parent=parent, + gcs_uri=src, + config=api_config, + ) + + options = [] + if response.recommendations: + options = [ + ModelGarden._extract_recommendation(recommendation) + for recommendation in response.recommendations + if recommendation.spec + ] + if config.filter_by_user_quota: + options = [ + option + for option in options + if option.get("user_quota_state") == "QUOTA_STATE_USER_HAS_QUOTA" + ] + elif response.specs: + options = [ + ModelGarden._extract_recommend_spec(spec) + for spec in response.specs + if spec + ] + + return ModelGarden._format_custom_deploy_options(options) diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 9bb4570d53..9c9866750c 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -133,6 +133,7 @@ from .common import _PurgeAgentEngineMemoriesRequestParameters from .common import _QueryAgentEngineRequestParameters from .common import _QueryAgentEngineRuntimeRevisionRequestParameters +from .common import _RecommendSpecRequestParameters from .common import _RestoreVersionRequestParameters from .common import _RetrieveAgentEngineMemoriesRequestParameters from .common import _RetrieveMemoryProfilesRequestParameters @@ -840,6 +841,9 @@ from .common import ListAgentEngineTasksResponse from .common import ListAgentEngineTasksResponseDict from .common import ListAgentEngineTasksResponseOrDict +from .common import ListCustomModelDeployOptionsConfig +from .common import ListCustomModelDeployOptionsConfigDict +from .common import ListCustomModelDeployOptionsConfigOrDict from .common import ListDatasetsResponse from .common import ListDatasetsResponseDict from .common import ListDatasetsResponseOrDict @@ -1199,6 +1203,7 @@ from .common import QueryReasoningEngineResponse from .common import QueryReasoningEngineResponseDict from .common import QueryReasoningEngineResponseOrDict +from .common import QuotaState from .common import RagContexts from .common import RagContextsContext from .common import RagContextsContextDict @@ -1420,6 +1425,18 @@ from .common import ReasoningEngineTrafficConfigTrafficSplitManualTarget from .common import ReasoningEngineTrafficConfigTrafficSplitManualTargetDict from .common import ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict +from .common import RecommendSpecConfig +from .common import RecommendSpecConfigDict +from .common import RecommendSpecConfigOrDict +from .common import RecommendSpecResponse +from .common import RecommendSpecResponseDict +from .common import RecommendSpecResponseMachineAndModelContainerSpec +from .common import RecommendSpecResponseMachineAndModelContainerSpecDict +from .common import RecommendSpecResponseMachineAndModelContainerSpecOrDict +from .common import RecommendSpecResponseOrDict +from .common import RecommendSpecResponseRecommendation +from .common import RecommendSpecResponseRecommendationDict +from .common import RecommendSpecResponseRecommendationOrDict from .common import RedTeamingAnalysisConfig from .common import RedTeamingAnalysisConfigDict from .common import RedTeamingAnalysisConfigOrDict @@ -3378,6 +3395,18 @@ "GetPublisherModelConfig", "GetPublisherModelConfigDict", "GetPublisherModelConfigOrDict", + "RecommendSpecConfig", + "RecommendSpecConfigDict", + "RecommendSpecConfigOrDict", + "RecommendSpecResponseMachineAndModelContainerSpec", + "RecommendSpecResponseMachineAndModelContainerSpecDict", + "RecommendSpecResponseMachineAndModelContainerSpecOrDict", + "RecommendSpecResponseRecommendation", + "RecommendSpecResponseRecommendationDict", + "RecommendSpecResponseRecommendationOrDict", + "RecommendSpecResponse", + "RecommendSpecResponseDict", + "RecommendSpecResponseOrDict", "PromptOptimizerConfig", "PromptOptimizerConfigDict", "PromptOptimizerConfigOrDict", @@ -3480,6 +3509,9 @@ "ListPublisherModelDeployOptionsConfig", "ListPublisherModelDeployOptionsConfigDict", "ListPublisherModelDeployOptionsConfigOrDict", + "ListCustomModelDeployOptionsConfig", + "ListCustomModelDeployOptionsConfigDict", + "ListCustomModelDeployOptionsConfigOrDict", "DeployOption", "DeployOptionDict", "DeployOptionOrDict", @@ -3508,6 +3540,7 @@ "LaunchStage", "OpenSourceCategory", "VersionState", + "QuotaState", "EvaluationItemType", "SamplingMethod", "EvaluationRunState", @@ -3668,6 +3701,7 @@ "_ListSkillRevisionsRequestParameters", "_ListPublisherModelsRequestParameters", "_GetPublisherModelRequestParameters", + "_RecommendSpecRequestParameters", "evals", "agent_engines", "prompts", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 8c1e7f688b..9732fea1ec 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -476,6 +476,17 @@ class VersionState(_common.CaseInSensitiveEnum): """Used to indicate the version is unstable.""" +class QuotaState(_common.CaseInSensitiveEnum): + """Output only. The user accelerator quota state.""" + + QUOTA_STATE_UNSPECIFIED = "QUOTA_STATE_UNSPECIFIED" + """Unspecified quota state. Quota information not available.""" + QUOTA_STATE_USER_HAS_QUOTA = "QUOTA_STATE_USER_HAS_QUOTA" + """User has enough accelerator quota for the machine type.""" + QUOTA_STATE_NO_USER_QUOTA = "QUOTA_STATE_NO_USER_QUOTA" + """User does not have enough accelerator quota for the machine type.""" + + class EvaluationItemType(_common.CaseInSensitiveEnum): """The type of the EvaluationItem.""" @@ -23558,6 +23569,151 @@ class _GetPublisherModelRequestParametersDict(TypedDict, total=False): ] +class RecommendSpecConfig(_common.BaseModel): + """Config for recommending spec.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + check_machine_availability: Optional[bool] = Field(default=None, description="""""") + check_user_quota: Optional[bool] = Field(default=None, description="""""") + + +class RecommendSpecConfigDict(TypedDict, total=False): + """Config for recommending spec.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + check_machine_availability: Optional[bool] + """""" + + check_user_quota: Optional[bool] + """""" + + +RecommendSpecConfigOrDict = Union[RecommendSpecConfig, RecommendSpecConfigDict] + + +class _RecommendSpecRequestParameters(_common.BaseModel): + """Parameters for recommending spec.""" + + parent: Optional[str] = Field(default=None, description="""""") + gcs_uri: Optional[str] = Field(default=None, description="""""") + config: Optional[RecommendSpecConfig] = Field(default=None, description="""""") + + +class _RecommendSpecRequestParametersDict(TypedDict, total=False): + """Parameters for recommending spec.""" + + parent: Optional[str] + """""" + + gcs_uri: Optional[str] + """""" + + config: Optional[RecommendSpecConfigDict] + """""" + + +_RecommendSpecRequestParametersOrDict = Union[ + _RecommendSpecRequestParameters, _RecommendSpecRequestParametersDict +] + + +class RecommendSpecResponseMachineAndModelContainerSpec(_common.BaseModel): + """A machine and model container spec.""" + + container_spec: Optional[ModelContainerSpec] = Field( + default=None, description="""Output only. The model container spec.""" + ) + machine_spec: Optional[MachineSpec] = Field( + default=None, description="""Output only. The machine spec.""" + ) + + +class RecommendSpecResponseMachineAndModelContainerSpecDict(TypedDict, total=False): + """A machine and model container spec.""" + + container_spec: Optional[ModelContainerSpecDict] + """Output only. The model container spec.""" + + machine_spec: Optional[MachineSpecDict] + """Output only. The machine spec.""" + + +RecommendSpecResponseMachineAndModelContainerSpecOrDict = Union[ + RecommendSpecResponseMachineAndModelContainerSpec, + RecommendSpecResponseMachineAndModelContainerSpecDict, +] + + +class RecommendSpecResponseRecommendation(_common.BaseModel): + """Recommendation of one deployment option for the given custom weights model in one region. Contains the machine and container spec, and user accelerator quota state.""" + + region: Optional[str] = Field( + default=None, description="""The region for the deployment spec (machine).""" + ) + spec: Optional[RecommendSpecResponseMachineAndModelContainerSpec] = Field( + default=None, + description="""Output only. The machine and model container specs.""", + ) + user_quota_state: Optional[QuotaState] = Field( + default=None, description="""Output only. The user accelerator quota state.""" + ) + + +class RecommendSpecResponseRecommendationDict(TypedDict, total=False): + """Recommendation of one deployment option for the given custom weights model in one region. Contains the machine and container spec, and user accelerator quota state.""" + + region: Optional[str] + """The region for the deployment spec (machine).""" + + spec: Optional[RecommendSpecResponseMachineAndModelContainerSpecDict] + """Output only. The machine and model container specs.""" + + user_quota_state: Optional[QuotaState] + """Output only. The user accelerator quota state.""" + + +RecommendSpecResponseRecommendationOrDict = Union[ + RecommendSpecResponseRecommendation, RecommendSpecResponseRecommendationDict +] + + +class RecommendSpecResponse(_common.BaseModel): + """Response for recommending spec.""" + + base_model: Optional[str] = Field( + default=None, + description="""Output only. The base model used to finetune the custom model.""", + ) + recommendations: Optional[list[RecommendSpecResponseRecommendation]] = Field( + default=None, + description="""Output only. Recommendations of deployment options for the given custom weights model.""", + ) + specs: Optional[list[RecommendSpecResponseMachineAndModelContainerSpec]] = Field( + default=None, + description="""Output only. The machine and model container specs.""", + ) + + +class RecommendSpecResponseDict(TypedDict, total=False): + """Response for recommending spec.""" + + base_model: Optional[str] + """Output only. The base model used to finetune the custom model.""" + + recommendations: Optional[list[RecommendSpecResponseRecommendationDict]] + """Output only. Recommendations of deployment options for the given custom weights model.""" + + specs: Optional[list[RecommendSpecResponseMachineAndModelContainerSpecDict]] + """Output only. The machine and model container specs.""" + + +RecommendSpecResponseOrDict = Union[RecommendSpecResponse, RecommendSpecResponseDict] + + class PromptOptimizerConfig(_common.BaseModel): """VAPO Prompt Optimizer Config.""" @@ -25666,6 +25822,32 @@ class ListPublisherModelDeployOptionsConfigDict(TypedDict, total=False): ] +class ListCustomModelDeployOptionsConfig(_common.BaseModel): + """Config for listing custom model deploy options.""" + + available_machines: Optional[bool] = Field( + default=True, description="""Whether to check machine availability.""" + ) + filter_by_user_quota: Optional[bool] = Field( + default=True, description="""Whether to filter by user quota.""" + ) + + +class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing custom model deploy options.""" + + available_machines: Optional[bool] + """Whether to check machine availability.""" + + filter_by_user_quota: Optional[bool] + """Whether to filter by user quota.""" + + +ListCustomModelDeployOptionsConfigOrDict = Union[ + ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict +] + + class DeployOption(_common.BaseModel): """A verified deploy option for a model.""" diff --git a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py index 9cebf7d6bb..cf5d8f5f15 100644 --- a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py @@ -102,6 +102,22 @@ def test_list_publisher_model_deploy_options_no_deploy_support(client): ) +def test_list_custom_model_deploy_options(client): + """Tests listing the recommended deploy options for a custom model. + + Exercises the RecommendSpec backend (distinct from the GetPublisherModel + path used by publisher models) and the human-readable string return type. + """ + options = client.model_garden.list_custom_model_deploy_options( + src="gs://gshuoy-mg-custom-model-replay/gemma3-1b", + config=types.ListCustomModelDeployOptionsConfig( + filter_by_user_quota=False, + ), + ) + assert isinstance(options, str) + assert "[Option 1" in options + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), @@ -147,3 +163,16 @@ async def test_list_publisher_model_deploy_options_async(client): assert len(options) > 0 assert isinstance(options[0], types.DeployOption) assert options[0].serving_container_image_uri + + +@pytest.mark.asyncio +async def test_list_custom_model_deploy_options_async(client): + """Tests listing the recommended deploy options for a custom model async.""" + options = await client.aio.model_garden.list_custom_model_deploy_options( + src="gs://gshuoy-mg-custom-model-replay/gemma3-1b", + config=types.ListCustomModelDeployOptionsConfig( + filter_by_user_quota=False, + ), + ) + assert isinstance(options, str) + assert "[Option 1" in options diff --git a/tests/unit/agentplatform/genai/test_genai_model_garden.py b/tests/unit/agentplatform/genai/test_genai_model_garden.py index 17b6fd937c..2790a7a1cd 100644 --- a/tests/unit/agentplatform/genai/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/test_genai_model_garden.py @@ -1129,3 +1129,223 @@ def test_list_publisher_model_deploy_options_concise_async(mock_async_client): assert isinstance(result, str) assert result.startswith("[Option 1: option-1]") + + +# ---- list_custom_model_deploy_options tests ---- + + +def test_list_custom_model_deploy_options_specs_path(mock_client): + """Tests the specs branch is extracted and formatted like the legacy SDK.""" + dummy_response = types.RecommendSpecResponse( + specs=[ + types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec( + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + ) + ) + ] + ) + + with mock.patch.object( + mock_client, "_recommend_spec", return_value=dummy_response + ): + result = mock_client.list_custom_model_deploy_options(src="gs://weights") + + expected = ( + "[Option 1]\n" + ' machine_type="g2-standard-12",\n' + ' accelerator_type="NVIDIA_L4",\n' + " accelerator_count=1" + ) + assert result == expected + + +def test_list_custom_model_deploy_options_request_config(mock_client): + """Tests the parent, gcs_uri and RecommendSpecConfig are passed through.""" + dummy_response = types.RecommendSpecResponse(specs=[]) + + with mock.patch.object( + mock_client, "_recommend_spec", return_value=dummy_response + ) as mock_recommend: + mock_client.list_custom_model_deploy_options(src="gs://weights") + + mock_recommend.assert_called_once_with( + parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", + gcs_uri="gs://weights", + config=types.RecommendSpecConfig( + check_machine_availability=True, + check_user_quota=True, + ), + ) + + +def test_list_custom_model_deploy_options_config_flags_forwarded(mock_client): + """Tests available_machines/filter_by_user_quota map to the API config.""" + dummy_response = types.RecommendSpecResponse(specs=[]) + + with mock.patch.object( + mock_client, "_recommend_spec", return_value=dummy_response + ) as mock_recommend: + mock_client.list_custom_model_deploy_options( + src="gs://weights", + config=types.ListCustomModelDeployOptionsConfig( + available_machines=False, + filter_by_user_quota=False, + ), + ) + + _, kwargs = mock_recommend.call_args + assert kwargs["config"] == types.RecommendSpecConfig( + check_machine_availability=False, + check_user_quota=False, + ) + + +def test_list_custom_model_deploy_options_recommendations_quota_filter( + mock_client, +): + """Tests recommendations are filtered to QUOTA_STATE_USER_HAS_QUOTA by default.""" + dummy_response = types.RecommendSpecResponse( + recommendations=[ + types.RecommendSpecResponseRecommendation( + region="us-central1", + spec=types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec( + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + ) + ), + user_quota_state="QUOTA_STATE_USER_HAS_QUOTA", + ), + types.RecommendSpecResponseRecommendation( + region="us-west1", + spec=types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec( + machine_type="a3-highgpu-8g", + accelerator_type="NVIDIA_H100_80GB", + accelerator_count=8, + ) + ), + user_quota_state="QUOTA_STATE_NO_USER_QUOTA", + ), + ] + ) + + with mock.patch.object( + mock_client, "_recommend_spec", return_value=dummy_response + ): + result = mock_client.list_custom_model_deploy_options(src="gs://weights") + + # Only the recommendation the user has quota for is kept, and region and + # user_quota_state are surfaced (legacy parity). + expected = ( + "[Option 1]\n" + ' machine_type="g2-standard-12",\n' + ' accelerator_type="NVIDIA_L4",\n' + " accelerator_count=1,\n" + ' region="us-central1",\n' + ' user_quota_state="QUOTA_STATE_USER_HAS_QUOTA"' + ) + assert result == expected + + +def test_list_custom_model_deploy_options_recommendations_no_quota_filter( + mock_client, +): + """Tests filter_by_user_quota=False keeps every recommendation.""" + dummy_response = types.RecommendSpecResponse( + recommendations=[ + types.RecommendSpecResponseRecommendation( + region="us-central1", + spec=types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec(machine_type="g2-standard-12") + ), + user_quota_state="QUOTA_STATE_USER_HAS_QUOTA", + ), + types.RecommendSpecResponseRecommendation( + region="us-west1", + spec=types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec(machine_type="a3-highgpu-8g") + ), + user_quota_state="QUOTA_STATE_NO_USER_QUOTA", + ), + ] + ) + + with mock.patch.object( + mock_client, "_recommend_spec", return_value=dummy_response + ): + result = mock_client.list_custom_model_deploy_options( + src="gs://weights", + config=types.ListCustomModelDeployOptionsConfig( + filter_by_user_quota=False + ), + ) + + assert "[Option 1]" in result + assert "[Option 2]" in result + assert "us-west1" in result + + +def test_list_custom_model_deploy_options_src_required_raises(mock_client): + """Tests an empty src raises ValueError, matching the legacy SDK guard.""" + with pytest.raises(ValueError, match="src must be specified"): + mock_client.list_custom_model_deploy_options(src="") + + +def test_list_custom_model_deploy_options_dict_config(mock_client): + """Tests a dict config is validated into ListCustomModelDeployOptionsConfig.""" + dummy_response = types.RecommendSpecResponse(specs=[]) + + with mock.patch.object( + mock_client, "_recommend_spec", return_value=dummy_response + ) as mock_recommend: + mock_client.list_custom_model_deploy_options( + src="gs://weights", + config={"available_machines": False, "filter_by_user_quota": False}, + ) + + _, kwargs = mock_recommend.call_args + assert kwargs["config"] == types.RecommendSpecConfig( + check_machine_availability=False, + check_user_quota=False, + ) + + +def test_list_custom_model_deploy_options_async(mock_async_client): + """Tests the async client returns the formatted deploy options string.""" + dummy_response = types.RecommendSpecResponse( + specs=[ + types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec( + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + ) + ) + ] + ) + + with mock.patch.object( + mock_async_client, + "_recommend_spec", + new=mock.AsyncMock(return_value=dummy_response), + ): + result = asyncio.run( + mock_async_client.list_custom_model_deploy_options(src="gs://weights") + ) + + assert isinstance(result, str) + assert result.startswith("[Option 1]") + assert 'machine_type="g2-standard-12"' in result + + +def test_list_custom_model_deploy_options_async_src_required_raises( + mock_async_client, +): + """Tests the async client also validates src.""" + with pytest.raises(ValueError, match="src must be specified"): + asyncio.run(mock_async_client.list_custom_model_deploy_options(src=""))