feat: Add support to create TensorboardExperiment (#909)

Author: sasha-gitg

google/cloud/aiplatform/__init__.py

@@ -47,7 +47,7 @@
     HyperparameterTuningJob,
 )
 from google.cloud.aiplatform.pipeline_jobs import PipelineJob
-from google.cloud.aiplatform.tensorboard import Tensorboard
+from google.cloud.aiplatform.tensorboard import Tensorboard, TensorboardExperiment
 from google.cloud.aiplatform.training_jobs import (
     CustomTrainingJob,
     CustomContainerTrainingJob,
@@ -105,8 +105,9 @@
     "Model",
     "PipelineJob",
     "TabularDataset",
+    "Tensorboard",
+    "TensorboardExperiment",
     "TextDataset",
     "TimeSeriesDataset",
     "VideoDataset",
-    "Tensorboard",
 )

google/cloud/aiplatform/tensorboard/__init__.py

@@ -15,7 +15,10 @@
 # limitations under the License.
 #
 
-from google.cloud.aiplatform.tensorboard.tensorboard_resource import Tensorboard
+from google.cloud.aiplatform.tensorboard.tensorboard_resource import (
+    Tensorboard,
+    TensorboardExperiment,
+)
 
 
-__all__ = ("Tensorboard",)
+__all__ = ("Tensorboard", "TensorboardExperiment")

google/cloud/aiplatform/tensorboard/tensorboard_resource.py

@@ -15,13 +15,16 @@
 # limitations under the License.
 #
 
-from typing import Optional, Sequence, Dict, Tuple
+from typing import Dict, List, Optional, Sequence, Tuple
 
 from google.auth import credentials as auth_credentials
 from google.protobuf import field_mask_pb2
 
 from google.cloud.aiplatform import base
 from google.cloud.aiplatform.compat.types import tensorboard as gca_tensorboard
+from google.cloud.aiplatform.compat.types import (
+    tensorboard_experiment as gca_tensorboard_experiment,
+)
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import utils
 
