fix: Remove the xprof dependency from Vertex Tensorboard uploader.

vertex-sdk-bot · copybara-github · commit 1336b85dab92 · 2025-03-28T09:41:54.000-07:00
PiperOrigin-RevId: 741553321
diff --git a/google/cloud/aiplatform/tensorboard/uploader.py b/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.
diff --git a/google/cloud/aiplatform/tensorboard/uploader_constants.py b/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,
     ]
 )
 

Original file line number	Diff line number	Diff line change
`@@ -13,8 +13,8 @@`
`13`	`13`	`from tensorboard.plugins.image import metadata as images_metadata`
`14`	`14`	`from tensorboard.plugins.scalar import metadata as scalar_metadata`
`15`	`15`	`from tensorboard.plugins.text import metadata as text_metadata`
`16`		`-from tensorboard_plugin_profile import profile_plugin`
`17`	`16`
	`17`	`+PROFILE_PLUGIN_NAME = "profile"`
`18`	`18`	`ALLOWED_PLUGINS = frozenset(`
`19`	`19`	`[`
`20`	`20`	`scalar_metadata.PLUGIN_NAME,`
`@@ -24,7 +24,7 @@`
`24`	`24`	`hparams_metadata.PLUGIN_NAME,`
`25`	`25`	`images_metadata.PLUGIN_NAME,`
`26`	`26`	`graphs_metadata.PLUGIN_NAME,`
`27`		`- profile_plugin.PLUGIN_NAME,`
	`27`	`+ PROFILE_PLUGIN_NAME,`
`28`	`28`	`]`
`29`	`29`	`)`
`30`	`30`