xprof-nightly 2.22.3a20251208__cp311-none-manylinux2014_x86_64.whl

xprof/standalone/plugin_event_multiplexer.py

@@ -0,0 +1,185 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A standalone version of Tensorboard's context and utils."""
+
+import collections
+import logging
+import os
+import re
+from typing import Iterator
+import urllib.parse
+
+from etils import epath
+
+from xprof.standalone import plugin_asset_util
+
+logger = logging.getLogger("tensorboard")
+
+
+def IsCloudPath(path: str) -> bool:
+  """Checks whether a given path is Cloud filesystem path."""
+  parsed_url = urllib.parse.urlparse(path)
+  return parsed_url.scheme and parsed_url.scheme not in ["file", ""]
+
+
+_ESCAPE_GLOB_CHARACTERS_REGEX = re.compile("([*?[])")
+
+
+def _EscapeGlobCharacters(path: str) -> str:
+  drive, path = os.path.splitdrive(path)
+  return "%s%s" % (drive, _ESCAPE_GLOB_CHARACTERS_REGEX.sub(r"[\1]", path))
+
+
+def ListRecursivelyViaGlobbing(
+    top: str,
+) -> Iterator[tuple[str, str]]:
+  """Recursively lists all files with the directory.
+
+  Based off of
+  tensorboard/backend/event_processing/io_wrapper.py:ListRecursivelyViaGlobbing
+
+  Args:
+    top: The top-level directory to traverse.
+
+  Yields:
+    (dir_name, file_paths) tuples.
+  """
+  fs = plugin_asset_util.get_fs_protocol(top)
+  current_glob_string = _EscapeGlobCharacters(top)
+  level = 0
+
+  while True:
+    logger.info("GlobAndListFiles: Starting to glob level %d", level)
+    glob = fs.glob(current_glob_string + "/*")
+    logger.info(
+        "GlobAndListFiles: %d files glob-ed at level %d", len(glob), level
+    )
+
+    if not glob:
+      return
+
+    pairs = collections.defaultdict(list)
+    for file_path in glob:
+      pairs[os.path.dirname(file_path)].append(file_path)
+    for dir_name, file_paths in pairs.items():
+      yield (dir_name, tuple(file_paths))
+
+    if len(pairs) == 1:
+      current_glob_string = list(pairs.keys())[0]
+
+    current_glob_string = os.path.join(current_glob_string, "*")
+    level += 1
+
+
+def ListRecursivelyViaWalking(
+    top: str,
+) -> Iterator[tuple[str, str]]:
+  for dir_path, _, filenames in plugin_asset_util.walk_with_fsspec(top):
+    yield (
+        dir_path,
+        tuple(os.path.join(dir_path, filename) for filename in filenames),
+    )
+
+
+def IsTensorFlowEventsFile(path: str):
+  if not path:
+    raise ValueError("Path must be a nonempty string")
+  return "tfevents" in os.path.basename(path)
+
+
+class Run:
+
+  def __init__(self, path):
+    self.run_name = path
+
+
+class EventMultiplexer:
+  """https://github.com/tensorflow/tensorboard/blob/a7178f4f622a786463d23ef645e0f16f6ea7a1cb/tensorboard/data/provider.py#L26."""
+
+  def __init__(self, logdir=""):
+    self._paths = {}
+    if logdir:
+      self.AddRunsFromDirectory(logdir)
+
+  def list_runs(self, *_, **__):  # pylint: disable=invalid-name,unused-argument
+    return self._paths.values()
+
+  def AddRun(self, path, name=None):
+    name = name or path
+    self._paths[name] = Run(name)
+    return self
+
+  def AddRunsFromDirectory(self, path: str, name: str = None):
+    path = str(epath.Path(path).expanduser())
+    logger.info("Starting AddRunsFromDirectory: %s", path)
+    for subdir in GetLogdirSubdirectories(path):
+      logger.info("Adding run from directory %s", subdir)
+      rpath = os.path.relpath(subdir, path)
+      subname = os.path.join(name, rpath) if name else rpath
+      self.AddRun(subdir, name=subname)
+    logger.info("Done with AddRunsFromDirectory: %s", path)
+    return self
+
+  def Reload(self):
+    pass
+
+DataProvider = EventMultiplexer
+
+
+def GetLogdirSubdirectories(
+    path: str,
+) -> Iterator[str]:
+  """Obtains all subdirectories with events files.
+
+  Based off of tensorboard/backend/event_processing/io_wrapper.py.
+
+  Args:
+    path: The top-level directory to traverse.
+
+  Returns:
+    Subdirectories containing event files.
+  """
+  fs = plugin_asset_util.get_fs_protocol(path)
+  if not fs.exists(path):
+    return ()
+
+  if not fs.isdir(path):
+    raise ValueError(
+        "GetLogdirSubdirectories: path exists and is not a directory, %s" % path
+    )
+
+  if IsCloudPath(path):
+    logger.info(
+        "GetLogdirSubdirectories: Starting to list directories via glob-ing."
+    )
+    traversal_method = ListRecursivelyViaGlobbing
+  else:
+    logger.info(
+        "GetLogdirSubdirectories: Starting to list directories via walking."
+    )
+    traversal_method = ListRecursivelyViaWalking
+
+  return (
+      subdir
+      for (subdir, files) in traversal_method(path)
+      if any(IsTensorFlowEventsFile(f) for f in files)
+  )
+
+
+class MultiplexerDataProvider:
+
+  def __init__(self, multiplexer, logdir):
+    self.multiplexer = multiplexer
+    self.logdir = logdir

xprof/standalone/tensorboard_shim.py

@@ -0,0 +1,31 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A shim to handle whether or not TensorBoard is available."""
+# pylint: disable=g-import-not-at-top,unused-import,line-too-long
+
+try:
+  from tensorboard import context
+  from tensorboard.backend.event_processing import data_provider
+  from tensorboard.backend.event_processing import plugin_asset_util
+  from tensorboard.backend.event_processing import plugin_event_multiplexer
+  from tensorboard.plugins import base_plugin
+  IS_TENSORBOARD_AVAILABLE = True
+except ImportError:
+  from xprof.standalone import base_plugin
+  from xprof.standalone import context
+  from xprof.standalone import data_provider
+  from xprof.standalone import plugin_asset_util
+  from xprof.standalone import plugin_event_multiplexer
+  IS_TENSORBOARD_AVAILABLE = False