xprof-nightly 2.22.3a20251208__cp311-none-manylinux2014_x86_64.whl

xprof/protobuf/trace_events_raw_pb2.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: plugin/xprof/protobuf/trace_events_raw.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'plugin/xprof/protobuf/trace_events_raw.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n,plugin/xprof/protobuf/trace_events_raw.proto\x12\x13tensorflow.profiler\"\xc0\x01\n\x07RawData\x12\x38\n\x0c\x64ma_activity\x18\x01 \x01(\x0b\x32 .tensorflow.profiler.DmaActivityH\x00\x12\x38\n\x04\x61rgs\x18\x02 \x01(\x0b\x32(.tensorflow.profiler.TraceEventArgumentsH\x00\x12\x35\n\x08tpu_data\x18\x03 \x01(\x0b\x32!.tensorflow.profiler.TpuTraceDataH\x00\x42\n\n\x08raw_data\"\xcf\x01\n\x0b\x44maActivity\x12\x19\n\x11start_time_cycles\x18\x01 \x01(\x04\x12\x17\n\x0f\x65nd_time_cycles\x18\x02 \x01(\x04\x12\x11\n\tkilobytes\x18\x04 \x01(\x04\x12\x14\n\x0cmesh_chip_id\x18\x05 \x01(\r\x12\x0f\n\x07\x63ore_id\x18\x0b \x01(\r\x12\x13\n\x0b\x64ma_address\x18\x06 \x01(\x04\x12\x11\n\tmulticast\x18\x08 \x01(\r\x12\x11\n\tsegmented\x18\t \x01(\r\x12\x11\n\ttemporary\x18\n \x01(\x04J\x04\x08\x03\x10\x04\"\xe6\x01\n\x13TraceEventArguments\x12>\n\x03\x61rg\x18\x01 \x03(\x0b\x32\x31.tensorflow.profiler.TraceEventArguments.Argument\x1a\x8e\x01\n\x08\x41rgument\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\tstr_value\x18\x02 \x01(\tH\x00\x12\x14\n\nuint_value\x18\x03 \x01(\x04H\x00\x12\x13\n\tint_value\x18\x05 \x01(\x03H\x00\x12\x16\n\x0c\x64ouble_value\x18\x04 \x01(\x01H\x00\x12\x13\n\tref_value\x18\x06 \x01(\x06H\x00\x42\x07\n\x05value\"\x1d\n\x0cTpuTraceData\x12\r\n\x05\x64ummy\x18\x01 \x01(\rB\x03\xf8\x01\x01')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'plugin.xprof.protobuf.trace_events_raw_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\370\001\001'
+  _globals['_RAWDATA']._serialized_start=70
+  _globals['_RAWDATA']._serialized_end=262
+  _globals['_DMAACTIVITY']._serialized_start=265
+  _globals['_DMAACTIVITY']._serialized_end=472
+  _globals['_TRACEEVENTARGUMENTS']._serialized_start=475
+  _globals['_TRACEEVENTARGUMENTS']._serialized_end=705
+  _globals['_TRACEEVENTARGUMENTS_ARGUMENT']._serialized_start=563
+  _globals['_TRACEEVENTARGUMENTS_ARGUMENT']._serialized_end=705
+  _globals['_TPUTRACEDATA']._serialized_start=707
+  _globals['_TPUTRACEDATA']._serialized_end=736
+# @@protoc_insertion_point(module_scope)

