fix: Remove the xprof dependency from Vertex Tensorboard uploader.

PiperOrigin-RevId: 741553321

Author: vertex-sdk-bot

google/cloud/aiplatform/tensorboard/uploader.py

@@ -25,6 +25,7 @@
 from typing import ContextManager, Dict, FrozenSet, Generator, Iterable, Optional, Tuple
 import uuid
 
+from google.api_core import exceptions
 from google.cloud import storage
 from google.cloud.aiplatform import base
 from google.cloud.aiplatform.compat.services import (
@@ -323,11 +324,16 @@ def _end_experiment_runs(self):
         for run_name in self._experiment_runs:
             if run_name:
                 logging.info("Ending run %s", run_name)
-                run = experiment_run_resource.ExperimentRun.get(
-                    project=self._project, location=self._location, run_name=run_name
-                )
-                if run:
-                    run.update_state(state=gca_execution.Execution.State.COMPLETE)
+                try:
+                    run = experiment_run_resource.ExperimentRun.get(
+                        project=self._project,
+                        location=self._location,
+                        run_name=run_name,
+                    )
+                    if run:
+                        run.update_state(state=gca_execution.Execution.State.COMPLETE)
+                except exceptions.MethodNotImplemented:
+                    logging.warning("Failed to end run %s", run_name)
 
     def start_uploading(self):
         """Blocks forever to continuously upload data from the logdir.

google/cloud/aiplatform/tensorboard/uploader_constants.py

@@ -13,8 +13,8 @@
 from tensorboard.plugins.image import metadata as images_metadata
 from tensorboard.plugins.scalar import metadata as scalar_metadata
 from tensorboard.plugins.text import metadata as text_metadata
-from tensorboard_plugin_profile import profile_plugin
 
+PROFILE_PLUGIN_NAME = "profile"
 ALLOWED_PLUGINS = frozenset(
     [
         scalar_metadata.PLUGIN_NAME,
@@ -24,7 +24,7 @@
         hparams_metadata.PLUGIN_NAME,
         images_metadata.PLUGIN_NAME,
         graphs_metadata.PLUGIN_NAME,
-        profile_plugin.PLUGIN_NAME,
+        PROFILE_PLUGIN_NAME,
     ]
 )