From 75d31cd03b1061a60483fea1a0d23850af0a5dbb Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 14 Jul 2026 11:33:26 -0700 Subject: [PATCH] feat: Onboard Vertex Model Garden to GenAI Python SDK: Add deploy_publisher_model support PiperOrigin-RevId: 947810037 --- agentplatform/_genai/model_garden.py | 627 ++++++++++++++ agentplatform/_genai/types/__init__.py | 72 ++ agentplatform/_genai/types/common.py | 657 +++++++++++++++ .../genai/replays/test_genai_model_garden.py | 27 + .../genai/test_genai_model_garden.py | 765 ++++++++++++++++++ 5 files changed, 2148 insertions(+) diff --git a/agentplatform/_genai/model_garden.py b/agentplatform/_genai/model_garden.py index 61d19be50f..2577fba646 100644 --- a/agentplatform/_genai/model_garden.py +++ b/agentplatform/_genai/model_garden.py @@ -32,6 +32,43 @@ logger = logging.getLogger("agentplatform_genai.modelgarden") +def _DeployRequestParameters_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(to_object, ["_url", "destination"], getv(from_object, ["destination"])) + + if getv(from_object, ["publisher_model_name"]) is not None: + setv( + to_object, + ["publisherModelName"], + getv(from_object, ["publisher_model_name"]), + ) + + if getv(from_object, ["hugging_face_model_id"]) is not None: + setv( + to_object, + ["huggingFaceModelId"], + getv(from_object, ["hugging_face_model_id"]), + ) + + if getv(from_object, ["custom_model"]) is not None: + setv(to_object, ["customModel"], getv(from_object, ["custom_model"])) + + if getv(from_object, ["model_config_val"]) is not None: + setv(to_object, ["modelConfig"], getv(from_object, ["model_config_val"])) + + if getv(from_object, ["endpoint_config"]) is not None: + setv(to_object, ["endpointConfig"], getv(from_object, ["endpoint_config"])) + + if getv(from_object, ["deploy_config"]) is not None: + setv(to_object, ["deployConfig"], getv(from_object, ["deploy_config"])) + + return to_object + + def _ExportPublisherModelConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -61,6 +98,19 @@ def _ExportPublisherModelRequestParameters_to_vertex( return to_object +def _GetDeployOperationParameters_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 _GetExportPublisherModelOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -571,12 +621,173 @@ def get_export_publisher_model_operation( self._api_client._verify_response(return_value) return return_value + def _deploy( + self, + *, + destination: str, + publisher_model_name: Optional[str] = None, + hugging_face_model_id: Optional[str] = None, + custom_model: Optional[types.DeployRequestCustomModelOrDict] = None, + model_config_val: Optional[types.DeployRequestModelConfigOrDict] = None, + endpoint_config: Optional[types.DeployRequestEndpointConfigOrDict] = None, + deploy_config: Optional[types.DeployRequestDeployConfigOrDict] = None, + config: Optional[types.DeployConfigOrDict] = None, + ) -> types.DeployModelOperation: + """ + Deploys a model (internal). + """ + + parameter_model = types._DeployRequestParameters( + destination=destination, + publisher_model_name=publisher_model_name, + hugging_face_model_id=hugging_face_model_id, + custom_model=custom_model, + model_config_val=model_config_val, + endpoint_config=endpoint_config, + deploy_config=deploy_config, + 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 = _DeployRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{destination}:deploy".format_map(request_url_dict) + else: + path = "{destination}:deploy" + + 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.DeployModelOperation._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_deploy_publisher_model_operation( + self, + *, + operation_name: str, + config: Optional[types.GetDeployOperationConfigOrDict] = None, + ) -> types.DeployModelOperation: + """ + Fetches the status of an in-flight ``deploy_publisher_model`` LRO. + """ + + parameter_model = types._GetDeployOperationParameters( + 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 = _GetDeployOperationParameters_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.DeployModelOperation._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 + # Deploy LRO polling defaults. 2h matches legacy vertexai.model_garden and + # the Vertex AI Console's one-click deployment timeout. + _DEFAULT_DEPLOY_TIMEOUT_SECONDS = 2 * 60 * 60 + _DEFAULT_DEPLOY_POLL_INTERVAL_SECONDS = 30 + @staticmethod def _build_filter_str( model_filter: Optional[str], @@ -893,6 +1104,134 @@ def _format_concise_deploy_options( blocks.append(header + "\n".join(lines)) return "\n\n".join(blocks) + @staticmethod + def _resolve_deploy_model_name(model: str) -> tuple[Optional[str], Optional[str]]: + """Returns the ``(publisher_model_name, hugging_face_model_id)`` pair for a deploy call. + + Exactly one element of the returned pair is set: HF model IDs (matched + by ``_is_hugging_face_model``) are sent as ``huggingFaceModelId`` and + lowercased for legacy parity; everything else is reconciled to the full + ``publishers/{pub}/models/{model}@{version}`` form. + """ + if ModelGarden._is_hugging_face_model(model): + return None, model.lower() + return ModelGarden._reconcile_model_name(model), None + + @staticmethod + def _validate_deploy_config( + config: types.DeployPublisherModelConfig, + ) -> None: + """Rejects deploy configs whose options would be silently dropped. + + Container overrides are only sent when ``serving_container_image_uri`` + is set (the backend has no way to patch a model's default container + piecemeal), so a config that sets only the command/args/env would + otherwise deploy successfully with none of the overrides applied. + + Raises: + ValueError: If any container override is set without + ``serving_container_image_uri``. + """ + if config.serving_container_image_uri: + return + dropped = [ + name + for name in ("container_command", "container_args", "container_variables") + if getattr(config, name) + ] + if dropped: + raise ValueError( + f"{', '.join(dropped)} require serving_container_image_uri to " + "also be set; container overrides are only applied on top of an " + "explicit serving container image." + ) + + @staticmethod + def _deploy_response_or_raise( + operation: types.DeployModelOperation, + ) -> types.DeployResponse: + """Unwraps a completed deploy LRO into its ``DeployResponse``. + + Shared by the sync and async surfaces so the two cannot drift as the + LRO shape evolves. + + Raises: + RuntimeError: If the LRO carries an error, or completed without an + endpoint resource name. + """ + if operation.error: + raise RuntimeError(f"Deploy failed: {operation.error}") + if not operation.response or not operation.response.endpoint: + raise RuntimeError( + f"Deploy completed but response has no endpoint: {operation!r}" + ) + return operation.response + + @staticmethod + def _build_container_spec( + config: types.DeployPublisherModelConfig, + ) -> Optional[types.ModelContainerSpec]: + """Returns a ``ModelContainerSpec`` when the user overrides the container, else None.""" + if not config.serving_container_image_uri: + return None + env = None + variables = config.container_variables + if variables: + env = [types.EnvVar(name=k, value=v) for k, v in variables.items()] + return types.ModelContainerSpec( + image_uri=config.serving_container_image_uri, + command=config.container_command, + args=config.container_args, + env=env, + ) + + @staticmethod + def _prepare_deploy_request( + config: types.DeployPublisherModelConfig, + ) -> tuple[ + types.DeployRequestModelConfig, + types.DeployRequestEndpointConfig, + types.DeployRequestDeployConfig, + ]: + """Translates ``DeployPublisherModelConfig`` to the three ``DeployRequest`` sub-messages.""" + model_config = types.DeployRequestModelConfig( + accept_eula=config.accept_eula, + model_display_name=config.model_display_name, + hugging_face_access_token=config.hugging_face_access_token, + container_spec=ModelGarden._build_container_spec(config), + ) + + endpoint_config = types.DeployRequestEndpointConfig( + endpoint_display_name=config.endpoint_display_name, + ) + disabled = config.dedicated_endpoint_disabled + if disabled is not None: + endpoint_config.dedicated_endpoint_enabled = not disabled + if config.enable_private_service_connect: + endpoint_config.private_service_connect_config = ( + types.PrivateServiceConnectConfig( + enable_private_service_connect=True, + project_allowlist=config.psc_project_allow_list, + ) + ) + + deploy_config = types.DeployRequestDeployConfig( + fast_tryout_enabled=config.fast_tryout_enabled, + ) + if config.machine_type or config.accelerator_type or config.accelerator_count: + deploy_config.dedicated_resources = types.DedicatedResources( + machine_spec=types.MachineSpec( + machine_type=config.machine_type, + accelerator_type=config.accelerator_type, + accelerator_count=config.accelerator_count, + ), + min_replica_count=config.min_replica_count, + max_replica_count=config.max_replica_count, + spot=config.spot, + ) + + return model_config, endpoint_config, deploy_config + def _list_all_publisher_models( self, api_config: types.ListPublisherModelsConfig, @@ -1081,6 +1420,86 @@ def list_publisher_model_deploy_options( return options + def deploy_publisher_model( + self, + *, + model: str, + config: Optional[types.DeployPublisherModelConfigOrDict] = None, + ) -> Union[types.DeployResponse, types.DeployModelOperation]: + """Deploys a Model Garden publisher model to a Vertex AI endpoint. + + Supports Google open models (e.g. ``'google/gemma3@gemma-3-12b-it'``), + partner publisher models (e.g. ``'ai21/jamba-large-1.6@001'``), and + Hugging Face model IDs (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). + + Args: + model: The publisher model to deploy. Accepts the full resource name + ``'publishers/{publisher}/models/{model}@{version}'``, a simplified + ``'{publisher}/{model}@{version}'`` (or without the ``@{version}``), + or a Hugging Face model ID ``'{organization}/{model}'``. Hugging + Face model IDs are lowercased before being sent, matching + ``vertexai.model_garden.OpenModel.deploy`` behavior. + config: Optional deployment configuration (machine shape, container + overrides, blocking behavior, poll timing). See + ``DeployPublisherModelConfig``. + + Returns: + When ``config.wait_for_completion`` is ``True`` (default), a + ``DeployResponse`` carrying the deployed ``endpoint`` and ``model`` + resource names (both ``str``). + When ``config.wait_for_completion`` is ``False``, the + ``DeployModelOperation`` for the caller to poll (via + ``get_deploy_publisher_model_operation`` or their own strategy). + + Raises: + ValueError: If ``model`` is not a valid publisher model name, or if + container overrides are set without ``serving_container_image_uri``. + 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 an endpoint resource name. + """ + if config is None: + config = types.DeployPublisherModelConfig() + elif isinstance(config, dict): + config = types.DeployPublisherModelConfig.model_validate(config) + ModelGarden._validate_deploy_config(config) + + publisher_model_name, hugging_face_model_id = ( + ModelGarden._resolve_deploy_model_name(model) + ) + model_config, endpoint_config, deploy_config = ( + ModelGarden._prepare_deploy_request(config) + ) + destination = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + + operation = self._deploy( + destination=destination, + publisher_model_name=publisher_model_name, + hugging_face_model_id=hugging_face_model_id, + model_config_val=model_config, + endpoint_config=endpoint_config, + deploy_config=deploy_config, + ) + if not config.wait_for_completion: + return operation + + operation = _operations_utils.await_operation( + operation_name=operation.name, + get_operation_fn=self.get_deploy_publisher_model_operation, + poll_interval=( + config.poll_interval_seconds + or ModelGarden._DEFAULT_DEPLOY_POLL_INTERVAL_SECONDS + ), + timeout_seconds=( + config.timeout_seconds or ModelGarden._DEFAULT_DEPLOY_TIMEOUT_SECONDS + ), + ) + return ModelGarden._deploy_response_or_raise(operation) + @staticmethod def _extract_recommend_spec(spec) -> dict[str, Any]: """Extracts machine spec fields from a single recommend-spec entry. @@ -1695,6 +2114,166 @@ async def get_export_publisher_model_operation( self._api_client._verify_response(return_value) return return_value + async def _deploy( + self, + *, + destination: str, + publisher_model_name: Optional[str] = None, + hugging_face_model_id: Optional[str] = None, + custom_model: Optional[types.DeployRequestCustomModelOrDict] = None, + model_config_val: Optional[types.DeployRequestModelConfigOrDict] = None, + endpoint_config: Optional[types.DeployRequestEndpointConfigOrDict] = None, + deploy_config: Optional[types.DeployRequestDeployConfigOrDict] = None, + config: Optional[types.DeployConfigOrDict] = None, + ) -> types.DeployModelOperation: + """ + Deploys a model (internal). + """ + + parameter_model = types._DeployRequestParameters( + destination=destination, + publisher_model_name=publisher_model_name, + hugging_face_model_id=hugging_face_model_id, + custom_model=custom_model, + model_config_val=model_config_val, + endpoint_config=endpoint_config, + deploy_config=deploy_config, + 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 = _DeployRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{destination}:deploy".format_map(request_url_dict) + else: + path = "{destination}:deploy" + + 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.DeployModelOperation._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_deploy_publisher_model_operation( + self, + *, + operation_name: str, + config: Optional[types.GetDeployOperationConfigOrDict] = None, + ) -> types.DeployModelOperation: + """ + Fetches the status of an in-flight ``deploy_publisher_model`` LRO. + """ + + parameter_model = types._GetDeployOperationParameters( + 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 = _GetDeployOperationParameters_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.DeployModelOperation._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, @@ -1880,6 +2459,54 @@ async def list_publisher_model_deploy_options( return options + async def deploy_publisher_model( + self, + *, + model: str, + config: Optional[types.DeployPublisherModelConfigOrDict] = None, + ) -> Union[types.DeployResponse, types.DeployModelOperation]: + """Async variant of ``ModelGarden.deploy_publisher_model``.""" + if config is None: + config = types.DeployPublisherModelConfig() + elif isinstance(config, dict): + config = types.DeployPublisherModelConfig.model_validate(config) + ModelGarden._validate_deploy_config(config) + + publisher_model_name, hugging_face_model_id = ( + ModelGarden._resolve_deploy_model_name(model) + ) + model_config, endpoint_config, deploy_config = ( + ModelGarden._prepare_deploy_request(config) + ) + destination = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + + operation = await self._deploy( + destination=destination, + publisher_model_name=publisher_model_name, + hugging_face_model_id=hugging_face_model_id, + model_config_val=model_config, + endpoint_config=endpoint_config, + deploy_config=deploy_config, + ) + if not config.wait_for_completion: + return operation + + operation = await _operations_utils.await_operation_async( + operation_name=operation.name, + get_operation_fn=self.get_deploy_publisher_model_operation, + poll_interval=( + config.poll_interval_seconds + or ModelGarden._DEFAULT_DEPLOY_POLL_INTERVAL_SECONDS + ), + timeout_seconds=( + config.timeout_seconds or ModelGarden._DEFAULT_DEPLOY_TIMEOUT_SECONDS + ), + ) + return ModelGarden._deploy_response_or_raise(operation) + async def list_custom_model_deploy_options( self, src: str, diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 6d719ddd84..37e64aaf1b 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -66,6 +66,7 @@ from .common import _DeleteSandboxEnvironmentSnapshotRequestParameters from .common import _DeleteSandboxEnvironmentTemplateRequestParameters from .common import _DeleteSkillRequestParameters +from .common import _DeployRequestParameters from .common import _EvaluateInstancesRequestParameters from .common import _ExecuteCodeAgentEngineSandboxRequestParameters from .common import _ExportPublisherModelRequestParameters @@ -93,6 +94,7 @@ from .common import _GetDatasetParameters from .common import _GetDatasetVersionParameters from .common import _GetDeleteAgentEngineRuntimeRevisionOperationParameters +from .common import _GetDeployOperationParameters from .common import _GetEvaluationExperimentParameters from .common import _GetEvaluationItemParameters from .common import _GetEvaluationMetricParameters @@ -501,9 +503,33 @@ from .common import DeleteSkillOperation from .common import DeleteSkillOperationDict from .common import DeleteSkillOperationOrDict +from .common import DeployConfig +from .common import DeployConfigDict +from .common import DeployConfigOrDict +from .common import DeployModelOperation +from .common import DeployModelOperationDict +from .common import DeployModelOperationOrDict from .common import DeployOption from .common import DeployOptionDict from .common import DeployOptionOrDict +from .common import DeployPublisherModelConfig +from .common import DeployPublisherModelConfigDict +from .common import DeployPublisherModelConfigOrDict +from .common import DeployRequestCustomModel +from .common import DeployRequestCustomModelDict +from .common import DeployRequestCustomModelOrDict +from .common import DeployRequestDeployConfig +from .common import DeployRequestDeployConfigDict +from .common import DeployRequestDeployConfigOrDict +from .common import DeployRequestEndpointConfig +from .common import DeployRequestEndpointConfigDict +from .common import DeployRequestEndpointConfigOrDict +from .common import DeployRequestModelConfig +from .common import DeployRequestModelConfigDict +from .common import DeployRequestModelConfigOrDict +from .common import DeployResponse +from .common import DeployResponseDict +from .common import DeployResponseOrDict from .common import DirectUploadSource from .common import DirectUploadSourceDict from .common import DirectUploadSourceOrDict @@ -746,6 +772,9 @@ from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfig from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfigDict from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict +from .common import GetDeployOperationConfig +from .common import GetDeployOperationConfigDict +from .common import GetDeployOperationConfigOrDict from .common import GetEvaluationExperimentConfig from .common import GetEvaluationExperimentConfigDict from .common import GetEvaluationExperimentConfigOrDict @@ -1186,6 +1215,9 @@ from .common import PredictSchemata from .common import PredictSchemataDict from .common import PredictSchemataOrDict +from .common import PrivateServiceConnectConfig +from .common import PrivateServiceConnectConfigDict +from .common import PrivateServiceConnectConfigOrDict from .common import Probe from .common import ProbeDict from .common import ProbeExecAction @@ -1227,6 +1259,10 @@ from .common import PromptVersionRefDict from .common import PromptVersionRefOrDict from .common import Protocol +from .common import PSCAutomationConfig +from .common import PSCAutomationConfigDict +from .common import PSCAutomationConfigOrDict +from .common import PscAutomationState from .common import PscInterfaceConfig from .common import PscInterfaceConfigDict from .common import PscInterfaceConfigOrDict @@ -3548,6 +3584,36 @@ "GetExportPublisherModelOperationConfig", "GetExportPublisherModelOperationConfigDict", "GetExportPublisherModelOperationConfigOrDict", + "DeployConfig", + "DeployConfigDict", + "DeployConfigOrDict", + "DeployRequestCustomModel", + "DeployRequestCustomModelDict", + "DeployRequestCustomModelOrDict", + "DeployRequestModelConfig", + "DeployRequestModelConfigDict", + "DeployRequestModelConfigOrDict", + "PSCAutomationConfig", + "PSCAutomationConfigDict", + "PSCAutomationConfigOrDict", + "PrivateServiceConnectConfig", + "PrivateServiceConnectConfigDict", + "PrivateServiceConnectConfigOrDict", + "DeployRequestEndpointConfig", + "DeployRequestEndpointConfigDict", + "DeployRequestEndpointConfigOrDict", + "DeployRequestDeployConfig", + "DeployRequestDeployConfigDict", + "DeployRequestDeployConfigOrDict", + "DeployResponse", + "DeployResponseDict", + "DeployResponseOrDict", + "DeployModelOperation", + "DeployModelOperationDict", + "DeployModelOperationOrDict", + "GetDeployOperationConfig", + "GetDeployOperationConfigDict", + "GetDeployOperationConfigOrDict", "CreateRuntimeFeedbackEntryConfig", "CreateRuntimeFeedbackEntryConfigDict", "CreateRuntimeFeedbackEntryConfigOrDict", @@ -3698,6 +3764,9 @@ "ExportOpenModelConfig", "ExportOpenModelConfigDict", "ExportOpenModelConfigOrDict", + "DeployPublisherModelConfig", + "DeployPublisherModelConfigDict", + "DeployPublisherModelConfigOrDict", "DeployOption", "DeployOptionDict", "DeployOptionOrDict", @@ -3727,6 +3796,7 @@ "OpenSourceCategory", "VersionState", "QuotaState", + "PscAutomationState", "FeedbackType", "EvaluationExperimentMergeStrategy", "EvaluationItemType", @@ -3896,6 +3966,8 @@ "_RecommendSpecRequestParameters", "_ExportPublisherModelRequestParameters", "_GetExportPublisherModelOperationParameters", + "_DeployRequestParameters", + "_GetDeployOperationParameters", "_CreateRuntimeFeedbackEntryRequestParameters", "_DeleteRuntimeFeedbackEntryRequestParameters", "_GetRuntimeFeedbackRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index e2ff847bac..3bdbf5fc0a 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -487,6 +487,17 @@ class QuotaState(_common.CaseInSensitiveEnum): """User does not have enough accelerator quota for the machine type.""" +class PscAutomationState(_common.CaseInSensitiveEnum): + """Output only. The state of the PSC service automation.""" + + PSC_AUTOMATION_STATE_UNSPECIFIED = "PSC_AUTOMATION_STATE_UNSPECIFIED" + """Should not be used.""" + PSC_AUTOMATION_STATE_SUCCESSFUL = "PSC_AUTOMATION_STATE_SUCCESSFUL" + """The PSC service automation is successful.""" + PSC_AUTOMATION_STATE_FAILED = "PSC_AUTOMATION_STATE_FAILED" + """The PSC service automation has failed.""" + + class FeedbackType(_common.CaseInSensitiveEnum): """The type of the feedback.""" @@ -24392,6 +24403,468 @@ class _GetExportPublisherModelOperationParametersDict(TypedDict, total=False): ] +class DeployConfig(_common.BaseModel): + """Config for deploying models.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class DeployConfigDict(TypedDict, total=False): + """Config for deploying models.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +DeployConfigOrDict = Union[DeployConfig, DeployConfigDict] + + +class DeployRequestCustomModel(_common.BaseModel): + """The custom model to deploy from model weights in a Google Cloud Storage URI or Model Registry model.""" + + gcs_uri: Optional[str] = Field( + default=None, + description="""Immutable. The Google Cloud Storage URI of the custom model, storing weights and config files (which can be used to infer the base model).""", + ) + model_id: Optional[str] = Field( + default=None, + description="""Optional. Deprecated. Use ModelConfig.model_user_id instead.""", + ) + + +class DeployRequestCustomModelDict(TypedDict, total=False): + """The custom model to deploy from model weights in a Google Cloud Storage URI or Model Registry model.""" + + gcs_uri: Optional[str] + """Immutable. The Google Cloud Storage URI of the custom model, storing weights and config files (which can be used to infer the base model).""" + + model_id: Optional[str] + """Optional. Deprecated. Use ModelConfig.model_user_id instead.""" + + +DeployRequestCustomModelOrDict = Union[ + DeployRequestCustomModel, DeployRequestCustomModelDict +] + + +class DeployRequestModelConfig(_common.BaseModel): + """The model config to use for the deployment.""" + + accept_eula: Optional[bool] = Field( + default=None, + description="""Optional. Whether the user accepts the End User License Agreement (EULA) for the model.""", + ) + container_spec: Optional[ModelContainerSpec] = Field( + default=None, + description="""Optional. The specification of the container that is to be used when deploying. If not set, the default container spec will be used.""", + ) + hugging_face_access_token: Optional[str] = Field( + default=None, + description="""Optional. The Hugging Face read access token used to access the model artifacts of gated models.""", + ) + hugging_face_cache_enabled: Optional[bool] = Field( + default=None, + description="""Optional. If true, the model will deploy with a cached version instead of directly downloading the model artifacts from Hugging Face. This is suitable for VPC-SC users with limited internet access.""", + ) + model_display_name: Optional[str] = Field( + default=None, + description="""Optional. The user-specified display name of the uploaded model. If not set, a default name will be used.""", + ) + model_user_id: Optional[str] = Field( + default=None, + description="""Optional. The ID to use for the uploaded Model, which will become the final component of the model resource name. When not provided, Vertex AI will generate a value for this ID. When Model Registry model is provided, this field will be ignored. This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`. The first character cannot be a number or hyphen.""", + ) + + +class DeployRequestModelConfigDict(TypedDict, total=False): + """The model config to use for the deployment.""" + + accept_eula: Optional[bool] + """Optional. Whether the user accepts the End User License Agreement (EULA) for the model.""" + + container_spec: Optional[ModelContainerSpecDict] + """Optional. The specification of the container that is to be used when deploying. If not set, the default container spec will be used.""" + + hugging_face_access_token: Optional[str] + """Optional. The Hugging Face read access token used to access the model artifacts of gated models.""" + + hugging_face_cache_enabled: Optional[bool] + """Optional. If true, the model will deploy with a cached version instead of directly downloading the model artifacts from Hugging Face. This is suitable for VPC-SC users with limited internet access.""" + + model_display_name: Optional[str] + """Optional. The user-specified display name of the uploaded model. If not set, a default name will be used.""" + + model_user_id: Optional[str] + """Optional. The ID to use for the uploaded Model, which will become the final component of the model resource name. When not provided, Vertex AI will generate a value for this ID. When Model Registry model is provided, this field will be ignored. This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`. The first character cannot be a number or hyphen.""" + + +DeployRequestModelConfigOrDict = Union[ + DeployRequestModelConfig, DeployRequestModelConfigDict +] + + +class PSCAutomationConfig(_common.BaseModel): + """PSC config that is used to automatically create PSC endpoints in the user projects.""" + + error_message: Optional[str] = Field( + default=None, + description="""Output only. Error message if the PSC service automation failed.""", + ) + forwarding_rule: Optional[str] = Field( + default=None, + description="""Output only. Forwarding rule created by the PSC service automation.""", + ) + ip_address: Optional[str] = Field( + default=None, + description="""Output only. IP address rule created by the PSC service automation.""", + ) + network: Optional[str] = Field( + default=None, + description="""Required. The full name of the Google Compute Engine [network](https://cloud.google.com/compute/docs/networks-and-firewalls#networks). [Format](https://cloud.google.com/compute/docs/reference/rest/v1/networks/get): `projects/{project}/global/networks/{network}`.""", + ) + project_id: Optional[str] = Field( + default=None, + description="""Required. Project id used to create forwarding rule.""", + ) + state: Optional[PscAutomationState] = Field( + default=None, + description="""Output only. The state of the PSC service automation.""", + ) + + +class PSCAutomationConfigDict(TypedDict, total=False): + """PSC config that is used to automatically create PSC endpoints in the user projects.""" + + error_message: Optional[str] + """Output only. Error message if the PSC service automation failed.""" + + forwarding_rule: Optional[str] + """Output only. Forwarding rule created by the PSC service automation.""" + + ip_address: Optional[str] + """Output only. IP address rule created by the PSC service automation.""" + + network: Optional[str] + """Required. The full name of the Google Compute Engine [network](https://cloud.google.com/compute/docs/networks-and-firewalls#networks). [Format](https://cloud.google.com/compute/docs/reference/rest/v1/networks/get): `projects/{project}/global/networks/{network}`.""" + + project_id: Optional[str] + """Required. Project id used to create forwarding rule.""" + + state: Optional[PscAutomationState] + """Output only. The state of the PSC service automation.""" + + +PSCAutomationConfigOrDict = Union[PSCAutomationConfig, PSCAutomationConfigDict] + + +class PrivateServiceConnectConfig(_common.BaseModel): + """Represents configuration for private service connect.""" + + enable_private_service_connect: Optional[bool] = Field( + default=None, + description="""Required. If true, expose the IndexEndpoint via private service connect.""", + ) + enable_secure_private_service_connect: Optional[bool] = Field( + default=None, + description="""Optional. If set to true, enable secure private service connect with IAM authorization. Otherwise, private service connect will be done without authorization. Note latency will be slightly increased if authorization is enabled.""", + ) + project_allowlist: Optional[list[str]] = Field( + default=None, + description="""A list of Projects from which the forwarding rule will target the service attachment.""", + ) + psc_automation_configs: Optional[list[PSCAutomationConfig]] = Field( + default=None, + description="""Optional. List of projects and networks where the PSC endpoints will be created. This field is used by Online Inference(Prediction) only.""", + ) + service_attachment: Optional[str] = Field( + default=None, + description="""Output only. The name of the generated service attachment resource. This is only populated if the endpoint is deployed with PrivateServiceConnect.""", + ) + + +class PrivateServiceConnectConfigDict(TypedDict, total=False): + """Represents configuration for private service connect.""" + + enable_private_service_connect: Optional[bool] + """Required. If true, expose the IndexEndpoint via private service connect.""" + + enable_secure_private_service_connect: Optional[bool] + """Optional. If set to true, enable secure private service connect with IAM authorization. Otherwise, private service connect will be done without authorization. Note latency will be slightly increased if authorization is enabled.""" + + project_allowlist: Optional[list[str]] + """A list of Projects from which the forwarding rule will target the service attachment.""" + + psc_automation_configs: Optional[list[PSCAutomationConfigDict]] + """Optional. List of projects and networks where the PSC endpoints will be created. This field is used by Online Inference(Prediction) only.""" + + service_attachment: Optional[str] + """Output only. The name of the generated service attachment resource. This is only populated if the endpoint is deployed with PrivateServiceConnect.""" + + +PrivateServiceConnectConfigOrDict = Union[ + PrivateServiceConnectConfig, PrivateServiceConnectConfigDict +] + + +class DeployRequestEndpointConfig(_common.BaseModel): + """The endpoint config to use for the deployment.""" + + dedicated_endpoint_disabled: Optional[bool] = Field( + default=None, + description="""Optional. By default, if dedicated endpoint is enabled and private service connect config is not set, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. If private service connect config is set, the endpoint will be exposed through private service connect. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon. If this field is set to true, the dedicated endpoint will be disabled and the deployed model will be exposed through the shared DNS {region}-aiplatform.googleapis.com.""", + ) + dedicated_endpoint_enabled: Optional[bool] = Field( + default=None, + description="""Optional. Deprecated. Use dedicated_endpoint_disabled instead. If true, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon.""", + ) + endpoint_display_name: Optional[str] = Field( + default=None, + description="""Optional. The user-specified display name of the endpoint. If not set, a default name will be used.""", + ) + endpoint_user_id: Optional[str] = Field( + default=None, + description="""Optional. Immutable. The ID to use for endpoint, which will become the final component of the endpoint resource name. If not provided, Vertex AI will generate a value for this ID. If the first character is a letter, this value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The last character must be a letter or number. If the first character is a number, this value may be up to 9 characters, and valid characters are `[0-9]` with no leading zeros. When using HTTP/JSON, this field is populated based on a query string argument, such as `?endpoint_id=12345`. This is the fallback for fields that are not included in either the URI or the body.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. The labels with user-defined metadata to organize your Endpoints. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", + ) + private_service_connect_config: Optional[PrivateServiceConnectConfig] = Field( + default=None, + description="""Optional. Configuration for private service connect. If set, the endpoint will be exposed through private service connect.""", + ) + + +class DeployRequestEndpointConfigDict(TypedDict, total=False): + """The endpoint config to use for the deployment.""" + + dedicated_endpoint_disabled: Optional[bool] + """Optional. By default, if dedicated endpoint is enabled and private service connect config is not set, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. If private service connect config is set, the endpoint will be exposed through private service connect. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon. If this field is set to true, the dedicated endpoint will be disabled and the deployed model will be exposed through the shared DNS {region}-aiplatform.googleapis.com.""" + + dedicated_endpoint_enabled: Optional[bool] + """Optional. Deprecated. Use dedicated_endpoint_disabled instead. If true, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon.""" + + endpoint_display_name: Optional[str] + """Optional. The user-specified display name of the endpoint. If not set, a default name will be used.""" + + endpoint_user_id: Optional[str] + """Optional. Immutable. The ID to use for endpoint, which will become the final component of the endpoint resource name. If not provided, Vertex AI will generate a value for this ID. If the first character is a letter, this value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The last character must be a letter or number. If the first character is a number, this value may be up to 9 characters, and valid characters are `[0-9]` with no leading zeros. When using HTTP/JSON, this field is populated based on a query string argument, such as `?endpoint_id=12345`. This is the fallback for fields that are not included in either the URI or the body.""" + + labels: Optional[dict[str, str]] + """Optional. The labels with user-defined metadata to organize your Endpoints. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" + + private_service_connect_config: Optional[PrivateServiceConnectConfigDict] + """Optional. Configuration for private service connect. If set, the endpoint will be exposed through private service connect.""" + + +DeployRequestEndpointConfigOrDict = Union[ + DeployRequestEndpointConfig, DeployRequestEndpointConfigDict +] + + +class DeployRequestDeployConfig(_common.BaseModel): + """The deploy config to use for the deployment.""" + + dedicated_resources: Optional[DedicatedResources] = Field( + default=None, + description="""Optional. The dedicated resources to use for the endpoint. If not set, the default resources will be used.""", + ) + fast_tryout_enabled: Optional[bool] = Field( + default=None, + description="""Optional. If true, enable the QMT fast tryout feature for this model if possible.""", + ) + system_labels: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. System labels for Model Garden deployments. These labels are managed by Google and for tracking purposes only.""", + ) + + +class DeployRequestDeployConfigDict(TypedDict, total=False): + """The deploy config to use for the deployment.""" + + dedicated_resources: Optional[DedicatedResourcesDict] + """Optional. The dedicated resources to use for the endpoint. If not set, the default resources will be used.""" + + fast_tryout_enabled: Optional[bool] + """Optional. If true, enable the QMT fast tryout feature for this model if possible.""" + + system_labels: Optional[dict[str, str]] + """Optional. System labels for Model Garden deployments. These labels are managed by Google and for tracking purposes only.""" + + +DeployRequestDeployConfigOrDict = Union[ + DeployRequestDeployConfig, DeployRequestDeployConfigDict +] + + +class _DeployRequestParameters(_common.BaseModel): + """Parameters for deployment.""" + + destination: Optional[str] = Field(default=None, description="""""") + publisher_model_name: Optional[str] = Field(default=None, description="""""") + hugging_face_model_id: Optional[str] = Field(default=None, description="""""") + custom_model: Optional[DeployRequestCustomModel] = Field( + default=None, description="""""" + ) + model_config_val: Optional[DeployRequestModelConfig] = Field( + default=None, description="""""" + ) + endpoint_config: Optional[DeployRequestEndpointConfig] = Field( + default=None, description="""""" + ) + deploy_config: Optional[DeployRequestDeployConfig] = Field( + default=None, description="""""" + ) + config: Optional[DeployConfig] = Field(default=None, description="""""") + + +class _DeployRequestParametersDict(TypedDict, total=False): + """Parameters for deployment.""" + + destination: Optional[str] + """""" + + publisher_model_name: Optional[str] + """""" + + hugging_face_model_id: Optional[str] + """""" + + custom_model: Optional[DeployRequestCustomModelDict] + """""" + + model_config_val: Optional[DeployRequestModelConfigDict] + """""" + + endpoint_config: Optional[DeployRequestEndpointConfigDict] + """""" + + deploy_config: Optional[DeployRequestDeployConfigDict] + """""" + + config: Optional[DeployConfigDict] + """""" + + +_DeployRequestParametersOrDict = Union[ + _DeployRequestParameters, _DeployRequestParametersDict +] + + +class DeployResponse(_common.BaseModel): + """Response for deployment.""" + + endpoint: Optional[str] = Field( + default=None, description="""The resource name of the deployed endpoint.""" + ) + model: Optional[str] = Field( + default=None, description="""The resource name of the deployed model.""" + ) + + +class DeployResponseDict(TypedDict, total=False): + """Response for deployment.""" + + endpoint: Optional[str] + """The resource name of the deployed endpoint.""" + + model: Optional[str] + """The resource name of the deployed model.""" + + +DeployResponseOrDict = Union[DeployResponse, DeployResponseDict] + + +class DeployModelOperation(_common.BaseModel): + """Operation that has a deploy response.""" + + 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[DeployResponse] = Field(default=None, description="""""") + + +class DeployModelOperationDict(TypedDict, total=False): + """Operation that has a deploy response.""" + + 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[DeployResponseDict] + """""" + + +DeployModelOperationOrDict = Union[DeployModelOperation, DeployModelOperationDict] + + +class GetDeployOperationConfig(_common.BaseModel): + """Config for ``get_deploy_publisher_model_operation``.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetDeployOperationConfigDict(TypedDict, total=False): + """Config for ``get_deploy_publisher_model_operation``.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetDeployOperationConfigOrDict = Union[ + GetDeployOperationConfig, GetDeployOperationConfigDict +] + + +class _GetDeployOperationParameters(_common.BaseModel): + """Parameters for polling a ``deploy`` operation.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetDeployOperationConfig] = Field(default=None, description="""""") + + +class _GetDeployOperationParametersDict(TypedDict, total=False): + """Parameters for polling a ``deploy`` operation.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetDeployOperationConfigDict] + """""" + + +_GetDeployOperationParametersOrDict = Union[ + _GetDeployOperationParameters, _GetDeployOperationParametersDict +] + + class CreateRuntimeFeedbackEntryConfig(_common.BaseModel): """Config for creating a Feedback Entry.""" @@ -27414,6 +27887,190 @@ class ExportOpenModelConfigDict(TypedDict, total=False): ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] +class DeployPublisherModelConfig(_common.BaseModel): + """Config for ``deploy_publisher_model``. + + Superset of options that apply to Google open, partner and Hugging Face + publisher models. Only fields relevant to the target model are honored; + the backend rejects unsupported fields with a clear error. + """ + + accept_eula: Optional[bool] = Field( + default=None, + description="""Whether to accept the model's End User License Agreement.""", + ) + hugging_face_access_token: Optional[str] = Field( + default=None, + description="""Hugging Face access token for gated HF models. See + https://huggingface.co/docs/hub/en/security-tokens.""", + ) + machine_type: Optional[str] = Field( + default=None, + description="""Machine type (e.g. ``'g2-standard-48'``). Leave unset for + automatic resources.""", + ) + min_replica_count: Optional[int] = Field( + default=1, description="""Minimum number of replicas.""" + ) + max_replica_count: Optional[int] = Field( + default=1, description="""Maximum number of replicas.""" + ) + accelerator_type: Optional[str] = Field( + default=None, description="""Accelerator type (e.g. ``'NVIDIA_L4'``).""" + ) + accelerator_count: Optional[int] = Field( + default=None, description="""Number of accelerators per replica.""" + ) + spot: Optional[bool] = Field(default=None, description="""Schedule on Spot VMs.""") + dedicated_endpoint_disabled: Optional[bool] = Field( + default=None, + description="""Set True to serve predictions via the shared endpoint DNS + instead of the dedicated endpoint DNS (default).""", + ) + fast_tryout_enabled: Optional[bool] = Field( + default=None, + description="""Use the fast-tryout deployment path (experimentation only, not + production). Only supported for select models and machine types.""", + ) + endpoint_display_name: Optional[str] = Field( + default=None, description="""Display name for the endpoint.""" + ) + model_display_name: Optional[str] = Field( + default=None, description="""Display name for the deployed model.""" + ) + serving_container_image_uri: Optional[str] = Field( + default=None, + description="""Custom serving container image URI overriding the model's + default container.""", + ) + container_command: Optional[list[str]] = Field( + default=None, description="""Serving container ENTRYPOINT override.""" + ) + container_args: Optional[list[str]] = Field( + default=None, description="""Serving container CMD override.""" + ) + container_variables: Optional[dict[str, str]] = Field( + default=None, description="""Environment variables for the serving container.""" + ) + enable_private_service_connect: Optional[bool] = Field( + default=None, description="""Enable Private Service Connect for the endpoint.""" + ) + psc_project_allow_list: Optional[list[str]] = Field( + default=None, + description="""Projects allowed to access the endpoint over Private Service + Connect. Only honored when ``enable_private_service_connect`` is True.""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to block on the deployment long-running operation. When + ``True`` (default), returns the ``DeployResponse`` (deployed endpoint + and model resource names) on completion. When ``False``, returns the + ``DeployModelOperation`` 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 deployment to complete + when ``wait_for_completion=True``. Defaults to 2 hours, matching the + Vertex AI Console one-click deployment timeout. Ignored when + ``wait_for_completion=False``.""", + ) + + +class DeployPublisherModelConfigDict(TypedDict, total=False): + """Config for ``deploy_publisher_model``. + + Superset of options that apply to Google open, partner and Hugging Face + publisher models. Only fields relevant to the target model are honored; + the backend rejects unsupported fields with a clear error. + """ + + accept_eula: Optional[bool] + """Whether to accept the model's End User License Agreement.""" + + hugging_face_access_token: Optional[str] + """Hugging Face access token for gated HF models. See + https://huggingface.co/docs/hub/en/security-tokens.""" + + machine_type: Optional[str] + """Machine type (e.g. ``'g2-standard-48'``). Leave unset for + automatic resources.""" + + min_replica_count: Optional[int] + """Minimum number of replicas.""" + + max_replica_count: Optional[int] + """Maximum number of replicas.""" + + accelerator_type: Optional[str] + """Accelerator type (e.g. ``'NVIDIA_L4'``).""" + + accelerator_count: Optional[int] + """Number of accelerators per replica.""" + + spot: Optional[bool] + """Schedule on Spot VMs.""" + + dedicated_endpoint_disabled: Optional[bool] + """Set True to serve predictions via the shared endpoint DNS + instead of the dedicated endpoint DNS (default).""" + + fast_tryout_enabled: Optional[bool] + """Use the fast-tryout deployment path (experimentation only, not + production). Only supported for select models and machine types.""" + + endpoint_display_name: Optional[str] + """Display name for the endpoint.""" + + model_display_name: Optional[str] + """Display name for the deployed model.""" + + serving_container_image_uri: Optional[str] + """Custom serving container image URI overriding the model's + default container.""" + + container_command: Optional[list[str]] + """Serving container ENTRYPOINT override.""" + + container_args: Optional[list[str]] + """Serving container CMD override.""" + + container_variables: Optional[dict[str, str]] + """Environment variables for the serving container.""" + + enable_private_service_connect: Optional[bool] + """Enable Private Service Connect for the endpoint.""" + + psc_project_allow_list: Optional[list[str]] + """Projects allowed to access the endpoint over Private Service + Connect. Only honored when ``enable_private_service_connect`` is True.""" + + wait_for_completion: Optional[bool] + """Whether to block on the deployment long-running operation. When + ``True`` (default), returns the ``DeployResponse`` (deployed endpoint + and model resource names) on completion. When ``False``, returns the + ``DeployModelOperation`` 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 deployment to complete + when ``wait_for_completion=True``. Defaults to 2 hours, matching the + Vertex AI Console one-click deployment timeout. Ignored when + ``wait_for_completion=False``.""" + + +DeployPublisherModelConfigOrDict = Union[ + DeployPublisherModelConfig, DeployPublisherModelConfigDict +] + + 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 0702fe1fdf..0cd7573fc6 100644 --- a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py @@ -252,6 +252,33 @@ def test_export_open_model_invalid_name_raises(client): ) +# TODO(gshuoy): restore the recorded deploy replays (open w/ machine +# config, Hugging Face, async) once they are re-recorded against the +# wait_for_completion LRO flow. The previous recordings polled a +# hand-rolled loop that no longer exists, and their recorded URL +# (``.../locations/{location}:deploy``) is not redacted by +# ``_redact_request_url``, so they only replayed under the exact project +# and location they were recorded in. Behavior stays covered by the mocked +# unit tests in ``tests/unit/agentplatform/genai/test_genai_model_garden.py``. + + +def test_deploy_publisher_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.deploy_publisher_model(model="not-a-valid-name") + + +def test_deploy_publisher_model_container_override_without_image_raises(client): + """Container overrides without an image URI are rejected client-side.""" + with pytest.raises(ValueError, match="require serving_container_image_uri"): + client.model_garden.deploy_publisher_model( + model="google/gemma2@gemma-2-2b-it", + config=types.DeployPublisherModelConfig( + container_variables={"MODEL_ID": "gemma-2-2b-it"}, + ), + ) + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), diff --git a/tests/unit/agentplatform/genai/test_genai_model_garden.py b/tests/unit/agentplatform/genai/test_genai_model_garden.py index f46752cebb..66666a51d2 100644 --- a/tests/unit/agentplatform/genai/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/test_genai_model_garden.py @@ -1770,3 +1770,768 @@ def test_export_open_model_async_invalid_name_raises(mock_async_client): ) ) m_rpc.assert_not_called() + + +# ---- deploy_publisher_model tests ------------------------------------- + + +_HF_MODEL = ( # mixed case to exercise lowering + "Meta-Llama/Llama-3.3-70B-Instruct" +) + +_TEST_DEPLOY_OPERATION = "projects/p/locations/us-central1/operations/1" +_TEST_ENDPOINT_NAME = "projects/p/locations/us-central1/endpoints/123" +_TEST_DEPLOYED_MODEL_NAME = "projects/p/locations/us-central1/models/456" + + +def _make_deploy_operation( + name=_TEST_DEPLOY_OPERATION, + *, + done=False, + endpoint=None, + model=None, + error=None, +): + """Builds a fake DeployModelOperation for use in tests.""" + op = types.DeployModelOperation(name=name, done=done) + if endpoint is not None or model is not None: + op.response = types.DeployResponse(endpoint=endpoint, model=model) + if error is not None: + op.error = error + return op + + +def _done_deploy_operation(): + """A completed deploy LRO carrying both resource names.""" + return _make_deploy_operation( + done=True, + endpoint=_TEST_ENDPOINT_NAME, + model=_TEST_DEPLOYED_MODEL_NAME, + ) + + +def _patch_deploy_rpc(mock_client, initial_op=None): + """Mocks the private ``_deploy`` RPC to return ``initial_op``.""" + return mock.patch.object( + mock_client, + "_deploy", + return_value=( + initial_op if initial_op is not None else _make_deploy_operation() + ), + ) + + +def _patch_async_deploy_rpc(mock_client, initial_op=None): + """Async variant of _patch_deploy_rpc.""" + return mock.patch.object( + mock_client, + "_deploy", + new=mock.AsyncMock( + return_value=( + initial_op if initial_op is not None else _make_deploy_operation() + ) + ), + ) + + +# --- Return type & blocking behavior ------------------------------------ + + +def test_deploy_publisher_model_returns_deploy_response_when_waiting( + mock_client, +): + """Default (wait_for_completion=True) blocks on the LRO and returns DeployResponse.""" + with ( + _patch_deploy_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ) as m_await, + ): + result = mock_client.deploy_publisher_model(model=_OPEN_MODEL) + m_await.assert_called_once() + assert isinstance(result, types.DeployResponse) + assert result.endpoint == _TEST_ENDPOINT_NAME + assert result.model == _TEST_DEPLOYED_MODEL_NAME + + +def test_deploy_publisher_model_does_not_return_gapic_endpoint(mock_client): + """The return value is a plain SDK type, never a GAPIC aiplatform.Endpoint.""" + with ( + _patch_deploy_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + result = mock_client.deploy_publisher_model(model=_OPEN_MODEL) + assert type(result).__module__.startswith( + "google.cloud.aiplatform.agentplatform" + ) + assert isinstance(result.endpoint, str) + assert isinstance(result.model, str) + + +def test_deploy_publisher_model_returns_operation_when_not_waiting(mock_client): + """wait_for_completion=False returns the operation immediately (no polling).""" + initial_op = _make_deploy_operation() + with ( + _patch_deploy_rpc(mock_client, initial_op=initial_op), + mock.patch.object( + model_garden._operations_utils, "await_operation" + ) as m_await, + ): + result = mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig(wait_for_completion=False), + ) + m_await.assert_not_called() + assert result is initial_op + + +def test_deploy_publisher_model_polls_via_public_getter(mock_client): + """``get_deploy_publisher_model_operation`` is used as the polling callback.""" + with ( + _patch_deploy_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ) as m_await, + ): + mock_client.deploy_publisher_model(model=_OPEN_MODEL) + kwargs = m_await.call_args.kwargs + assert ( + kwargs["get_operation_fn"] + == mock_client.get_deploy_publisher_model_operation + ) + assert kwargs["operation_name"] == _TEST_DEPLOY_OPERATION + + +def test_deploy_publisher_model_uses_default_poll_and_timeout(mock_client): + """No override in config -> ModelGarden defaults are forwarded to the poller.""" + with ( + _patch_deploy_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ) as m_await, + ): + mock_client.deploy_publisher_model(model=_OPEN_MODEL) + kwargs = m_await.call_args.kwargs + assert ( + kwargs["poll_interval"] == 30 + ) # _DEFAULT_DEPLOY_POLL_INTERVAL_SECONDS + assert ( + kwargs["timeout_seconds"] == 2 * 60 * 60 + ) # _DEFAULT_DEPLOY_TIMEOUT_SECONDS + + +def test_deploy_publisher_model_config_overrides_poll_and_timeout(mock_client): + """Config values override the defaults on the polling call.""" + with ( + _patch_deploy_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ) as m_await, + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + 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_deploy_publisher_model_raises_on_operation_error(mock_client): + """A backend error on the LRO surfaces as RuntimeError.""" + failed_op = _make_deploy_operation( + done=True, error={"message": "quota exceeded"} + ) + with ( + _patch_deploy_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=failed_op, + ), + ): + with pytest.raises(RuntimeError, match="Deploy failed"): + mock_client.deploy_publisher_model(model=_OPEN_MODEL) + + +def test_deploy_publisher_model_raises_when_response_missing_endpoint( + mock_client, +): + """A done LRO with no endpoint surfaces as RuntimeError.""" + done_op = _make_deploy_operation(done=True) # no response set + with ( + _patch_deploy_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ), + ): + with pytest.raises(RuntimeError, match="no endpoint"): + mock_client.deploy_publisher_model(model=_OPEN_MODEL) + + +# --- Model-name resolution --------------------------------------------- + + +def test_deploy_publisher_model_simplified_name_reconciled(mock_client): + """Simplified '{publisher}/{model}@{version}' -> full resource name.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model(model=_OPEN_MODEL) + kwargs = m_rpc.call_args.kwargs + assert ( + kwargs["publisher_model_name"] + == "publishers/google/models/gemma3@gemma-3-12b-it" + ) + assert kwargs["hugging_face_model_id"] is None + assert kwargs["destination"] == ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" + ) + + +def test_deploy_publisher_model_full_resource_name_preserved(mock_client): + """A full 'publishers/.../models/...@version' is passed through unchanged.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model="publishers/google/models/gemma3@gemma-3-12b-it" + ) + assert ( + m_rpc.call_args.kwargs["publisher_model_name"] + == "publishers/google/models/gemma3@gemma-3-12b-it" + ) + + +def test_deploy_publisher_model_hugging_face_uses_hf_model_id(mock_client): + """HF IDs are lowercased and routed to huggingFaceModelId (legacy parity).""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model(model=_HF_MODEL) + kwargs = m_rpc.call_args.kwargs + assert kwargs["publisher_model_name"] is None + assert ( + kwargs["hugging_face_model_id"] == "meta-llama/llama-3.3-70b-instruct" + ) + + +def test_deploy_publisher_model_invalid_name_raises_before_rpc(mock_client): + """Malformed model names raise ValueError; no _deploy call and no polling.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, "await_operation" + ) as m_await, + ): + with pytest.raises(ValueError, match="not a valid publisher model name"): + mock_client.deploy_publisher_model(model="not-a-valid-name") + m_rpc.assert_not_called() + m_await.assert_not_called() + + +# --- Config validation -------------------------------------------------- + + +@pytest.mark.parametrize( + "override", + [ + {"container_command": ["python", "-m", "server"]}, + {"container_args": ["--tensor-parallel-size=2"]}, + {"container_variables": {"MODEL_ID": "gemma-3-12b-it"}}, + ], +) +def test_deploy_publisher_model_container_override_without_image_raises( + mock_client, override +): + """Container overrides without serving_container_image_uri fail loudly.""" + with _patch_deploy_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="require serving_container_image_uri"): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig(**override), + ) + m_rpc.assert_not_called() + + +def test_deploy_publisher_model_container_override_error_names_fields( + mock_client, +): + """The ValueError names every offending field so the fix is obvious.""" + with _patch_deploy_rpc(mock_client): + with pytest.raises(ValueError) as excinfo: + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + container_command=["python"], + container_variables={"FOO": "bar"}, + ), + ) + message = str(excinfo.value) + assert "container_command" in message + assert "container_variables" in message + + +# --- Config translation to DeployRequest sub-messages ------------------- + + +def test_deploy_publisher_model_default_config_no_optional_blocks(mock_client): + """Empty config -> no DedicatedResources, no ContainerSpec, no PSC.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model(model=_OPEN_MODEL) + kwargs = m_rpc.call_args.kwargs + assert kwargs["deploy_config"].dedicated_resources is None + assert kwargs["model_config_val"].container_spec is None + assert kwargs["endpoint_config"].dedicated_endpoint_enabled is None + assert kwargs["endpoint_config"].private_service_connect_config is None + + +def test_deploy_publisher_model_dict_config_validated(mock_client): + """A plain dict is validated into DeployPublisherModelConfig.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config={"machine_type": "g2-standard-12", "accept_eula": True}, + ) + kwargs = m_rpc.call_args.kwargs + assert kwargs["model_config_val"].accept_eula is True + assert ( + kwargs["deploy_config"].dedicated_resources.machine_spec.machine_type + == "g2-standard-12" + ) + + +def test_deploy_publisher_model_dict_config_honors_wait_for_completion( + mock_client, +): + """wait_for_completion=False survives dict validation.""" + initial_op = _make_deploy_operation() + with ( + _patch_deploy_rpc(mock_client, initial_op=initial_op), + mock.patch.object( + model_garden._operations_utils, "await_operation" + ) as m_await, + ): + result = mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config={"wait_for_completion": False}, + ) + m_await.assert_not_called() + assert result is initial_op + + +def test_deploy_publisher_model_machine_replica_and_spot_forwarded(mock_client): + """machine_type/accelerator/replicas/spot populate DedicatedResources.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + min_replica_count=2, + max_replica_count=5, + spot=True, + ), + ) + dr = m_rpc.call_args.kwargs["deploy_config"].dedicated_resources + assert dr.machine_spec.machine_type == "g2-standard-12" + assert dr.machine_spec.accelerator_type == "NVIDIA_L4" + assert dr.machine_spec.accelerator_count == 1 + assert dr.min_replica_count == 2 + assert dr.max_replica_count == 5 + assert dr.spot is True + + +def test_deploy_publisher_model_accelerator_only_triggers_dedicated_resources( + mock_client, +): + """A single machine-shape field (accelerator_type) is enough to send DedicatedResources.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig(accelerator_type="NVIDIA_L4"), + ) + dr = m_rpc.call_args.kwargs["deploy_config"].dedicated_resources + assert dr is not None + assert dr.machine_spec.accelerator_type == "NVIDIA_L4" + + +def test_deploy_publisher_model_dedicated_endpoint_disabled_inverts( + mock_client, +): + """dedicated_endpoint_disabled=True maps to dedicated_endpoint_enabled=False on wire.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + dedicated_endpoint_disabled=True + ), + ) + assert ( + m_rpc.call_args.kwargs["endpoint_config"].dedicated_endpoint_enabled + is False + ) + + +def test_deploy_publisher_model_dedicated_endpoint_disabled_false_maps_to_true( + mock_client, +): + """dedicated_endpoint_disabled=False maps to dedicated_endpoint_enabled=True.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + dedicated_endpoint_disabled=False + ), + ) + assert ( + m_rpc.call_args.kwargs["endpoint_config"].dedicated_endpoint_enabled + is True + ) + + +def test_deploy_publisher_model_private_service_connect(mock_client): + """PSC fields populate PrivateServiceConnectConfig.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + enable_private_service_connect=True, + psc_project_allow_list=["project-a", "project-b"], + ), + ) + psc = m_rpc.call_args.kwargs[ + "endpoint_config" + ].private_service_connect_config + assert psc.enable_private_service_connect is True + assert psc.project_allowlist == ["project-a", "project-b"] + + +def test_deploy_publisher_model_private_service_connect_without_allowlist( + mock_client, +): + """PSC without an allowlist is allowed (backend applies its own default).""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + enable_private_service_connect=True, + ), + ) + psc = m_rpc.call_args.kwargs[ + "endpoint_config" + ].private_service_connect_config + assert psc.enable_private_service_connect is True + assert psc.project_allowlist is None + + +def test_deploy_publisher_model_custom_container_spec(mock_client): + """image_uri + command + args + variables all populate ModelContainerSpec.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vllm", + container_command=["python", "-m", "vllm.entrypoints.api_server"], + container_args=["--tensor-parallel-size=2"], + container_variables={"MODEL_ID": "gemma-3-12b-it"}, + ), + ) + spec = m_rpc.call_args.kwargs["model_config_val"].container_spec + assert spec.image_uri == "us-docker.pkg.dev/vertex-ai/vllm" + assert spec.command == ["python", "-m", "vllm.entrypoints.api_server"] + assert spec.args == ["--tensor-parallel-size=2"] + assert spec.env == [types.EnvVar(name="MODEL_ID", value="gemma-3-12b-it")] + + +def test_deploy_publisher_model_container_image_without_variables_leaves_env_none( + mock_client, +): + """serving_container_image_uri alone -> container_spec set, env stays None.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + serving_container_image_uri="us-docker.pkg.dev/vertex-ai/vllm", + ), + ) + spec = m_rpc.call_args.kwargs["model_config_val"].container_spec + assert spec.image_uri == "us-docker.pkg.dev/vertex-ai/vllm" + assert spec.env is None + + +def test_deploy_publisher_model_display_names_and_fast_tryout(mock_client): + """endpoint/model display names and fast_tryout_enabled land on the right sub-messages.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + endpoint_display_name="my-endpoint", + model_display_name="my-model", + fast_tryout_enabled=True, + ), + ) + kwargs = m_rpc.call_args.kwargs + assert kwargs["endpoint_config"].endpoint_display_name == "my-endpoint" + assert kwargs["model_config_val"].model_display_name == "my-model" + assert kwargs["deploy_config"].fast_tryout_enabled is True + + +def test_deploy_publisher_model_forwards_eula_and_hf_token(mock_client): + """accept_eula and hugging_face_access_token land on model_config.""" + with ( + _patch_deploy_rpc(mock_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=_done_deploy_operation(), + ), + ): + mock_client.deploy_publisher_model( + model=_HF_MODEL, + config=types.DeployPublisherModelConfig( + accept_eula=True, + hugging_face_access_token="hf_token_xyz", + ), + ) + model_config = m_rpc.call_args.kwargs["model_config_val"] + assert model_config.accept_eula is True + assert model_config.hugging_face_access_token == "hf_token_xyz" + + +# --- Async surface parity ---------------------------------------------- + + +def test_deploy_publisher_model_async_returns_deploy_response_when_waiting( + mock_async_client, +): + """Async default (wait=True) awaits the LRO and returns DeployResponse.""" + with ( + _patch_async_deploy_rpc(mock_async_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation_async", + new=mock.AsyncMock(return_value=_done_deploy_operation()), + ), + ): + result = asyncio.run( + mock_async_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + machine_type="g2-standard-12" + ), + ) + ) + assert isinstance(result, types.DeployResponse) + assert result.endpoint == _TEST_ENDPOINT_NAME + assert result.model == _TEST_DEPLOYED_MODEL_NAME + + +def test_deploy_publisher_model_async_returns_operation_when_not_waiting( + mock_async_client, +): + """Async wait=False returns the operation immediately.""" + initial_op = _make_deploy_operation() + with ( + _patch_async_deploy_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.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig(wait_for_completion=False), + ) + ) + m_await.assert_not_called() + assert result is initial_op + + +def test_deploy_publisher_model_async_hugging_face(mock_async_client): + """Async HF path also lowercases and routes to hugging_face_model_id.""" + with ( + _patch_async_deploy_rpc(mock_async_client) as m_rpc, + mock.patch.object( + model_garden._operations_utils, + "await_operation_async", + new=mock.AsyncMock(return_value=_done_deploy_operation()), + ), + ): + asyncio.run(mock_async_client.deploy_publisher_model(model=_HF_MODEL)) + kwargs = m_rpc.call_args.kwargs + assert kwargs["publisher_model_name"] is None + assert ( + kwargs["hugging_face_model_id"] == "meta-llama/llama-3.3-70b-instruct" + ) + + +def test_deploy_publisher_model_async_invalid_name_raises(mock_async_client): + """Async invalid model name raises ValueError before any RPC or polling.""" + with _patch_async_deploy_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="not a valid publisher model name"): + asyncio.run( + mock_async_client.deploy_publisher_model(model="not-a-valid-name") + ) + m_rpc.assert_not_called() + + +def test_deploy_publisher_model_async_container_override_without_image_raises( + mock_async_client, +): + """Async container-override validation matches the sync surface.""" + with _patch_async_deploy_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="require serving_container_image_uri"): + asyncio.run( + mock_async_client.deploy_publisher_model( + model=_OPEN_MODEL, + config=types.DeployPublisherModelConfig( + container_args=["--flag"] + ), + ) + ) + m_rpc.assert_not_called() + + +def test_deploy_publisher_model_async_raises_on_operation_error( + mock_async_client, +): + """Async backend error on the LRO surfaces as RuntimeError.""" + failed_op = _make_deploy_operation( + done=True, error={"message": "quota exceeded"} + ) + with ( + _patch_async_deploy_rpc(mock_async_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation_async", + new=mock.AsyncMock(return_value=failed_op), + ), + ): + with pytest.raises(RuntimeError, match="Deploy failed"): + asyncio.run(mock_async_client.deploy_publisher_model(model=_OPEN_MODEL)) + + +# --- Private (pre-GA) surface ------------------------------------------ +# +# No direct unit tests: `agentplatform/private/**` is excluded from the +# buildable ":aiplatform" library (CORE_LIBRARY_EXCLUDES) and is not +# importable today -- `private/types/evals.py` references +# `genai_types.RubricVerdict`, which does not exist in the vendored +# `google.genai`, so importing the package raises AttributeError at HEAD +# independently of Model Garden. +# +# The private deploy surface cannot drift from the public one regardless: +# both are generated from the same `model_garden.py.in` + `deploy.py` +# config, and `generator_diff_test_private_vertex_sdk` fails if the +# checked-in private files stop matching generator output.