From 6f8d3d1ed89f0aa6f2f0418ae752185104196c63 Mon Sep 17 00:00:00 2001 From: Morgan Du Date: Tue, 20 Jul 2021 17:29:55 -0700 Subject: [PATCH] feat: add tensorboard resource management (#539) * feat: add tensorboard resource management * addressing feedbacks * fix: correct doc string to format --- google/cloud/aiplatform/__init__.py | 2 + .../cloud/aiplatform/compat/types/__init__.py | 3 +- .../cloud/aiplatform/tensorboard/__init__.py | 5 + .../aiplatform/tensorboard/tensorboard.py | 286 +++++++++++++++++ tests/unit/aiplatform/test_tensorboard.py | 290 ++++++++++++++++++ 5 files changed, 585 insertions(+), 1 deletion(-) create mode 100644 google/cloud/aiplatform/tensorboard/tensorboard.py create mode 100644 tests/unit/aiplatform/test_tensorboard.py diff --git a/google/cloud/aiplatform/__init__.py b/google/cloud/aiplatform/__init__.py index ddf9c4e6e7..17c6952f35 100644 --- a/google/cloud/aiplatform/__init__.py +++ b/google/cloud/aiplatform/__init__.py @@ -36,6 +36,7 @@ CustomJob, HyperparameterTuningJob, ) +from google.cloud.aiplatform.tensorboard import Tensorboard from google.cloud.aiplatform.training_jobs import ( CustomTrainingJob, CustomContainerTrainingJob, @@ -91,4 +92,5 @@ "TextDataset", "TimeSeriesDataset", "VideoDataset", + "Tensorboard", ) diff --git a/google/cloud/aiplatform/compat/types/__init__.py b/google/cloud/aiplatform/compat/types/__init__.py index 93345cb5d8..f9ded28f12 100644 --- a/google/cloud/aiplatform/compat/types/__init__.py +++ b/google/cloud/aiplatform/compat/types/__init__.py @@ -53,7 +53,7 @@ study as study_v1beta1, training_pipeline as training_pipeline_v1beta1, metadata_service as metadata_service_v1beta1, - tensorboard_service as tensorboard_service_v1beta1, + tensorboard as tensorboard_v1beta1, tensorboard_data as tensorboard_data_v1beta1, tensorboard_experiment as tensorboard_experiment_v1beta1, tensorboard_run as tensorboard_run_v1beta1, @@ -167,6 +167,7 @@ specialist_pool_service_v1beta1, training_pipeline_v1beta1, metadata_service_v1beta1, + tensorboard_v1beta1, tensorboard_service_v1beta1, tensorboard_data_v1beta1, tensorboard_experiment_v1beta1, diff --git a/google/cloud/aiplatform/tensorboard/__init__.py b/google/cloud/aiplatform/tensorboard/__init__.py index a6fbe4122f..93c48cd46c 100644 --- a/google/cloud/aiplatform/tensorboard/__init__.py +++ b/google/cloud/aiplatform/tensorboard/__init__.py @@ -14,3 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # + +from google.cloud.aiplatform.tensorboard.tensorboard import Tensorboard + + +__all__ = ("Tensorboard",) diff --git a/google/cloud/aiplatform/tensorboard/tensorboard.py b/google/cloud/aiplatform/tensorboard/tensorboard.py new file mode 100644 index 0000000000..4a9598c5aa --- /dev/null +++ b/google/cloud/aiplatform/tensorboard/tensorboard.py @@ -0,0 +1,286 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Optional, Sequence, Dict, Tuple + +from google.auth import credentials as auth_credentials + +from google.cloud.aiplatform import base +from google.cloud.aiplatform import compat +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform import utils + + +from google.cloud.aiplatform.compat.types import tensorboard_v1beta1 as gca_tensorboard + +from google.protobuf import field_mask_pb2 + +_LOGGER = base.Logger(__name__) + + +class Tensorboard(base.VertexAiResourceNounWithFutureManager): + """Managed tensorboard resource for Vertex AI.""" + + client_class = utils.TensorboardClientWithOverride + _is_client_prediction_client = False + _resource_noun = "tensorboards" + _getter_method = "get_tensorboard" + _list_method = "list_tensorboards" + _delete_method = "delete_tensorboard" + + def __init__( + self, + tensorboard_name: str, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing managed tensorboard given a tensorboard name or ID. + + Args: + tensorboard_name (str): + Required. A fully-qualified tensorboard resource name or tensorboard ID. + Example: "projects/123/locations/us-central1/tensorboards/456" or + "456" when project and location are initialized or passed. + project (str): + Optional. Project to retrieve tensorboard from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve tensorboard from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retreive this Tensorboard. Overrides + credentials set in aiplatform.init. + """ + + super().__init__( + project=project, + location=location, + credentials=credentials, + resource_name=tensorboard_name, + ) + self._gca_resource = self._get_gca_resource(resource_name=tensorboard_name) + + @classmethod + def create( + cls, + display_name: str, + description: Optional[str] = None, + labels: Optional[Dict[str, str]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + encryption_spec_key_name: Optional[str] = None, + ) -> "Tensorboard": + """Creates a new tensorboard. + + Example Usage: + + tb = aiplatform.Tensorboard.create( + display_name='my display name', + description='my description', + labels={ + 'key1': 'value1', + 'key2': 'value2' + } + ) + + Args: + display_name (str): + Required. The user-defined name of the Tensorboard. + The name can be up to 128 characters long and can be consist + of any UTF-8 characters. + description (str): + Optional. Description of this Tensorboard. + labels (Dict[str, str]): + Optional. Labels with user-defined metadata to organize your Tensorboards. + Label keys and values can be no longer than 64 characters + (Unicode codepoints), can only contain lowercase letters, numeric + characters, underscores and dashes. International characters are allowed. + No more than 64 user labels can be associated with one Tensorboard + (System labels are excluded). + See https://goo.gl/xmQnxf for more information and examples of labels. + System reserved label keys are prefixed with "aiplatform.googleapis.com/" + and are immutable. + project (str): + Optional. Project to upload this model to. Overrides project set in + aiplatform.init. + location (str): + Optional. Location to upload this model to. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to upload this model. Overrides + credentials set in aiplatform.init. + request_metadata (Sequence[Tuple[str, str]]): + Optional. Strings which should be sent along with the request as metadata. + encryption_spec_key_name (str): + Optional. Cloud KMS resource identifier of the customer + managed encryption key used to protect the tensorboard. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. + + If set, this Tensorboard and all sub-resources of this Tensorboard will be secured by this key. + + Overrides encryption_spec_key_name set in aiplatform.init. + + Returns: + tensorboard (Tensorboard): + Instantiated representation of the managed tensorboard resource. + """ + + utils.validate_display_name(display_name) + + api_client = cls._instantiate_client(location=location, credentials=credentials) + + parent = initializer.global_config.common_location_path( + project=project, location=location + ) + + encryption_spec = initializer.global_config.get_encryption_spec( + encryption_spec_key_name=encryption_spec_key_name, + select_version=compat.V1BETA1, + ) + + gapic_tensorboard = gca_tensorboard.Tensorboard( + display_name=display_name, + description=description, + labels=labels, + encryption_spec=encryption_spec, + ) + + create_tensorboard_lro = api_client.create_tensorboard( + parent=parent, tensorboard=gapic_tensorboard, metadata=request_metadata + ) + + _LOGGER.log_create_with_lro(cls, create_tensorboard_lro) + + created_tensorboard = create_tensorboard_lro.result() + + _LOGGER.log_create_complete(cls, created_tensorboard, "tb") + + return cls( + tensorboard_name=created_tensorboard.name, + project=project or initializer.global_config.project, + location=location or initializer.global_config.location, + credentials=credentials, + ) + + def update( + self, + display_name: Optional[str] = None, + description: Optional[str] = None, + labels: Optional[Dict[str, str]] = None, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + encryption_spec_key_name: Optional[str] = None, + ) -> "Tensorboard": + """Updates an existing tensorboard. + + Example Usage: + + tb = aiplatform.Tensorboard(tensorboard_name='123456') + tb.update( + display_name='update my display name', + description='update my description', + ) + + Args: + display_name (str): + Optional. User-defined name of the Tensorboard. + The name can be up to 128 characters long and can be consist + of any UTF-8 characters. + description (str): + Optional. Description of this Tensorboard. + labels (Dict[str, str]): + Optional. Labels with user-defined metadata to organize your Tensorboards. + Label keys and values can be no longer than 64 characters + (Unicode codepoints), can only contain lowercase letters, numeric + characters, underscores and dashes. International characters are allowed. + No more than 64 user labels can be associated with one Tensorboard + (System labels are excluded). + See https://goo.gl/xmQnxf for more information and examples of labels. + System reserved label keys are prefixed with "aiplatform.googleapis.com/" + and are immutable. + request_metadata (Sequence[Tuple[str, str]]): + Optional. Strings which should be sent along with the request as metadata. + encryption_spec_key_name (str): + Optional. Cloud KMS resource identifier of the customer + managed encryption key used to protect the tensorboard. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. + + If set, this Tensorboard and all sub-resources of this Tensorboard will be secured by this key. + + Overrides encryption_spec_key_name set in aiplatform.init. + + Returns: + tensorboard (Tensorboard): + The managed tensorboard resource. + """ + update_mask = list() + + if display_name: + utils.validate_display_name(display_name) + update_mask.append("display_name") + + if description: + update_mask.append("description") + + if labels: + update_mask.append("labels") + + encryption_spec = None + if encryption_spec_key_name: + encryption_spec = initializer.global_config.get_encryption_spec( + encryption_spec_key_name=encryption_spec_key_name, + select_version=compat.V1BETA1, + ) + update_mask.append("encryption_spec") + + update_mask = field_mask_pb2.FieldMask(paths=update_mask) + + gapic_tensorboard = gca_tensorboard.Tensorboard( + name=self.resource_name, + display_name=display_name, + description=description, + labels=labels, + encryption_spec=encryption_spec, + ) + + _LOGGER.log_action_start_against_resource( + "Updating", "tensorboard", self, + ) + + update_tensorboard_lro = self.api_client.update_tensorboard( + tensorboard=gapic_tensorboard, + update_mask=update_mask, + metadata=request_metadata, + ) + + _LOGGER.log_action_started_against_resource_with_lro( + "Update", "tensorboard", self.__class__, update_tensorboard_lro + ) + + update_tensorboard_lro.result() + + _LOGGER.log_action_completed_against_resource("tensorboard", "updated", self) + + return self diff --git a/tests/unit/aiplatform/test_tensorboard.py b/tests/unit/aiplatform/test_tensorboard.py new file mode 100644 index 0000000000..60f85683df --- /dev/null +++ b/tests/unit/aiplatform/test_tensorboard.py @@ -0,0 +1,290 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os + +import pytest + +from unittest import mock +from unittest.mock import patch +from importlib import reload + +from google.api_core import operation +from google.auth.exceptions import GoogleAuthError +from google.auth import credentials as auth_credentials + +from google.cloud import aiplatform + +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform import tensorboard + +from google.cloud.aiplatform_v1beta1.services.tensorboard_service import ( + client as tensorboard_service_client, +) + +from google.cloud.aiplatform_v1beta1.types import ( + tensorboard as gca_tensorboard, + tensorboard_service as gca_tensorboard_service, + encryption_spec as gca_encryption_spec, +) + +from google.protobuf import field_mask_pb2 + +# project +_TEST_PROJECT = "test-project" +_TEST_LOCATION = "us-central1" +_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" +_TEST_ALT_PROJECT = "test-project_alt" + +_TEST_ALT_LOCATION = "europe-west4" +_TEST_INVALID_LOCATION = "us-central2" + +# tensorboard +_TEST_ID = "1028944691210842416" +_TEST_DISPLAY_NAME = "my_tensorboard_1234" +_TEST_DISPLAY_NAME_UPDATE = "my_tensorboard_1234_update" + +_TEST_NAME = ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/tensorboards/{_TEST_ID}" +) +_TEST_ALT_NAME = ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_ALT_LOCATION}/tensorboards/{_TEST_ID}" +) +_TEST_INVALID_NAME = f"prj/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/{_TEST_ID}" + +# request_metadata +_TEST_REQUEST_METADATA = () + +# CMEK encryption +_TEST_ENCRYPTION_KEY_NAME = "key_1234" +_TEST_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec( + kms_key_name=_TEST_ENCRYPTION_KEY_NAME +) + + +@pytest.fixture +def get_tensorboard_mock(): + with patch.object( + tensorboard_service_client.TensorboardServiceClient, "get_tensorboard" + ) as get_tensorboard_mock: + get_tensorboard_mock.return_value = gca_tensorboard.Tensorboard( + name=_TEST_NAME, + display_name=_TEST_DISPLAY_NAME, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + yield get_tensorboard_mock + + +@pytest.fixture +def create_tensorboard_mock(): + with patch.object( + tensorboard_service_client.TensorboardServiceClient, "create_tensorboard" + ) as create_tensorboard_mock: + create_tensorboard_lro_mock = mock.Mock(operation.Operation) + create_tensorboard_lro_mock.result.return_value = gca_tensorboard.Tensorboard( + name=_TEST_NAME, + display_name=_TEST_DISPLAY_NAME, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + create_tensorboard_mock.return_value = create_tensorboard_lro_mock + yield create_tensorboard_mock + + +@pytest.fixture +def update_tensorboard_mock(): + with patch.object( + tensorboard_service_client.TensorboardServiceClient, "update_tensorboard" + ) as update_tensorboard_mock: + update_tensorboard_lro_mock = mock.Mock(operation.Operation) + update_tensorboard_lro_mock.result.return_value = gca_tensorboard.Tensorboard( + name=_TEST_NAME, + display_name=_TEST_DISPLAY_NAME_UPDATE, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + update_tensorboard_mock.return_value = update_tensorboard_lro_mock + yield update_tensorboard_mock + + +@pytest.fixture +def delete_tensorboard_mock(): + with mock.patch.object( + tensorboard_service_client.TensorboardServiceClient, "delete_tensorboard" + ) as delete_tensorboard_mock: + delete_tensorboard_lro_mock = mock.Mock(operation.Operation) + delete_tensorboard_lro_mock.result.return_value = gca_tensorboard_service.DeleteTensorboardRequest( + name=_TEST_NAME, + ) + delete_tensorboard_mock.return_value = delete_tensorboard_lro_mock + yield delete_tensorboard_mock + + +class TestTensorboard: + def setup_method(self): + reload(initializer) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_init_tensorboard(self, get_tensorboard_mock): + aiplatform.init(project=_TEST_PROJECT) + tensorboard.Tensorboard(tensorboard_name=_TEST_NAME) + get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + + def test_init_tensorboard_with_id_only_with_project_and_location( + self, get_tensorboard_mock + ): + aiplatform.init(project=_TEST_PROJECT) + tensorboard.Tensorboard( + tensorboard_name=_TEST_ID, project=_TEST_PROJECT, location=_TEST_LOCATION + ) + get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + + def test_init_tensorboard_with_project_and_location(self, get_tensorboard_mock): + aiplatform.init(project=_TEST_PROJECT) + tensorboard.Tensorboard( + tensorboard_name=_TEST_NAME, project=_TEST_PROJECT, location=_TEST_LOCATION + ) + get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + + def test_init_tensorboard_with_alt_project_and_location(self, get_tensorboard_mock): + aiplatform.init(project=_TEST_PROJECT) + tensorboard.Tensorboard( + tensorboard_name=_TEST_NAME, + project=_TEST_ALT_PROJECT, + location=_TEST_LOCATION, + ) + get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + + def test_init_tensorboard_with_alt_location(self, get_tensorboard_mock): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_ALT_LOCATION) + tensorboard.Tensorboard(tensorboard_name=_TEST_NAME,) + get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + + def test_init_tensorboard_with_project_and_alt_location(self): + aiplatform.init(project=_TEST_PROJECT) + with pytest.raises(RuntimeError): + tensorboard.Tensorboard( + tensorboard_name=_TEST_NAME, + project=_TEST_PROJECT, + location=_TEST_ALT_LOCATION, + ) + + @patch.dict( + os.environ, {"GOOGLE_CLOUD_PROJECT": "", "GOOGLE_APPLICATION_CREDENTIALS": ""} + ) + def test_init_tensorboard_with_id_only_without_project_or_location(self): + with pytest.raises(GoogleAuthError): + tensorboard.Tensorboard( + tensorboard_name=_TEST_ID, + credentials=auth_credentials.AnonymousCredentials(), + ) + + def test_init_tensorboard_with_location_override(self, get_tensorboard_mock): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + tensorboard.Tensorboard(tensorboard_name=_TEST_ID, location=_TEST_ALT_LOCATION) + get_tensorboard_mock.assert_called_once_with(name=_TEST_ALT_NAME) + + @pytest.mark.usefixtures("get_tensorboard_mock") + def test_init_tensorboard_with_invalid_name(self): + with pytest.raises(ValueError): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + tensorboard.Tensorboard(tensorboard_name=_TEST_INVALID_NAME) + + @pytest.mark.usefixtures("get_tensorboard_mock") + def test_create_tensorboard_with_default_encryption_key( + self, create_tensorboard_mock + ): + aiplatform.init( + project=_TEST_PROJECT, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + ) + + tensorboard.Tensorboard.create(display_name=_TEST_DISPLAY_NAME,) + + expected_tensorboard = gca_tensorboard.Tensorboard( + display_name=_TEST_DISPLAY_NAME, encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + + create_tensorboard_mock.assert_called_once_with( + parent=_TEST_PARENT, + tensorboard=expected_tensorboard, + metadata=_TEST_REQUEST_METADATA, + ) + + @pytest.mark.usefixtures("get_tensorboard_mock") + def test_create_tensorboard(self, create_tensorboard_mock): + + aiplatform.init(project=_TEST_PROJECT,) + + tensorboard.Tensorboard.create( + display_name=_TEST_DISPLAY_NAME, + encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + ) + + expected_tensorboard = gca_tensorboard.Tensorboard( + display_name=_TEST_DISPLAY_NAME, encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + + create_tensorboard_mock.assert_called_once_with( + parent=_TEST_PARENT, + tensorboard=expected_tensorboard, + metadata=_TEST_REQUEST_METADATA, + ) + + @pytest.mark.usefixtures("get_tensorboard_mock") + def test_delete_tensorboard(self, delete_tensorboard_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_tensorboard = tensorboard.Tensorboard(tensorboard_name=_TEST_NAME) + + my_tensorboard.delete() + + delete_tensorboard_mock.assert_called_once_with( + name=my_tensorboard.resource_name + ) + + @pytest.mark.usefixtures("get_tensorboard_mock") + def test_update_tensorboard_display_name(self, update_tensorboard_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_tensorboard = tensorboard.Tensorboard(tensorboard_name=_TEST_NAME) + my_tensorboard.update(display_name=_TEST_DISPLAY_NAME_UPDATE) + + expected_tensorboard = gca_tensorboard.Tensorboard( + name=_TEST_NAME, display_name=_TEST_DISPLAY_NAME_UPDATE, + ) + update_tensorboard_mock.assert_called_once_with( + update_mask=field_mask_pb2.FieldMask(paths=["display_name"]), + tensorboard=expected_tensorboard, + metadata=_TEST_REQUEST_METADATA, + ) + + @pytest.mark.usefixtures("get_tensorboard_mock") + def test_update_tensorboard_encryption_spec(self, update_tensorboard_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_tensorboard = tensorboard.Tensorboard(tensorboard_name=_TEST_NAME) + my_tensorboard.update(encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME) + + expected_tensorboard = gca_tensorboard.Tensorboard( + name=_TEST_NAME, encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + update_tensorboard_mock.assert_called_once_with( + update_mask=field_mask_pb2.FieldMask(paths=["encryption_spec"]), + tensorboard=expected_tensorboard, + metadata=_TEST_REQUEST_METADATA, + )