xprof/protobuf/trace_filter_config_pb2.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: plugin/xprof/protobuf/trace_filter_config.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'plugin/xprof/protobuf/trace_filter_config.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n/plugin/xprof/protobuf/trace_filter_config.proto\x12\x13tensorflow.profiler\"\xc0\x02\n\x10TraceEventFilter\x12\x12\n\nfield_name\x18\x05 \x01(\t\x12=\n\x05op_id\x18\x06 \x01(\x0e\x32..tensorflow.profiler.TraceEventFilter.Operator\x12\x10\n\x08negation\x18\x07 \x01(\x08\x12\x13\n\tstr_value\x18\x08 \x01(\tH\x00\x12\x15\n\x0bregex_value\x18\t \x01(\tH\x00\x12\x14\n\nuint_value\x18\n \x01(\x04H\x00\x12\x13\n\tint_value\x18\x0c \x01(\x03H\x00\x12\x16\n\x0c\x64ouble_value\x18\x0b \x01(\x01H\x00\"O\n\x08Operator\x12\t\n\x05OP_EQ\x10\x00\x12\t\n\x05OP_LT\x10\x01\x12\t\n\x05OP_GT\x10\x02\x12\t\n\x05OP_LE\x10\x03\x12\t\n\x05OP_GE\x10\x04\x12\x0c\n\x08OP_REGEX\x10\x05\x42\x07\n\x05value\"\xe3\x01\n\x11TraceFilterConfig\x12\x16\n\x0e\x64\x65vice_regexes\x18\x01 \x03(\t\x12\x18\n\x10resource_regexes\x18\x02 \x03(\t\x12\x42\n\x13trace_event_filters\x18\x03 \x03(\x0b\x32%.tensorflow.profiler.TraceEventFilter\x12\x46\n\x17trace_event_arg_filters\x18\x04 \x03(\x0b\x32%.tensorflow.profiler.TraceEventFilter\x12\x10\n\x08negation\x18\x05 \x01(\x08\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'plugin.xprof.protobuf.trace_filter_config_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_TRACEEVENTFILTER']._serialized_start=73
+  _globals['_TRACEEVENTFILTER']._serialized_end=393
+  _globals['_TRACEEVENTFILTER_OPERATOR']._serialized_start=305
+  _globals['_TRACEEVENTFILTER_OPERATOR']._serialized_end=384
+  _globals['_TRACEFILTERCONFIG']._serialized_start=396
+  _globals['_TRACEFILTERCONFIG']._serialized_end=623
+# @@protoc_insertion_point(module_scope)

xprof/server.py

