Merge branch 'main' into ivanmkc--package-folders

docs/aiplatform_v1/services.rst

@@ -16,4 +16,5 @@ Services for Google Cloud Aiplatform v1 API
     pipeline_service
     prediction_service
     specialist_pool_service
+    tensorboard_service
     vizier_service

docs/aiplatform_v1/tensorboard_service.rst

@@ -0,0 +1,10 @@
+TensorboardService
+------------------------------------
+
+.. automodule:: google.cloud.aiplatform_v1.services.tensorboard_service
+    :members:
+    :inherited-members:
+
+.. automodule:: google.cloud.aiplatform_v1.services.tensorboard_service.pagers
+    :members:
+    :inherited-members:

google/cloud/aiplatform/compat/__init__.py

@@ -77,13 +77,14 @@
     types.specialist_pool = types.specialist_pool_v1beta1
     types.specialist_pool_service = types.specialist_pool_service_v1beta1
     types.study = types.study_v1beta1
-    types.training_pipeline = types.training_pipeline_v1beta1
+    types.tensorboard = types.tensorboard_v1beta1
     types.tensorboard_service = types.tensorboard_service_v1beta1
     types.tensorboard_data = types.tensorboard_data_v1beta1
     types.tensorboard_experiment = types.tensorboard_experiment_v1beta1
     types.tensorboard_run = types.tensorboard_run_v1beta1
     types.tensorboard_service = types.tensorboard_service_v1beta1
     types.tensorboard_time_series = types.tensorboard_time_series_v1beta1
+    types.training_pipeline = types.training_pipeline_v1beta1
 
 if DEFAULT_VERSION == V1:
 
@@ -135,6 +136,13 @@
     types.specialist_pool = types.specialist_pool_v1
     types.specialist_pool_service = types.specialist_pool_service_v1
     types.study = types.study_v1
+    types.tensorboard = types.tensorboard_v1
+    types.tensorboard_service = types.tensorboard_service_v1
+    types.tensorboard_data = types.tensorboard_data_v1
+    types.tensorboard_experiment = types.tensorboard_experiment_v1
+    types.tensorboard_run = types.tensorboard_run_v1
+    types.tensorboard_service = types.tensorboard_service_v1
+    types.tensorboard_time_series = types.tensorboard_time_series_v1
     types.training_pipeline = types.training_pipeline_v1
 
 __all__ = (

google/cloud/aiplatform/compat/services/__init__.py

@@ -67,6 +67,9 @@
 from google.cloud.aiplatform_v1.services.specialist_pool_service import (
     client as specialist_pool_service_client_v1,
 )
+from google.cloud.aiplatform_v1.services.tensorboard_service import (
+    client as tensorboard_service_client_v1,
+)
 
 __all__ = (
     # v1
@@ -78,6 +81,7 @@
     pipeline_service_client_v1,
     prediction_service_client_v1,
     specialist_pool_service_client_v1,
+    tensorboard_service_client_v1,
     # v1beta1
     dataset_service_client_v1beta1,
     endpoint_service_client_v1beta1,

google/cloud/aiplatform/compat/types/__init__.py

@@ -57,13 +57,13 @@
     specialist_pool as specialist_pool_v1beta1,
     specialist_pool_service as specialist_pool_service_v1beta1,
     study as study_v1beta1,
-    training_pipeline as training_pipeline_v1beta1,
     tensorboard as tensorboard_v1beta1,
     tensorboard_data as tensorboard_data_v1beta1,
     tensorboard_experiment as tensorboard_experiment_v1beta1,
     tensorboard_run as tensorboard_run_v1beta1,
     tensorboard_service as tensorboard_service_v1beta1,
     tensorboard_time_series as tensorboard_time_series_v1beta1,
+    training_pipeline as training_pipeline_v1beta1,
 )
 from google.cloud.aiplatform_v1.types import (
     accelerator_type as accelerator_type_v1,
@@ -107,6 +107,12 @@
     specialist_pool as specialist_pool_v1,
     specialist_pool_service as specialist_pool_service_v1,
     study as study_v1,
+    tensorboard as tensorboard_v1,
+    tensorboard_data as tensorboard_data_v1,
+    tensorboard_experiment as tensorboard_experiment_v1,
+    tensorboard_run as tensorboard_run_v1,
+    tensorboard_service as tensorboard_service_v1,
+    tensorboard_time_series as tensorboard_time_series_v1,
     training_pipeline as training_pipeline_v1,
 )
 
@@ -152,6 +158,12 @@
     prediction_service_v1,
     specialist_pool_v1,
     specialist_pool_service_v1,
+    tensorboard_v1,
+    tensorboard_data_v1,
+    tensorboard_experiment_v1,
+    tensorboard_run_v1,
+    tensorboard_service_v1,
+    tensorboard_time_series_v1,
     training_pipeline_v1,
     # v1beta1
     accelerator_type_v1beta1,
@@ -194,13 +206,11 @@
     prediction_service_v1beta1,
     specialist_pool_v1beta1,
     specialist_pool_service_v1beta1,
-    training_pipeline_v1beta1,
-    metadata_service_v1beta1,
     tensorboard_v1beta1,
-    tensorboard_service_v1beta1,
     tensorboard_data_v1beta1,
     tensorboard_experiment_v1beta1,
     tensorboard_run_v1beta1,
     tensorboard_service_v1beta1,
     tensorboard_time_series_v1beta1,
+    training_pipeline_v1beta1,
 )

