diff --git a/agentplatform/_genai/model_garden.py b/agentplatform/_genai/model_garden.py index 511f958367..61d19be50f 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,162 @@ 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 + + # Fallbacks for ``ExportOpenModelConfig`` when the caller does not + # override them. 2h matches the legacy SDK's blocking ``.export()`` and is + # generous enough for large open weights (e.g. Gemma 3 27B). + _DEFAULT_EXPORT_TIMEOUT_SECONDS = 2 * 60 * 60 + _DEFAULT_EXPORT_POLL_INTERVAL_SECONDS = 30 + @staticmethod def _build_filter_str( model_filter: Optional[str], @@ -1027,6 +1227,91 @@ def list_custom_model_deploy_options( return self._format_custom_deploy_options(options) + def export_open_model( + self, + *, + model: str, + output_gcs_uri: 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}'``. Hugging Face + model IDs are not supported. + output_gcs_uri: Cloud Storage URI prefix to write the model weights + to (e.g. ``'gs://my-bucket/gemma-weights/'``). + config: Optional export configuration (blocking behavior, poll + timing). See ``ExportOpenModelConfig``. + + Returns: + When ``config.wait_for_completion`` is ``True`` (default), the + Cloud Storage URI (``str``) where the weights were written. + When ``config.wait_for_completion`` is ``False``, the + ``ExportModelOperation`` for the caller to poll (via + ``get_export_publisher_model_operation`` or their own strategy). + + Raises: + ValueError: If ``output_gcs_uri`` is empty, if ``model`` is not a + valid Google publisher model resource name, or if ``model`` looks + like a Hugging Face model ID. + TimeoutError: If ``wait_for_completion=True`` and the LRO does not + complete within ``config.timeout_seconds`` (default 2 hours). + RuntimeError: If the LRO completes with an error, or completes + successfully but without a ``destination_uri``. + """ + if not output_gcs_uri: + raise ValueError("output_gcs_uri must be a non-empty Cloud Storage URI.") + if ModelGarden._is_hugging_face_model(model): + raise ValueError( + f"export_open_model does not support Hugging Face model IDs " + f"(got {model!r}); only Google publisher open models are " + f"exportable via this API." + ) + if config is None: + config = types.ExportOpenModelConfig() + elif isinstance(config, dict): + config = types.ExportOpenModelConfig.model_validate(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=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=( + config.poll_interval_seconds + or ModelGarden._DEFAULT_EXPORT_POLL_INTERVAL_SECONDS + ), + timeout_seconds=( + config.timeout_seconds or ModelGarden._DEFAULT_EXPORT_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 destination_uri: " + f"{operation!r}" + ) + return operation.response.destination_uri + class AsyncModelGarden(_api_module.BaseModule): """Model Garden module.""" @@ -1256,6 +1541,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 +1951,61 @@ 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, + output_gcs_uri: str, + config: Optional[types.ExportOpenModelConfigOrDict] = None, + ) -> Union[str, types.ExportModelOperation]: + """Async variant of ``ModelGarden.export_open_model``.""" + if not output_gcs_uri: + raise ValueError("output_gcs_uri must be a non-empty Cloud Storage URI.") + if ModelGarden._is_hugging_face_model(model): + raise ValueError( + f"export_open_model does not support Hugging Face model IDs " + f"(got {model!r}); only Google publisher open models are " + f"exportable via this API." + ) + if config is None: + config = types.ExportOpenModelConfig() + elif isinstance(config, dict): + config = types.ExportOpenModelConfig.model_validate(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=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=( + config.poll_interval_seconds + or ModelGarden._DEFAULT_EXPORT_POLL_INTERVAL_SECONDS + ), + timeout_seconds=( + config.timeout_seconds or ModelGarden._DEFAULT_EXPORT_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 destination_uri: " + f"{operation!r}" + ) + return operation.response.destination_uri diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index e7612f28d1..b0fa3d2b3c 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -67,6 +67,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 @@ -96,6 +97,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 @@ -622,6 +624,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 @@ -739,6 +753,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 @@ -3499,6 +3516,18 @@ "RecommendSpecResponse", "RecommendSpecResponseDict", "RecommendSpecResponseOrDict", + "ExportPublisherModelConfig", + "ExportPublisherModelConfigDict", + "ExportPublisherModelConfigOrDict", + "ExportPublisherModelResponse", + "ExportPublisherModelResponseDict", + "ExportPublisherModelResponseOrDict", + "ExportModelOperation", + "ExportModelOperationDict", + "ExportModelOperationOrDict", + "GetExportPublisherModelOperationConfig", + "GetExportPublisherModelOperationConfigDict", + "GetExportPublisherModelOperationConfigOrDict", "CreateRuntimeFeedbackEntryConfig", "CreateRuntimeFeedbackEntryConfigDict", "CreateRuntimeFeedbackEntryConfigOrDict", @@ -3646,6 +3675,9 @@ "ListCustomModelDeployOptionsConfig", "ListCustomModelDeployOptionsConfigDict", "ListCustomModelDeployOptionsConfigOrDict", + "ExportOpenModelConfig", + "ExportOpenModelConfigDict", + "ExportOpenModelConfigOrDict", "DeployOption", "DeployOptionDict", "DeployOptionOrDict", @@ -3840,6 +3872,8 @@ "_ListPublisherModelsRequestParameters", "_GetPublisherModelRequestParameters", "_RecommendSpecRequestParameters", + "_ExportPublisherModelRequestParameters", + "_GetExportPublisherModelOperationParameters", "_CreateRuntimeFeedbackEntryRequestParameters", "_DeleteRuntimeFeedbackEntryRequestParameters", "_GetRuntimeFeedbackRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index cd6fb70b6c..60223bacdc 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -24029,6 +24029,193 @@ 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. + + Fields are re-declared as ``SdkFieldPatch`` (both are already in the + discovery-generated class) so the SDK's dependency on ``destination_uri`` + is visible in one place and proto drift is caught at codegen time + instead of at first user call. + """ + + destination_uri: Optional[str] = Field( + default=None, + description="""Cloud Storage URI where the exported weights were written.""", + ) + publisher_model: Optional[str] = Field( + default=None, + description="""Resource name of the publisher model that was exported.""", + ) + + +class ExportPublisherModelResponseDict(TypedDict, total=False): + """Response for the ``ExportPublisherModel`` RPC. + + Fields are re-declared as ``SdkFieldPatch`` (both are already in the + discovery-generated class) so the SDK's dependency on ``destination_uri`` + is visible in one place and proto drift is caught at codegen time + instead of at first user call. + """ + + destination_uri: Optional[str] + """Cloud Storage URI where the exported weights were written.""" + + publisher_model: Optional[str] + """Resource name of the publisher model that was exported.""" + + +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.""" @@ -27004,6 +27191,53 @@ class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): ] +class ExportOpenModelConfig(_common.BaseModel): + """Config for ``export_open_model``.""" + + 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.""", + ) + poll_interval_seconds: Optional[float] = Field( + default=None, + description="""Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""", + ) + timeout_seconds: Optional[float] = Field( + default=None, + description="""Total wall-clock seconds to wait for the export to complete + when ``wait_for_completion=True``. Defaults to 2 hours to + accommodate large model weights. Ignored when + ``wait_for_completion=False``.""", + ) + + +class ExportOpenModelConfigDict(TypedDict, total=False): + """Config for ``export_open_model``.""" + + 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.""" + + poll_interval_seconds: Optional[float] + """Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""" + + timeout_seconds: Optional[float] + """Total wall-clock seconds to wait for the export to complete + when ``wait_for_completion=True``. Defaults to 2 hours to + accommodate large model weights. Ignored when + ``wait_for_completion=False``.""" + + +ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] + + class DeployOption(_common.BaseModel): """A verified deploy option for a model.""" diff --git a/noxfile.py b/noxfile.py index e6c97e2d5f..b43a081d26 100644 --- a/noxfile.py +++ b/noxfile.py @@ -105,11 +105,11 @@ "unit_ray", "unit_langchain", "unit_ag2", - "unit_llama_index", + # "unit_llama_index", "unit_agentplatform_adk", "unit_agentplatform_langchain", "unit_agentplatform_ag2", - "unit_agentplatform_llama_index", + # "unit_agentplatform_llama_index", "unit_agentplatform_a2a", "unit_a2a", "system", 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..0702fe1fdf 100644 --- a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py @@ -36,11 +36,11 @@ 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( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=False, - model_filter="timesfm", - ) - ) + 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() @@ -83,9 +83,9 @@ def test_list_publisher_model_deploy_options_concise(client): def test_list_publisher_model_deploy_options_hugging_face(client): """Tests deploy options for a Hugging Face model. - Exercises the distinct GetPublisherModel request path where - is_hugging_face_model=True is sent. - """ + 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" ) @@ -143,7 +143,9 @@ 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), + config=types.ListCustomModelDeployOptionsConfig( + filter_by_user_quota=False + ), ) assert isinstance(options, str) assert "[Option 1]" in options @@ -161,6 +163,95 @@ def test_list_custom_model_deploy_options_dict_config(client): 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 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", + 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", + output_gcs_uri="gs://mg-sdk-model-export-testing/gemma-2-2b-it-no-wait/", + config=types.ExportOpenModelConfig(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", + output_gcs_uri=( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-dict-no-wait/" + ), + config={"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", + output_gcs_uri="gs://mg-sdk-model-export-testing/gemma-2-2b-it-poll/", + config=types.ExportOpenModelConfig(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_empty_uri_raises(client): + """Empty output_gcs_uri is rejected client-side; no RPC needed.""" + with pytest.raises(ValueError, match="output_gcs_uri must be a non-empty"): + client.model_garden.export_open_model( + model="google/gemma3@gemma-3-12b-it", output_gcs_uri="" + ) + + +def test_export_open_model_hf_model_id_raises(client): + """Hugging Face model IDs are rejected client-side; no RPC needed.""" + with pytest.raises(ValueError, match="does not support Hugging Face"): + client.model_garden.export_open_model( + model="meta-llama/Llama-3.3-70B-Instruct", + output_gcs_uri="gs://mg-sdk-model-export-testing/hf-should-not-record/", + ) + + +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", output_gcs_uri="gs://b/out/" + ) + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), @@ -220,3 +311,34 @@ async def test_list_custom_model_deploy_options_async(client): 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", + output_gcs_uri=( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-async-no-wait/" + ), + config=types.ExportOpenModelConfig(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", + output_gcs_uri=( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-async-poll/" + ), + config=types.ExportOpenModelConfig(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..f46752cebb 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 @@ -219,8 +222,8 @@ def test_list_deployable_models_hf_filter_string(mock_client): def test_list_models_excludes_hf_via_server_filter(mock_client): - """Tests that is_hf_wildcard(false) is sent to exclude HF models server-side.""" - dummy_response = types.ListPublisherModelsResponse( + """Tests that is_hf_wildcard(false) is sent to exclude HF models server-side.""" + dummy_response = types.ListPublisherModelsResponse( publisher_models=[ types.PublisherModel( name="publishers/google/models/gemma-2b", version_id="001" @@ -228,24 +231,24 @@ def test_list_models_excludes_hf_via_server_filter(mock_client): ] ) - with mock.patch.object( + with mock.patch.object( 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 - ) + models = mock_client.list_models( + config=types.ListModelGardenModelsConfig( + include_hugging_face_models=False ) + ) - api_config = mock_list.call_args.kwargs.get("config") - assert "is_hf_wildcard(false)" in api_config.filter - assert len(models) == 1 - assert models[0] == "google/gemma-2b@001" + api_config = mock_list.call_args.kwargs.get("config") + assert "is_hf_wildcard(false)" in api_config.filter + assert len(models) == 1 + assert models[0] == "google/gemma-2b@001" def test_list_models_with_hf(mock_client): - """Tests list_models with include_hugging_face_models=True.""" - dummy_response = types.ListPublisherModelsResponse( + """Tests list_models with include_hugging_face_models=True.""" + dummy_response = types.ListPublisherModelsResponse( publisher_models=[ types.PublisherModel( name="publishers/google/models/gemma-2b", version_id="001" @@ -256,27 +259,27 @@ def test_list_models_with_hf(mock_client): ] ) - with mock.patch.object( + with mock.patch.object( mock_client, "_list_publisher_models", return_value=dummy_response ): - models = mock_client.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=True - ) + models = mock_client.list_models( + config=types.ListModelGardenModelsConfig( + include_hugging_face_models=True ) + ) - assert len(models) == 2 - assert models[0] == "google/gemma-2b" - assert models[1] == "meta/llama-3" + assert len(models) == 2 + assert models[0] == "google/gemma-2b" + assert models[1] == "meta/llama-3" def test_list_models_includes_non_deployable(mock_client): - """Tests that list_models includes models without deploy configs. + """Tests that list_models includes models without deploy configs. Unlike list_deployable_models, list_models should return ALL models regardless of whether they have multi_deploy_vertex configurations. """ - dummy_response = types.ListPublisherModelsResponse( + dummy_response = types.ListPublisherModelsResponse( publisher_models=[ _make_deployable_model("publishers/google/models/gemma-2b"), # This model has no deploy config -- should still be included @@ -286,18 +289,18 @@ def test_list_models_includes_non_deployable(mock_client): ] ) - with mock.patch.object( + with mock.patch.object( mock_client, "_list_publisher_models", return_value=dummy_response ): - models = mock_client.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=False - ) + models = mock_client.list_models( + config=types.ListModelGardenModelsConfig( + include_hugging_face_models=False ) + ) - assert len(models) == 2 - assert models[0] == "google/gemma-2b@001" - assert models[1] == "google/bert-base@001" + assert len(models) == 2 + assert models[0] == "google/gemma-2b@001" + assert models[1] == "google/bert-base@001" def test_list_models_default_config(mock_client): @@ -419,8 +422,10 @@ 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 - ) + model_filter="model.v2+", + include_hugging_face_models=False, + deployable_only=False, + ) # re.escape turns '.' into '\\.' and '+' into '\\+' assert r"model\.v2\+" in result @@ -435,8 +440,8 @@ def _make_deploy_option( accelerator_count=1, image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/vllm", ): - """Builds a single PublisherModelCallToActionDeploy option.""" - return types.PublisherModelCallToActionDeploy( + """Builds a single PublisherModelCallToActionDeploy option.""" + return types.PublisherModelCallToActionDeploy( deploy_task_name=deploy_task_name, dedicated_resources=types.DedicatedResources( machine_spec=types.MachineSpec( @@ -454,13 +459,13 @@ def _make_model_with_deploy_options( ): """Builds a PublisherModel exposing the given deploy options.""" return types.PublisherModel( - name=name, - supported_actions=types.PublisherModelCallToAction( - multi_deploy_vertex=types.PublisherModelCallToActionDeployVertex( - multi_deploy_vertex=options, - ) - ), - ) + name=name, + supported_actions=types.PublisherModelCallToAction( + multi_deploy_vertex=types.PublisherModelCallToActionDeployVertex( + multi_deploy_vertex=options, + ) + ), + ) def test_list_publisher_model_deploy_options_basic(mock_client): @@ -468,8 +473,8 @@ def test_list_publisher_model_deploy_options_basic(mock_client): dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ) as mock_get: + 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" @@ -497,12 +502,10 @@ def test_list_publisher_model_deploy_options_basic(mock_client): 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( - [ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - ] - ) + 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 @@ -516,12 +519,10 @@ def test_list_publisher_model_deploy_options_multiple(mock_client): 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( - [ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - ] - ) + 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 @@ -564,12 +565,10 @@ def test_list_publisher_model_deploy_options_accelerator_type_filter( 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( - [ - _make_deploy_option(deploy_task_name="vllm", image_uri="docker/vllm"), - _make_deploy_option(deploy_task_name="tgi", image_uri="docker/tgi"), - ] - ) + 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 @@ -637,13 +636,11 @@ def test_list_publisher_model_deploy_options_accelerator_type_filter_list( 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( - [ - _make_deploy_option(deploy_task_name="vllm", image_uri="docker/vllm"), - _make_deploy_option(deploy_task_name="tgi", image_uri="docker/tgi"), - _make_deploy_option(deploy_task_name="sglang", image_uri="docker/sglang"), - ] - ) + 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"), + _make_deploy_option(deploy_task_name="sglang", image_uri="docker/sglang"), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -678,12 +675,10 @@ def test_matches_filter(): 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( - [ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - ] - ) + 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 @@ -820,8 +815,8 @@ def test_list_publisher_model_deploy_options_hugging_face_model(mock_client): dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ) as mock_get: + 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" ) @@ -840,10 +835,10 @@ def test_list_publisher_model_deploy_options_async(mock_async_client): dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) with mock.patch.object( - mock_async_client, - "_get_publisher_model", - new=mock.AsyncMock(return_value=dummy_model), - ): + 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" @@ -862,10 +857,10 @@ def test_list_publisher_model_deploy_options_async_no_deploy_support_raises( dummy_model = types.PublisherModel(name="publishers/google/models/bert-base") with mock.patch.object( - mock_async_client, - "_get_publisher_model", - new=mock.AsyncMock(return_value=dummy_model), - ): + 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( @@ -881,12 +876,12 @@ def _make_deploy_option_no_accelerator( ): """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) - ), - container_spec=types.ModelContainerSpec(image_uri=image_uri), - ) + deploy_task_name=deploy_task_name, + dedicated_resources=types.DedicatedResources( + machine_spec=types.MachineSpec(machine_type=machine_type) + ), + container_spec=types.ModelContainerSpec(image_uri=image_uri), + ) def test_list_publisher_model_deploy_options_no_accelerator_defaults( @@ -940,12 +935,10 @@ def test_list_publisher_model_deploy_options_accelerator_filter_excludes_unspeci mock_client, ): """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"), - ] - ) + 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 @@ -1017,24 +1010,22 @@ def test_list_publisher_model_deploy_options_concise_returns_string( 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( - [ - _make_deploy_option( - deploy_task_name="g2", - machine_type="g2-standard-12", - accelerator_type="NVIDIA_L4", - accelerator_count=1, - image_uri="docker/vllm", - ), - _make_deploy_option( - deploy_task_name="a3", - machine_type="a3-highgpu-8g", - accelerator_type="NVIDIA_H100_80GB", - accelerator_count=8, - image_uri="docker/hexllm", - ), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option( + deploy_task_name="g2", + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + image_uri="docker/vllm", + ), + _make_deploy_option( + deploy_task_name="a3", + machine_type="a3-highgpu-8g", + accelerator_type="NVIDIA_H100_80GB", + accelerator_count=8, + image_uri="docker/hexllm", + ), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -1062,12 +1053,10 @@ def test_list_publisher_model_deploy_options_concise_multiple(mock_client): 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( - [ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - ] - ) + 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 @@ -1087,39 +1076,35 @@ def test_list_publisher_model_deploy_options_concise_with_filter(mock_client): def test_format_concise_deploy_options_omits_none_fields(): """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, - ) - ] + 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," - ) + '[Option 1]\n machine_type="g2-standard-12",\n accelerator_count=1,' + ) 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( - [ - _make_deploy_option( - deploy_task_name="option-1", - machine_type="g2-standard-12", - accelerator_type="NVIDIA_L4", - accelerator_count=1, - image_uri="docker/vllm", - ) - ] - ) + 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_async_client, - "_get_publisher_model", - new=mock.AsyncMock(return_value=dummy_model), - ): + 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", @@ -1399,3 +1384,389 @@ def test_list_custom_model_deploy_options_async_src_required_raises( """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.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ) as m_await, + ): + result = mock_client.export_open_model( + model=_OPEN_MODEL, 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.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ) as m_await, + ): + mock_client.export_open_model( + model=_OPEN_MODEL, 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_uses_default_poll_and_timeout(mock_client): + """No override in config -> ModelGarden defaults are forwarded to the poller.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ) as m_await, + ): + mock_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="gs://b/out/" + ) + kwargs = m_await.call_args.kwargs + assert ( + kwargs["poll_interval"] == 30 + ) # _DEFAULT_EXPORT_POLL_INTERVAL_SECONDS + assert ( + kwargs["timeout_seconds"] == 2 * 60 * 60 + ) # _DEFAULT_EXPORT_TIMEOUT_SECONDS + + +def test_export_open_model_config_overrides_poll_and_timeout(mock_client): + """Config values override the defaults on the polling call.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ) as m_await, + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig( + poll_interval_seconds=5, + timeout_seconds=600, + ), + ) + kwargs = m_await.call_args.kwargs + assert kwargs["poll_interval"] == 5 + assert kwargs["timeout_seconds"] == 600 + + +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.object( + 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, output_gcs_uri="gs://b/out/" + ) + + +def test_export_open_model_raises_when_response_missing_destination_uri( + mock_client, +): + """A done LRO with no destination_uri surfaces as RuntimeError.""" + done_op = _make_export_operation(done=True) # no destination_uri set + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ), + ): + with pytest.raises(RuntimeError, match="destination_uri"): + mock_client.export_open_model( + model=_OPEN_MODEL, 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.object( + model_garden._operations_utils, "await_operation" + ) as m_await, + ): + result = mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + m_await.assert_not_called() + assert result is initial_op + + +# --- Input validation ------------------------------------------------- + + +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 a non-empty"): + mock_client.export_open_model(model=_OPEN_MODEL, output_gcs_uri="") + m_rpc.assert_not_called() + + +def test_export_open_model_hf_model_id_raises(mock_client): + """Hugging Face model IDs are rejected client-side with a targeted message.""" + with _patch_export_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="does not support Hugging Face"): + mock_client.export_open_model( + model="meta-llama/Llama-3.3-70B-Instruct", + output_gcs_uri="gs://b/out/", + ) + m_rpc.assert_not_called() + + +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", output_gcs_uri="gs://b/out/" + ) + m_rpc.assert_not_called() + + +def test_export_open_model_dict_config_validated(mock_client): + """A plain dict is validated into ExportOpenModelConfig.""" + initial_op = _make_export_operation() + with ( + _patch_export_rpc(mock_client, initial_op=initial_op), + mock.patch.object( + model_garden._operations_utils, "await_operation" + ) as m_await, + ): + result = mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config={"wait_for_completion": False}, + ) + m_await.assert_not_called() + assert result is initial_op + + +# --- Model-name resolution --------------------------------------------- + + +def test_export_open_model_simplified_name_reconciled(mock_client): + """Simplified '{publisher}/{model}@{version}' -> full resource name.""" + with ( + _patch_export_rpc(mock_client) as m_rpc, + mock.patch.object(model_garden._operations_utils, "await_operation"), + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig(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.object(model_garden._operations_utils, "await_operation"), + ): + mock_client.export_open_model( + model="publishers/google/models/gemma3@gemma-3-12b-it", + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + assert ( + m_rpc.call_args.kwargs["name"] + == "publishers/google/models/gemma3@gemma-3-12b-it" + ) + + +# --- Destination forwarded to RPC config -------------------------------- + + +def test_export_open_model_forwards_output_gcs_uri(mock_client): + """Top-level output_gcs_uri becomes destination.output_uri_prefix on the RPC.""" + with ( + _patch_export_rpc(mock_client) as m_rpc, + mock.patch.object(model_garden._operations_utils, "await_operation"), + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://my-bucket/weights/", + config=types.ExportOpenModelConfig(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.object( + 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, 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.object( + model_garden._operations_utils, "await_operation_async" + ) as m_await, + ): + result = asyncio.run( + mock_async_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + ) + m_await.assert_not_called() + assert result is initial_op + + +def test_export_open_model_async_empty_uri_raises(mock_async_client): + """Async missing output_gcs_uri 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 must be a non-empty"): + asyncio.run( + mock_async_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="" + ) + ) + m_rpc.assert_not_called() + + +def test_export_open_model_async_hf_model_id_raises(mock_async_client): + """Async Hugging Face model IDs are rejected client-side.""" + with _patch_async_export_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="does not support Hugging Face"): + asyncio.run( + mock_async_client.export_open_model( + model="meta-llama/Llama-3.3-70B-Instruct", + output_gcs_uri="gs://b/out/", + ) + ) + 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", output_gcs_uri="gs://b/out/" + ) + ) + m_rpc.assert_not_called() diff --git a/tests/unit/agentplatform/genai/test_model_garden.py b/tests/unit/agentplatform/genai/test_model_garden.py new file mode 100644 index 0000000000..e820283234 --- /dev/null +++ b/tests/unit/agentplatform/genai/test_model_garden.py @@ -0,0 +1,198 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for stateless GenAI SDK Model Garden extension.""" + +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from unittest import mock +from agentplatform._genai import model_garden +from agentplatform._genai import types +from google.genai import client +import pytest + +_TEST_PROJECT = "test-project" +_TEST_LOCATION = "us-central1" + + +@pytest.fixture +def mock_client(): + mock_api_client = mock.Mock(spec=client.Client) + mock_api_client.project = _TEST_PROJECT + mock_api_client.location = _TEST_LOCATION + mock_api_client.vertexai = True + + return model_garden.ModelGarden(mock_api_client) + + +def test_list_deployable_models_mocked(mock_client): + dummy_response = types.ListPublisherModelsResponse( + publisher_models=[ + types.PublisherModel(name="publishers/google/models/gemma-2b"), + types.PublisherModel(name="publishers/hf-google/models/gemma-2b"), + ] + ) + + with mock.patch.object( + mock_client, "_list_publisher_models", return_value=dummy_response + ) as mock_list: + models = mock_client.list_deployable_models( + config=types.ListDeployableModelsConfig( + include_hugging_face_models=False + ) + ) + + mock_list.assert_called_once() + # Check that client-side Hugging Face filtering worked + assert len(models) == 1 + assert models[0].name == "publishers/google/models/gemma-2b" + + +def test_list_open_model_deploy_options_mocked(mock_client): + dummy_model = types.PublisherModel( + name="publishers/google/models/gemma-2b", + supported_actions=types.PublisherModelCallToAction( + multi_deploy_vertex=types.PublisherModelCallToActionDeployVertex( + multi_deploy_vertex=[ + types.PublisherModelCallToActionDeploy( + deploy_task_name="option-1", + dedicated_resources=types.DedicatedResources( + machine_spec=types.MachineSpec( + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + ) + ), + container_spec=types.ModelContainerSpec( + image_uri="us-docker.pkg.dev/vllm-serve" + ), + ) + ] + ) + ), + ) + + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ) as mock_get: + options = mock_client.list_open_model_deploy_options( + model="publishers/google/models/gemma-2b" + ) + + mock_get.assert_called_once_with( + name="publishers/google/models/gemma-2b", + config=types.GetPublisherModelConfig( + include_equivalent_model_garden_model_deployment_configs=True + ), + ) + assert len(options) == 1 + assert options[0].option_name == "option-1" + assert options[0].machine_type == "g2-standard-12" + + +def test_list_partner_model_deploy_options_mocked(mock_client): + dummy_model = types.PublisherModel( + name="publishers/meta/models/llama-3-8b", + supported_actions=types.PublisherModelCallToAction( + multi_deploy_vertex=types.PublisherModelCallToActionDeployVertex( + multi_deploy_vertex=[ + types.PublisherModelCallToActionDeploy( + deploy_task_name="partner-option", + dedicated_resources=types.DedicatedResources( + machine_spec=types.MachineSpec( + machine_type="g2-standard-24", + accelerator_type="NVIDIA_L4", + accelerator_count=2, + ) + ), + container_spec=types.ModelContainerSpec( + image_uri="us-docker.pkg.dev/partner-serve" + ), + ) + ] + ) + ), + ) + + with mock.patch.object( + mock_client, "_get_publisher_model", return_value=dummy_model + ) as mock_get: + options = mock_client.list_partner_model_deploy_options( + model="publishers/meta/models/llama-3-8b" + ) + + mock_get.assert_called_once_with( + name="publishers/meta/models/llama-3-8b", + config=types.GetPublisherModelConfig( + include_equivalent_model_garden_model_deployment_configs=True + ), + ) + assert len(options) == 1 + assert options[0].option_name == "partner-option" + assert options[0].machine_type == "g2-standard-24" + + +def test_list_custom_model_deploy_options_mocked(mock_client): + 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 + ) as mock_recommend: + options_str = mock_client.list_custom_model_deploy_options( + src="gs://weights-path" + ) + + mock_recommend.assert_called_once() + assert "Option 1" in options_str + assert 'machine_type="g2-standard-12"' in options_str + + +def test_deploy_open_model_mocked(mock_client): + dummy_operation = types.DeployModelOperation( + name="projects/123/locations/us-central1/operations/456" + ) + + with mock.patch.object( + mock_client, "_deploy", return_value=dummy_operation + ) as mock_deploy: + op = mock_client.deploy_open_model( + model="publishers/google/models/gemma-2b", + config=types.DeployOpenModelConfig( + machine_type="g2-standard-8", + min_replica_count=1, + max_replica_count=2, + accelerator_type="NVIDIA_L4", + accelerator_count=1, + ), + ) + + mock_deploy.assert_called_once() + call_kwargs = mock_deploy.call_args[1] + assert ( + call_kwargs["publisher_model_name"] + == "publishers/google/models/gemma-2b" + ) + assert call_kwargs["hugging_face_model_id"] is None + assert op.name == "projects/123/locations/us-central1/operations/456"