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

PiperOrigin-RevId: 697830649

Author: KCFindstr

google/cloud/aiplatform/preview/models.py

@@ -582,6 +582,7 @@ def deploy(
         deployment_resource_pool: Optional[DeploymentResourcePool] = None,
         disable_container_logging: bool = False,
         fast_tryout_enabled: bool = False,
+        system_labels: Optional[Dict[str, str]] = None,
     ) -> None:
         """Deploys a Model to the Endpoint.
 
@@ -667,6 +668,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 and machine types. Defaults to False.
+            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()
@@ -706,6 +710,7 @@ def deploy(
             deployment_resource_pool=deployment_resource_pool,
             disable_container_logging=disable_container_logging,
             fast_tryout_enabled=fast_tryout_enabled,
+            system_labels=system_labels,
         )
 
     @base.optional_sync()
@@ -730,6 +735,7 @@ def _deploy(
         deployment_resource_pool: Optional[DeploymentResourcePool] = None,
         disable_container_logging: bool = False,
         fast_tryout_enabled: bool = False,
+        system_labels: Optional[Dict[str, str]] = None,
     ) -> None:
         """Deploys a Model to the Endpoint.
 
@@ -809,6 +815,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 and machine types. Defaults to False.
+            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(
@@ -838,6 +847,7 @@ def _deploy(
             deployment_resource_pool=deployment_resource_pool,
             disable_container_logging=disable_container_logging,
             fast_tryout_enabled=fast_tryout_enabled,
+            system_labels=system_labels,
         )
 
         _LOGGER.log_action_completed_against_resource("model", "deployed", self)
@@ -869,6 +879,7 @@ def _deploy_call(
         deployment_resource_pool: Optional[DeploymentResourcePool] = None,
         disable_container_logging: bool = False,
         fast_tryout_enabled: bool = False,
+        system_labels: Optional[Dict[str, str]] = None,
     ) -> None:
         """Helper method to deploy model to endpoint.
 
@@ -955,6 +966,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 and machine types. Defaults to False.
+            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
@@ -990,6 +1004,9 @@ def _deploy_call(
                 enable_container_logging=not 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
@@ -1074,6 +1091,9 @@ def _deploy_call(
                 enable_container_logging=not 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
@@ -1359,6 +1379,7 @@ def deploy(
         deployment_resource_pool: Optional[DeploymentResourcePool] = None,
         disable_container_logging: bool = False,
         fast_tryout_enabled: bool = False,
+        system_labels: Optional[Dict[str, str]] = None,
     ) -> Union[Endpoint, models.PrivateEndpoint]:
         """Deploys model to endpoint.
 
@@ -1465,6 +1486,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 and machine types. Defaults to False.
+            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, models.PrivateEndpoint]):
@@ -1521,6 +1545,7 @@ def deploy(
             deployment_resource_pool=deployment_resource_pool,
             disable_container_logging=disable_container_logging,
             fast_tryout_enabled=fast_tryout_enabled,
+            system_labels=system_labels,
         )
 
     @base.optional_sync(return_input_arg="endpoint", bind_future_to_self=False)
@@ -1547,6 +1572,7 @@ def _deploy(
         deployment_resource_pool: Optional[DeploymentResourcePool] = None,
         disable_container_logging: bool = False,
         fast_tryout_enabled: bool = False,
+        system_labels: Optional[Dict[str, str]] = None,
     ) -> Union[Endpoint, models.PrivateEndpoint]:
         """Deploys model to endpoint.
 
@@ -1644,6 +1670,9 @@ def _deploy(
             fast_tryout_enabled (bool):
               Optional. Whether to enable fast deployment. Defaults to False.
               Useful for quick experiments. Not for production workloads.
+            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, models.PrivateEndpoint]):
@@ -1696,6 +1725,7 @@ def _deploy(
             deployment_resource_pool=deployment_resource_pool,
             disable_container_logging=disable_container_logging,
             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

@@ -2161,6 +2161,53 @@ def test_deploy_with_fast_tryout_enabled(self, deploy_model_mock, sync):
             timeout=None,
         )
 
+    @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock")
+    @pytest.mark.parametrize("sync", [True, False])
+    def test_preview_deploy_with_system_labels(self, preview_deploy_model_mock, sync):
+        test_endpoint = preview_models.Endpoint(_TEST_ENDPOINT_NAME)
+        test_model = preview_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_v1beta1.MachineSpec(
+            machine_type=_TEST_MACHINE_TYPE,
+            accelerator_type=_TEST_ACCELERATOR_TYPE,
+            accelerator_count=_TEST_ACCELERATOR_COUNT,
+        )
+        expected_dedicated_resources = gca_machine_resources_v1beta1.DedicatedResources(
+            machine_spec=expected_machine_spec,
+            min_replica_count=1,
+            max_replica_count=1,
+        )
+        expected_deployed_model = gca_endpoint_v1beta1.DeployedModel(
+            dedicated_resources=expected_dedicated_resources,
+            model=test_model.resource_name,
+            display_name=None,
+            faster_deployment_config=gca_endpoint_v1beta1.FasterDeploymentConfig(),
+            enable_container_logging=True,
+            system_labels=_TEST_LABELS,
+        )
+        preview_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

@@ -501,6 +501,8 @@
     "aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle"
 )
 
+_TEST_LABELS = {"label1": "value1", "label2": "value2"}
+
 
 @pytest.fixture
 def mock_model():
@@ -2627,6 +2629,56 @@ def test_deploy_with_fast_tryout_enabled(self, 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_preview_deploy_with_system_labels(self, preview_deploy_model_mock, sync):
+        test_model = models.Model(_TEST_ID).preview
+        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_v1beta1.MachineSpec(
+            machine_type=_TEST_MACHINE_TYPE,
+            accelerator_type=_TEST_ACCELERATOR_TYPE,
+            accelerator_count=_TEST_ACCELERATOR_COUNT,
+        )
+        expected_dedicated_resources = gca_machine_resources_v1beta1.DedicatedResources(
+            machine_spec=expected_machine_spec,
+            min_replica_count=1,
+            max_replica_count=1,
+        )
+        expected_deployed_model = gca_endpoint_v1beta1.DeployedModel(
+            dedicated_resources=expected_dedicated_resources,
+            model=test_model.resource_name,
+            display_name=None,
+            enable_container_logging=True,
+            faster_deployment_config=gca_endpoint_v1beta1.FasterDeploymentConfig(),
+            system_labels=_TEST_LABELS,
+        )
+        preview_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",