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",