diff --git a/agentplatform/_genai/model_garden.py b/agentplatform/_genai/model_garden.py index ee9cebac2d..8c78cef833 100644 --- a/agentplatform/_genai/model_garden.py +++ b/agentplatform/_genai/model_garden.py @@ -712,7 +712,7 @@ def get_deploy_publisher_model_operation( config: Optional[types.GetDeployOperationConfigOrDict] = None, ) -> types.DeployModelOperation: """ - Fetches the status of an in-flight ``deploy_publisher_model`` LRO. + Fetches the status of an in-flight ``deploy_publisher_model`` or ``deploy_custom_model`` LRO. """ parameter_model = types._GetDeployOperationParameters( @@ -1232,6 +1232,94 @@ def _prepare_deploy_request( return model_config, endpoint_config, deploy_config + @staticmethod + def _validate_custom_deploy_config( + config: types.DeployCustomModelConfig, + ) -> None: + """Rejects a partially-specified machine shape for a custom model deploy. + + A publisher model can fall back to its own recommended machine spec, so + a partial override there is still deployable. A custom model has no such + fallback: the backend needs the machine type and the accelerator pair + together or not at all. Matches the all-or-none check in + ``vertexai.model_garden.CustomModel.deploy``, including its truthiness + test: ``accelerator_count=0`` counts as unset, since an accelerator type + with zero accelerators is not a deployable shape. + + Raises: + ValueError: If some but not all of ``machine_type``, + ``accelerator_type`` and ``accelerator_count`` are set. + """ + fields = ("machine_type", "accelerator_type", "accelerator_count") + set_fields = [name for name in fields if getattr(config, name)] + if set_fields and len(set_fields) != len(fields): + missing = [name for name in fields if name not in set_fields] + raise ValueError( + f"{', '.join(fields)} must all be provided or all be omitted; " + f"got {', '.join(set_fields)} without {', '.join(missing)}." + ) + + @staticmethod + def _prepare_custom_deploy_request( + config: types.DeployCustomModelConfig, + ) -> tuple[ + types.DeployRequestModelConfig, + types.DeployRequestEndpointConfig, + types.DeployRequestDeployConfig, + ]: + """Translates ``DeployCustomModelConfig`` to the three ``DeployRequest`` sub-messages. + + Mirrors ``vertexai.model_garden.CustomModel.deploy`` field for field: + replica counts and reservation affinity are sent even without a machine + shape, and PSC is sent only with a non-empty project allowlist. + """ + model_config = types.DeployRequestModelConfig( + model_display_name=config.model_display_name, + ) + + endpoint_config = types.DeployRequestEndpointConfig( + endpoint_display_name=config.endpoint_display_name, + ) + if config.enable_private_service_connect and config.psc_project_allow_list: + endpoint_config.private_service_connect_config = ( + types.PrivateServiceConnectConfig( + enable_private_service_connect=True, + project_allowlist=config.psc_project_allow_list, + ) + ) + + reservation_affinity = None + if config.reservation_affinity_type or ( + config.reservation_affinity_key and config.reservation_affinity_values + ): + reservation_affinity = types.ReservationAffinity( + reservation_affinity_type=config.reservation_affinity_type, + ) + if config.reservation_affinity_key and config.reservation_affinity_values: + reservation_affinity.key = config.reservation_affinity_key + reservation_affinity.values = config.reservation_affinity_values + + machine_spec = None + if config.machine_type or reservation_affinity: + machine_spec = types.MachineSpec( + machine_type=config.machine_type, + accelerator_type=config.accelerator_type, + accelerator_count=config.accelerator_count, + reservation_affinity=reservation_affinity, + ) + + deploy_config = types.DeployRequestDeployConfig( + system_labels=config.system_labels or None, + ) + if machine_spec or config.min_replica_count or config.max_replica_count: + deploy_config.dedicated_resources = types.DedicatedResources( + machine_spec=machine_spec, + min_replica_count=config.min_replica_count or None, + max_replica_count=config.max_replica_count or None, + ) + + return model_config, endpoint_config, deploy_config + def _list_all_publisher_models( self, api_config: types.ListPublisherModelsConfig, @@ -1500,6 +1588,81 @@ def deploy_publisher_model( ) return ModelGarden._deploy_response_or_raise(operation) + def deploy_custom_model( + self, + *, + src: str, + config: Optional[types.DeployCustomModelConfigOrDict] = None, + ) -> Union[types.DeployResponse, types.DeployModelOperation]: + """Deploys a custom model from its weights in Cloud Storage. + + Use this for weights you produced yourself (a tuning job checkpoint, a + converted open model) rather than a model served from a Model Garden + publisher. For the latter, use ``deploy_publisher_model``. + + Args: + src: The Cloud Storage URI holding the model weights and config files, + e.g. ``'gs://my-bucket/weights/'``. The base model is inferred from + the config files found there. + config: Optional deployment configuration (machine shape, display + names, Private Service Connect, blocking behavior, poll timing). See + ``DeployCustomModelConfig``. + + 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 ``src`` is empty, or if the machine shape is only + partially specified (see ``DeployCustomModelConfig.machine_type``). + 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 not src: + raise ValueError("src must be specified.") + if config is None: + config = types.DeployCustomModelConfig() + elif isinstance(config, dict): + config = types.DeployCustomModelConfig.model_validate(config) + ModelGarden._validate_custom_deploy_config(config) + + model_config, endpoint_config, deploy_config = ( + ModelGarden._prepare_custom_deploy_request(config) + ) + destination = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + + operation = self._deploy( + destination=destination, + custom_model=types.DeployRequestCustomModel(gcs_uri=src), + 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. @@ -2207,7 +2370,7 @@ async def get_deploy_publisher_model_operation( config: Optional[types.GetDeployOperationConfigOrDict] = None, ) -> types.DeployModelOperation: """ - Fetches the status of an in-flight ``deploy_publisher_model`` LRO. + Fetches the status of an in-flight ``deploy_publisher_model`` or ``deploy_custom_model`` LRO. """ parameter_model = types._GetDeployOperationParameters( @@ -2507,6 +2670,52 @@ async def deploy_publisher_model( ) return ModelGarden._deploy_response_or_raise(operation) + async def deploy_custom_model( + self, + *, + src: str, + config: Optional[types.DeployCustomModelConfigOrDict] = None, + ) -> Union[types.DeployResponse, types.DeployModelOperation]: + """Async variant of ``ModelGarden.deploy_custom_model``.""" + if not src: + raise ValueError("src must be specified.") + if config is None: + config = types.DeployCustomModelConfig() + elif isinstance(config, dict): + config = types.DeployCustomModelConfig.model_validate(config) + ModelGarden._validate_custom_deploy_config(config) + + model_config, endpoint_config, deploy_config = ( + ModelGarden._prepare_custom_deploy_request(config) + ) + destination = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + + operation = await self._deploy( + destination=destination, + custom_model=types.DeployRequestCustomModel(gcs_uri=src), + 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 4c0d33c86e..394f649938 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -538,6 +538,9 @@ from .common import DeployConfig from .common import DeployConfigDict from .common import DeployConfigOrDict +from .common import DeployCustomModelConfig +from .common import DeployCustomModelConfigDict +from .common import DeployCustomModelConfigOrDict from .common import DeployedModel from .common import DeployedModelDict from .common import DeployedModelOrDict @@ -4360,6 +4363,9 @@ "DeployPublisherModelConfig", "DeployPublisherModelConfigDict", "DeployPublisherModelConfigOrDict", + "DeployCustomModelConfig", + "DeployCustomModelConfigDict", + "DeployCustomModelConfigOrDict", "DeployOption", "DeployOptionDict", "DeployOptionOrDict", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 8895b9feef..49657d375f 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -31797,6 +31797,187 @@ class DeployPublisherModelConfigDict(TypedDict, total=False): ] +class DeployCustomModelConfig(_common.BaseModel): + """Config for ``deploy_custom_model``. + + Deliberately narrower than ``DeployPublisherModelConfig``: a custom model + has no publisher, so there is no EULA to accept, no Hugging Face artifacts + to authenticate against, and no publisher-supplied default container for a + container override to patch. + """ + + machine_type: Optional[str] = Field( + default=None, + description="""Machine type (e.g. ``'a3-highgpu-8g'``). Must be set together + with ``accelerator_type`` and ``accelerator_count``, or left unset for + automatic resources.""", + ) + min_replica_count: Optional[int] = Field( + default=1, + description="""Minimum number of replicas. ``0`` is treated as unset and not + sent, as in ``CustomModel.deploy``.""", + ) + max_replica_count: Optional[int] = Field( + default=1, + description="""Maximum number of replicas. ``0`` is treated as unset and not + sent, as in ``CustomModel.deploy``.""", + ) + accelerator_type: Optional[str] = Field( + default=None, + description="""Accelerator type (e.g. ``'NVIDIA_H100_80GB'``). Must be set + together with ``machine_type`` and ``accelerator_count``.""", + ) + accelerator_count: Optional[int] = Field( + default=None, + description="""Number of accelerators per replica; must be positive. Must be set + together with ``machine_type`` and ``accelerator_type``. ``0`` counts as + unset for that check, as in ``CustomModel.deploy``.""", + ) + 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.""" + ) + enable_private_service_connect: Optional[bool] = Field( + default=None, + description="""Enable Private Service Connect for the endpoint. Only sent + together with a non-empty ``psc_project_allow_list``.""", + ) + psc_project_allow_list: Optional[list[str]] = Field( + default=None, + description="""Projects allowed to access the endpoint over Private Service + Connect. Only sent 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``.""", + ) + reservation_affinity_type: Optional[str] = Field( + default=None, + description="""Reservation affinity type. One of ``NO_RESERVATION``, + ``ANY_RESERVATION`` or ``SPECIFIC_RESERVATION``.""", + ) + reservation_affinity_key: Optional[str] = Field( + default=None, + description="""Label key of the reservation resource. To target a + ``SPECIFIC_RESERVATION`` by name, use + ``compute.googleapis.com/reservation-name``. Only sent together with + ``reservation_affinity_values``.""", + ) + reservation_affinity_values: Optional[list[str]] = Field( + default=None, + description="""Label values of the reservation resource, as full reservation + resource names + (``projects/{project}/zones/{zone}/reservations/{name}``). Only sent + together with ``reservation_affinity_key``.""", + ) + system_labels: Optional[dict[str, str]] = Field( + default=None, description="""System labels for Model Garden deployments.""" + ) + + +class DeployCustomModelConfigDict(TypedDict, total=False): + """Config for ``deploy_custom_model``. + + Deliberately narrower than ``DeployPublisherModelConfig``: a custom model + has no publisher, so there is no EULA to accept, no Hugging Face artifacts + to authenticate against, and no publisher-supplied default container for a + container override to patch. + """ + + machine_type: Optional[str] + """Machine type (e.g. ``'a3-highgpu-8g'``). Must be set together + with ``accelerator_type`` and ``accelerator_count``, or left unset for + automatic resources.""" + + min_replica_count: Optional[int] + """Minimum number of replicas. ``0`` is treated as unset and not + sent, as in ``CustomModel.deploy``.""" + + max_replica_count: Optional[int] + """Maximum number of replicas. ``0`` is treated as unset and not + sent, as in ``CustomModel.deploy``.""" + + accelerator_type: Optional[str] + """Accelerator type (e.g. ``'NVIDIA_H100_80GB'``). Must be set + together with ``machine_type`` and ``accelerator_count``.""" + + accelerator_count: Optional[int] + """Number of accelerators per replica; must be positive. Must be set + together with ``machine_type`` and ``accelerator_type``. ``0`` counts as + unset for that check, as in ``CustomModel.deploy``.""" + + endpoint_display_name: Optional[str] + """Display name for the endpoint.""" + + model_display_name: Optional[str] + """Display name for the deployed model.""" + + enable_private_service_connect: Optional[bool] + """Enable Private Service Connect for the endpoint. Only sent + together with a non-empty ``psc_project_allow_list``.""" + + psc_project_allow_list: Optional[list[str]] + """Projects allowed to access the endpoint over Private Service + Connect. Only sent 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``.""" + + reservation_affinity_type: Optional[str] + """Reservation affinity type. One of ``NO_RESERVATION``, + ``ANY_RESERVATION`` or ``SPECIFIC_RESERVATION``.""" + + reservation_affinity_key: Optional[str] + """Label key of the reservation resource. To target a + ``SPECIFIC_RESERVATION`` by name, use + ``compute.googleapis.com/reservation-name``. Only sent together with + ``reservation_affinity_values``.""" + + reservation_affinity_values: Optional[list[str]] + """Label values of the reservation resource, as full reservation + resource names + (``projects/{project}/zones/{zone}/reservations/{name}``). Only sent + together with ``reservation_affinity_key``.""" + + system_labels: Optional[dict[str, str]] + """System labels for Model Garden deployments.""" + + +DeployCustomModelConfigOrDict = Union[ + DeployCustomModelConfig, DeployCustomModelConfigDict +] + + 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 b6a7c8fb46..b122588248 100644 --- a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py @@ -410,6 +410,120 @@ def test_deploy_publisher_model_container_args_without_image_raises(client): ) +def test_deploy_custom_model_empty_src_raises(client): + """An empty src is rejected client-side; no RPC needed.""" + with pytest.raises(ValueError, match="src must be specified"): + client.model_garden.deploy_custom_model(src="") + + +def test_deploy_custom_model_partial_machine_shape_raises(client): + """A half-specified machine shape is rejected before any RPC. + + Validation runs after the src check but before the deploy call, so a + regression here would surface as an unrecorded interaction rather than a + silent partial deploy. + """ + with pytest.raises(ValueError, match="must all be provided or all be"): + client.model_garden.deploy_custom_model( + src=_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig(machine_type="g2-standard-12"), + ) + + +def test_deploy_custom_model_accelerator_without_machine_type_raises(client): + """Accelerators without a machine type hit the same all-or-none check.""" + with pytest.raises(ValueError, match="must all be provided or all be"): + client.model_garden.deploy_custom_model( + src=_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + accelerator_type="NVIDIA_L4", + accelerator_count=1, + ), + ) + + +def test_get_deploy_operation_for_custom_model(client): + """The deploy operation getter also polls custom-model deploy LROs. + + Chains a wait_for_completion=False deploy with an explicit poll via + ``get_deploy_publisher_model_operation`` -- the same `{operationName}` GET + serves both deploy methods. + """ + op = client.model_garden.deploy_custom_model( + src=_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + wait_for_completion=False, + ), + ) + polled = client.model_garden.get_deploy_publisher_model_operation( + operation_name=op.name, + ) + assert isinstance(polled, types.DeployModelOperation) + assert polled.name == op.name + + +def test_deploy_custom_model_no_wait_returns_operation(client): + """wait_for_completion=False returns the DeployModelOperation immediately. + + Records only the initial Deploy POST. Also pins the wire mapping: ``src`` + must serialize to ``customModel.gcs_uri`` and the machine shape to + ``deployConfig.dedicated_resources`` exactly as recorded, or the replay + framework's request-body comparison fails. + """ + operation = client.model_garden.deploy_custom_model( + src=_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + min_replica_count=1, + max_replica_count=1, + wait_for_completion=False, + ), + ) + assert isinstance(operation, types.DeployModelOperation) + assert "/operations/" in operation.name + # The non-blocking path must not resolve the LRO. + assert operation.done is not True + + +def test_deploy_custom_model_default_config_no_wait(client): + """Without a machine shape, only the 1/1 replica counts are sent. + + Matches ``CustomModel.deploy`` and confirms the backend accepts + ``dedicated_resources`` without a ``machine_spec``. + """ + operation = client.model_garden.deploy_custom_model( + src=_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig(wait_for_completion=False), + ) + assert isinstance(operation, types.DeployModelOperation) + assert operation.name + + +def test_deploy_custom_model_no_wait_dict_config(client): + """The config may be passed as a plain dict; wait_for_completion=False path. + + Also pins ``endpoint_display_name`` and ``system_labels`` on the wire. + """ + operation = client.model_garden.deploy_custom_model( + src=_CUSTOM_MODEL_SRC, + config={ + "machine_type": "g2-standard-12", + "accelerator_type": "NVIDIA_L4", + "accelerator_count": 1, + "endpoint_display_name": "sdk-replay-custom-model-dict", + "system_labels": {"source": "sdk-replay"}, + "wait_for_completion": False, + }, + ) + assert isinstance(operation, types.DeployModelOperation) + assert operation.name + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), @@ -542,3 +656,37 @@ async def test_deploy_publisher_model_async_container_override_raises(client): container_args=["--tensor-parallel-size=2"], ), ) + + +@pytest.mark.asyncio +async def test_deploy_custom_model_async_empty_src_raises(client): + """Async empty src is rejected client-side; no RPC needed.""" + with pytest.raises(ValueError, match="src must be specified"): + await client.aio.model_garden.deploy_custom_model(src="") + + +@pytest.mark.asyncio +async def test_deploy_custom_model_async_partial_machine_shape_raises(client): + """Async machine-shape validation matches the sync surface.""" + with pytest.raises(ValueError, match="must all be provided or all be"): + await client.aio.model_garden.deploy_custom_model( + src=_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig(machine_type="g2-standard-12"), + ) + + +@pytest.mark.asyncio +async def test_deploy_custom_model_async_no_wait_returns_operation(client): + """Async wait_for_completion=False returns the DeployModelOperation immediately.""" + operation = await client.aio.model_garden.deploy_custom_model( + src=_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + wait_for_completion=False, + ), + ) + assert isinstance(operation, types.DeployModelOperation) + assert "/operations/" in operation.name + assert operation.done is not True diff --git a/tests/unit/agentplatform/genai/test_genai_model_garden.py b/tests/unit/agentplatform/genai/test_genai_model_garden.py index ed727d0af8..f5ebbb5218 100644 --- a/tests/unit/agentplatform/genai/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/test_genai_model_garden.py @@ -2520,6 +2520,627 @@ def test_deploy_publisher_model_async_raises_on_operation_error( asyncio.run(mock_async_client.deploy_publisher_model(model=_OPEN_MODEL)) +# ---- deploy_custom_model tests ---------------------------------------- + + +_TEST_CUSTOM_MODEL_SRC = "gs://tuning-job-output/node-0/checkpoints/final" + + +# --- Return type & blocking behavior ------------------------------------ + + +def test_deploy_custom_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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + 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_custom_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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig(wait_for_completion=False), + ) + m_await.assert_not_called() + assert result is initial_op + + +def test_deploy_custom_model_polls_via_public_getter(mock_client): + """Custom deploys reuse ``get_deploy_publisher_model_operation`` to poll.""" + 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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + 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_custom_model_uses_default_poll_and_timeout(mock_client): + """No override in config -> ModelGarden deploy defaults reach 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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + kwargs = m_await.call_args.kwargs + assert kwargs["poll_interval"] == 30 + assert kwargs["timeout_seconds"] == 2 * 60 * 60 + + +def test_deploy_custom_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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + 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_custom_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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + + +def test_deploy_custom_model_raises_when_response_missing_endpoint(mock_client): + """An LRO that completes without an endpoint is reported, not returned.""" + 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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + + +# --- Request shaping ----------------------------------------------------- + + +def test_deploy_custom_model_sends_gcs_uri_as_custom_model(mock_client): + """``src`` rides on ``custom_model.gcs_uri``, not as a publisher model.""" + 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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + kwargs = m_rpc.call_args.kwargs + assert kwargs["custom_model"].gcs_uri == _TEST_CUSTOM_MODEL_SRC + assert "publisher_model_name" not in kwargs + assert "hugging_face_model_id" not in kwargs + assert kwargs["destination"] == ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" + ) + + +def test_deploy_custom_model_default_config_sends_only_replica_counts( + mock_client, +): + """An empty config sends only the 1/1 replica counts, as CustomModel does.""" + 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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + kwargs = m_rpc.call_args.kwargs + resources = kwargs["deploy_config"].dedicated_resources + assert resources.machine_spec is None + assert resources.min_replica_count == 1 + assert resources.max_replica_count == 1 + assert kwargs["deploy_config"].system_labels is None + assert kwargs["endpoint_config"].private_service_connect_config is None + + +def test_deploy_custom_model_replica_counts_without_machine_shape(mock_client): + """Replica counts are sent even when the machine shape is left automatic.""" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + min_replica_count=2, + max_replica_count=5, + ), + ) + resources = m_rpc.call_args.kwargs["deploy_config"].dedicated_resources + assert resources.machine_spec is None + assert resources.min_replica_count == 2 + assert resources.max_replica_count == 5 + + +def test_deploy_custom_model_specific_reservation_forwarded(mock_client): + """Reservation type, key and values land on the machine spec.""" + reservation = "projects/p/zones/us-central1-a/reservations/r" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + machine_type="a3-highgpu-8g", + accelerator_type="NVIDIA_H100_80GB", + accelerator_count=8, + reservation_affinity_type="SPECIFIC_RESERVATION", + reservation_affinity_key="compute.googleapis.com/reservation-name", + reservation_affinity_values=[reservation], + ), + ) + machine_spec = m_rpc.call_args.kwargs[ + "deploy_config" + ].dedicated_resources.machine_spec + assert machine_spec.machine_type == "a3-highgpu-8g" + affinity = machine_spec.reservation_affinity + assert affinity.reservation_affinity_type == "SPECIFIC_RESERVATION" + assert affinity.key == "compute.googleapis.com/reservation-name" + assert affinity.values == [reservation] + + +def test_deploy_custom_model_reservation_without_machine_shape(mock_client): + """Reservation affinity is sent even when the machine shape is automatic.""" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + reservation_affinity_type="ANY_RESERVATION", + ), + ) + machine_spec = m_rpc.call_args.kwargs[ + "deploy_config" + ].dedicated_resources.machine_spec + assert machine_spec.machine_type is None + affinity = machine_spec.reservation_affinity + assert affinity.reservation_affinity_type == "ANY_RESERVATION" + assert affinity.key is None + assert affinity.values is None + + +def test_deploy_custom_model_reservation_key_without_values_dropped( + mock_client, +): + """A reservation key without values is not sent, as in CustomModel.deploy.""" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + reservation_affinity_type="SPECIFIC_RESERVATION", + reservation_affinity_key="compute.googleapis.com/reservation-name", + ), + ) + affinity = m_rpc.call_args.kwargs[ + "deploy_config" + ].dedicated_resources.machine_spec.reservation_affinity + assert affinity.reservation_affinity_type == "SPECIFIC_RESERVATION" + assert affinity.key is None + assert affinity.values is None + + +def test_deploy_custom_model_system_labels_forwarded(mock_client): + """System labels land on the deploy 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + system_labels={"source": "notebook"}, + ), + ) + assert m_rpc.call_args.kwargs["deploy_config"].system_labels == { + "source": "notebook" + } + + +def test_deploy_custom_model_machine_shape_forwarded(mock_client): + """A complete machine shape becomes dedicated resources with replica counts.""" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + machine_type="a3-highgpu-8g", + accelerator_type="NVIDIA_H100_80GB", + accelerator_count=8, + min_replica_count=2, + max_replica_count=4, + ), + ) + resources = m_rpc.call_args.kwargs["deploy_config"].dedicated_resources + assert resources.machine_spec.machine_type == "a3-highgpu-8g" + assert resources.machine_spec.accelerator_type == "NVIDIA_H100_80GB" + assert resources.machine_spec.accelerator_count == 8 + assert resources.min_replica_count == 2 + assert resources.max_replica_count == 4 + + +def test_deploy_custom_model_display_names_forwarded(mock_client): + """Endpoint and model display names land on their respective 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + endpoint_display_name="my-endpoint", + model_display_name="my-model", + ), + ) + 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" + + +def test_deploy_custom_model_private_service_connect(mock_client): + """PSC config is built only when explicitly enabled, carrying the allowlist.""" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + enable_private_service_connect=True, + psc_project_allow_list=["proj-a", "proj-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 == ["proj-a", "proj-b"] + + +def test_deploy_custom_model_psc_without_allow_list_not_sent(mock_client): + """PSC enabled with no allowlist is not sent, as in CustomModel.deploy.""" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + enable_private_service_connect=True, + ), + ) + assert ( + m_rpc.call_args.kwargs["endpoint_config"].private_service_connect_config + is None + ) + + +def test_deploy_custom_model_sends_no_publisher_only_fields(mock_client): + """A custom model carries no EULA, HF token or container override.""" + 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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + model_config = m_rpc.call_args.kwargs["model_config_val"] + assert model_config.accept_eula is None + assert model_config.hugging_face_access_token is None + assert model_config.container_spec is None + + +def test_deploy_custom_model_dict_config_validated(mock_client): + """A plain dict config is coerced to DeployCustomModelConfig.""" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config={"model_display_name": "from-dict"}, + ) + assert ( + m_rpc.call_args.kwargs["model_config_val"].model_display_name + == "from-dict" + ) + + +# --- Client-side validation ---------------------------------------------- + + +def test_deploy_custom_model_empty_src_raises_before_rpc(mock_client): + """An empty src is rejected locally, with no RPC issued.""" + with _patch_deploy_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="src must be specified"): + mock_client.deploy_custom_model(src="") + m_rpc.assert_not_called() + + +def test_deploy_custom_model_partial_machine_shape_raises_before_rpc( + mock_client, +): + """Machine type without the accelerator pair is rejected locally.""" + with _patch_deploy_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="must all be provided or all be"): + mock_client.deploy_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig(machine_type="a3-highgpu-8g"), + ) + m_rpc.assert_not_called() + + +def test_deploy_custom_model_partial_machine_shape_error_names_missing_fields( + mock_client, +): + """The error tells the caller which fields are still missing.""" + with pytest.raises(ValueError) as exc_info: + mock_client.deploy_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + machine_type="a3-highgpu-8g", + accelerator_type="NVIDIA_H100_80GB", + ), + ) + message = str(exc_info.value) + assert "accelerator_count" in message + assert "machine_type, accelerator_type" in message + + +def test_deploy_custom_model_accelerator_only_raises(mock_client): + """Accelerators without a machine type are rejected too (not silently sent).""" + with _patch_deploy_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="must all be provided or all be"): + mock_client.deploy_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + accelerator_type="NVIDIA_H100_80GB", + accelerator_count=8, + ), + ) + m_rpc.assert_not_called() + + +def test_deploy_custom_model_zero_accelerator_count_raises(mock_client): + """accelerator_count=0 counts as unset (as in CustomModel), so it is rejected.""" + with _patch_deploy_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="without accelerator_count"): + mock_client.deploy_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + machine_type="a3-highgpu-8g", + accelerator_type="NVIDIA_H100_80GB", + accelerator_count=0, + ), + ) + m_rpc.assert_not_called() + + +def test_deploy_custom_model_zero_replica_count_not_sent(mock_client): + """A replica count of 0 is treated as unset, as in CustomModel.deploy.""" + 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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + min_replica_count=0, + max_replica_count=3, + ), + ) + resources = m_rpc.call_args.kwargs["deploy_config"].dedicated_resources + assert resources.min_replica_count is None + assert resources.max_replica_count == 3 + + +# --- Async surface parity ------------------------------------------------ + + +def test_deploy_custom_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) as m_rpc, + 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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + ) + assert isinstance(result, types.DeployResponse) + assert result.endpoint == _TEST_ENDPOINT_NAME + assert ( + m_rpc.call_args.kwargs["custom_model"].gcs_uri == _TEST_CUSTOM_MODEL_SRC + ) + + +def test_deploy_custom_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_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig(wait_for_completion=False), + ) + ) + m_await.assert_not_called() + assert result is initial_op + + +def test_deploy_custom_model_async_partial_machine_shape_raises( + mock_async_client, +): + """Async machine-shape validation matches the sync surface.""" + with _patch_async_deploy_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="must all be provided or all be"): + asyncio.run( + mock_async_client.deploy_custom_model( + src=_TEST_CUSTOM_MODEL_SRC, + config=types.DeployCustomModelConfig( + machine_type="a3-highgpu-8g" + ), + ) + ) + m_rpc.assert_not_called() + + +def test_deploy_custom_model_async_empty_src_raises(mock_async_client): + """Async empty src is rejected locally, with no RPC issued.""" + with _patch_async_deploy_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="src must be specified"): + asyncio.run(mock_async_client.deploy_custom_model(src="")) + m_rpc.assert_not_called() + + +def test_deploy_custom_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_custom_model(src=_TEST_CUSTOM_MODEL_SRC) + ) + + # --- Private (pre-GA) surface ------------------------------------------ # # No direct unit tests: `agentplatform/private/**` is excluded from the