fix: Important the correct constants and use v1 for tensorboard experiments (#905)

mkovalski · web-flow · commit 48c2bf1ea2fa · 2021-12-14T08:50:51.000-05:00
* fix: Important the correct constants and use v1 for tensorboard experiments

* fix: fixing linting errors
diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tensorboard_api.py b/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tensorboard_api.py
@@ -29,11 +29,10 @@
 from google.api_core import exceptions
 from google.cloud import aiplatform
 from google.cloud import storage
+from google.cloud.aiplatform.constants import base as constants
 from google.cloud.aiplatform.utils import TensorboardClientWithOverride
 from google.cloud.aiplatform.tensorboard import uploader_utils
-from google.cloud.aiplatform.compat.types import (
-    tensorboard_experiment_v1beta1 as tensorboard_experiment,
-)
+from google.cloud.aiplatform.compat.types import tensorboard_experiment
 from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
 from google.cloud.aiplatform import training_utils
 
@@ -42,9 +41,8 @@
 
 def _get_api_client() -> TensorboardClientWithOverride:
     """Creates an Tensorboard API client."""
-    aiplatform.constants.API_BASE_PATH = (
-        training_utils.environment_variables.tensorboard_api_uri
-    )
+    constants.API_BASE_PATH = training_utils.environment_variables.tensorboard_api_uri
+
     m = re.match(
         "projects/.*/locations/(.*)/tensorboards/.*",
         training_utils.environment_variables.tensorboard_resource_name,
