Skip to content

Commit

Permalink
feat: Expose system_labels field to model deployment APIs in Vertex p…
Browse files Browse the repository at this point in the history
…ublic preview SDK

PiperOrigin-RevId: 697830649
KCFindstr authored and copybara-github committed Nov 19, 2024
1 parent 6a22bef commit 13cede4
Showing 3 changed files with 129 additions and 0 deletions.
30 changes: 30 additions & 0 deletions google/cloud/aiplatform/preview/models.py
Original file line number Diff line number Diff line change
@@ -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)
47 changes: 47 additions & 0 deletions tests/unit/aiplatform/test_endpoints.py
Original file line number Diff line number Diff line change
@@ -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):
52 changes: 52 additions & 0 deletions tests/unit/aiplatform/test_models.py
Original file line number Diff line number Diff line change
@@ -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",

0 comments on commit 13cede4

Please sign in to comment.