From 22741fd1c4eed52dfaebc3f586b255000f350705 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 21 Jul 2026 09:12:31 -0700 Subject: [PATCH] feat: Onboard Vertex Model Garden to GenAI Python SDK: Add export_open_model support PiperOrigin-RevId: 951526334 --- agentplatform/_genai/model_garden.py | 466 +++++ agentplatform/_genai/types/__init__.py | 34 + agentplatform/_genai/types/common.py | 208 +++ .../genai/replays/test_genai_model_garden.py | 320 ++-- .../genai/test_genai_model_garden.py | 1575 ++++++++++------- 5 files changed, 1879 insertions(+), 724 deletions(-) diff --git a/agentplatform/_genai/model_garden.py b/agentplatform/_genai/model_garden.py index 511f958367..823cf46cfe 100644 --- a/agentplatform/_genai/model_garden.py +++ b/agentplatform/_genai/model_garden.py @@ -22,14 +22,58 @@ from google.genai import _api_module from google.genai import _common +from google.genai import types as genai_types from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv +from . import _operations_utils from . import types logger = logging.getLogger("agentplatform_genai.modelgarden") +def _ExportPublisherModelConfig_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, ["destination"]) is not None: + setv(parent_object, ["destination"], getv(from_object, ["destination"])) + + return to_object + + +def _ExportPublisherModelRequestParameters_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, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _ExportPublisherModelConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _GetExportPublisherModelOperationParameters_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, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) + + return to_object + + def _GetPublisherModelConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -377,6 +421,159 @@ def _recommend_spec( self._api_client._verify_response(return_value) return return_value + def _export_publisher_model( + self, + *, + parent: str, + name: str, + config: Optional[types.ExportPublisherModelConfigOrDict] = None, + ) -> types.ExportModelOperation: + """ + Exports a publisher model (internal). + """ + + parameter_model = types._ExportPublisherModelRequestParameters( + parent=parent, + name=name, + 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 = _ExportPublisherModelRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/{name}:export".format_map(request_url_dict) + else: + path = "{parent}/{name}:export" + + 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.ExportModelOperation._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 + + def get_export_publisher_model_operation( + self, + *, + operation_name: str, + config: Optional[types.GetExportPublisherModelOperationConfigOrDict] = None, + ) -> types.ExportModelOperation: + """ + Fetches the status of an in-flight ``export_open_model`` LRO. + """ + + parameter_model = types._GetExportPublisherModelOperationParameters( + operation_name=operation_name, + 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 = _GetExportPublisherModelOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + 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("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ExportModelOperation._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 + + _EXPORT_POLL_TIMEOUT_SECONDS = 2 * 60 * 60 # 2 hours; matches legacy SDK. + _EXPORT_POLL_INTERVAL_SECONDS = 30 + @staticmethod def _build_filter_str( model_filter: Optional[str], @@ -1027,6 +1224,74 @@ def list_custom_model_deploy_options( return self._format_custom_deploy_options(options) + def export_open_model( + self, + *, + model: str, + config: Optional[types.ExportOpenModelConfigOrDict] = None, + ) -> Union[str, types.ExportModelOperation]: + """Exports Google open model weights to a Cloud Storage bucket. + + Args: + model: The publisher model to export. Accepts a full resource name + ``'publishers/{publisher}/models/{model}@{version}'`` or the + simplified form ``'{publisher}/{model}@{version}'``. + config: Export configuration. Must specify ``output_gcs_uri``. + + Returns: + The Cloud Storage URI (``str``) where the weights were written, when + ``config.wait_for_completion`` is ``True`` (default). + The ``ExportModelOperation`` for the caller to poll, when + ``config.wait_for_completion`` is ``False``. + + Raises: + ValueError: If ``model`` is not a valid publisher model name, or if + ``config`` is missing or its ``output_gcs_uri`` is empty. + TimeoutError: If ``wait_for_completion`` is ``True`` and the LRO does + not complete within 2 hours. + RuntimeError: If the LRO completes with an error, or completes + successfully but without a ``destinationUri``. + """ + if config is None: + raise ValueError( + "config with output_gcs_uri is required to export a model." + ) + if isinstance(config, dict): + config = types.ExportOpenModelConfig.model_validate(config) + if not config.output_gcs_uri: + raise ValueError("output_gcs_uri must be specified in config.") + + publisher_model_name = ModelGarden._reconcile_model_name(model) + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + operation = self._export_publisher_model( + parent=parent, + name=publisher_model_name, + config=types.ExportPublisherModelConfig( + destination=genai_types.GcsDestination( + output_uri_prefix=config.output_gcs_uri, + ), + ), + ) + if not config.wait_for_completion: + return operation + + operation = _operations_utils.await_operation( + operation_name=operation.name, + get_operation_fn=self.get_export_publisher_model_operation, + poll_interval=ModelGarden._EXPORT_POLL_INTERVAL_SECONDS, + timeout_seconds=ModelGarden._EXPORT_POLL_TIMEOUT_SECONDS, + ) + if operation.error: + raise RuntimeError(f"Export failed: {operation.error}") + if not operation.response or not operation.response.destination_uri: + raise RuntimeError( + f"Export completed but response has no destinationUri: {operation!r}" + ) + return operation.response.destination_uri + class AsyncModelGarden(_api_module.BaseModule): """Model Garden module.""" @@ -1256,6 +1521,160 @@ async def _recommend_spec( self._api_client._verify_response(return_value) return return_value + async def _export_publisher_model( + self, + *, + parent: str, + name: str, + config: Optional[types.ExportPublisherModelConfigOrDict] = None, + ) -> types.ExportModelOperation: + """ + Exports a publisher model (internal). + """ + + parameter_model = types._ExportPublisherModelRequestParameters( + parent=parent, + name=name, + 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 = _ExportPublisherModelRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/{name}:export".format_map(request_url_dict) + else: + path = "{parent}/{name}:export" + + 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.ExportModelOperation._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 get_export_publisher_model_operation( + self, + *, + operation_name: str, + config: Optional[types.GetExportPublisherModelOperationConfigOrDict] = None, + ) -> types.ExportModelOperation: + """ + Fetches the status of an in-flight ``export_open_model`` LRO. + """ + + parameter_model = types._GetExportPublisherModelOperationParameters( + operation_name=operation_name, + 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 = _GetExportPublisherModelOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + 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( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ExportModelOperation._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, @@ -1512,3 +1931,50 @@ async def list_custom_model_deploy_options( raise ValueError("No deploy options found.") return ModelGarden._format_custom_deploy_options(options) + + async def export_open_model( + self, + *, + model: str, + config: Optional[types.ExportOpenModelConfigOrDict] = None, + ) -> Union[str, types.ExportModelOperation]: + """Async variant of ``ModelGarden.export_open_model``.""" + if config is None: + raise ValueError( + "config with output_gcs_uri is required to export a model." + ) + if isinstance(config, dict): + config = types.ExportOpenModelConfig.model_validate(config) + if not config.output_gcs_uri: + raise ValueError("output_gcs_uri must be specified in config.") + + publisher_model_name = ModelGarden._reconcile_model_name(model) + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + operation = await self._export_publisher_model( + parent=parent, + name=publisher_model_name, + config=types.ExportPublisherModelConfig( + destination=genai_types.GcsDestination( + output_uri_prefix=config.output_gcs_uri, + ), + ), + ) + if not config.wait_for_completion: + return operation + + operation = await _operations_utils.await_operation_async( + operation_name=operation.name, + get_operation_fn=self.get_export_publisher_model_operation, + poll_interval=ModelGarden._EXPORT_POLL_INTERVAL_SECONDS, + timeout_seconds=ModelGarden._EXPORT_POLL_TIMEOUT_SECONDS, + ) + if operation.error: + raise RuntimeError(f"Export failed: {operation.error}") + if not operation.response or not operation.response.destination_uri: + raise RuntimeError( + f"Export completed but response has no destinationUri: {operation!r}" + ) + return operation.response.destination_uri diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index a731e74a0c..adf2d397f6 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -66,6 +66,7 @@ from .common import _DeleteSkillRequestParameters from .common import _EvaluateInstancesRequestParameters from .common import _ExecuteCodeAgentEngineSandboxRequestParameters +from .common import _ExportPublisherModelRequestParameters from .common import _GenerateAgentEngineMemoriesRequestParameters from .common import _GenerateInstanceRubricsRequest from .common import _GenerateLossClustersParameters @@ -94,6 +95,7 @@ from .common import _GetEvaluationMetricParameters from .common import _GetEvaluationRunParameters from .common import _GetEvaluationSetParameters +from .common import _GetExportPublisherModelOperationParameters from .common import _GetImportFilesOperationParameters from .common import _GetMultimodalDatasetOperationParameters from .common import _GetMultimodalDatasetParameters @@ -616,6 +618,18 @@ from .common import ExecuteSandboxEnvironmentResponse from .common import ExecuteSandboxEnvironmentResponseDict from .common import ExecuteSandboxEnvironmentResponseOrDict +from .common import ExportModelOperation +from .common import ExportModelOperationDict +from .common import ExportModelOperationOrDict +from .common import ExportOpenModelConfig +from .common import ExportOpenModelConfigDict +from .common import ExportOpenModelConfigOrDict +from .common import ExportPublisherModelConfig +from .common import ExportPublisherModelConfigDict +from .common import ExportPublisherModelConfigOrDict +from .common import ExportPublisherModelResponse +from .common import ExportPublisherModelResponseDict +from .common import ExportPublisherModelResponseOrDict from .common import FailedRubric from .common import FailedRubricDict from .common import FailedRubricOrDict @@ -730,6 +744,9 @@ from .common import GetEvaluationSetConfig from .common import GetEvaluationSetConfigDict from .common import GetEvaluationSetConfigOrDict +from .common import GetExportPublisherModelOperationConfig +from .common import GetExportPublisherModelOperationConfigDict +from .common import GetExportPublisherModelOperationConfigOrDict from .common import GetImportFilesOperationConfig from .common import GetImportFilesOperationConfigDict from .common import GetImportFilesOperationConfigOrDict @@ -3469,6 +3486,18 @@ "RecommendSpecResponse", "RecommendSpecResponseDict", "RecommendSpecResponseOrDict", + "ExportPublisherModelConfig", + "ExportPublisherModelConfigDict", + "ExportPublisherModelConfigOrDict", + "ExportPublisherModelResponse", + "ExportPublisherModelResponseDict", + "ExportPublisherModelResponseOrDict", + "ExportModelOperation", + "ExportModelOperationDict", + "ExportModelOperationOrDict", + "GetExportPublisherModelOperationConfig", + "GetExportPublisherModelOperationConfigDict", + "GetExportPublisherModelOperationConfigOrDict", "CreateRuntimeFeedbackEntryConfig", "CreateRuntimeFeedbackEntryConfigDict", "CreateRuntimeFeedbackEntryConfigOrDict", @@ -3619,6 +3648,9 @@ "ListCustomModelDeployOptionsConfig", "ListCustomModelDeployOptionsConfigDict", "ListCustomModelDeployOptionsConfigOrDict", + "ExportOpenModelConfig", + "ExportOpenModelConfigDict", + "ExportOpenModelConfigOrDict", "DeployOption", "DeployOptionDict", "DeployOptionOrDict", @@ -3810,6 +3842,8 @@ "_ListPublisherModelsRequestParameters", "_GetPublisherModelRequestParameters", "_RecommendSpecRequestParameters", + "_ExportPublisherModelRequestParameters", + "_GetExportPublisherModelOperationParameters", "_CreateRuntimeFeedbackEntryRequestParameters", "_DeleteRuntimeFeedbackEntryRequestParameters", "_GetRuntimeFeedbackRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index dd57c62b1e..12851051fd 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -23755,6 +23755,180 @@ class RecommendSpecResponseDict(TypedDict, total=False): RecommendSpecResponseOrDict = Union[RecommendSpecResponse, RecommendSpecResponseDict] +class ExportPublisherModelConfig(_common.BaseModel): + """RPC-level config for ``export_publisher_model``.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + destination: Optional[genai_types.GcsDestination] = Field( + default=None, description="""""" + ) + + +class ExportPublisherModelConfigDict(TypedDict, total=False): + """RPC-level config for ``export_publisher_model``.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + destination: Optional[genai_types.GcsDestination] + """""" + + +ExportPublisherModelConfigOrDict = Union[ + ExportPublisherModelConfig, ExportPublisherModelConfigDict +] + + +class _ExportPublisherModelRequestParameters(_common.BaseModel): + """Parameters for ``export_publisher_model``.""" + + parent: Optional[str] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") + config: Optional[ExportPublisherModelConfig] = Field( + default=None, description="""""" + ) + + +class _ExportPublisherModelRequestParametersDict(TypedDict, total=False): + """Parameters for ``export_publisher_model``.""" + + parent: Optional[str] + """""" + + name: Optional[str] + """""" + + config: Optional[ExportPublisherModelConfigDict] + """""" + + +_ExportPublisherModelRequestParametersOrDict = Union[ + _ExportPublisherModelRequestParameters, _ExportPublisherModelRequestParametersDict +] + + +class ExportPublisherModelResponse(_common.BaseModel): + """Response for the ``ExportPublisherModel`` RPC.""" + + destination_uri: Optional[str] = Field( + default=None, description="""The destination uri of the model weights.""" + ) + publisher_model: Optional[str] = Field( + default=None, + description="""The name of the PublisherModel resource. Format: `publishers/{publisher}/models/{publisher_model}@{version_id}`""", + ) + + +class ExportPublisherModelResponseDict(TypedDict, total=False): + """Response for the ``ExportPublisherModel`` RPC.""" + + destination_uri: Optional[str] + """The destination uri of the model weights.""" + + publisher_model: Optional[str] + """The name of the PublisherModel resource. Format: `publishers/{publisher}/models/{publisher_model}@{version_id}`""" + + +ExportPublisherModelResponseOrDict = Union[ + ExportPublisherModelResponse, ExportPublisherModelResponseDict +] + + +class ExportModelOperation(_common.BaseModel): + """Long-running operation returned by ``ExportPublisherModel``.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[ExportPublisherModelResponse] = Field( + default=None, description="""""" + ) + + +class ExportModelOperationDict(TypedDict, total=False): + """Long-running operation returned by ``ExportPublisherModel``.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + response: Optional[ExportPublisherModelResponseDict] + """""" + + +ExportModelOperationOrDict = Union[ExportModelOperation, ExportModelOperationDict] + + +class GetExportPublisherModelOperationConfig(_common.BaseModel): + """Config for ``get_export_publisher_model_operation``.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetExportPublisherModelOperationConfigDict(TypedDict, total=False): + """Config for ``get_export_publisher_model_operation``.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetExportPublisherModelOperationConfigOrDict = Union[ + GetExportPublisherModelOperationConfig, GetExportPublisherModelOperationConfigDict +] + + +class _GetExportPublisherModelOperationParameters(_common.BaseModel): + """Parameters for polling an ``export_publisher_model`` operation.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetExportPublisherModelOperationConfig] = Field( + default=None, description="""""" + ) + + +class _GetExportPublisherModelOperationParametersDict(TypedDict, total=False): + """Parameters for polling an ``export_publisher_model`` operation.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetExportPublisherModelOperationConfigDict] + """""" + + +_GetExportPublisherModelOperationParametersOrDict = Union[ + _GetExportPublisherModelOperationParameters, + _GetExportPublisherModelOperationParametersDict, +] + + class CreateRuntimeFeedbackEntryConfig(_common.BaseModel): """Config for creating a Feedback Entry.""" @@ -26797,6 +26971,40 @@ class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): ] +class ExportOpenModelConfig(_common.BaseModel): + """Config for ``export_open_model``.""" + + output_gcs_uri: Optional[str] = Field( + default=None, + description="""Required. The Cloud Storage URI prefix to export the model + weights to (e.g. ``'gs://my-bucket/gemma-weights/'``).""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to block on the export long-running operation. When + ``True`` (default), returns the destination URI on completion. When + ``False``, returns the ``ExportModelOperation`` for the caller to + poll.""", + ) + + +class ExportOpenModelConfigDict(TypedDict, total=False): + """Config for ``export_open_model``.""" + + output_gcs_uri: Optional[str] + """Required. The Cloud Storage URI prefix to export the model + weights to (e.g. ``'gs://my-bucket/gemma-weights/'``).""" + + wait_for_completion: Optional[bool] + """Whether to block on the export long-running operation. When + ``True`` (default), returns the destination URI on completion. When + ``False``, returns the ``ExportModelOperation`` for the caller to + poll.""" + + +ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] + + 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 828976c517..a65667bc7d 100644 --- a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py @@ -34,72 +34,72 @@ def test_list_deployable_models(client): def test_list_models(client): - """Tests listing all baseline models in Model Garden.""" - models = client.model_garden.list_models( + """Tests listing all baseline models in Model Garden.""" + models = client.model_garden.list_models( config=types.ListModelGardenModelsConfig( include_hugging_face_models=False, model_filter="timesfm", ) ) - assert len(models) > 0 - assert isinstance(models[0], str) - assert "timesfm" in models[0].lower() + assert len(models) > 0 + assert isinstance(models[0], str) + assert "timesfm" in models[0].lower() def test_list_publisher_model_deploy_options(client): - """Tests listing the verified deploy options for an open model.""" - options = client.model_garden.list_publisher_model_deploy_options( - model="google/gemma3@gemma-3-12b-it" - ) - assert len(options) > 0 - assert isinstance(options[0], types.DeployOption) - # Every verified deploy option exposes a serving container image. - assert options[0].serving_container_image_uri + """Tests listing the verified deploy options for an open model.""" + options = client.model_garden.list_publisher_model_deploy_options( + model="google/gemma3@gemma-3-12b-it" + ) + assert len(options) > 0 + assert isinstance(options[0], types.DeployOption) + # Every verified deploy option exposes a serving container image. + assert options[0].serving_container_image_uri def test_list_publisher_model_deploy_options_with_filter(client): - """Tests filtering an open model's deploy options by accelerator type.""" - options = client.model_garden.list_publisher_model_deploy_options( - model="google/gemma3@gemma-3-12b-it", - config=types.ListPublisherModelDeployOptionsConfig( - accelerator_type_filter="NVIDIA" - ), - ) - assert len(options) > 0 - for option in options: - assert "NVIDIA" in (option.accelerator_type or "") + """Tests filtering an open model's deploy options by accelerator type.""" + options = client.model_garden.list_publisher_model_deploy_options( + model="google/gemma3@gemma-3-12b-it", + config=types.ListPublisherModelDeployOptionsConfig( + accelerator_type_filter="NVIDIA" + ), + ) + assert len(options) > 0 + for option in options: + assert "NVIDIA" in (option.accelerator_type or "") def test_list_publisher_model_deploy_options_concise(client): - """Tests the concise (human-readable string) output for an open model.""" - options = client.model_garden.list_publisher_model_deploy_options( - model="google/gemma3@gemma-3-12b-it", - config=types.ListPublisherModelDeployOptionsConfig(concise=True), - ) - assert isinstance(options, str) - assert "[Option 1" in options + """Tests the concise (human-readable string) output for an open model.""" + options = client.model_garden.list_publisher_model_deploy_options( + model="google/gemma3@gemma-3-12b-it", + config=types.ListPublisherModelDeployOptionsConfig(concise=True), + ) + assert isinstance(options, str) + assert "[Option 1" in options def test_list_publisher_model_deploy_options_hugging_face(client): - """Tests deploy options for a Hugging Face model. + """Tests deploy options for a Hugging Face model. Exercises the distinct GetPublisherModel request path where is_hugging_face_model=True is sent. """ - options = client.model_garden.list_publisher_model_deploy_options( - model="codellama/codellama-7b-hf" - ) - assert len(options) > 0 - assert isinstance(options[0], types.DeployOption) - assert options[0].serving_container_image_uri + options = client.model_garden.list_publisher_model_deploy_options( + model="codellama/codellama-7b-hf" + ) + assert len(options) > 0 + assert isinstance(options[0], types.DeployOption) + assert options[0].serving_container_image_uri def test_list_publisher_model_deploy_options_no_deploy_support(client): - """Tests a model with no verified deployment config raises ValueError.""" - with pytest.raises(ValueError, match="does not support deployment"): - client.model_garden.list_publisher_model_deploy_options( - model="google/gemini-embedding-001@default" - ) + """Tests a model with no verified deployment config raises ValueError.""" + with pytest.raises(ValueError, match="does not support deployment"): + client.model_garden.list_publisher_model_deploy_options( + model="google/gemini-embedding-001@default" + ) # A GCS folder holding a custom model (config.json + a weights file) that @@ -108,57 +108,144 @@ def test_list_publisher_model_deploy_options_no_deploy_support(client): def test_list_custom_model_deploy_options(client): - """Default config: machine availability + user-quota filter. - - Exercises the RecommendSpec backend (distinct from the GetPublisherModel - path used by publisher models) and the human-readable string return type. - With ``check_machine_availability`` defaulting to True, the response is - the per-region ``recommendations`` list, so each option carries a - ``region``. - """ - options = client.model_garden.list_custom_model_deploy_options( - src=_CUSTOM_MODEL_SRC, - ) - assert isinstance(options, str) - assert "[Option 1]" in options - assert "region=" in options + """Default config: machine availability + user-quota filter. + + Exercises the RecommendSpec backend (distinct from the GetPublisherModel + path used by publisher models) and the human-readable string return type. + With ``check_machine_availability`` defaulting to True, the response is + the per-region ``recommendations`` list, so each option carries a + ``region``. + """ + options = client.model_garden.list_custom_model_deploy_options( + src=_CUSTOM_MODEL_SRC, + ) + assert isinstance(options, str) + assert "[Option 1]" in options + assert "region=" in options def test_list_custom_model_deploy_options_check_machine_availability_false( client, ): - """check_machine_availability=False -> RecommendSpec 'specs' path (no region field).""" - options = client.model_garden.list_custom_model_deploy_options( - src=_CUSTOM_MODEL_SRC, - config=types.ListCustomModelDeployOptionsConfig( - check_machine_availability=False - ), - ) - assert isinstance(options, str) - assert "[Option 1]" in options - assert "region=" not in options + """check_machine_availability=False -> RecommendSpec 'specs' path (no region field).""" + options = client.model_garden.list_custom_model_deploy_options( + src=_CUSTOM_MODEL_SRC, + config=types.ListCustomModelDeployOptionsConfig( + check_machine_availability=False + ), + ) + assert isinstance(options, str) + assert "[Option 1]" in options + assert "region=" not in options def test_list_custom_model_deploy_options_no_user_quota_filter(client): - """filter_by_user_quota=False -> recommendations returned without quota filtering.""" - options = client.model_garden.list_custom_model_deploy_options( - src=_CUSTOM_MODEL_SRC, - config=types.ListCustomModelDeployOptionsConfig(filter_by_user_quota=False), - ) - assert isinstance(options, str) - assert "[Option 1]" in options - assert "region=" in options + """filter_by_user_quota=False -> recommendations returned without quota filtering.""" + options = client.model_garden.list_custom_model_deploy_options( + src=_CUSTOM_MODEL_SRC, + config=types.ListCustomModelDeployOptionsConfig(filter_by_user_quota=False), + ) + assert isinstance(options, str) + assert "[Option 1]" in options + assert "region=" in options def test_list_custom_model_deploy_options_dict_config(client): - """The config may be passed as a plain dict instead of the typed config.""" - options = client.model_garden.list_custom_model_deploy_options( - src=_CUSTOM_MODEL_SRC, - config={"check_machine_availability": False}, - ) - assert isinstance(options, str) - assert "[Option 1]" in options - assert "region=" not in options + """The config may be passed as a plain dict instead of the typed config.""" + options = client.model_garden.list_custom_model_deploy_options( + src=_CUSTOM_MODEL_SRC, + config={"check_machine_availability": False}, + ) + assert isinstance(options, str) + assert "[Option 1]" in options + assert "region=" not in options + + +def test_export_open_model_wait_for_completion(client): + """Default wait_for_completion=True path: blocks on the LRO, returns URI. + + End-to-end coverage of the polling loop -- records the initial POST plus + every GET on the operation until it's done, so recording this takes as + long as the export itself (~2h for Gemma-2-2B). Once recorded, replaying + is instant. Skip this filter when iterating on unrelated changes. + """ + destination_uri = client.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + config=types.ExportOpenModelConfig( + output_gcs_uri="gs://mg-sdk-model-export-testing/gemma-2-2b-it/", + ), + ) + assert isinstance(destination_uri, str) + assert destination_uri.startswith("gs://") + + +def test_export_open_model_no_wait_returns_operation(client): + """wait_for_completion=False returns the ExportModelOperation immediately. + + Records only the initial ExportPublisherModel POST (no LRO polling), so this + is cheap to record (seconds). + """ + operation = client.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + config=types.ExportOpenModelConfig( + output_gcs_uri=("gs://mg-sdk-model-export-testing/gemma-2-2b-it-no-wait/"), + wait_for_completion=False, + ), + ) + assert isinstance(operation, types.ExportModelOperation) + assert operation.name + assert "/operations/" in operation.name + + +def test_export_open_model_no_wait_dict_config(client): + """The config may be passed as a plain dict; wait_for_completion=False path.""" + operation = client.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + config={ + "output_gcs_uri": ( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-dict-no-wait/" + ), + "wait_for_completion": False, + }, + ) + assert isinstance(operation, types.ExportModelOperation) + assert operation.name + + +def test_get_export_publisher_model_operation(client): + """Public getter returns an ExportModelOperation for a running export LRO. + + Chains a wait_for_completion=False export with an explicit poll via the + public getter -- exercises the same call users doing their own polling + (custom backoff, cancellation, etc.) will make. + """ + op = client.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + config=types.ExportOpenModelConfig( + output_gcs_uri="gs://mg-sdk-model-export-testing/gemma-2-2b-it-poll/", + wait_for_completion=False, + ), + ) + polled = client.model_garden.get_export_publisher_model_operation( + operation_name=op.name, + ) + assert isinstance(polled, types.ExportModelOperation) + assert polled.name == op.name + + +def test_export_open_model_missing_config_raises(client): + """No config is rejected client-side (output_gcs_uri is required); no RPC needed.""" + with pytest.raises(ValueError, match="output_gcs_uri is required"): + client.model_garden.export_open_model(model="google/gemma3@gemma-3-12b-it") + + +def test_export_open_model_invalid_name_raises(client): + """Invalid model name is rejected client-side; no RPC needed.""" + with pytest.raises(ValueError, match="not a valid publisher model name"): + client.model_garden.export_open_model( + model="not-a-valid-name", + config=types.ExportOpenModelConfig(output_gcs_uri="gs://b/out/"), + ) pytestmark = pytest_helper.setup( @@ -199,24 +286,59 @@ async def test_list_models_async(client): @pytest.mark.asyncio async def test_list_publisher_model_deploy_options_async(client): - """Tests listing the deploy options for an open model asynchronously.""" - options = await client.aio.model_garden.list_publisher_model_deploy_options( - model="google/gemma3@gemma-3-12b-it" - ) - assert len(options) > 0 - assert isinstance(options[0], types.DeployOption) - assert options[0].serving_container_image_uri + """Tests listing the deploy options for an open model asynchronously.""" + options = await client.aio.model_garden.list_publisher_model_deploy_options( + model="google/gemma3@gemma-3-12b-it" + ) + 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=_CUSTOM_MODEL_SRC, - config=types.ListCustomModelDeployOptionsConfig( - filter_by_user_quota=False, - ), - ) - assert isinstance(options, str) - assert "[Option 1]" in options - assert "region=" in options + """Tests listing the recommended deploy options for a custom model async.""" + options = await client.aio.model_garden.list_custom_model_deploy_options( + src=_CUSTOM_MODEL_SRC, + config=types.ListCustomModelDeployOptionsConfig( + filter_by_user_quota=False, + ), + ) + assert isinstance(options, str) + assert "[Option 1]" in options + assert "region=" in options + + +@pytest.mark.asyncio +async def test_export_open_model_async_no_wait_returns_operation(client): + """Async wait_for_completion=False returns the ExportModelOperation immediately.""" + operation = await client.aio.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + config=types.ExportOpenModelConfig( + output_gcs_uri=( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-async-no-wait/" + ), + wait_for_completion=False, + ), + ) + assert isinstance(operation, types.ExportModelOperation) + assert operation.name + + +@pytest.mark.asyncio +async def test_get_export_publisher_model_operation_async(client): + """Public async getter returns an ExportModelOperation for an in-flight LRO.""" + op = await client.aio.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + config=types.ExportOpenModelConfig( + output_gcs_uri=( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-async-poll/" + ), + wait_for_completion=False, + ), + ) + polled = await client.aio.model_garden.get_export_publisher_model_operation( + operation_name=op.name, + ) + assert isinstance(polled, types.ExportModelOperation) + assert polled.name == op.name diff --git a/tests/unit/agentplatform/genai/test_genai_model_garden.py b/tests/unit/agentplatform/genai/test_genai_model_garden.py index bd39dd7e82..d1d8f46647 100644 --- a/tests/unit/agentplatform/genai/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/test_genai_model_garden.py @@ -16,13 +16,16 @@ import asyncio from unittest import mock -from agentplatform._genai import model_garden +from agentplatform._genai import ( + model_garden, +) # pylint: disable=unused-import from agentplatform._genai import types from google.genai import client import pytest _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" +_OPEN_MODEL = "google/gemma3@gemma-3-12b-it" @pytest.fixture @@ -232,9 +235,7 @@ def test_list_models_excludes_hf_via_server_filter(mock_client): mock_client, "_list_publisher_models", return_value=dummy_response ) as mock_list: models = mock_client.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=False - ) + config=types.ListModelGardenModelsConfig(include_hugging_face_models=False) ) api_config = mock_list.call_args.kwargs.get("config") @@ -260,9 +261,7 @@ def test_list_models_with_hf(mock_client): mock_client, "_list_publisher_models", return_value=dummy_response ): models = mock_client.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=True - ) + config=types.ListModelGardenModelsConfig(include_hugging_face_models=True) ) assert len(models) == 2 @@ -290,9 +289,7 @@ def test_list_models_includes_non_deployable(mock_client): mock_client, "_list_publisher_models", return_value=dummy_response ): models = mock_client.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=False - ) + config=types.ListModelGardenModelsConfig(include_hugging_face_models=False) ) assert len(models) == 2 @@ -416,13 +413,15 @@ def test_build_filter_str_with_model_filter(): def test_build_filter_str_escapes_special_chars(): - """Tests that special regex characters in model_filter are escaped.""" - build_filter = model_garden.ModelGarden._build_filter_str - result = build_filter( - model_filter="model.v2+", include_hugging_face_models=False, deployable_only=False + """Tests that special regex characters in model_filter are escaped.""" + build_filter = model_garden.ModelGarden._build_filter_str + result = build_filter( + model_filter="model.v2+", + include_hugging_face_models=False, + deployable_only=False, ) - # re.escape turns '.' into '\\.' and '+' into '\\+' - assert r"model\.v2\+" in result + # re.escape turns '.' into '\\.' and '+' into '\\+' + assert r"model\.v2\+" in result # ---- list_publisher_model_deploy_options tests ---- @@ -449,11 +448,9 @@ def _make_deploy_option( ) -def _make_model_with_deploy_options( - options, name="publishers/google/models/gemma-2b" -): - """Builds a PublisherModel exposing the given deploy options.""" - return types.PublisherModel( +def _make_model_with_deploy_options(options, name="publishers/google/models/gemma-2b"): + """Builds a PublisherModel exposing the given deploy options.""" + return types.PublisherModel( name=name, supported_actions=types.PublisherModelCallToAction( multi_deploy_vertex=types.PublisherModelCallToActionDeployVertex( @@ -464,180 +461,186 @@ def _make_model_with_deploy_options( def test_list_publisher_model_deploy_options_basic(mock_client): - """Tests extraction of a single deploy option into a DeployOption.""" - dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) + """Tests extraction of a single deploy option into a DeployOption.""" + dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) - with mock.patch.object( + with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model ) as mock_get: - # A simplified name is reconciled to the full publisher resource name. - options = mock_client.list_publisher_model_deploy_options( - model="google/gemma3@gemma-3-12b-it" - ) + # A simplified name is reconciled to the full publisher resource name. + options = mock_client.list_publisher_model_deploy_options( + model="google/gemma3@gemma-3-12b-it" + ) - # The GetPublisherModel call must use the reconciled name, flag the model - # as non-Hugging-Face (it has an @version), and request equivalent - # deployment configs. - mock_get.assert_called_once_with( - name="publishers/google/models/gemma3@gemma-3-12b-it", - config=types.GetPublisherModelConfig( - is_hugging_face_model=False, - include_equivalent_model_garden_model_deployment_configs=True, - ), - ) - assert len(options) == 1 - assert isinstance(options[0], types.DeployOption) - assert options[0].option_name == "option-1" - assert options[0].machine_type == "g2-standard-12" - # accelerator_type is returned as the enum's string value (legacy parity). - assert options[0].accelerator_type == "NVIDIA_L4" - assert options[0].accelerator_count == 1 - assert "vllm" in options[0].serving_container_image_uri + # The GetPublisherModel call must use the reconciled name, flag the model + # as non-Hugging-Face (it has an @version), and request equivalent + # deployment configs. + mock_get.assert_called_once_with( + name="publishers/google/models/gemma3@gemma-3-12b-it", + config=types.GetPublisherModelConfig( + is_hugging_face_model=False, + include_equivalent_model_garden_model_deployment_configs=True, + ), + ) + assert len(options) == 1 + assert isinstance(options[0], types.DeployOption) + assert options[0].option_name == "option-1" + assert options[0].machine_type == "g2-standard-12" + # accelerator_type is returned as the enum's string value (legacy parity). + assert options[0].accelerator_type == "NVIDIA_L4" + assert options[0].accelerator_count == 1 + assert "vllm" in options[0].serving_container_image_uri def test_list_publisher_model_deploy_options_multiple(mock_client): - """Tests that all deploy options are returned when no filters are set.""" - dummy_model = _make_model_with_deploy_options( + """Tests that all deploy options are returned when no filters are set.""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), ] ) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001" - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001" + ) - assert [o.option_name for o in options] == ["g2", "a3"] + assert [o.option_name for o in options] == ["g2", "a3"] def test_list_publisher_model_deploy_options_machine_type_filter(mock_client): - """Tests machine_type_filter is a case-insensitive substring match.""" - dummy_model = _make_model_with_deploy_options( + """Tests machine_type_filter is a case-insensitive substring match.""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), ] ) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - machine_type_filter="G2" - ), - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + machine_type_filter="G2" + ), + ) - assert len(options) == 1 - assert options[0].option_name == "g2" + assert len(options) == 1 + assert options[0].option_name == "g2" def test_list_publisher_model_deploy_options_accelerator_type_filter( mock_client, ): - """Tests accelerator_type_filter is a case-insensitive substring match.""" - dummy_model = _make_model_with_deploy_options([ - _make_deploy_option(deploy_task_name="l4", accelerator_type="NVIDIA_L4"), - _make_deploy_option( - deploy_task_name="h100", accelerator_type="NVIDIA_H100_80GB" - ), - ]) - - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - accelerator_type_filter="h100" - ), + """Tests accelerator_type_filter is a case-insensitive substring match.""" + dummy_model = _make_model_with_deploy_options( + [ + _make_deploy_option(deploy_task_name="l4", accelerator_type="NVIDIA_L4"), + _make_deploy_option( + deploy_task_name="h100", accelerator_type="NVIDIA_H100_80GB" + ), + ] ) - assert len(options) == 1 - assert options[0].option_name == "h100" + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + accelerator_type_filter="h100" + ), + ) + + assert len(options) == 1 + assert options[0].option_name == "h100" def test_list_publisher_model_deploy_options_image_uri_filter(mock_client): - """Tests serving_container_image_uri_filter is a case-insensitive match.""" - dummy_model = _make_model_with_deploy_options( + """Tests serving_container_image_uri_filter is a case-insensitive match.""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option(deploy_task_name="vllm", image_uri="docker/vllm"), _make_deploy_option(deploy_task_name="tgi", image_uri="docker/tgi"), ] ) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - serving_container_image_uri_filter="VLLM" - ), - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + serving_container_image_uri_filter="VLLM" + ), + ) - assert len(options) == 1 - assert options[0].option_name == "vllm" + assert len(options) == 1 + assert options[0].option_name == "vllm" def test_list_publisher_model_deploy_options_machine_type_filter_list( mock_client, ): - """Tests a list of keywords matches options containing ANY of them (legacy parity).""" - dummy_model = _make_model_with_deploy_options([ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - _make_deploy_option(deploy_task_name="n1", machine_type="n1-standard-8"), - ]) - - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - machine_type_filter=["A3", "n1"] - ), + """Tests a list of keywords matches options containing ANY of them (legacy parity).""" + dummy_model = _make_model_with_deploy_options( + [ + _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), + _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), + _make_deploy_option(deploy_task_name="n1", machine_type="n1-standard-8"), + ] ) - assert [o.option_name for o in options] == ["a3", "n1"] + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + machine_type_filter=["A3", "n1"] + ), + ) + + assert [o.option_name for o in options] == ["a3", "n1"] def test_list_publisher_model_deploy_options_accelerator_type_filter_list( mock_client, ): - """Tests a list accelerator filter matches options containing ANY keyword.""" - dummy_model = _make_model_with_deploy_options([ - _make_deploy_option(deploy_task_name="l4", accelerator_type="NVIDIA_L4"), - _make_deploy_option( - deploy_task_name="t4", accelerator_type="NVIDIA_TESLA_T4" - ), - _make_deploy_option( - deploy_task_name="h100", accelerator_type="NVIDIA_H100_80GB" - ), - ]) - - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - accelerator_type_filter=["T4", "L4"] - ), + """Tests a list accelerator filter matches options containing ANY keyword.""" + dummy_model = _make_model_with_deploy_options( + [ + _make_deploy_option(deploy_task_name="l4", accelerator_type="NVIDIA_L4"), + _make_deploy_option( + deploy_task_name="t4", accelerator_type="NVIDIA_TESLA_T4" + ), + _make_deploy_option( + deploy_task_name="h100", accelerator_type="NVIDIA_H100_80GB" + ), + ] ) - assert [o.option_name for o in options] == ["l4", "t4"] + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + accelerator_type_filter=["T4", "L4"] + ), + ) + + assert [o.option_name for o in options] == ["l4", "t4"] def test_list_publisher_model_deploy_options_image_uri_filter_list(mock_client): - """Tests a list image-uri filter matches any keyword.""" - dummy_model = _make_model_with_deploy_options( + """Tests a list image-uri filter matches any keyword.""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option(deploy_task_name="vllm", image_uri="docker/vllm"), _make_deploy_option(deploy_task_name="tgi", image_uri="docker/tgi"), @@ -645,101 +648,101 @@ def test_list_publisher_model_deploy_options_image_uri_filter_list(mock_client): ] ) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - serving_container_image_uri_filter=["vllm", "tgi"] - ), - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + serving_container_image_uri_filter=["vllm", "tgi"] + ), + ) - assert [o.option_name for o in options] == ["vllm", "tgi"] + assert [o.option_name for o in options] == ["vllm", "tgi"] def test_matches_filter(): - """Tests the keyword-filter helper (single keyword, list, None, misses).""" - matches = model_garden.ModelGarden._matches_filter - # No filter -> always matches. - assert matches("g2-standard-12", None) is True - assert matches(None, None) is True - # Single keyword, case-insensitive substring. - assert matches("g2-standard-12", "G2") is True - assert matches("g2-standard-12", "n1") is False - # List of keywords -> match if ANY is contained. - assert matches("a3-highgpu-8g", ["n1", "a3"]) is True - assert matches("a3-highgpu-8g", ["n1", "g2"]) is False - # A missing (None) value never matches a non-empty filter. - assert matches(None, "g2") is False - # An empty list filter behaves like "no filter". - assert matches("anything", []) is True + """Tests the keyword-filter helper (single keyword, list, None, misses).""" + matches = model_garden.ModelGarden._matches_filter + # No filter -> always matches. + assert matches("g2-standard-12", None) is True + assert matches(None, None) is True + # Single keyword, case-insensitive substring. + assert matches("g2-standard-12", "G2") is True + assert matches("g2-standard-12", "n1") is False + # List of keywords -> match if ANY is contained. + assert matches("a3-highgpu-8g", ["n1", "a3"]) is True + assert matches("a3-highgpu-8g", ["n1", "g2"]) is False + # A missing (None) value never matches a non-empty filter. + assert matches(None, "g2") is False + # An empty list filter behaves like "no filter". + assert matches("anything", []) is True def test_list_publisher_model_deploy_options_dict_config(mock_client): - """Tests config passed as a dict is validated and applied.""" - dummy_model = _make_model_with_deploy_options( + """Tests config passed as a dict is validated and applied.""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), ] ) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config={"machine_type_filter": "a3"}, - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config={"machine_type_filter": "a3"}, + ) - assert len(options) == 1 - assert options[0].option_name == "a3" + assert len(options) == 1 + assert options[0].option_name == "a3" def test_list_publisher_model_deploy_options_default_config(mock_client): - """Tests config=None returns all options.""" - dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) - - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001" - ) + """Tests config=None returns all options.""" + dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) - assert len(options) == 1 + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001" + ) + + assert len(options) == 1 def test_list_publisher_model_deploy_options_no_deploy_support_raises( mock_client, ): - """Tests ValueError when the model does not support deployment (legacy parity).""" - dummy_model = types.PublisherModel(name="publishers/google/models/bert-base") + """Tests ValueError when the model does not support deployment (legacy parity).""" + dummy_model = types.PublisherModel(name="publishers/google/models/bert-base") - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - with pytest.raises(ValueError, match="does not support deployment"): - mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/bert-base@001" - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + with pytest.raises(ValueError, match="does not support deployment"): + mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/bert-base@001" + ) def test_list_publisher_model_deploy_options_no_match_raises(mock_client): - """Tests ValueError when filters exclude every option (legacy parity).""" - dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) + """Tests ValueError when filters exclude every option (legacy parity).""" + dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - with pytest.raises(ValueError, match="No deploy options found."): - mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - machine_type_filter="does-not-exist" - ), - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + with pytest.raises(ValueError, match="No deploy options found."): + mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + machine_type_filter="does-not-exist" + ), + ) def test_reconcile_model_name_simplified_with_version(): @@ -776,102 +779,102 @@ def test_reconcile_model_name_lowercases(): def test_reconcile_model_name_invalid_raises(): - """Tests an invalid name raises ValueError.""" - reconcile = model_garden.ModelGarden._reconcile_model_name - with pytest.raises(ValueError, match="not a valid publisher model name"): - reconcile("invalid-name-without-slash") + """Tests an invalid name raises ValueError.""" + reconcile = model_garden.ModelGarden._reconcile_model_name + with pytest.raises(ValueError, match="not a valid publisher model name"): + reconcile("invalid-name-without-slash") def test_reconcile_model_name_model_registry_raises(): - """Tests a Model Registry resource name raises ValueError. - - Without an explicit guard, ``projects/.../locations/.../models/...`` would - match the simplified ``{publisher}/{model}`` regex and be silently mangled - into ``publishers/projects/models/.../locations/.../models/...``. The guard - rejects it loudly with the same ``not a valid publisher model name`` - message used for any other unsupported input. - """ - reconcile = model_garden.ModelGarden._reconcile_model_name - for name in ( - "projects/123/locations/us-central1/models/456", - "projects/my-project/locations/europe-west1/models/9876543210@1", - "projects/p/locations/l/models/m", - ): - with pytest.raises(ValueError, match="not a valid publisher model name"): - reconcile(name) + """Tests a Model Registry resource name raises ValueError. + + Without an explicit guard, ``projects/.../locations/.../models/...`` would + match the simplified ``{publisher}/{model}`` regex and be silently mangled + into ``publishers/projects/models/.../locations/.../models/...``. The guard + rejects it loudly with the same ``not a valid publisher model name`` + message used for any other unsupported input. + """ + reconcile = model_garden.ModelGarden._reconcile_model_name + for name in ( + "projects/123/locations/us-central1/models/456", + "projects/my-project/locations/europe-west1/models/9876543210@1", + "projects/p/locations/l/models/m", + ): + with pytest.raises(ValueError, match="not a valid publisher model name"): + reconcile(name) def test_is_hugging_face_model(): - """Tests the Hugging Face model heuristic.""" - is_hf = model_garden.ModelGarden._is_hugging_face_model - # Bare org/model (single slash, no @version) -> Hugging Face. - assert is_hf("meta-llama/Llama-3.3-70B-Instruct") is True - # Simplified native names without @version also match (handled - # correctly by _reconcile_model_name). - assert is_hf("google/gemma3") is True - # Names with @version or a publishers/ prefix are not Hugging Face. - assert is_hf("google/gemma3@gemma-3-12b-it") is False - assert is_hf("publishers/google/models/gemma3@gemma-3-12b-it") is False - assert is_hf("publishers/hf-meta-llama/models/llama-3.3") is False + """Tests the Hugging Face model heuristic.""" + is_hf = model_garden.ModelGarden._is_hugging_face_model + # Bare org/model (single slash, no @version) -> Hugging Face. + assert is_hf("meta-llama/Llama-3.3-70B-Instruct") is True + # Simplified native names without @version also match (handled + # correctly by _reconcile_model_name). + assert is_hf("google/gemma3") is True + # Names with @version or a publishers/ prefix are not Hugging Face. + assert is_hf("google/gemma3@gemma-3-12b-it") is False + assert is_hf("publishers/google/models/gemma3@gemma-3-12b-it") is False + assert is_hf("publishers/hf-meta-llama/models/llama-3.3") is False def test_list_publisher_model_deploy_options_hugging_face_model(mock_client): - """Tests an HF model name sends is_hugging_face_model=True (legacy parity).""" - dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) + """Tests an HF model name sends is_hugging_face_model=True (legacy parity).""" + dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) - with mock.patch.object( + with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model ) as mock_get: - mock_client.list_publisher_model_deploy_options( - model="meta-llama/Llama-3.3-70B-Instruct" - ) + mock_client.list_publisher_model_deploy_options( + model="meta-llama/Llama-3.3-70B-Instruct" + ) - mock_get.assert_called_once_with( - name="publishers/meta-llama/models/llama-3.3-70b-instruct", - config=types.GetPublisherModelConfig( - is_hugging_face_model=True, - include_equivalent_model_garden_model_deployment_configs=True, - ), - ) + mock_get.assert_called_once_with( + name="publishers/meta-llama/models/llama-3.3-70b-instruct", + config=types.GetPublisherModelConfig( + is_hugging_face_model=True, + include_equivalent_model_garden_model_deployment_configs=True, + ), + ) def test_list_publisher_model_deploy_options_async(mock_async_client): - """Tests the async client returns deploy options.""" - dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) + """Tests the async client returns deploy options.""" + dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) - with mock.patch.object( + with mock.patch.object( mock_async_client, "_get_publisher_model", new=mock.AsyncMock(return_value=dummy_model), ): - options = asyncio.run( - mock_async_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001" + options = asyncio.run( + mock_async_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001" + ) ) - ) - assert len(options) == 1 - assert options[0].option_name == "option-1" - assert options[0].machine_type == "g2-standard-12" + assert len(options) == 1 + assert options[0].option_name == "option-1" + assert options[0].machine_type == "g2-standard-12" def test_list_publisher_model_deploy_options_async_no_deploy_support_raises( mock_async_client, ): - """Tests the async client raises when deployment is unsupported.""" - dummy_model = types.PublisherModel(name="publishers/google/models/bert-base") + """Tests the async client raises when deployment is unsupported.""" + dummy_model = types.PublisherModel(name="publishers/google/models/bert-base") - with mock.patch.object( + with mock.patch.object( mock_async_client, "_get_publisher_model", new=mock.AsyncMock(return_value=dummy_model), ): - with pytest.raises(ValueError, match="does not support deployment"): - asyncio.run( - mock_async_client.list_publisher_model_deploy_options( - model="publishers/google/models/bert-base@001" - ) - ) + with pytest.raises(ValueError, match="does not support deployment"): + asyncio.run( + mock_async_client.list_publisher_model_deploy_options( + model="publishers/google/models/bert-base@001" + ) + ) def _make_deploy_option_no_accelerator( @@ -879,8 +882,8 @@ def _make_deploy_option_no_accelerator( machine_type="ct5lp-hightpu-1t", image_uri="docker/hexllm", ): - """Builds a deploy option whose machine has no GPU accelerator (e.g. TPU).""" - return types.PublisherModelCallToActionDeploy( + """Builds a deploy option whose machine has no GPU accelerator (e.g. TPU).""" + return types.PublisherModelCallToActionDeploy( deploy_task_name=deploy_task_name, dedicated_resources=types.DedicatedResources( machine_spec=types.MachineSpec(machine_type=machine_type) @@ -892,72 +895,72 @@ def _make_deploy_option_no_accelerator( def test_list_publisher_model_deploy_options_no_accelerator_defaults( mock_client, ): - """Tests no-accelerator machines report UNSPECIFIED/0 (legacy proto parity).""" - dummy_model = _make_model_with_deploy_options( - [_make_deploy_option_no_accelerator()] - ) - - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001" + """Tests no-accelerator machines report UNSPECIFIED/0 (legacy proto parity).""" + dummy_model = _make_model_with_deploy_options( + [_make_deploy_option_no_accelerator()] ) - assert options[0].machine_type == "ct5lp-hightpu-1t" - # Over gRPC the legacy SDK surfaces these proto3 defaults; we match them. - assert options[0].accelerator_type == "ACCELERATOR_TYPE_UNSPECIFIED" - assert options[0].accelerator_count == 0 + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001" + ) + + assert options[0].machine_type == "ct5lp-hightpu-1t" + # Over gRPC the legacy SDK surfaces these proto3 defaults; we match them. + assert options[0].accelerator_type == "ACCELERATOR_TYPE_UNSPECIFIED" + assert options[0].accelerator_count == 0 def test_list_publisher_model_deploy_options_no_accelerator_concise( mock_client, ): - """Tests concise output for a no-accelerator machine matches legacy.""" - dummy_model = _make_model_with_deploy_options( - [_make_deploy_option_no_accelerator()] - ) - - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - result = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig(concise=True), + """Tests concise output for a no-accelerator machine matches legacy.""" + dummy_model = _make_model_with_deploy_options( + [_make_deploy_option_no_accelerator()] ) - expected = ( - "[Option 1: tpu]\n" - ' serving_container_image_uri="docker/hexllm",\n' - ' machine_type="ct5lp-hightpu-1t",\n' - ' accelerator_type="ACCELERATOR_TYPE_UNSPECIFIED",\n' - " accelerator_count=0," - ) - assert result == expected + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + result = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig(concise=True), + ) + + expected = ( + "[Option 1: tpu]\n" + ' serving_container_image_uri="docker/hexllm",\n' + ' machine_type="ct5lp-hightpu-1t",\n' + ' accelerator_type="ACCELERATOR_TYPE_UNSPECIFIED",\n' + " accelerator_count=0," + ) + assert result == expected def test_list_publisher_model_deploy_options_accelerator_filter_excludes_unspecified( mock_client, ): - """Tests accelerator_type_filter excludes no-accelerator options (legacy parity).""" - dummy_model = _make_model_with_deploy_options( + """Tests accelerator_type_filter excludes no-accelerator options (legacy parity).""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option_no_accelerator(deploy_task_name="tpu"), _make_deploy_option(deploy_task_name="gpu", accelerator_type="NVIDIA_L4"), ] ) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - options = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - accelerator_type_filter="NVIDIA" - ), - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + options = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + accelerator_type_filter="NVIDIA" + ), + ) - assert [o.option_name for o in options] == ["gpu"] + assert [o.option_name for o in options] == ["gpu"] # ---- concise option tests ---- @@ -966,58 +969,60 @@ def test_list_publisher_model_deploy_options_accelerator_filter_excludes_unspeci def test_list_publisher_model_deploy_options_not_concise_returns_list( mock_client, ): - """Tests that without concise the method returns a list of DeployOption.""" - dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) - - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - result = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig(concise=False), - ) + """Tests that without concise the method returns a list of DeployOption.""" + dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) - assert isinstance(result, list) - assert isinstance(result[0], types.DeployOption) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + result = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig(concise=False), + ) + + assert isinstance(result, list) + assert isinstance(result[0], types.DeployOption) def test_list_publisher_model_deploy_options_concise_returns_string( mock_client, ): - """Tests concise=True returns the legacy-format human-readable string.""" - dummy_model = _make_model_with_deploy_options([ - _make_deploy_option( - deploy_task_name="option-1", - machine_type="g2-standard-12", - accelerator_type="NVIDIA_L4", - accelerator_count=1, - image_uri="docker/vllm", - ) - ]) - - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - result = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig(concise=True), + """Tests concise=True returns the legacy-format human-readable string.""" + dummy_model = _make_model_with_deploy_options( + [ + _make_deploy_option( + deploy_task_name="option-1", + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + image_uri="docker/vllm", + ) + ] ) - # Matches the legacy SDK's concise formatting exactly. - expected = ( - "[Option 1: option-1]\n" - ' serving_container_image_uri="docker/vllm",\n' - ' machine_type="g2-standard-12",\n' - ' accelerator_type="NVIDIA_L4",\n' - " accelerator_count=1," - ) - assert isinstance(result, str) - assert result == expected + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + result = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig(concise=True), + ) + + # Matches the legacy SDK's concise formatting exactly. + expected = ( + "[Option 1: option-1]\n" + ' serving_container_image_uri="docker/vllm",\n' + ' machine_type="g2-standard-12",\n' + ' accelerator_type="NVIDIA_L4",\n' + " accelerator_count=1," + ) + assert isinstance(result, str) + assert result == expected def test_list_publisher_model_deploy_options_concise_multiple(mock_client): - """Tests concise formatting of multiple options separated by a blank line.""" - dummy_model = _make_model_with_deploy_options( + """Tests concise formatting of multiple options separated by a blank line.""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option( deploy_task_name="g2", @@ -1036,74 +1041,72 @@ def test_list_publisher_model_deploy_options_concise_multiple(mock_client): ] ) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - result = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig(concise=True), - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + result = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig(concise=True), + ) - expected = ( - "[Option 1: g2]\n" - ' serving_container_image_uri="docker/vllm",\n' - ' machine_type="g2-standard-12",\n' - ' accelerator_type="NVIDIA_L4",\n' - " accelerator_count=1," - "\n\n" - "[Option 2: a3]\n" - ' serving_container_image_uri="docker/hexllm",\n' - ' machine_type="a3-highgpu-8g",\n' - ' accelerator_type="NVIDIA_H100_80GB",\n' - " accelerator_count=8," - ) - assert result == expected + expected = ( + "[Option 1: g2]\n" + ' serving_container_image_uri="docker/vllm",\n' + ' machine_type="g2-standard-12",\n' + ' accelerator_type="NVIDIA_L4",\n' + " accelerator_count=1," + "\n\n" + "[Option 2: a3]\n" + ' serving_container_image_uri="docker/hexllm",\n' + ' machine_type="a3-highgpu-8g",\n' + ' accelerator_type="NVIDIA_H100_80GB",\n' + " accelerator_count=8," + ) + assert result == expected def test_list_publisher_model_deploy_options_concise_with_filter(mock_client): - """Tests concise output honors filters before formatting.""" - dummy_model = _make_model_with_deploy_options( + """Tests concise output honors filters before formatting.""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), ] ) - with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ): - result = mock_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig( - machine_type_filter="a3", concise=True - ), - ) + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ): + result = mock_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig( + machine_type_filter="a3", concise=True + ), + ) - assert result.startswith("[Option 1: a3]") - assert "a3-highgpu-8g" in result - assert "g2-standard-12" not in result + assert result.startswith("[Option 1: a3]") + assert "a3-highgpu-8g" in result + assert "g2-standard-12" not in result def test_format_concise_deploy_options_omits_none_fields(): - """Tests that None fields are omitted and the header has no name if unset.""" - options = [ + """Tests that None fields are omitted and the header has no name if unset.""" + options = [ types.DeployOption( machine_type="g2-standard-12", accelerator_count=1, ) ] - result = model_garden.ModelGarden._format_concise_deploy_options(options) - expected = ( - "[Option 1]\n" - ' machine_type="g2-standard-12",\n' - " accelerator_count=1," + result = model_garden.ModelGarden._format_concise_deploy_options(options) + expected = ( + '[Option 1]\n machine_type="g2-standard-12",\n accelerator_count=1,' ) - assert result == expected + assert result == expected def test_list_publisher_model_deploy_options_concise_async(mock_async_client): - """Tests the async client returns a concise string when requested.""" - dummy_model = _make_model_with_deploy_options( + """Tests the async client returns a concise string when requested.""" + dummy_model = _make_model_with_deploy_options( [ _make_deploy_option( deploy_task_name="option-1", @@ -1115,287 +1118,609 @@ def test_list_publisher_model_deploy_options_concise_async(mock_async_client): ] ) - with mock.patch.object( + with mock.patch.object( mock_async_client, "_get_publisher_model", new=mock.AsyncMock(return_value=dummy_model), ): - result = asyncio.run( - mock_async_client.list_publisher_model_deploy_options( - model="publishers/google/models/gemma-2b@001", - config=types.ListPublisherModelDeployOptionsConfig(concise=True), + result = asyncio.run( + mock_async_client.list_publisher_model_deploy_options( + model="publishers/google/models/gemma-2b@001", + config=types.ListPublisherModelDeployOptionsConfig(concise=True), + ) ) - ) - assert isinstance(result, str) - assert result.startswith("[Option 1: option-1]") + assert isinstance(result, str) + assert result.startswith("[Option 1: option-1]") 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" + """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, + ) + ) + ] ) - assert result == expected + + 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=[ - types.RecommendSpecResponseMachineAndModelContainerSpec( - machine_spec=types.MachineSpec(machine_type="g2-standard-12") - ) - ] - ) - - 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, - ), + """Tests the parent, gcs_uri and RecommendSpecConfig are passed through.""" + dummy_response = types.RecommendSpecResponse( + specs=[ + types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec(machine_type="g2-standard-12") + ) + ] ) + 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 check_machine_availability/filter_by_user_quota map to the API config.""" - dummy_response = types.RecommendSpecResponse( - specs=[ - types.RecommendSpecResponseMachineAndModelContainerSpec( - machine_spec=types.MachineSpec(machine_type="g2-standard-12") - ) - ] - ) - - 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( - check_machine_availability=False, - filter_by_user_quota=False, - ), + """Tests check_machine_availability/filter_by_user_quota map to the API config.""" + dummy_response = types.RecommendSpecResponse( + specs=[ + types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec(machine_type="g2-standard-12") + ) + ] ) - _, kwargs = mock_recommend.call_args - assert kwargs["config"] == types.RecommendSpecConfig( - check_machine_availability=False, - check_user_quota=False, - ) + 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( + check_machine_availability=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") - - 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"' + """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", + ), + ] ) - assert result == expected + + 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,\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 - ), + """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", + ), + ] ) - assert "[Option 1]" in result - assert "[Option 2]" in result - assert "us-west1" in result + 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="") + """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_empty_options_raises(mock_client): - """Tests an empty API response raises, matching the publisher-model sibling.""" - dummy_response = types.RecommendSpecResponse(specs=[]) + """Tests an empty API response raises, matching the publisher-model sibling.""" + dummy_response = types.RecommendSpecResponse(specs=[]) - with mock.patch.object( - mock_client, "_recommend_spec", return_value=dummy_response - ): - with pytest.raises(ValueError, match="No deploy options found"): - mock_client.list_custom_model_deploy_options(src="gs://weights") + with mock.patch.object(mock_client, "_recommend_spec", return_value=dummy_response): + with pytest.raises(ValueError, match="No deploy options found"): + mock_client.list_custom_model_deploy_options(src="gs://weights") def test_list_custom_model_deploy_options_quota_filter_empty_raises( mock_client, ): - """Tests filter_by_user_quota dropping every recommendation raises.""" - 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_NO_USER_QUOTA", - ), - ] - ) - - with mock.patch.object( - mock_client, "_recommend_spec", return_value=dummy_response - ): - with pytest.raises(ValueError, match="No deploy options found"): - mock_client.list_custom_model_deploy_options(src="gs://weights") + """Tests filter_by_user_quota dropping every recommendation raises.""" + 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_NO_USER_QUOTA", + ), + ] + ) + + with mock.patch.object(mock_client, "_recommend_spec", return_value=dummy_response): + with pytest.raises(ValueError, match="No deploy options found"): + mock_client.list_custom_model_deploy_options(src="gs://weights") def test_list_custom_model_deploy_options_dict_config(mock_client): - """Tests a dict config is validated into ListCustomModelDeployOptionsConfig.""" - dummy_response = types.RecommendSpecResponse( - specs=[ - types.RecommendSpecResponseMachineAndModelContainerSpec( - machine_spec=types.MachineSpec(machine_type="g2-standard-12") - ) - ] - ) - - 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={ - "check_machine_availability": False, - "filter_by_user_quota": False, - }, + """Tests a dict config is validated into ListCustomModelDeployOptionsConfig.""" + dummy_response = types.RecommendSpecResponse( + specs=[ + types.RecommendSpecResponseMachineAndModelContainerSpec( + machine_spec=types.MachineSpec(machine_type="g2-standard-12") + ) + ] ) - _, kwargs = mock_recommend.call_args - assert kwargs["config"] == types.RecommendSpecConfig( - check_machine_availability=False, - check_user_quota=False, - ) + 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={ + "check_machine_availability": 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") + """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, + ) + ) + ] ) - assert isinstance(result, str) - assert result.startswith("[Option 1]") - assert 'machine_type="g2-standard-12"' in result + 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="")) + """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="")) + + +# ===================================================================== +# export_open_model +# ===================================================================== + +_TEST_EXPORT_OPERATION = "projects/p/locations/us-central1/operations/exp-1" +_TEST_EXPORT_URI = "gs://my-bucket/gemma-weights/exported/" + + +def _make_export_operation( + name=_TEST_EXPORT_OPERATION, *, done=False, destination_uri=None, error=None +): + """Builds a fake ExportModelOperation for use in tests.""" + op = types.ExportModelOperation(name=name, done=done) + if destination_uri is not None: + op.response = types.ExportPublisherModelResponse( + destination_uri=destination_uri + ) + if error is not None: + op.error = error + return op + + +def _patch_export_rpc(mock_client, initial_op=None): + """Mocks ``_export_publisher_model`` to return ``initial_op``.""" + return mock.patch.object( + mock_client, + "_export_publisher_model", + return_value=initial_op if initial_op is not None else _make_export_operation(), + ) + + +def _patch_async_export_rpc(mock_client, initial_op=None): + """Async variant of _patch_export_rpc.""" + return mock.patch.object( + mock_client, + "_export_publisher_model", + new=mock.AsyncMock( + return_value=( + initial_op if initial_op is not None else _make_export_operation() + ) + ), + ) + + +# --- Return type: wait_for_completion=True (default) ------------------- + + +def test_export_open_model_returns_destination_uri_when_waiting(mock_client): + """Default (wait_for_completion=True) blocks on LRO and returns the URI string.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation", + return_value=done_op, + ) as m_await, + ): + result = mock_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig(output_gcs_uri="gs://b/out/"), + ) + m_await.assert_called_once() + assert result == _TEST_EXPORT_URI + + +def test_export_open_model_polls_via_public_getter(mock_client): + """``get_export_publisher_model_operation`` is used as the polling callback.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation", + return_value=done_op, + ) as m_await, + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig(output_gcs_uri="gs://b/out/"), + ) + kwargs = m_await.call_args.kwargs + assert ( + kwargs["get_operation_fn"] + == mock_client.get_export_publisher_model_operation + ) + assert kwargs["operation_name"] == _TEST_EXPORT_OPERATION + + +def test_export_open_model_raises_on_operation_error(mock_client): + """A backend error on the LRO surfaces as RuntimeError.""" + failed_op = _make_export_operation(done=True, error={"message": "EULA required"}) + with ( + _patch_export_rpc(mock_client), + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation", + return_value=failed_op, + ), + ): + with pytest.raises(RuntimeError, match="Export failed"): + mock_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig(output_gcs_uri="gs://b/out/"), + ) + + +def test_export_open_model_raises_when_response_missing_destination_uri( + mock_client, +): + """A done LRO with no destinationUri surfaces as RuntimeError.""" + done_op = _make_export_operation(done=True) # no destination_uri set + with ( + _patch_export_rpc(mock_client), + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation", + return_value=done_op, + ), + ): + with pytest.raises(RuntimeError, match="destinationUri"): + mock_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig(output_gcs_uri="gs://b/out/"), + ) + + +# --- Return type: wait_for_completion=False ---------------------------- + + +def test_export_open_model_returns_operation_when_not_waiting(mock_client): + """wait_for_completion=False returns the operation immediately (no polling).""" + initial_op = _make_export_operation() + with ( + _patch_export_rpc(mock_client, initial_op=initial_op), + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation" + ) as m_await, + ): + result = mock_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig( + output_gcs_uri="gs://b/out/", wait_for_completion=False + ), + ) + m_await.assert_not_called() + assert result is initial_op + + +# --- Config validation ------------------------------------------------- + + +def test_export_open_model_missing_config_raises(mock_client): + """Omitting config raises ValueError before any RPC or polling.""" + with _patch_export_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="output_gcs_uri is required"): + mock_client.export_open_model(model=_OPEN_MODEL) + m_rpc.assert_not_called() + + +def test_export_open_model_empty_output_gcs_uri_raises(mock_client): + """An empty output_gcs_uri raises before any RPC or polling.""" + with _patch_export_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="output_gcs_uri must be specified"): + mock_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig(output_gcs_uri=""), + ) + m_rpc.assert_not_called() + + +def test_export_open_model_dict_config_validated(mock_client): + """A plain dict is validated into ExportOpenModelConfig.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client) as m_rpc, + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation", + return_value=done_op, + ), + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + config={"output_gcs_uri": "gs://b/out/"}, + ) + assert ( + m_rpc.call_args.kwargs["config"].destination.output_uri_prefix + == "gs://b/out/" + ) + + +# --- Model-name resolution --------------------------------------------- + + +def test_export_open_model_simplified_name_reconciled(mock_client): + """Simplified '{publisher}/{model}@{version}' -> full resource name.""" + initial_op = _make_export_operation() + with ( + _patch_export_rpc(mock_client, initial_op=initial_op) as m_rpc, + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation" + ), + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig( + output_gcs_uri="gs://b/out/", wait_for_completion=False + ), + ) + kwargs = m_rpc.call_args.kwargs + assert kwargs["name"] == "publishers/google/models/gemma3@gemma-3-12b-it" + assert ( + kwargs["parent"] == f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" + ) + + +def test_export_open_model_full_resource_name_preserved(mock_client): + """A full 'publishers/.../models/...@version' is passed through unchanged.""" + with ( + _patch_export_rpc(mock_client) as m_rpc, + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation" + ), + ): + mock_client.export_open_model( + model="publishers/google/models/gemma3@gemma-3-12b-it", + config=types.ExportOpenModelConfig( + output_gcs_uri="gs://b/out/", wait_for_completion=False + ), + ) + assert ( + m_rpc.call_args.kwargs["name"] + == "publishers/google/models/gemma3@gemma-3-12b-it" + ) + + +def test_export_open_model_invalid_name_raises_before_rpc(mock_client): + """Malformed model names raise ValueError; no RPC and no polling.""" + with _patch_export_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="not a valid publisher model name"): + mock_client.export_open_model( + model="not-a-valid-name", + config=types.ExportOpenModelConfig(output_gcs_uri="gs://b/out/"), + ) + m_rpc.assert_not_called() + + +# --- Destination forwarded to RPC config -------------------------------- + + +def test_export_open_model_forwards_output_gcs_uri(mock_client): + """output_gcs_uri from public config becomes destination.output_uri_prefix.""" + with ( + _patch_export_rpc(mock_client) as m_rpc, + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation" + ), + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig( + output_gcs_uri="gs://my-bucket/weights/", wait_for_completion=False + ), + ) + assert ( + m_rpc.call_args.kwargs["config"].destination.output_uri_prefix + == "gs://my-bucket/weights/" + ) + + +# --- Async surface parity ---------------------------------------------- + + +def test_export_open_model_async_returns_destination_uri_when_waiting( + mock_async_client, +): + """Async default (wait=True) awaits the LRO and returns the URI string.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_async_export_rpc(mock_async_client), + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation_async", + new=mock.AsyncMock(return_value=done_op), + ), + ): + result = asyncio.run( + mock_async_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig(output_gcs_uri="gs://b/out/"), + ) + ) + assert result == _TEST_EXPORT_URI + + +def test_export_open_model_async_returns_operation_when_not_waiting( + mock_async_client, +): + """Async wait=False returns the operation immediately.""" + initial_op = _make_export_operation() + with ( + _patch_async_export_rpc(mock_async_client, initial_op=initial_op), + mock.patch( + "google.cloud.aiplatform.agentplatform._genai.model_garden._operations_utils.await_operation_async" + ) as m_await, + ): + result = asyncio.run( + mock_async_client.export_open_model( + model=_OPEN_MODEL, + config=types.ExportOpenModelConfig( + output_gcs_uri="gs://b/out/", wait_for_completion=False + ), + ) + ) + m_await.assert_not_called() + assert result is initial_op + + +def test_export_open_model_async_missing_config_raises(mock_async_client): + """Async missing config raises ValueError before any RPC or polling.""" + with _patch_async_export_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="output_gcs_uri is required"): + asyncio.run(mock_async_client.export_open_model(model=_OPEN_MODEL)) + m_rpc.assert_not_called() + + +def test_export_open_model_async_invalid_name_raises(mock_async_client): + """Async invalid model name raises ValueError before any RPC or polling.""" + with _patch_async_export_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="not a valid publisher model name"): + asyncio.run( + mock_async_client.export_open_model( + model="not-a-valid-name", + config=types.ExportOpenModelConfig(output_gcs_uri="gs://b/out/"), + ) + ) + m_rpc.assert_not_called()