@@ -0,0 +1,319 @@
+# 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.
+# ==============================================================================
+"""Utilities to start up a standalone webserver."""
+
+import argparse
+import collections
+import dataclasses
+import socket
+import sys
+from typing import Optional
+
+from cheroot import wsgi
+from etils import epath
+
+from xprof import profile_plugin_loader
+from xprof.standalone import base_plugin
+from xprof.standalone import plugin_event_multiplexer
+from xprof.convert import _pywrap_profiler_plugin
+
+DataProvider = plugin_event_multiplexer.DataProvider
+TBContext = base_plugin.TBContext
+ProfilePluginLoader = profile_plugin_loader.ProfilePluginLoader
+
+
+_DEFAULT_GRPC_PORT = 50051
+
+
+@dataclasses.dataclass(frozen=True)
+class ServerConfig:
+  """Configuration parameters for launching the XProf server.
+
+  This dataclass holds all the settings required to initialize and run the XProf
+  profiling server, including network ports, log locations, and feature flags.
+  """
+
+  logdir: Optional[str]
+  port: int
+  grpc_port: int
+  worker_service_address: str
+  hide_capture_profile_button: bool
+
+
+def make_wsgi_app(plugin):
+  """Create a WSGI application for the standalone server."""
+
+  apps = plugin.get_plugin_apps()
+
+  prefix = "/data/plugin/profile"
+
+  def application(environ, start_response):
+    path = environ["PATH_INFO"]
+    if path.startswith(prefix):
+      path = path[len(prefix) :]
+    if path != "/" and path.endswith("/"):
+      path = path[:-1]
+    handler = apps.get(path, plugin.default_handler)
+    return handler(environ, start_response)
+
+  return application
+
+
+def run_server(plugin, host, port):
+  """Starts a webserver for the standalone server."""
+
+  app = make_wsgi_app(plugin)
+
+  server = wsgi.Server((host, port), app)
+
+  try:
+    print(f"XProf at http://localhost:{port}/ (Press CTRL+C to quit)")
+    server.start()
+  except KeyboardInterrupt:
+    server.stop()
+
+
+def _get_wildcard_address(port) -> str:
+  """Returns a wildcard address for the port in question.
+
+  This will attempt to follow the best practice of calling
+  getaddrinfo() with a null host and AI_PASSIVE to request a
+  server-side socket wildcard address. If that succeeds, this
+  returns the first IPv6 address found, or if none, then returns
+  the first IPv4 address. If that fails, then this returns the
+  hardcoded address "::" if socket.has_ipv6 is True, else
+  "0.0.0.0".
+
+  Args:
+    port: The port number.
+
+  Returns:
+    The wildcard address.
+  """
+  fallback_address = "::" if socket.has_ipv6 else "0.0.0.0"
+  if hasattr(socket, "AI_PASSIVE"):
+    try:
+      addrinfos = socket.getaddrinfo(
+          None,
+          port,
+          socket.AF_UNSPEC,
+          socket.SOCK_STREAM,
+          socket.IPPROTO_TCP,
+          socket.AI_PASSIVE,
+      )
+    except socket.gaierror:
+      return fallback_address
+    addrs_by_family = collections.defaultdict(list)
+    for family, _, _, _, sockaddr in addrinfos:
+      # Format of the "sockaddr" socket address varies by address family,
+      # but [0] is always the IP address portion.
+      addrs_by_family[family].append(sockaddr[0])
+    if hasattr(socket, "AF_INET6") and addrs_by_family[socket.AF_INET6]:
+      return addrs_by_family[socket.AF_INET6][0]
+    if hasattr(socket, "AF_INET") and addrs_by_family[socket.AF_INET]:
+      return addrs_by_family[socket.AF_INET][0]
+  return fallback_address
+
+
+def _launch_server(
+    config: ServerConfig,
+):
+  """Initializes and launches the main XProf server.
+
+  This function sets up the necessary components for the XProf server based on
+  the provided configuration. It starts the gRPC worker service if distributed
+  processing is enabled, creates the TensorBoard context, loads the profile
+  plugin, and finally starts the web server to handle HTTP requests.
+
+  Args:
+    config: The ServerConfig object containing all server settings.
+  """
+  _pywrap_profiler_plugin.initialize_stubs(config.worker_service_address)
+  _pywrap_profiler_plugin.start_grpc_server(config.grpc_port)
+
+  context = TBContext(
+      config.logdir, DataProvider(config.logdir), TBContext.Flags(False)
+  )
+  context.hide_capture_profile_button = config.hide_capture_profile_button
+  loader = ProfilePluginLoader()
+  plugin = loader.load(context)
+  run_server(plugin, _get_wildcard_address(config.port), config.port)
+
+
+def get_abs_path(logdir: str) -> str:
+  """Gets the absolute path for a given log directory string.
+
+  This function correctly handles both Google Cloud Storage (GCS) paths and
+  local filesystem paths.
+
+  - GCS paths (e.g., "gs://bucket/log") are returned as is.
+  - Local filesystem paths (e.g., "~/logs", "log", ".") are made absolute.
+
+  Args:
+      logdir: The path string.
+
+  Returns:
+      The corresponding absolute path as a string.
+  """
+  if logdir.startswith("gs://"):
+    return logdir
+
+  return str(epath.Path(logdir).expanduser().resolve())
+
+
+def _create_argument_parser() -> argparse.ArgumentParser:
+  """Creates and configures the argument parser for the XProf server CLI.
+
+  This function sets up argparse to handle command-line flags for specifying
+  the log directory, server port, and other operational modes.
+
+  Returns:
+    The configured argument parser.
+  """
+  parser = argparse.ArgumentParser(
+      prog="xprof",
+      description="Launch the XProf profiling server.",
+      formatter_class=argparse.RawDescriptionHelpFormatter,
+      epilog=(
+          "Examples:\n"
+          "\txprof ~/jax/profile-logs -p 8080\n"
+          "\txprof --logdir ~/jax/profile-logs -p 8080"
+      ),
+  )
+
+  logdir_group = parser.add_mutually_exclusive_group(required=False)
+
+  logdir_group.add_argument(
+      "-l",
+      "--logdir",
+      dest="logdir_opt",
+      metavar="<logdir>",
+      type=str,
+      help="The directory where profile files will be stored.",
+  )
+
+  logdir_group.add_argument(
+      "logdir_pos",
+      nargs="?",
+      metavar="logdir",
+      type=str,
+      default=None,
+      help="Positional argument for the profile log directory.",
+  )
+
+  parser.add_argument(
+      "-p",
+      "--port",
+      metavar="<port>",
+      type=int,
+      default=8791,
+      help="The port number for the server (default: %(default)s).",
+  )
+
+  parser.add_argument(
+      "-hcpb",
+      "--hide_capture_profile_button",
+      action="store_true",
+      default=False,
+      help="Hides the 'Capture Profile' button in the UI.",
+  )
+
+  parser.add_argument(
+      "-wsa",
+      "--worker_service_address",
+      type=str,
+      default=None,
+      help=(
+          "A comma-separated list of worker service addresses (IPs or FQDNs)"
+          " with their gRPC ports, used in distributed profiling. Example:"
+          " 'worker-a.project.internal:50051,worker-b.project.internal:50051'."
+          " If not provided, it will use 0.0.0.0 with the gRPC port."
+      ),
+  )
+
+  parser.add_argument(
+      "-gp",
+      "--grpc_port",
+      type=int,
+      default=_DEFAULT_GRPC_PORT,
+      help=(
+          "The port for the gRPC server, which runs alongside the main HTTP"
+          " server for distributed profiling. This must be different from the"
+          " main server port (--port)."
+      ),
+  )
+  return parser
+
+
+def main() -> int:
+  """Parses command-line arguments and launches the XProf server.
+
+  This is the main entry point for the XProf server application. It parses
+  command-line arguments, creates a ServerConfig, and then launches the
+  server.
+
+  Returns:
+    An exit code, 0 for success and non-zero for errors.
+  """
+  parser = _create_argument_parser()
+  try:
+    args = parser.parse_args()
+  except SystemExit as e:
+    return e.code
+
+  logdir = (
+      get_abs_path(args.logdir_opt or args.logdir_pos)
+      if args.logdir_opt or args.logdir_pos
+      else None
+  )
+
+  worker_service_address = args.worker_service_address
+  if worker_service_address is None:
+    worker_service_address = f"0.0.0.0:{args.grpc_port}"
+
+  config = ServerConfig(
+      logdir=logdir,
+      port=args.port,
+      grpc_port=args.grpc_port,
+      worker_service_address=worker_service_address,
+      hide_capture_profile_button=args.hide_capture_profile_button,
+  )
+
+  print("Attempting to start XProf server:")
+  print(f"  Log Directory: {logdir}")
+  print(f"  Port: {config.port}")
+  print(f"  Worker Service Address: {config.worker_service_address}")
+  print(f"  Hide Capture Button: {config.hide_capture_profile_button}")
+
+  if logdir and not epath.Path(logdir).exists():
+    print(
+        f"Error: Log directory '{logdir}' does not exist or is not a"
+        " directory.",
+        file=sys.stderr,
+    )
+    return 1
+
+  if config.port == config.grpc_port:
+    print(
+        "Error: The main server port (--port) and the gRPC port (--grpc_port)"
+        " must be different.",
+        file=sys.stderr,
+    )
+    return 1
+
+  _launch_server(
+      config,
+  )
+  return 0

