From 5b3dba96038e008786f538e90e733266ee765eb7 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 26 Oct 2026 13:01:11 -0700 Subject: [PATCH] fix: Model Garden deploy_publisher_model now honors dedicated_endpoint_disabled and spot without a machine type PiperOrigin-RevId: 1007487713 --- agentplatform/_genai/model_garden.py | 12 +++-- agentplatform/_genai/types/common.py | 9 +++- .../genai/test_genai_model_garden.py | 51 ++++++++++++++----- 3 files changed, 55 insertions(+), 17 deletions(-) diff --git a/agentplatform/_genai/model_garden.py b/agentplatform/_genai/model_garden.py index 70d8cbade7..3db508eb01 100644 --- a/agentplatform/_genai/model_garden.py +++ b/agentplatform/_genai/model_garden.py @@ -1573,9 +1573,11 @@ def _prepare_deploy_request( 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 + # Send the current field, as vertexai.model_garden does. The deprecated + # dedicated_endpoint_enabled cannot express "off": False is its proto3 + # default, so the backend would keep its default (dedicated endpoint on). + if config.dedicated_endpoint_disabled: + endpoint_config.dedicated_endpoint_disabled = True if config.enable_private_service_connect: endpoint_config.private_service_connect_config = ( types.PrivateServiceConnectConfig( @@ -1598,6 +1600,10 @@ def _prepare_deploy_request( max_replica_count=config.max_replica_count, spot=config.spot, ) + elif config.spot: + # Without a machine shape the backend deploys the model's default + # verified config and applies spot to it, as with vertexai.model_garden. + deploy_config.dedicated_resources = types.DedicatedResources(spot=True) return model_config, endpoint_config, deploy_config diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 7caa3c0147..9a59cdf042 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -32128,7 +32128,11 @@ class DeployPublisherModelConfig(_common.BaseModel): accelerator_count: Optional[int] = Field( default=None, description="""Number of accelerators per replica.""" ) - spot: Optional[bool] = Field(default=None, description="""Schedule on Spot VMs.""") + spot: Optional[bool] = Field( + default=None, + description="""Schedule on Spot VMs. Without ``machine_type``, applies to the + model's default machine.""", + ) dedicated_endpoint_disabled: Optional[bool] = Field( default=None, description="""Set True to serve predictions via the shared endpoint DNS @@ -32217,7 +32221,8 @@ class DeployPublisherModelConfigDict(TypedDict, total=False): """Number of accelerators per replica.""" spot: Optional[bool] - """Schedule on Spot VMs.""" + """Schedule on Spot VMs. Without ``machine_type``, applies to the + model's default machine.""" dedicated_endpoint_disabled: Optional[bool] """Set True to serve predictions via the shared endpoint DNS diff --git a/tests/unit/agentplatform/genai/test_genai_model_garden.py b/tests/unit/agentplatform/genai/test_genai_model_garden.py index 82a92e1d06..fc00909140 100644 --- a/tests/unit/agentplatform/genai/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/test_genai_model_garden.py @@ -2120,6 +2120,7 @@ def test_deploy_publisher_model_default_config_no_optional_blocks(mock_client): 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"].dedicated_endpoint_disabled is None assert kwargs["endpoint_config"].private_service_connect_config is None @@ -2215,10 +2216,38 @@ def test_deploy_publisher_model_accelerator_only_triggers_dedicated_resources( assert dr.machine_spec.accelerator_type == "NVIDIA_L4" -def test_deploy_publisher_model_dedicated_endpoint_disabled_inverts( +def test_deploy_publisher_model_spot_without_machine_shape_forwarded( mock_client, ): - """dedicated_endpoint_disabled=True maps to dedicated_endpoint_enabled=False on wire.""" + """spot alone is sent without a machine spec; the backend applies it to the default machine.""" + 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(spot=True), + ) + dr = m_rpc.call_args.kwargs["deploy_config"].dedicated_resources + assert dr is not None + assert dr.spot is True + assert dr.machine_spec is None + assert dr.min_replica_count is None + assert dr.max_replica_count is None + + +def test_deploy_publisher_model_dedicated_endpoint_disabled_forwarded( + mock_client, +): + """dedicated_endpoint_disabled=True is sent as-is, not as the deprecated field. + + dedicated_endpoint_enabled=False is the proto3 default, so the backend + cannot tell it from unset and would keep the dedicated endpoint on. + """ with ( _patch_deploy_rpc(mock_client) as m_rpc, mock.patch.object( @@ -2233,16 +2262,15 @@ def test_deploy_publisher_model_dedicated_endpoint_disabled_inverts( dedicated_endpoint_disabled=True ), ) - assert ( - m_rpc.call_args.kwargs["endpoint_config"].dedicated_endpoint_enabled - is False - ) + endpoint_config = m_rpc.call_args.kwargs["endpoint_config"] + assert endpoint_config.dedicated_endpoint_disabled is True + assert endpoint_config.dedicated_endpoint_enabled is None -def test_deploy_publisher_model_dedicated_endpoint_disabled_false_maps_to_true( +def test_deploy_publisher_model_dedicated_endpoint_disabled_false_not_sent( mock_client, ): - """dedicated_endpoint_disabled=False maps to dedicated_endpoint_enabled=True.""" + """dedicated_endpoint_disabled=False sends neither field (server default).""" with ( _patch_deploy_rpc(mock_client) as m_rpc, mock.patch.object( @@ -2257,10 +2285,9 @@ def test_deploy_publisher_model_dedicated_endpoint_disabled_false_maps_to_true( dedicated_endpoint_disabled=False ), ) - assert ( - m_rpc.call_args.kwargs["endpoint_config"].dedicated_endpoint_enabled - is True - ) + endpoint_config = m_rpc.call_args.kwargs["endpoint_config"] + assert endpoint_config.dedicated_endpoint_disabled is None + assert endpoint_config.dedicated_endpoint_enabled is None def test_deploy_publisher_model_private_service_connect(mock_client):