diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index 661868af17..e79e308693 100644 --- a/google/cloud/aiplatform/models.py +++ b/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) diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py index 0c41bc70df..ee8b222c1c 100644 --- a/tests/unit/aiplatform/test_endpoints.py +++ b/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): diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py index 7ce51ef605..e5ab109d42 100644 --- a/tests/unit/aiplatform/test_models.py +++ b/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",