xprof/standalone/base_plugin.py

@@ -0,0 +1,52 @@
+# 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.
+# ==============================================================================
+"""Base plugin classes and context for the TensorBoard free plugin."""
+
+from typing import Any
+
+from xprof.standalone import plugin_event_multiplexer
+
+
+class TBPlugin:
+  pass
+
+
+class FrontendMetadata:
+  def __init__(self, es_module_path: str):
+    self.es_module_path = es_module_path
+
+
+class TBLoader:
+  pass
+
+
+class TBContext():
+  """https://github.com/tensorflow/tensorboard/blob/a7178f4f622a786463d23ef645e0f16f6ea7a1cb/tensorboard/plugins/base_plugin.py#L229."""
+
+  class Flags():
+    def __init__(self, master_tpu_unsecure_channel):
+      self.master_tpu_unsecure_channel = master_tpu_unsecure_channel
+
+  def __init__(
+      self,
+      logdir: str,
+      data_provider: plugin_event_multiplexer.DataProvider,
+      flags: dict[str, Any],
+      multiplexer=None,
+  ):
+    self.logdir = logdir
+    self.data_provider = data_provider
+    self.flags = flags
+    self.multiplexer = multiplexer

xprof/standalone/context.py

@@ -0,0 +1,22 @@
+# 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."""
+
+
+class RequestContext:
+  """Overload of tensorboard/context.py:RequestContext."""
+
+  def __init__(self):
+    pass

xprof/standalone/data_provider.py