google/cloud/aiplatform/jobs.py

@@ -35,12 +35,10 @@
     batch_prediction_job as gca_bp_job_compat,
     completion_stats as gca_completion_stats,
     custom_job as gca_custom_job_compat,
-    custom_job_v1beta1 as gca_custom_job_v1beta1,
     explanation as gca_explanation_compat,
     io as gca_io_compat,
     job_state as gca_job_state,
     hyperparameter_tuning_job as gca_hyperparameter_tuning_job_compat,
-    hyperparameter_tuning_job_v1beta1 as gca_hyperparameter_tuning_job_v1beta1,
     machine_resources as gca_machine_resources_compat,
     study as gca_study_compat,
 )
@@ -1388,17 +1386,11 @@ def run(
             self._gca_resource.job_spec.enable_web_access = enable_web_access
 
         if tensorboard:
-            v1beta1_gca_resource = gca_custom_job_v1beta1.CustomJob()
-            v1beta1_gca_resource._pb.MergeFromString(
-                self._gca_resource._pb.SerializeToString()
-            )
-            self._gca_resource = v1beta1_gca_resource
             self._gca_resource.job_spec.tensorboard = tensorboard
 
         _LOGGER.log_create_with_lro(self.__class__)
 
-        version = "v1beta1" if tensorboard else "v1"
-        self._gca_resource = self.api_client.select_version(version).create_custom_job(
+        self._gca_resource = self.api_client.create_custom_job(
             parent=self._parent, custom_job=self._gca_resource
         )
 
@@ -1773,21 +1765,11 @@ def run(
             self._gca_resource.trial_job_spec.enable_web_access = enable_web_access
 
         if tensorboard:
-            v1beta1_gca_resource = (
-                gca_hyperparameter_tuning_job_v1beta1.HyperparameterTuningJob()
-            )
-            v1beta1_gca_resource._pb.MergeFromString(
-                self._gca_resource._pb.SerializeToString()
-            )
-            self._gca_resource = v1beta1_gca_resource
             self._gca_resource.trial_job_spec.tensorboard = tensorboard
 
         _LOGGER.log_create_with_lro(self.__class__)
 
-        version = "v1beta1" if tensorboard else "v1"
-        self._gca_resource = self.api_client.select_version(
-            version
-        ).create_hyperparameter_tuning_job(
+        self._gca_resource = self.api_client.create_hyperparameter_tuning_job(
             parent=self._parent, hyperparameter_tuning_job=self._gca_resource
         )
 

google/cloud/aiplatform/tensorboard/tensorboard.py

@@ -18,17 +18,13 @@
 from typing import Optional, Sequence, Dict, 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 import compat
+from google.cloud.aiplatform.compat.types import tensorboard as gca_tensorboard
 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__)
 
 
@@ -156,8 +152,7 @@ def create(
         )
 
         encryption_spec = initializer.global_config.get_encryption_spec(
-            encryption_spec_key_name=encryption_spec_key_name,
-            select_version=compat.V1BETA1,
+            encryption_spec_key_name=encryption_spec_key_name
         )
 
         gapic_tensorboard = gca_tensorboard.Tensorboard(
@@ -254,7 +249,6 @@ def update(
         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")
 

google/cloud/aiplatform/tensorboard/uploader_main.py

@@ -28,8 +28,10 @@
 from tensorboard.plugins.image import metadata as images_metadata
 from tensorboard.plugins.graph import metadata as graphs_metadata
 
+from google.api_core import exceptions
 from google.cloud import storage
 from google.cloud import aiplatform
+from google.cloud.aiplatform import jobs
 from google.cloud.aiplatform.tensorboard import uploader
 from google.cloud.aiplatform.utils import TensorboardClientWithOverride
 
@@ -123,9 +125,14 @@ def main(argv):
             exitcode=0,
         )
 
+    experiment_name = FLAGS.experiment_name
+    experiment_display_name = get_experiment_display_name_with_override(
+        experiment_name, FLAGS.experiment_display_name, project_id, region
+    )
+
     tb_uploader = uploader.TensorBoardUploader(
-        experiment_name=FLAGS.experiment_name,
-        experiment_display_name=FLAGS.experiment_display_name,
+        experiment_name=experiment_name,
+        experiment_display_name=experiment_display_name,
         tensorboard_resource_name=tensorboard.name,
         blob_storage_bucket=blob_storage_bucket,
         blob_storage_folder=blob_storage_folder,
@@ -149,6 +156,19 @@ def main(argv):
     tb_uploader.start_uploading()
 
 
+def get_experiment_display_name_with_override(
+    experiment_name, experiment_display_name, project_id, region
+):
+    if experiment_name.isdecimal() and not experiment_display_name:
+        try:
+            return jobs.CustomJob.get(
+                resource_name=experiment_name, project=project_id, location=region,
+            ).display_name
+        except exceptions.NotFound:
+            return experiment_display_name
+    return experiment_display_name
+
+
 def flags_parser(args):
     # Plumbs the flags defined in this file to the main module, mostly for the
     # console script wrapper tb-gcp-uploader.

google/cloud/aiplatform/utils/__init__.py

@@ -51,6 +51,7 @@
     model_service_client_v1,
     pipeline_service_client_v1,
     prediction_service_client_v1,
+    tensorboard_service_client_v1,
 )
 
 from google.cloud.aiplatform.compat.types import (
@@ -67,6 +68,7 @@
     pipeline_service_client_v1beta1.PipelineServiceClient,
     job_service_client_v1beta1.JobServiceClient,
     metadata_service_client_v1beta1.MetadataServiceClient,
+    tensorboard_service_client_v1beta1.TensorboardServiceClient,
     # v1
     dataset_service_client_v1.DatasetServiceClient,
     endpoint_service_client_v1.EndpointServiceClient,
@@ -75,6 +77,7 @@
     prediction_service_client_v1.PredictionServiceClient,
     pipeline_service_client_v1.PipelineServiceClient,
     job_service_client_v1.JobServiceClient,
+    tensorboard_service_client_v1.TensorboardServiceClient,
 )
 
 RESOURCE_NAME_PATTERN = re.compile(
@@ -506,8 +509,9 @@ class MetadataClientWithOverride(ClientWithOverride):
 
 class TensorboardClientWithOverride(ClientWithOverride):
     _is_temporary = False
-    _default_version = compat.V1BETA1
+    _default_version = compat.DEFAULT_VERSION
     _version_map = (
+        (compat.V1, tensorboard_service_client_v1.TensorboardServiceClient),
         (compat.V1BETA1, tensorboard_service_client_v1beta1.TensorboardServiceClient),
     )
 

google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/types/automl_tables.py

@@ -55,11 +55,13 @@ class AutoMlTablesInputs(proto.Message):
             Required when optimization_objective is
             "maximize-precision-at-recall". Must be between 0 and 1,
             inclusive.
+
             This field is a member of `oneof`_ ``additional_optimization_objective_config``.
         optimization_objective_precision_value (float):
             Required when optimization_objective is
             "maximize-recall-at-precision". Must be between 0 and 1,
             inclusive.
+
             This field is a member of `oneof`_ ``additional_optimization_objective_config``.
         prediction_type (str):
             The type of prediction the Model is to

google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/types/automl_tables.py

@@ -55,11 +55,13 @@ class AutoMlTablesInputs(proto.Message):
             Required when optimization_objective is
             "maximize-precision-at-recall". Must be between 0 and 1,
             inclusive.
+
             This field is a member of `oneof`_ ``additional_optimization_objective_config``.
         optimization_objective_precision_value (float):
             Required when optimization_objective is
             "maximize-recall-at-precision". Must be between 0 and 1,
             inclusive.
+
             This field is a member of `oneof`_ ``additional_optimization_objective_config``.
         prediction_type (str):
             The type of prediction the Model is to