diff --git a/tests/unit/vertexai/model_garden/test_model_garden.py b/tests/unit/vertexai/model_garden/test_model_garden.py index 1e0ee250cb..0fa7cb565e 100644 --- a/tests/unit/vertexai/model_garden/test_model_garden.py +++ b/tests/unit/vertexai/model_garden/test_model_garden.py @@ -1695,6 +1695,29 @@ def test_deploy_custom_model_with_reservation_success(self, deploy_mock): ) ) + def test_deploy_custom_model_with_system_labels_success(self, deploy_mock): + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) + model = model_garden_preview.CustomModel(gcs_uri=_TEST_GCS_URI) + model.deploy(system_labels={"test-key": "test-value"}) + deploy_mock.assert_called_once_with( + types.DeployRequest( + destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", + custom_model=types.DeployRequest.CustomModel( + gcs_uri=_TEST_GCS_URI, + ), + deploy_config=types.DeployRequest.DeployConfig( + dedicated_resources=types.DedicatedResources( + min_replica_count=1, + max_replica_count=1, + ), + system_labels={"test-key": "test-value"}, + ), + ) + ) + @pytest.mark.parametrize("filter_by_user_quota", [True, False]) def test_list_deploy_options_with_recommendations(self, filter_by_user_quota): """Tests list_deploy_options when recommend_spec returns recommendations.""" diff --git a/vertexai/model_garden/_model_garden.py b/vertexai/model_garden/_model_garden.py index e67557bc30..98e828bf40 100644 --- a/vertexai/model_garden/_model_garden.py +++ b/vertexai/model_garden/_model_garden.py @@ -1120,6 +1120,7 @@ def deploy( reservation_affinity_type: Optional[str] = None, reservation_affinity_key: Optional[str] = None, reservation_affinity_values: Optional[List[str]] = None, + system_labels: Optional[Dict[str, str]] = None, endpoint_display_name: Optional[str] = None, model_display_name: Optional[str] = None, enable_private_service_connect: bool = False, @@ -1164,6 +1165,7 @@ def deploy( name of the reservation. Format: 'projects/{project_id_or_number}/zones/{zone}/reservations/{reservation_name}' + system_labels (Dict[str, str]): Optional. System labels for Model Garden deployments. endpoint_display_name: The display name of the created endpoint. model_display_name: The display name of the custom model. enable_private_service_connect (bool): Whether to enable private service @@ -1186,6 +1188,7 @@ def deploy( reservation_affinity_type=reservation_affinity_type, reservation_affinity_key=reservation_affinity_key, reservation_affinity_values=reservation_affinity_values, + system_labels=system_labels, endpoint_display_name=endpoint_display_name, model_display_name=model_display_name, enable_private_service_connect=enable_private_service_connect, @@ -1211,6 +1214,7 @@ def _deploy_gcs_uri( reservation_affinity_type: Optional[str] = None, reservation_affinity_key: Optional[str] = None, reservation_affinity_values: Optional[List[str]] = None, + system_labels: Optional[Dict[str, str]] = None, endpoint_display_name: Optional[str] = None, model_display_name: Optional[str] = None, deploy_request_timeout: Optional[float] = None, @@ -1255,6 +1259,7 @@ def _deploy_gcs_uri( name of the reservation. Format: 'projects/{project_id_or_number}/zones/{zone}/reservations/{reservation_name}' + system_labels (Dict[str, str]): Optional. System labels for Model Garden deployments. endpoint_display_name: The display name of the created endpoint. model_display_name: The display name of the custom model. enable_private_service_connect (bool): Whether to enable private service @@ -1289,6 +1294,8 @@ def has_all_or_none_values(var1, var2, var3) -> bool: request.endpoint_config.endpoint_display_name = endpoint_display_name if model_display_name: request.model_config.model_display_name = model_display_name + if system_labels: + request.deploy_config.system_labels = system_labels if enable_private_service_connect and psc_project_allow_list: request.endpoint_config.private_service_connect_config = (