@@ -0,0 +1,32 @@
+# 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."""
+
+
+class MultiplexerDataProvider:
+  """Data provider that reads data from a multiplexer."""
+
+  def __init__(self, multiplexer, logdir):
+    """Creates a new MultiplexerDataProvider.
+
+    Args:
+      multiplexer: A multiplexer object.
+      logdir: The log directory.
+    """
+    self.multiplexer = multiplexer
+    self.logdir = logdir
+
+  def list_runs(self, *args, **kwargs):  # pylint: disable=invalid-name
+    return self.multiplexer.list_runs(*args, **kwargs)

xprof/standalone/plugin_asset_util.py

@@ -0,0 +1,131 @@
+# 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.
+# ==============================================================================
+"""Utility functions for plugin assets."""
+
+import os
+import urllib.parse
+import fsspec
+
+_PLUGINS_DIR = "plugins"
+# In the future, if S3 support is needed, "s3://" can be added here.
+_PROTOCOL_PREFIXES = ("gs://",)
+
+
+def _get_protocol_from_parsed_url_scheme(scheme: str) -> str:
+  """Returns the protocol from the parsed URL scheme.
+
+  If the scheme is not recognized, returns "file" as the default protocol.
+
+  Args:
+    scheme: The parsed URL scheme.
+
+  Returns:
+    The protocol from the parsed URL scheme.
+  """
+  if not scheme:
+    return "file"
+  for prefix in _PROTOCOL_PREFIXES:
+    if prefix.startswith(scheme):
+      return scheme
+  return "file"
+
+
+def get_fs_protocol(path):
+  string_path = str(path)
+  parsed_url = urllib.parse.urlparse(string_path)
+  protocol = _get_protocol_from_parsed_url_scheme(parsed_url.scheme)
+  return fsspec.filesystem(protocol)
+
+
+def PluginDirectory(logdir, plugin_name):  # pylint: disable=invalid-name
+  return os.path.join(logdir, _PLUGINS_DIR, plugin_name)
+
+
+def walk_with_fsspec(top):
+  """Mimics os.walk using fsspec."""
+  try:
+    parsed_url = urllib.parse.urlparse(top)
+    protocol = parsed_url.scheme or "file"
+    fs = fsspec.filesystem(protocol)
+
+    # Use a queue for breadth-first traversal
+    queue = [top]
+
+    while queue:
+      current_dir = queue.pop(0)
+      dirnames = []
+      filenames = []
+
+      try:
+        items = fs.listdir(current_dir)
+      except FileNotFoundError:
+        continue  # skip if the directory is not found.
+
+      for item in items:
+        full_path = os.path.join(current_dir, item["name"])
+        if item["type"] == "directory":
+          dirnames.append(item["name"])
+          queue.append(full_path)
+        else:
+          filenames.append(item["name"])
+
+      yield (current_dir, dirnames, filenames)
+
+  except OSError as e:
+    print(f"An error occurred: {e}")
+    return
+
+
+def _reconstruct_path(base_path: str, item_path: str) -> str:
+  """Prepends a protocol to the item path if the base path has one.
+
+  This handles inconsistencies in fsspec's listdir output:
+  - GCS: returns paths without the 'gs://' protocol.
+  - Local: returns absolute paths.
+  This function reconstructs the full path by adding the protocol if needed.
+
+  Args:
+    base_path: The base path that may have a protocol (e.g., 'gs://...').
+    item_path: The path to an item, which may be missing the protocol.
+
+  Returns:
+    The item path, with the protocol prepended if necessary.
+  """
+  prefix = next((p for p in _PROTOCOL_PREFIXES if base_path.startswith(p)), "")
+  return prefix + item_path
+
+
+def iterate_directory_with_fsspec(plugin_dir):
+  """Replaces upath.UPath(plugin_dir).iterdir() with fsspec."""
+  try:
+    fs = get_fs_protocol(plugin_dir)
+    for item in fs.listdir(plugin_dir):
+      yield _reconstruct_path(plugin_dir, item["name"])
+  except FileNotFoundError:
+    # Handle cases where the directory doesn't exist.
+    print(f"Directory not found: {plugin_dir}")
+    return
+
+
+def ListAssets(logdir, plugin_name):  # pylint: disable=invalid-name
+  plugin_dir = PluginDirectory(logdir, plugin_name)
+  try:
+    # Strip trailing slashes, which listdir() includes for some filesystems.
+    return [
+        x.rstrip("/")[len(plugin_dir) + 1 :]
+        for x in iterate_directory_with_fsspec(plugin_dir)
+    ]
+  except FileNotFoundError:
+    return []