Skip to content

feat: Add system_labels parameter to Model Garden deploy methods. #6541

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Add system_labels parameter to Model Garden deploy methods. #6541

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions tests/unit/vertexai/model_garden/test_model_garden.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
7 changes: 7 additions & 0 deletions vertexai/model_garden/_model_garden.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 = (
Expand Down
Loading