feat: Expose system_labels field to model deployment APIs in Vertex Python SDK

PiperOrigin-RevId: 697852093

Author: KCFindstr

google/cloud/aiplatform/models.py

@@ -1291,6 +1291,7 @@ def deploy(
  reservation_affinity_values: Optional[List[str]] = None,
  spot: bool = False,
  fast_tryout_enabled: bool = False,
+ system_labels: Optional[Dict[str, str]] = None,
  ) -> None:
  """Deploys a Model to the Endpoint.
 
@@ -1403,6 +1404,9 @@ def deploy(
  If True, model will be deployed using faster deployment path.
  Useful for quick experiments. Not for production workloads. Only
  available for most popular models with certain machine types.
+ system_labels (Dict[str, str]):
+ Optional. System labels to apply to Model Garden deployments.
+ System labels are managed by Google for internal use only.
  """
  self._sync_gca_resource_if_skipped()
 
@@ -1447,6 +1451,7 @@ def deploy(
  disable_container_logging=disable_container_logging,
  deployment_resource_pool=deployment_resource_pool,
  fast_tryout_enabled=fast_tryout_enabled,
+ system_labels=system_labels,
  )
 
  @base.optional_sync()
@@ -1477,6 +1482,7 @@ def _deploy(
  disable_container_logging: bool = False,
  deployment_resource_pool: Optional[DeploymentResourcePool] = None,
  fast_tryout_enabled: bool = False,
+ system_labels: Optional[Dict[str, str]] = None,
  ) -> None:
  """Deploys a Model to the Endpoint.
 
@@ -1583,6 +1589,9 @@ def _deploy(
  If True, model will be deployed using faster deployment path.
  Useful for quick experiments. Not for production workloads. Only
  available for most popular models with certain machine types.
+ system_labels (Dict[str, str]):
+ Optional. System labels to apply to Model Garden deployments.
+ System labels are managed by Google for internal use only.
  """
  _LOGGER.log_action_start_against_resource(
  f"Deploying Model {model.resource_name} to", "", self
@@ -1617,6 +1626,7 @@ def _deploy(
  disable_container_logging=disable_container_logging,
  deployment_resource_pool=deployment_resource_pool,
  fast_tryout_enabled=fast_tryout_enabled,
+ system_labels=system_labels,
  )
 
  _LOGGER.log_action_completed_against_resource("model", "deployed", self)
@@ -1654,6 +1664,7 @@ def _deploy_call(
  disable_container_logging: bool = False,
  deployment_resource_pool: Optional[DeploymentResourcePool] = None,
  fast_tryout_enabled: bool = False,
+ system_labels: Optional[Dict[str, str]] = None,
  ) -> None:
  """Helper method to deploy model to endpoint.
 
@@ -1767,6 +1778,9 @@ def _deploy_call(
  If True, model will be deployed using faster deployment path.
  Useful for quick experiments. Not for production workloads. Only
  available for most popular models with certain machine types.
+ system_labels (Dict[str, str]):
+ Optional. System labels to apply to Model Garden deployments.
+ System labels are managed by Google for internal use only.
 
  Raises:
  ValueError: If only `accelerator_type` or `accelerator_count` is specified.
@@ -1788,6 +1802,9 @@ def _deploy_call(
  disable_container_logging=disable_container_logging,
  )
 
+ if system_labels:
+ deployed_model.system_labels = system_labels
+
  supports_shared_resources = (
  gca_model_compat.Model.DeploymentResourcesType.SHARED_RESOURCES
  in model.supported_deployment_resources_types
@@ -1847,6 +1864,9 @@ def _deploy_call(
  disable_container_logging=disable_container_logging,
  )
 
+ if system_labels:
+ deployed_model.system_labels = system_labels
+
  supports_automatic_resources = (
  gca_model_compat.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES
  in model.supported_deployment_resources_types
@@ -3911,6 +3931,7 @@ def deploy(
  reservation_affinity_key: Optional[str] = None,
  reservation_affinity_values: Optional[List[str]] = None,
  spot: bool = False,
+ system_labels: Optional[Dict[str, str]] = None,
  ) -> None:
  """Deploys a Model to the PrivateEndpoint.
 
@@ -4026,6 +4047,9 @@ def deploy(
  Format: 'projects/{project_id_or_number}/zones/{zone}/reservations/{reservation_name}'
  spot (bool):
  Optional. Whether to schedule the deployment workload on spot VMs.
+ system_labels (Dict[str, str]):
+ Optional. System labels to apply to Model Garden deployments.
+ System labels are managed by Google for internal use only.
  """
 
  if self.network:
@@ -4070,6 +4094,7 @@ def deploy(
  sync=sync,
  spot=spot,
  disable_container_logging=disable_container_logging,
+ system_labels=system_labels,
  )
 
  def update(
@@ -5133,6 +5158,7 @@ def deploy(
  reservation_affinity_values: Optional[List[str]] = None,
  spot: bool = False,
  fast_tryout_enabled: bool = False,
+ system_labels: Optional[Dict[str, str]] = None,
  ) -> Union[Endpoint, PrivateEndpoint]:
  """Deploys model to endpoint. Endpoint will be created if unspecified.
 
@@ -5267,6 +5293,9 @@ def deploy(
  If True, model will be deployed using faster deployment path.
  Useful for quick experiments. Not for production workloads. Only
  available for most popular models with certain machine types.
+ system_labels (Dict[str, str]):
+ Optional. System labels to apply to Model Garden deployments.
+ System labels are managed by Google for internal use only.
 
  Returns:
  endpoint (Union[Endpoint, PrivateEndpoint]):
@@ -5336,6 +5365,7 @@ def deploy(
  private_service_connect_config=private_service_connect_config,
  deployment_resource_pool=deployment_resource_pool,
  fast_tryout_enabled=fast_tryout_enabled,
+ system_labels=system_labels,
  )
 
  @base.optional_sync(return_input_arg="endpoint", bind_future_to_self=False)
@@ -5371,6 +5401,7 @@ def _deploy(
  ] = None,
  deployment_resource_pool: Optional[DeploymentResourcePool] = None,
  fast_tryout_enabled: bool = False,
+ system_labels: Optional[Dict[str, str]] = None,
  ) -> Union[Endpoint, PrivateEndpoint]:
  """Deploys model to endpoint. Endpoint will be created if unspecified.
 
@@ -5498,6 +5529,9 @@ def _deploy(
  If True, model will be deployed using faster deployment path.
  Useful for quick experiments. Not for production workloads. Only
  available for most popular models with certain machine types.
+ system_labels (Dict[str, str]):
+ Optional. System labels to apply to Model Garden deployments.
+ System labels are managed by Google for internal use only.
 
  Returns:
  endpoint (Union[Endpoint, PrivateEndpoint]):
@@ -5557,6 +5591,7 @@ def _deploy(
  disable_container_logging=disable_container_logging,
  deployment_resource_pool=deployment_resource_pool,
  fast_tryout_enabled=fast_tryout_enabled,
+ system_labels=system_labels,
  )
 
  _LOGGER.log_action_completed_against_resource("model", "deployed", endpoint)

tests/unit/aiplatform/test_endpoints.py

@@ -2208,6 +2208,51 @@ def test_preview_deploy_with_system_labels(self, preview_deploy_model_mock, sync
  timeout=None,
  )
 
+ @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock")
+ @pytest.mark.parametrize("sync", [True, False])
+ def test_deploy_with_system_labels(self, deploy_model_mock, sync):
+ test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
+ test_model = models.Model(_TEST_ID)
+ test_model._gca_resource.supported_deployment_resources_types.append(
+ aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES,
+ )
+
+ test_endpoint.deploy(
+ model=test_model,
+ sync=sync,
+ deploy_request_timeout=None,
+ machine_type=_TEST_MACHINE_TYPE,
+ accelerator_type=_TEST_ACCELERATOR_TYPE,
+ accelerator_count=_TEST_ACCELERATOR_COUNT,
+ system_labels=_TEST_LABELS,
+ )
+ if not sync:
+ test_endpoint.wait()
+
+ expected_machine_spec = gca_machine_resources.MachineSpec(
+ machine_type=_TEST_MACHINE_TYPE,
+ accelerator_type=_TEST_ACCELERATOR_TYPE,
+ accelerator_count=_TEST_ACCELERATOR_COUNT,
+ )
+ expected_dedicated_resources = gca_machine_resources.DedicatedResources(
+ machine_spec=expected_machine_spec,
+ min_replica_count=1,
+ max_replica_count=1,
+ )
+ expected_deployed_model = gca_endpoint.DeployedModel(
+ dedicated_resources=expected_dedicated_resources,
+ model=test_model.resource_name,
+ display_name=None,
+ system_labels=_TEST_LABELS,
+ )
+ deploy_model_mock.assert_called_once_with(
+ endpoint=test_endpoint.resource_name,
+ deployed_model=expected_deployed_model,
+ traffic_split={"0": 100},
+ metadata=(),
+ timeout=None,
+ )
+
  @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock", "get_drp_mock")
  @pytest.mark.parametrize("sync", [True, False])
  def test_deploy_with_deployment_resource_pool(self, deploy_model_mock, sync):

tests/unit/aiplatform/test_models.py

@@ -2679,6 +2679,54 @@ def test_preview_deploy_with_system_labels(self, preview_deploy_model_mock, sync
  timeout=None,
  )
 
+ @pytest.mark.usefixtures(
+ "get_model_mock",
+ "create_endpoint_mock",
+ "get_endpoint_mock",
+ )
+ @pytest.mark.parametrize("sync", [True, False])
+ def test_deploy_with_system_labels(self, deploy_model_mock, sync):
+ test_model = models.Model(_TEST_ID)
+ test_model._gca_resource.supported_deployment_resources_types.append(
+ aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES
+ )
+
+ test_endpoint = test_model.deploy(
+ machine_type=_TEST_MACHINE_TYPE,
+ accelerator_type=_TEST_ACCELERATOR_TYPE,
+ accelerator_count=_TEST_ACCELERATOR_COUNT,
+ sync=sync,
+ deploy_request_timeout=None,
+ system_labels=_TEST_LABELS,
+ )
+
+ if not sync:
+ test_endpoint.wait()
+
+ expected_machine_spec = gca_machine_resources.MachineSpec(
+ machine_type=_TEST_MACHINE_TYPE,
+ accelerator_type=_TEST_ACCELERATOR_TYPE,
+ accelerator_count=_TEST_ACCELERATOR_COUNT,
+ )
+ expected_dedicated_resources = gca_machine_resources.DedicatedResources(
+ machine_spec=expected_machine_spec,
+ min_replica_count=1,
+ max_replica_count=1,
+ )
+ expected_deployed_model = gca_endpoint.DeployedModel(
+ dedicated_resources=expected_dedicated_resources,
+ model=test_model.resource_name,
+ display_name=None,
+ system_labels=_TEST_LABELS,
+ )
+ deploy_model_mock.assert_called_once_with(
+ endpoint=test_endpoint.resource_name,
+ deployed_model=expected_deployed_model,
+ traffic_split={"0": 100},
+ metadata=(),
+ timeout=None,
+ )
+
  @pytest.mark.usefixtures(
  "get_model_mock",
  "preview_get_drp_mock",