@@ -176,12 +179,7 @@ def create(
 
         _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,
-        )
+        return cls(tensorboard_name=created_tensorboard.name, credentials=credentials,)
 
     def update(
         self,
@@ -233,8 +231,7 @@ def update(
                 Overrides encryption_spec_key_name set in aiplatform.init.
 
         Returns:
-            tensorboard (Tensorboard):
-                The managed tensorboard resource.
+            Tensorboard: The managed tensorboard resource.
         """
         update_mask = list()
 
@@ -285,3 +282,240 @@ def update(
         _LOGGER.log_action_completed_against_resource("tensorboard", "updated", self)
 
         return self
+
+
+class TensorboardExperiment(_TensorboardServiceResource):
+    """Managed tensorboard resource for Vertex AI."""
+
+    _resource_noun = "experiments"
+    _getter_method = "get_tensorboard_experiment"
+    _list_method = "list_tensorboard_experiments"
+    _delete_method = "delete_tensorboard_experiment"
+    _parse_resource_name_method = "parse_tensorboard_experiment_path"
+    _format_resource_name_method = "tensorboard_experiment_path"
+
+    def __init__(
+        self,
+        tensorboard_experiment_name: str,
+        tensorboard_id: Optional[str] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ):
+        """Retrieves an existing tensorboard experiment given a tensorboard experiment name or ID.
+
+        Example Usage:
+
+            tb_exp = aiplatform.TensorboardExperiment(
+                tensorboard_experiment_name= "projects/123/locations/us-central1/tensorboards/456/experiments/678"
+            )
+
+            tb_exp = aiplatform.TensorboardExperiment(
+                tensorboard_experiment_name= "678"
+                tensorboard_id = "456"
+            )
+
+        Args:
+            tensorboard_experiment_name (str):
+                Required. A fully-qualified tensorboard experiment resource name or resource ID.
+                Example: "projects/123/locations/us-central1/tensorboards/456/experiments/678" or
+                "678" when tensorboard_id is passed and project and location are initialized or passed.
+            tensorboard_id (str):
+                Optional. A tensorboard resource ID.
+            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 retrieve this Tensorboard. Overrides
+                credentials set in aiplatform.init.
+        """
+
+        super().__init__(
+            project=project,
+            location=location,
+            credentials=credentials,
+            resource_name=tensorboard_experiment_name,
+        )
+        self._gca_resource = self._get_gca_resource(
+            resource_name=tensorboard_experiment_name,
+            parent_resource_name_fields={Tensorboard._resource_noun: tensorboard_id}
+            if tensorboard_id
+            else tensorboard_id,
+        )
+
+    @classmethod
+    def create(
+        cls,
+        tensorboard_experiment_id: str,
+        tensorboard_name: str,
+        display_name: Optional[str] = None,
+        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: Sequence[Tuple[str, str]] = (),
+    ) -> "TensorboardExperiment":
+        """Creates a new TensorboardExperiment.
+
+        Example Usage:
+
+            tb = aiplatform.TensorboardExperiment.create(
+                tensorboard_experiment_id='my-experiment'
+                tensorboard_id='456'
+                display_name='my display name',
+                description='my description',
+                labels={
+                    'key1': 'value1',
+                    'key2': 'value2'
+                }
+            )
+
+        Args:
+            tensorboard_experiment_id (str):
+                Required. The ID to use for the Tensorboard experiment,
+                which will become the final component of the Tensorboard
+                experiment's resource name.
+
+                This value should be 1-128 characters, and valid
+                characters are /[a-z][0-9]-/.
+
+                This corresponds to the ``tensorboard_experiment_id`` field
+                on the ``request`` instance; if ``request`` is provided, this
+                should not be set.
+            tensorboard_name (str):
+                Required. The resource name or ID of the Tensorboard to create
+                the TensorboardExperiment in. Format of resource name:
+                ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
+            display_name (str):
+                Optional. The user-defined name of the Tensorboard Experiment.
+                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 Experiment.
+            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.
+        Returns:
+            TensorboardExperiment: The TensorboardExperiment resource.
+        """
+
+        if display_name:
+            utils.validate_display_name(display_name)
+
+        if labels:
+            utils.validate_labels(labels)
+
+        api_client = cls._instantiate_client(location=location, credentials=credentials)
+
+        parent = utils.full_resource_name(
+            resource_name=tensorboard_name,
+            resource_noun=Tensorboard._resource_noun,
+            parse_resource_name_method=Tensorboard._parse_resource_name,
+            format_resource_name_method=Tensorboard._format_resource_name,
+            project=project,
+            location=location,
+        )
+
+        gapic_tensorboard_experiment = gca_tensorboard_experiment.TensorboardExperiment(
+            display_name=display_name, description=description, labels=labels,
+        )
+
+        _LOGGER.log_create_with_lro(cls)
+
+        tensorboard_experiment = api_client.create_tensorboard_experiment(
+            parent=parent,
+            tensorboard_experiment=gapic_tensorboard_experiment,
+            tensorboard_experiment_id=tensorboard_experiment_id,
+            metadata=request_metadata,
+        )
+
+        _LOGGER.log_create_complete(cls, tensorboard_experiment, "tb experiment")
+
+        return cls(
+            tensorboard_experiment_name=tensorboard_experiment.name,
+            credentials=credentials,
+        )
+
+    @classmethod
+    def list(
+        cls,
+        tensorboard_name: str,
+        filter: Optional[str] = None,
+        order_by: Optional[str] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> List["TensorboardExperiment"]:
+        """List TensorboardExperiemnts in a Tensorboard resource.
+
+        Example Usage:
+
+            aiplatform.TensorboardExperiment.list(
+                tensorboard_name='projects/my-project/locations/us-central1/tensorboards/123'
+            )
+
+        Args:
+            tensorboard_name(str):
+                Required. The resource name or resource ID of the
+                Tensorboard to list
+                TensorboardExperiments. Format, if resource name:
+                'projects/{project}/locations/{location}/tensorboards/{tensorboard}'
+            filter (str):
+                Optional. An expression for filtering the results of the request.
+                For field names both snake_case and camelCase are supported.
+            order_by (str):
+                Optional. A comma-separated list of fields to order by, sorted in
+                ascending order. Use "desc" after a field name for descending.
+                Supported fields: `display_name`, `create_time`, `update_time`
+            project (str):
+                Optional. Project to retrieve list from. If not set, project
+                set in aiplatform.init will be used.
+            location (str):
+                Optional. Location to retrieve list from. If not set, location
+                set in aiplatform.init will be used.
+            credentials (auth_credentials.Credentials):
+                Optional. Custom credentials to use to retrieve list. Overrides
+                credentials set in aiplatform.init.
+        Returns:
+            List[TensorboardExperiment] - A list of TensorboardExperiments
+        """
+
+        parent = utils.full_resource_name(
+            resource_name=tensorboard_name,
+            resource_noun=Tensorboard._resource_noun,
+            parse_resource_name_method=Tensorboard._parse_resource_name,
+            format_resource_name_method=Tensorboard._format_resource_name,
+            project=project,
+            location=location,
+        )
+
+        return super()._list(
+            filter=filter,
+            order_by=order_by,
+            project=project,
+            location=location,
+            credentials=credentials,
+            parent=parent,
+        )

tests/system/aiplatform/test_tensorboard.py

@@ -42,3 +42,25 @@ def test_create_and_get_tensorboard(self, shared_state):
         list_tb = aiplatform.Tensorboard.list()
 
         assert len(list_tb) > 0
+
+        tb_experiment = aiplatform.TensorboardExperiment.create(
+            tensorboard_experiment_id="vertex-sdk-e2e-test-experiment",
+            tensorboard_name=tb.resource_name,
+            display_name=self._make_display_name("tensorboard_experiment"),
+            description="Vertex SDK Integration test.",
+            labels={"test": "labels"},
+        )
+
+        shared_state["resources"].append(tb_experiment)
+
+        get_tb_experiment = aiplatform.TensorboardExperiment(
+            tb_experiment.resource_name
+        )
+
+        assert tb_experiment.resource_name == get_tb_experiment.resource_name
+
+        list_tb_experiment = aiplatform.TensorboardExperiment.list(
+            tensorboard_name=tb.resource_name
+        )
+
+        assert len(list_tb_experiment) > 0