From 98a07dd614063cb6a4c55c9024893874d3c95a1f Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 2 Apr 2024 17:03:32 -0700 Subject: [PATCH] feat: Added GA support for PersistentResource management PiperOrigin-RevId: 621337736 --- .../aiplatform/compat/services/__init__.py | 4 + .../cloud/aiplatform/compat/types/__init__.py | 7 +- .../{preview => }/persistent_resource.py | 12 +- .../preview/vertex_ray/util/_gapic_utils.py | 2 +- google/cloud/aiplatform/utils/__init__.py | 7 +- .../aiplatform/test_persistent_resource.py | 126 +++++++++--------- vertexai/resources/preview/__init__.py | 3 - 7 files changed, 84 insertions(+), 77 deletions(-) rename google/cloud/aiplatform/{preview => }/persistent_resource.py (97%) diff --git a/google/cloud/aiplatform/compat/services/__init__.py b/google/cloud/aiplatform/compat/services/__init__.py index ce7665be9c..2ee0dd4015 100644 --- a/google/cloud/aiplatform/compat/services/__init__.py +++ b/google/cloud/aiplatform/compat/services/__init__.py @@ -124,6 +124,9 @@ from google.cloud.aiplatform_v1.services.model_service import ( client as model_service_client_v1, ) +from google.cloud.aiplatform_v1.services.persistent_resource_service import ( + client as persistent_resource_service_client_v1, +) from google.cloud.aiplatform_v1.services.pipeline_service import ( client as pipeline_service_client_v1, ) @@ -160,6 +163,7 @@ metadata_service_client_v1, model_garden_service_client_v1, model_service_client_v1, + persistent_resource_service_client_v1, pipeline_service_client_v1, prediction_service_client_v1, prediction_service_async_client_v1, diff --git a/google/cloud/aiplatform/compat/types/__init__.py b/google/cloud/aiplatform/compat/types/__init__.py index d36312b025..81c5dd9b4c 100644 --- a/google/cloud/aiplatform/compat/types/__init__.py +++ b/google/cloud/aiplatform/compat/types/__init__.py @@ -156,6 +156,8 @@ model_service as model_service_v1, model_monitoring as model_monitoring_v1, operation as operation_v1, + persistent_resource as persistent_resource_v1, + persistent_resource_service as persistent_resource_service_v1, pipeline_failure_policy as pipeline_failure_policy_v1, pipeline_job as pipeline_job_v1, pipeline_service as pipeline_service_v1, @@ -230,7 +232,8 @@ model_service_v1, model_monitoring_v1, operation_v1, - persistent_resource_v1beta1, + persistent_resource_v1, + persistent_resource_service_v1, pipeline_failure_policy_v1, pipeline_job_v1, pipeline_service_v1, @@ -306,6 +309,8 @@ model_service_v1beta1, model_monitoring_v1beta1, operation_v1beta1, + persistent_resource_v1beta1, + persistent_resource_service_v1beta1, pipeline_failure_policy_v1beta1, pipeline_job_v1beta1, pipeline_service_v1beta1, diff --git a/google/cloud/aiplatform/preview/persistent_resource.py b/google/cloud/aiplatform/persistent_resource.py similarity index 97% rename from google/cloud/aiplatform/preview/persistent_resource.py rename to google/cloud/aiplatform/persistent_resource.py index 0823af0db4..0bd6dbe404 100644 --- a/google/cloud/aiplatform/preview/persistent_resource.py +++ b/google/cloud/aiplatform/persistent_resource.py @@ -23,13 +23,11 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.services import ( - persistent_resource_service_client_v1beta1 as persistent_resource_service_client_compat, + persistent_resource_service_client_v1 as persistent_resource_service_client_compat, ) from google.cloud.aiplatform.compat.types import ( - persistent_resource_v1beta1 as gca_persistent_resource_compat, -) -from google.cloud.aiplatform_v1beta1.types import ( - encryption_spec as gca_encryption_spec, + encryption_spec as gca_encryption_spec_compat, + persistent_resource_v1 as gca_persistent_resource_compat, ) from google.protobuf import timestamp_pb2 # type: ignore from google.rpc import status_pb2 # type: ignore @@ -194,7 +192,7 @@ def create( This corresponds to the ``persistent_resource_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - resource_pools (MutableSequence[google.cloud.aiplatform_v1beta1.types.ResourcePool]): + resource_pools (MutableSequence[google.cloud.aiplatform_v1.types.ResourcePool]): Required. The list of resource pools to create for the PersistentResource. display_name (str): @@ -294,7 +292,7 @@ def create( if kms_key_name: gca_persistent_resource.encryption_spec = ( - gca_encryption_spec.EncryptionSpec(kms_key_name=kms_key_name) + gca_encryption_spec_compat.EncryptionSpec(kms_key_name=kms_key_name) ) if service_account: diff --git a/google/cloud/aiplatform/preview/vertex_ray/util/_gapic_utils.py b/google/cloud/aiplatform/preview/vertex_ray/util/_gapic_utils.py index aabfbd1171..3a15f039f0 100644 --- a/google/cloud/aiplatform/preview/vertex_ray/util/_gapic_utils.py +++ b/google/cloud/aiplatform/preview/vertex_ray/util/_gapic_utils.py @@ -47,7 +47,7 @@ def create_persistent_resource_client(): return initializer.global_config.create_client( client_class=PersistentResourceClientWithOverride, appended_gapic_version="vertex_ray", - ) + ).select_version("v1beta1") def polling_delay(num_attempts: int, time_scale: float) -> datetime.timedelta: diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py index f7c4ee725e..0a9aa2b684 100644 --- a/google/cloud/aiplatform/utils/__init__.py +++ b/google/cloud/aiplatform/utils/__init__.py @@ -79,6 +79,7 @@ schedule_service_client_v1, tensorboard_service_client_v1, vizier_service_client_v1, + persistent_resource_service_client_v1, ) from google.cloud.aiplatform.compat.types import ( @@ -739,8 +740,12 @@ class ModelGardenClientWithOverride(ClientWithOverride): class PersistentResourceClientWithOverride(ClientWithOverride): _is_temporary = True - _default_version = compat.V1BETA1 + _default_version = compat.DEFAULT_VERSION _version_map = ( + ( + compat.V1, + persistent_resource_service_client_v1.PersistentResourceServiceClient, + ), ( compat.V1BETA1, persistent_resource_service_client_v1beta1.PersistentResourceServiceClient, diff --git a/tests/unit/aiplatform/test_persistent_resource.py b/tests/unit/aiplatform/test_persistent_resource.py index b3480a6c0a..14421e9066 100644 --- a/tests/unit/aiplatform/test_persistent_resource.py +++ b/tests/unit/aiplatform/test_persistent_resource.py @@ -22,14 +22,14 @@ from google.api_core import operation as ga_operation from google.cloud import aiplatform from google.cloud.aiplatform.compat.services import ( - persistent_resource_service_client_v1beta1, + persistent_resource_service_client_v1, ) -from google.cloud.aiplatform.compat.types import encryption_spec_v1beta1 +from google.cloud.aiplatform.compat.types import encryption_spec_v1 from google.cloud.aiplatform.compat.types import ( - persistent_resource_service_v1beta1, + persistent_resource_service_v1, ) -from google.cloud.aiplatform.compat.types import persistent_resource_v1beta1 -from google.cloud.aiplatform.preview import persistent_resource +from google.cloud.aiplatform.compat.types import persistent_resource_v1 +from google.cloud.aiplatform import persistent_resource import constants as test_constants import pytest @@ -50,7 +50,7 @@ _TEST_KEY_NAME = test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME _TEST_SERVICE_ACCOUNT = test_constants.ProjectConstants._TEST_SERVICE_ACCOUNT -_TEST_PERSISTENT_RESOURCE_PROTO = persistent_resource_v1beta1.PersistentResource( +_TEST_PERSISTENT_RESOURCE_PROTO = persistent_resource_v1.PersistentResource( name=_TEST_PERSISTENT_RESOURCE_ID, resource_pools=[ test_constants.PersistentResourceConstants._TEST_RESOURCE_POOL, @@ -60,7 +60,7 @@ def _get_persistent_resource_proto( state=None, name=None, error=None -) -> persistent_resource_v1beta1.PersistentResource: +) -> persistent_resource_v1.PersistentResource: persistent_resource_proto = copy.deepcopy(_TEST_PERSISTENT_RESOURCE_PROTO) persistent_resource_proto.name = name persistent_resource_proto.state = state @@ -75,49 +75,49 @@ def _get_resource_name(name=None, project=_TEST_PROJECT, location=_TEST_LOCATION @pytest.fixture -def create_preview_persistent_resource_mock(): +def create_persistent_resource_mock(): with mock.patch.object( - (persistent_resource_service_client_v1beta1.PersistentResourceServiceClient), + (persistent_resource_service_client_v1.PersistentResourceServiceClient), "create_persistent_resource", - ) as create_preview_persistent_resource_mock: + ) as create_persistent_resource_mock: create_lro = mock.Mock(ga_operation.Operation) create_lro.result.return_value = None - create_preview_persistent_resource_mock.return_value = create_lro - yield create_preview_persistent_resource_mock + create_persistent_resource_mock.return_value = create_lro + yield create_persistent_resource_mock @pytest.fixture -def get_preview_persistent_resource_mock(): +def get_persistent_resource_mock(): with mock.patch.object( - (persistent_resource_service_client_v1beta1.PersistentResourceServiceClient), + (persistent_resource_service_client_v1.PersistentResourceServiceClient), "get_persistent_resource", - ) as get_preview_persistent_resource_mock: - get_preview_persistent_resource_mock.side_effect = [ + ) as get_persistent_resource_mock: + get_persistent_resource_mock.side_effect = [ _get_persistent_resource_proto( name=_TEST_PERSISTENT_RESOURCE_ID, - state=(persistent_resource_v1beta1.PersistentResource.State.RUNNING), + state=(persistent_resource_v1.PersistentResource.State.RUNNING), ), ] - yield get_preview_persistent_resource_mock + yield get_persistent_resource_mock _TEST_LIST_RESOURCE_1 = _get_persistent_resource_proto( name="resource_1", - state=(persistent_resource_v1beta1.PersistentResource.State.RUNNING), + state=(persistent_resource_v1.PersistentResource.State.RUNNING), ) _TEST_LIST_RESOURCE_2 = _get_persistent_resource_proto( name="resource_2", - state=(persistent_resource_v1beta1.PersistentResource.State.PROVISIONING), + state=(persistent_resource_v1.PersistentResource.State.PROVISIONING), ) _TEST_LIST_RESOURCE_3 = _get_persistent_resource_proto( name="resource_3", - state=(persistent_resource_v1beta1.PersistentResource.State.STOPPING), + state=(persistent_resource_v1.PersistentResource.State.STOPPING), ) _TEST_LIST_RESOURCE_4 = _get_persistent_resource_proto( name="resource_4", - state=(persistent_resource_v1beta1.PersistentResource.State.ERROR), + state=(persistent_resource_v1.PersistentResource.State.ERROR), ) _TEST_PERSISTENT_RESOURCE_LIST = [ @@ -129,30 +129,28 @@ def get_preview_persistent_resource_mock(): @pytest.fixture -def list_preview_persistent_resources_mock(): +def list_persistent_resources_mock(): with mock.patch.object( - (persistent_resource_service_client_v1beta1.PersistentResourceServiceClient), + (persistent_resource_service_client_v1.PersistentResourceServiceClient), "list_persistent_resources", - ) as list_preview_persistent_resources_mock: - list_preview_persistent_resources_mock.return_value = ( - _TEST_PERSISTENT_RESOURCE_LIST - ) + ) as list_persistent_resources_mock: + list_persistent_resources_mock.return_value = _TEST_PERSISTENT_RESOURCE_LIST - yield list_preview_persistent_resources_mock + yield list_persistent_resources_mock @pytest.fixture -def delete_preview_persistent_resource_mock(): +def delete_persistent_resource_mock(): with mock.patch.object( - (persistent_resource_service_client_v1beta1.PersistentResourceServiceClient), + (persistent_resource_service_client_v1.PersistentResourceServiceClient), "delete_persistent_resource", - ) as delete_preview_persistent_resource_mock: + ) as delete_persistent_resource_mock: delete_lro = mock.Mock(ga_operation.Operation) delete_lro.result.return_value = ( - persistent_resource_service_v1beta1.DeletePersistentResourceRequest() + persistent_resource_service_v1.DeletePersistentResourceRequest() ) - delete_preview_persistent_resource_mock.return_value = delete_lro - yield delete_preview_persistent_resource_mock + delete_persistent_resource_mock.return_value = delete_lro + yield delete_persistent_resource_mock @pytest.mark.usefixtures("google_auth_mock") @@ -168,8 +166,8 @@ def teardown_method(self): @pytest.mark.parametrize("sync", [True, False]) def test_create_persistent_resource( self, - create_preview_persistent_resource_mock, - get_preview_persistent_resource_mock, + create_persistent_resource_mock, + get_persistent_resource_mock, sync, ): my_test_resource = persistent_resource.PersistentResource.create( @@ -194,15 +192,15 @@ def test_create_persistent_resource( ) expected_persistent_resource_arg.labels = _TEST_LABELS - create_preview_persistent_resource_mock.assert_called_once_with( + create_persistent_resource_mock.assert_called_once_with( parent=_TEST_PARENT, persistent_resource_id=_TEST_PERSISTENT_RESOURCE_ID, persistent_resource=expected_persistent_resource_arg, timeout=None, ) - get_preview_persistent_resource_mock.assert_called_once() - _, mock_kwargs = get_preview_persistent_resource_mock.call_args + get_persistent_resource_mock.assert_called_once() + _, mock_kwargs = get_persistent_resource_mock.call_args assert mock_kwargs["name"] == _get_resource_name( name=_TEST_PERSISTENT_RESOURCE_ID ) @@ -210,8 +208,8 @@ def test_create_persistent_resource( @pytest.mark.parametrize("sync", [True, False]) def test_create_persistent_resource_with_network( self, - create_preview_persistent_resource_mock, - get_preview_persistent_resource_mock, + create_persistent_resource_mock, + get_persistent_resource_mock, sync, ): my_test_resource = persistent_resource.PersistentResource.create( @@ -234,14 +232,14 @@ def test_create_persistent_resource_with_network( expected_persistent_resource_arg.network = _TEST_NETWORK expected_persistent_resource_arg.reserved_ip_ranges = _TEST_RESERVED_IP_RANGES - create_preview_persistent_resource_mock.assert_called_once_with( + create_persistent_resource_mock.assert_called_once_with( parent=_TEST_PARENT, persistent_resource_id=_TEST_PERSISTENT_RESOURCE_ID, persistent_resource=expected_persistent_resource_arg, timeout=None, ) - get_preview_persistent_resource_mock.assert_called_once() - _, mock_kwargs = get_preview_persistent_resource_mock.call_args + get_persistent_resource_mock.assert_called_once() + _, mock_kwargs = get_persistent_resource_mock.call_args assert mock_kwargs["name"] == _get_resource_name( name=_TEST_PERSISTENT_RESOURCE_ID ) @@ -249,8 +247,8 @@ def test_create_persistent_resource_with_network( @pytest.mark.parametrize("sync", [True, False]) def test_create_persistent_resource_with_kms_key( self, - create_preview_persistent_resource_mock, - get_preview_persistent_resource_mock, + create_persistent_resource_mock, + get_persistent_resource_mock, sync, ): my_test_resource = persistent_resource.PersistentResource.create( @@ -270,17 +268,17 @@ def test_create_persistent_resource_with_kms_key( ) expected_persistent_resource_arg.encryption_spec = ( - encryption_spec_v1beta1.EncryptionSpec(kms_key_name=_TEST_KEY_NAME) + encryption_spec_v1.EncryptionSpec(kms_key_name=_TEST_KEY_NAME) ) - create_preview_persistent_resource_mock.assert_called_once_with( + create_persistent_resource_mock.assert_called_once_with( parent=_TEST_PARENT, persistent_resource_id=_TEST_PERSISTENT_RESOURCE_ID, persistent_resource=expected_persistent_resource_arg, timeout=None, ) - get_preview_persistent_resource_mock.assert_called_once() - _, mock_kwargs = get_preview_persistent_resource_mock.call_args + get_persistent_resource_mock.assert_called_once() + _, mock_kwargs = get_persistent_resource_mock.call_args assert mock_kwargs["name"] == _get_resource_name( name=_TEST_PERSISTENT_RESOURCE_ID ) @@ -288,8 +286,8 @@ def test_create_persistent_resource_with_kms_key( @pytest.mark.parametrize("sync", [True, False]) def test_create_persistent_resource_with_service_account( self, - create_preview_persistent_resource_mock, - get_preview_persistent_resource_mock, + create_persistent_resource_mock, + get_persistent_resource_mock, sync, ): my_test_resource = persistent_resource.PersistentResource.create( @@ -308,31 +306,31 @@ def test_create_persistent_resource_with_service_account( name=_TEST_PERSISTENT_RESOURCE_ID, ) - service_account_spec = persistent_resource_v1beta1.ServiceAccountSpec( + service_account_spec = persistent_resource_v1.ServiceAccountSpec( enable_custom_service_account=True, service_account=_TEST_SERVICE_ACCOUNT ) expected_persistent_resource_arg.resource_runtime_spec = ( - persistent_resource_v1beta1.ResourceRuntimeSpec( + persistent_resource_v1.ResourceRuntimeSpec( service_account_spec=service_account_spec ) ) - create_preview_persistent_resource_mock.assert_called_once_with( + create_persistent_resource_mock.assert_called_once_with( parent=_TEST_PARENT, persistent_resource_id=_TEST_PERSISTENT_RESOURCE_ID, persistent_resource=expected_persistent_resource_arg, timeout=None, ) - get_preview_persistent_resource_mock.assert_called_once() - _, mock_kwargs = get_preview_persistent_resource_mock.call_args + get_persistent_resource_mock.assert_called_once() + _, mock_kwargs = get_persistent_resource_mock.call_args assert mock_kwargs["name"] == _get_resource_name( name=_TEST_PERSISTENT_RESOURCE_ID ) - def test_list_persistent_resources(self, list_preview_persistent_resources_mock): + def test_list_persistent_resources(self, list_persistent_resources_mock): resource_list = persistent_resource.PersistentResource.list() - list_preview_persistent_resources_mock.assert_called_once() + list_persistent_resources_mock.assert_called_once() assert len(resource_list) == len(_TEST_PERSISTENT_RESOURCE_LIST) for i in range(len(resource_list)): @@ -345,8 +343,8 @@ def test_list_persistent_resources(self, list_preview_persistent_resources_mock) @pytest.mark.parametrize("sync", [True, False]) def test_delete_persistent_resource( self, - get_preview_persistent_resource_mock, - delete_preview_persistent_resource_mock, + get_persistent_resource_mock, + delete_persistent_resource_mock, sync, ): test_resource = persistent_resource.PersistentResource( @@ -357,7 +355,7 @@ def test_delete_persistent_resource( if not sync: test_resource.wait() - get_preview_persistent_resource_mock.assert_called_once() - delete_preview_persistent_resource_mock.assert_called_once_with( + get_persistent_resource_mock.assert_called_once() + delete_persistent_resource_mock.assert_called_once_with( name=_TEST_PERSISTENT_RESOURCE_ID, ) diff --git a/vertexai/resources/preview/__init__.py b/vertexai/resources/preview/__init__.py index c23b8a11f4..b4f8a3842d 100644 --- a/vertexai/resources/preview/__init__.py +++ b/vertexai/resources/preview/__init__.py @@ -26,9 +26,6 @@ Endpoint, Model, ) -from google.cloud.aiplatform.preview.persistent_resource import ( - PersistentResource, -) from google.cloud.aiplatform.preview.featurestore.entity_type import ( EntityType, )