# Copyright 2019 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.
# ==============================================================================
"""The TensorBoard plugin for performance profiling."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
import concurrent.futures
import gzip
import json
import logging
import os
import re
import sys
import threading
import time
from typing import Any, TextIO, TypedDict

from etils import epath
import etils.epath.backend
import six
from werkzeug import wrappers

from xprof import profile_io
from xprof import version
from xprof.convert import counter_extractor
from xprof.convert import raw_to_tool_data as convert
from xprof.standalone.tensorboard_shim import base_plugin
from xprof.standalone.tensorboard_shim import plugin_asset_util
from xprof.convert import _pywrap_profiler_plugin


logger = logging.getLogger('tensorboard.plugins.profile')
logger.setLevel(logging.INFO)
if not logger.handlers:
  handler = logging.StreamHandler(sys.stderr)
  formatter = logging.Formatter(
      '%(levelname)s %(asctime)s [%(filename)s:%(lineno)d] %(message)s'
  )
  handler.setFormatter(formatter)
  logger.addHandler(handler)
  logger.propagate = False

try:
  import tensorflow.compat.v2 as tf  # pylint: disable=g-import-not-at-top # pytype: disable=import-error

  tf.enable_v2_behavior()
except ImportError:
  logger.info(
      'Disabling some remote capture features as tensorflow is not available'
  )
  tf = None

# The prefix of routes provided by this plugin.
TB_NAME = 'plugins'
PLUGIN_NAME = 'profile'

BASE_ROUTE = '/'
INDEX_JS_ROUTE = '/index.js'
INDEX_HTML_ROUTE = '/index.html'
BUNDLE_JS_ROUTE = '/bundle.js'
STYLES_CSS_ROUTE = '/styles.css'
MATERIALICONS_WOFF2_ROUTE = '/materialicons.woff2'
TRACE_VIEWER_INDEX_HTML_ROUTE = '/trace_viewer_index.html'
TRACE_VIEWER_INDEX_JS_ROUTE = '/trace_viewer_index.js'
TRACE_VIEWER_V2_JS_ROUTE = '/trace_viewer_v2.js'
TRACE_VIEWER_V2_WASM_ROUTE = '/trace_viewer_v2.wasm'
ZONE_JS_ROUTE = '/zone.js'
DATA_ROUTE = '/data'
DATA_CSV_ROUTE = '/data_csv'
VERSION_ROUTE = '/version'
RUNS_ROUTE = '/runs'
RUN_TOOLS_ROUTE = '/run_tools'
HOSTS_ROUTE = '/hosts'
HLO_MODULE_LIST_ROUTE = '/module_list'
CAPTURE_ROUTE = '/capture_profile'
LOCAL_ROUTE = '/local'
CONFIG_ROUTE = '/config'
CACHE_VERSION_FILE = 'cache_version.txt'
GENERATE_CACHE_ROUTE = '/generate_cache'

# Suffixes of "^, #, @" symbols represent different input data formats for the
# same tool.
# 1) '^': data generate from XPlane.
# 2) '#': data is in gzip format.
# 3) '@': data generate from proto, or tracetable for streaming trace viewer.
# 4) no suffix: data is in json format, ready to feed to frontend.
TOOLS = {
    'xplane': 'xplane.pb',
    'hlo_proto': 'hlo_proto.pb',
}

ALL_HOSTS = 'ALL_HOSTS'

HostMetadata = TypedDict('HostMetadata', {'hostname': str})

_EXTENSION_TO_TOOL = {extension: tool for tool, extension in TOOLS.items()}
_EXTENSION_TO_TOOL['xplane.riegeli'] = 'xplane'

_FILENAME_RE = re.compile(
    r"""
    (?:            # Start optional non-capturing group for the host.
      (.*)         #   Capture group 1: The host name.
      \.           #   A literal dot.
    )?             # End optional non-capturing group.
    (              # Start capture group 2: The tool extension.
    """
    + '|'.join(re.escape(v) for v in _EXTENSION_TO_TOOL.keys())
    + r"""
    )              # End capture group 2.
    """,
    re.VERBOSE,
)


# Tools that can be generated from xplane end with ^.
XPLANE_TOOLS = [
    'trace_viewer',  # non-streaming before TF 2.13
    'trace_viewer@',  # streaming since TF 2.14
    'overview_page',
    'input_pipeline_analyzer',
    'framework_op_stats',
    'kernel_stats',
    'memory_profile',
    'pod_viewer',
    'op_profile',
    'hlo_stats',
    'roofline_model',
    'inference_profile',
    'memory_viewer',
    'graph_viewer',
    'megascale_stats',
    'perf_counters',
    'utilization_viewer',
    'smart_suggestion',
]

XPLANE_TOOLS_SET = frozenset(XPLANE_TOOLS)
DEFAULT_CACHE_TOOLS = ('overview_page', 'trace_viewer@')

# XPlane generated tools that support all host mode.
XPLANE_TOOLS_ALL_HOSTS_SUPPORTED = frozenset([
    'input_pipeline_analyzer',
    'framework_op_stats',
    'kernel_stats',
    'overview_page',
    'pod_viewer',
    'megascale_stats',
])

# XPlane generated tools that only support all host mode.
XPLANE_TOOLS_ALL_HOSTS_ONLY = frozenset(
    ['overview_page', 'pod_viewer', 'smart_suggestion']
)

# Rate limiter constants, the GCS quota defined below
# https://cloud.google.com/storage/quotas#rate-quotas.
# currently set to 1000 request per minute.
# TODO(kcai): The assumption on the average number of subdirs is not
# always true. If this is not sufficient, we can consider a token-based
# approach that counts the number of subdirs after calling iterdir.
MAX_GCS_REQUESTS = 1000
LIMIT_WINDOW_SECONDS = 60
AVERAGE_SUBDIR_NUMBER = 10


def use_xplane(tool: str) -> bool:
  return tool in XPLANE_TOOLS


# HLO generated tools.
HLO_TOOLS = frozenset(['graph_viewer', 'memory_viewer'])


def use_hlo(tool: str) -> bool:
  return tool in HLO_TOOLS


def make_filename(host: str, tool: str) -> str:
  """Returns the name of the file containing data for the given host and tool.

  Args:
    host: Name of the host that produced the profile data, e.g., 'localhost'.
    tool: Name of the tool, e.g., 'trace_viewer'.

  Returns:
    The host name concatenated with the tool-specific extension, e.g.,
    'localhost.trace'.
  """
  filename = str(host) + '.' if host else ''
  if use_hlo(tool):
    tool = 'hlo_proto'
  elif use_xplane(tool):
    tool = 'xplane'
  return filename + TOOLS[tool]


def _parse_filename(filename: str) -> tuple[str | None, str | None]:
  """Returns the host and tool encoded in a filename in the run directory.

  Args:
    filename: Name of a file in the run directory. The name might encode a host
      and tool, e.g., 'host.tracetable', 'host.domain.op_profile.json', or just
      a tool, e.g., 'trace', 'tensorflow_stats.pb'.

  Returns:
    A tuple (host, tool) containing the names of the host and tool, e.g.,
    ('localhost', 'trace_viewer'). Either of the tuple's components can be None.
  """
  m = _FILENAME_RE.fullmatch(filename)
  if m is None:
    return filename, None
  return m.group(1), _EXTENSION_TO_TOOL[m.group(2)]


def _get_hosts(filenames: Sequence[str]) -> set[str]:
  """Parses a sequence of filenames and returns the set of hosts.

  Args:
    filenames: A sequence of filenames (just basenames, no directory).

  Returns:
    A set of host names encoded in the filenames.
  """
  hosts = set()
  for name in filenames:
    host, _ = _parse_filename(name)
    if host:
      hosts.add(host)
  return hosts


def _get_tools(filenames: Sequence[str], profile_run_dir: str) -> set[str]:
  """Parses a sequence of filenames and returns the set of tools.

  If xplane is present in the repository, add tools that can be generated by
  xplane if we don't have a file for the tool.

  Args:
    filenames: A sequence of filenames.
    profile_run_dir: The run directory of the profile.

  Returns:
    A set of tool names encoded in the filenames.
  """
  tools = set()
  found = set()
  xplane_filenames = []
  for name in filenames:
    _, tool = _parse_filename(name)
    if tool == 'xplane':
      xplane_filenames.append(os.path.join(profile_run_dir, name))
      continue
    elif tool == 'hlo_proto':
      continue
    elif tool:
      tools.add(tool)
      if tool[-1] in ('@'):
        found.add(tool[:-1])
      else:
        found.add(tool)
  # profile_run_dir might be empty, like in cloud AI use case.
  if not profile_run_dir:
    if xplane_filenames:
      for item in XPLANE_TOOLS:
        if item[:-1] not in found:
          tools.add(item)
  else:
    try:
      if xplane_filenames:
        return set(convert.xspace_to_tool_names(xplane_filenames))
    except AttributeError:
      logger.warning('XPlane converters are available after Tensorflow 2.4')
  return tools


@wrappers.Request.application
def version_route(_: wrappers.Request) -> wrappers.Response:
  return respond(version.__version__, 'text/plain')


def respond(
    body: Any,
    content_type: str,
    code: int = 200,
    content_encoding: tuple[str, str] | None = None,
    extra_headers: Mapping[str, str] | None = None,
) -> wrappers.Response:
  """Create a Werkzeug response, handling JSON serialization and CSP.

  Args:
    body: For JSON responses, a JSON-serializable object; otherwise, a raw
      `bytes` string or Unicode `str` (which will be encoded as UTF-8).
    content_type: Response content-type (`str`); use `application/json` to
      automatically serialize structures.
    code: HTTP status code (`int`).
    content_encoding: Response Content-Encoding header ('str'); e.g. 'gzip'. If
      the content type is not set, The data would be compressed and the content
      encoding would be set to gzip.
    extra_headers: Optional additional headers to add to the response.

  Returns:
    A `werkzeug.wrappers.Response` object.
  """
  if content_type == 'application/json' and isinstance(
      body, (dict, list, set, tuple)
  ):
    body = json.dumps(body, sort_keys=True)
  if not isinstance(body, bytes):
    body = body.encode('utf-8')
  csp_parts = {
      'default-src': ["'self'"],
      'script-src': [
          "'self'",
          "'unsafe-eval'",
          "'unsafe-inline'",
          'https://www.gstatic.com',
      ],
      'object-src': ["'none'"],
      'style-src': [
          "'self'",
          "'unsafe-inline'",
          'https://fonts.googleapis.com',
          'https://www.gstatic.com',
      ],
      'font-src': [
          "'self'",
          'https://fonts.googleapis.com',
          'https://fonts.gstatic.com',
          'data:',
      ],
      'connect-src': [
          "'self'",
          'data:',
          'www.gstatic.com',
      ],
      'img-src': [
          "'self'",
          'blob:',
          'data:',
      ],
      'frame-src': [
          "'self'",
          'https://ui.perfetto.dev',
      ],
      'script-src-elem': [
          "'self'",
          "'unsafe-inline'",
          # Remember to restrict on integrity when importing from jsdelivr
          # Whitelist this domain to support hlo_graph_dumper html format
          'https://cdn.jsdelivr.net/npm/',
          'https://www.gstatic.com',
      ],
  }
  csp = ';'.join((' '.join([k] + v) for (k, v) in csp_parts.items()))
  headers = [
      ('Content-Security-Policy', csp),
      ('X-Content-Type-Options', 'nosniff'),
  ]
  if content_encoding:
    headers.append(('Content-Encoding', content_encoding))
  else:
    headers.append(('Content-Encoding', 'gzip'))
    body = gzip.compress(body)

  if extra_headers is not None:
    headers.extend(extra_headers.items())

  return wrappers.Response(
      body, content_type=content_type, status=code, headers=headers
  )


def _plugin_assets(
    session_dir: str, runs: Sequence[str], plugin_name: str
) -> dict[str, Sequence[str]]:
  result = {}
  for run in runs:
    run_path = _tb_run_directory(session_dir, run)
    assets = plugin_asset_util.ListAssets(run_path, plugin_name)
    result[run] = assets
  return result


def _tb_run_directory(session_dir: str, run: str) -> str:
  """Returns the TensorBoard run directory for a TensorBoard run name.

  This helper returns the TensorBoard-level run directory (the one that would)
  contain tfevents files) for a given TensorBoard run name (aka the relative
  path from the session_dir root to this directory). For the root run '.'
  this is the bare session_dir path; for all other runs this is the
  session_dir joined with the run name.

  Args:
    session_dir: the TensorBoard log directory root path
    run: the TensorBoard run name, e.g. '.' or 'train'

  Returns:
    The TensorBoard run directory path, e.g. my/session_dir or
    my/session_dir/train.
  """
  return session_dir if run == '.' else os.path.join(session_dir, run)


def hosts_from_xplane_filenames(
    filenames: Sequence[str], tool: str
) -> Sequence[str]:
  """Converts a sequence of filenames to a list of host names given a tool.

  Args:
    filenames: A sequence of filenames.
    tool: A string representing the profiling tool.

  Returns:
    A list of hostnames.
  """
  hosts = _get_hosts(filenames)
  if len(hosts) > 1:
    if tool in XPLANE_TOOLS_ALL_HOSTS_ONLY:
      hosts = [ALL_HOSTS]
    elif tool in XPLANE_TOOLS_ALL_HOSTS_SUPPORTED:
      hosts.add(ALL_HOSTS)
  return sorted(hosts)


def _get_bool_arg(
    args: Mapping[str, Any], arg_name: str, default: bool
) -> bool:
  """Gets a boolean argument from a request.

  Args:
    args: The werkzeug request arguments.
    arg_name: The name of the argument.
    default: The default value if the argument is not present.

  Returns:
    The boolean value of the argument.
  """
  arg_str = args.get(arg_name)
  if arg_str is None:
    return default
  return arg_str.lower() == 'true'


def _generate_csv_filename(request: wrappers.Request) -> str:
  """Generates a sanitized filename for the CSV export."""
  tool = request.args.get('tag', 'data')
  run = request.args.get('run', '')
  host = request.args.get('host', '')

  safe_tool = re.sub(r'[^a-zA-Z0-9_\-]', '_', tool)
  safe_run = re.sub(r'[^a-zA-Z0-9_\-]', '_', run)

  host_suffix = host.split('-')[-1] if host else ''
  safe_host_suffix = re.sub(r'[^a-zA-Z0-9_\-]', '_', host_suffix)

  filename_parts = (p for p in [safe_tool, safe_run, safe_host_suffix] if p)
  return '_'.join(filename_parts) + '.csv'


class ToolsCache:
  """A cache for tool lists based on file content hashes or mtimes.

  Attributes:
    CACHE_FILE_NAME: The name of the cache file.
    CACHE_VERSION: The version of the cache format.
  """

  CACHE_FILE_NAME = '.cached_tools.json'
  CACHE_VERSION = 1

  def __init__(
      self, profile_run_dir: epath.Path, fs: profile_io.ProfileFileSystem
  ):
    """Initializes the ToolsCache.

    Args:
      profile_run_dir: The directory containing the profile run data.
      fs: The file system object to use for file operations.
    """
    self._profile_run_dir = profile_run_dir
    self._cache_file = self._profile_run_dir / self.CACHE_FILE_NAME
    self._fs = fs
    logger.info('ToolsCache initialized for %s', self._cache_file)

  def load(self) -> list[str] | None:
    """Loads the cached list of tools if the cache is valid.

    The cache is valid if the cache file exists, the version matches, and
    the file states (hashes/mtimes) of the XPlane files have not changed.

    Returns:
      A list of tool names if the cache is valid, otherwise None.
    """
    cached_data = self._fs.read_json(str(self._cache_file))
    if cached_data is None:
      return None

    if cached_data.get('version') != self.CACHE_VERSION:
      logger.info(
          'ToolsCache invalid: version mismatch, expected %s, got %s.'
          ' Invalidating %s',
          self.CACHE_VERSION,
          cached_data.get('version'),
          self._cache_file,
      )
      self.invalidate()
      return None

    current_files = self._fs.get_xplane_file_states(str(self._profile_run_dir))
    if current_files is None:
      logger.info(
          'ToolsCache invalid: could not determine current file states.'
          ' Invalidating %s',
          self._cache_file,
      )
      self.invalidate()
      return None

    if cached_data.get('files') != current_files:
      logger.info(
          'ToolsCache invalid: file states differ. Invalidating %s',
          self._cache_file,
      )
      self.invalidate()
      return None

    logger.info('ToolsCache hit: %s', self._cache_file)
    return cached_data.get('tools')

  def save(self, tools: Sequence[str]) -> None:
    """Saves the list of tools and the current file states to the cache file.

    Args:
      tools: The list of tool names to cache.
    """
    current_files_for_cache = self._fs.get_xplane_file_states(
        str(self._profile_run_dir)
    )
    if current_files_for_cache is None:
      logger.warning(
          'ToolsCache not saved: could not get file states %s', self._cache_file
      )
      return

    new_cache_data = {
        'version': self.CACHE_VERSION,
        'files': current_files_for_cache,
        'tools': tools,
    }
    self._fs.write_json(str(self._cache_file), new_cache_data)

  def invalidate(self) -> None:
    """Deletes the cache file, forcing regeneration on the next load."""
    self._fs.delete_file(str(self._cache_file))


class _TfProfiler:
  """A helper class to encapsulate all TensorFlow-dependent profiler logic."""

  def __init__(self, tf_module):
    if not tf_module:
      raise ImportError('TensorFlow module is not available.')
    self.tf = tf_module

  def _get_worker_list(self, cluster_resolver) -> str:
    """Parses TPU workers list from the cluster resolver."""
    cluster_spec = cluster_resolver.cluster_spec()
    task_indices = cluster_spec.task_indices('worker')
    worker_list = [
        cluster_spec.task_address('worker', i).replace(':8470', ':8466')
        for i in task_indices
    ]
    return ','.join(worker_list)

  def resolve_tpu_name(
      self, tpu_name: str, worker_list: str
  ) -> tuple[str, str, str]:
    """Resolves a TPU name to its master IP, service address, and worker list.

    Args:
      tpu_name: The name of the TPU to resolve.
      worker_list: A comma-separated list of worker addresses.

    Returns:
      A tuple containing (service_addr, worker_list, master_ip).
    """
    try:
      resolver = self.tf.distribute.cluster_resolver.TPUClusterResolver(
          tpu_name
      )
      master_grpc_addr = resolver.get_master()
    except RuntimeError as err:
      # Propagate error to be handled by the caller.
      raise RuntimeError(
          f'Error initializing TPUClusterResolver: {err}'
      ) from err
    except (ValueError, TypeError) as e:
      # Handle cases where the TPU name is invalid.
      raise ValueError(f'No TPU found with the name: {tpu_name}') from e

    if not worker_list:
      worker_list = self._get_worker_list(resolver)

    # TPU cluster resolver always returns port 8470. Replace it with 8466
    # on which profiler service is running.
    master_ip = master_grpc_addr.replace('grpc://', '').replace(':8470', '')
    service_addr = f'{master_ip}:8466'
    return service_addr, worker_list, master_ip


def _logging_wrapper(
    app: Callable[
        [
            Mapping[str, Any],
            Callable[[str, Sequence[tuple[str, str]], Any | None], Any],
        ],
        Iterable[bytes],
    ],
    clock: Callable[[], float] = time.time,
    writer: TextIO = sys.stderr,
) -> Callable[
    [
        Mapping[str, Any],
        Callable[[str, Sequence[tuple[str, str]], Any | None], Any],
    ],
    Iterable[bytes],
]:
  """Wraps a WSGI application to log request timing in GKE-friendly JSON format.

  Args:
    app: The WSGI application to wrap.
    clock: A function that returns the current time. Defaults to `time.time`.
    writer: A file-like object to write logs to. Defaults to `sys.stderr`.

  Returns:
    A new WSGI application that logs request timing.
  """

  def wrapper(
      environ: Mapping[str, Any],
      start_response: Callable[
          [str, Sequence[tuple[str, str]], Any | None], Any
      ],
  ) -> Iterable[bytes]:
    start_time = clock()
    request_method = environ.get('REQUEST_METHOD', '')
    # Reconstruct the URL as best as possible from WSGI environ.
    url_path = environ.get('SCRIPT_NAME', '') + environ.get('PATH_INFO', '')
    query_string = environ.get('QUERY_STRING')
    request_url = f'{url_path}?{query_string}' if query_string else url_path

    # container for status code captured from start_response
    status_info = {'code': 200}  # Default to 200 if not captured (unlikely)

    def custom_start_response(
        status: str,
        headers: Sequence[tuple[str, str]],
        exc_info: Any | None = None,
    ) -> Any:
      try:
        status_code_str, *_ = status.split(' ', 1)
        status_info['code'] = int(status_code_str)
      except (ValueError, IndexError):
        logger.warning(
            'Failed to parse status code from %r', status, exc_info=True
        )
      return start_response(status, headers, exc_info)

    response = app(environ, custom_start_response)
    try:
      yield from response
    finally:
      duration = clock() - start_time
      log_entry = {
          'httpRequest': {
              'requestMethod': request_method,
              'requestUrl': request_url,
              'status': status_info['code'],
              'latency': f'{duration:.9f}s',
          }
      }
      # Write raw JSON to stderr for GKE/Cloud Logging to pick up as structured
      # payload. We use sys.stderr instead of the logging module because the
      # latter would add text prefixes (timestamp, level, etc.) that break the
      # JSON parsing required for the "httpRequest" field to be recognized by
      # GKE.
      writer.write(json.dumps(log_entry) + '\n')
      writer.flush()
      if hasattr(response, 'close'):
        response.close()

  return wrapper


class ProfilePlugin(base_plugin.TBPlugin):
  """Profile Plugin for TensorBoard."""

  plugin_name = PLUGIN_NAME

  def __init__(
      self,
      context,
      *,
      epath_module: Any = epath,
      xspace_to_tool_data_fn: Callable[
          [Sequence[epath.Path], str, dict[str, Any]],
          tuple[bytes | str | None, str],
      ] = convert.xspace_to_tool_data,
      version_module: Any = version,
      cache_generation_executor: concurrent.futures.Executor | None = None,
  ):
    """Constructs a profiler plugin for TensorBoard.

    This plugin adds handlers for performance-related frontends.
    Args:
      context: A base_plugin.TBContext instance.
      epath_module: The epath module to use, can be injected for testing.
      xspace_to_tool_data_fn: Function to convert xspace to tool data.
      version_module: The version module to use, can be injected for testing.
      cache_generation_executor: A `concurrent.futures.Executor` instance for
        async cache generation. If None, a default executor is created.
    """
    self.logdir = context.logdir
    self.data_provider = context.data_provider
    self.master_tpu_unsecure_channel = context.flags.master_tpu_unsecure_channel
    self.hide_capture_profile_button = getattr(
        context, 'hide_capture_profile_button', False
    )
    self.enable_tab_name_label = getattr(
        context, 'enable_tab_name_label', False
    )
    self.src_prefix = getattr(context, 'src_prefix', '')
    self._epath = epath_module
    self._xspace_to_tool_data = xspace_to_tool_data_fn
    self._version = version_module

    # Whether the plugin is active. This is an expensive computation, so we
    # compute this asynchronously and cache positive results indefinitely.
    self._is_active = False
    # Lock to ensure at most one thread computes _is_active at a time.
    self._is_active_lock = threading.Lock()
    # Lock to protect access to _run_to_profile_run_dir.
    self._run_dir_cache_lock = threading.Lock()
    # Cache to map profile run name to corresponding tensorboard dir name
    self._run_to_profile_run_dir = {}
    self._tf_profiler = _TfProfiler(tf) if tf else None
    # Limit to 1 worker to prevent potential Out-of-Memory (OOM) errors.
    # Cache generation, especially for tools like trace viewer, can be
    # memory-intensive when processing large XPlane files. Running multiple
    # cache generation tasks in parallel increases the risk of excessive
    # memory consumption.
    self._cache_generation_pool = cache_generation_executor or (
        concurrent.futures.ThreadPoolExecutor(
            max_workers=1, thread_name_prefix='XprofCacheGen'
        )
    )

  def _get_xplane_basenames(self, path_str: str) -> Sequence[str]:
    """Returns a list of .xplane.pb base filenames in the given path."""
    fs = profile_io.get_file_system(path_str, self._epath)
    return fs.get_xplane_basenames(path_str)

  def _dir_has_xplane_files(self, path_str: str) -> bool:
    """Checks if the directory contains any .xplane.pb files."""
    fs = profile_io.get_file_system(path_str, self._epath)
    return fs.dir_has_xplane_files(path_str)

  def _get_all_basenames(self, path_str: str) -> Sequence[str]:
    """Returns a list of all base filenames in the given path."""
    fs = profile_io.get_file_system(path_str, self._epath)
    return fs.get_all_basenames(path_str)

  def is_active(self) -> bool:
    """Whether this plugin is active and has any profile data to show.

    Returns:
      Whether any run has profile data.
    """
    if not self._is_active:
      self._is_active = any(self.generate_runs())
    return self._is_active

  def _does_tool_support_multi_hosts_processing(self, tool: str) -> bool:
    """Returns true if the tool supports multi-hosts processing."""
    return tool == 'trace_viewer@' or tool == 'trace_viewer'

  def get_plugin_apps(
      self,
  ) -> dict[str, Callable[[wrappers.Request], wrappers.Response]]:
    return {
        route: _logging_wrapper(app)
        for route, app in {
            BASE_ROUTE: self.default_handler,
            INDEX_JS_ROUTE: self.static_file_route,
            INDEX_HTML_ROUTE: self.static_file_route,
            BUNDLE_JS_ROUTE: self.static_file_route,
            STYLES_CSS_ROUTE: self.static_file_route,
            MATERIALICONS_WOFF2_ROUTE: self.static_file_route,
            TRACE_VIEWER_INDEX_HTML_ROUTE: self.static_file_route,
            TRACE_VIEWER_INDEX_JS_ROUTE: self.static_file_route,
            TRACE_VIEWER_V2_JS_ROUTE: self.static_file_route,
            TRACE_VIEWER_V2_WASM_ROUTE: self.static_file_route,
            ZONE_JS_ROUTE: self.static_file_route,
            RUNS_ROUTE: self.runs_route,
            RUN_TOOLS_ROUTE: self.run_tools_route,
            HOSTS_ROUTE: self.hosts_route,
            DATA_ROUTE: self.data_route,
            DATA_CSV_ROUTE: self.data_csv_route,
            VERSION_ROUTE: version_route,
            HLO_MODULE_LIST_ROUTE: self.hlo_module_list_route,
            CAPTURE_ROUTE: self.capture_route,
            LOCAL_ROUTE: self.default_handler,
            CONFIG_ROUTE: self.config_route,
            GENERATE_CACHE_ROUTE: self.generate_cache_route,
        }.items()
    }  # pytype: disable=bad-return-type

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def default_handler(self, _: wrappers.Request) -> wrappers.Response:
    contents = self._read_static_file_impl('index.html')
    return respond(contents, 'text/html')

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def config_route(self, _: wrappers.Request) -> wrappers.Response:
    # pytype: enable=wrong-arg-types
    """Returns UI configuration details."""
    config_data = {
        'hideCaptureProfileButton': self.hide_capture_profile_button,
        'srcPathPrefix': self.src_prefix,
        'enableTabNameLabel': self.enable_tab_name_label,
    }
    logger.info('config_route: %s', config_data)
    return respond(config_data, 'application/json')

  def frontend_metadata(self):
    return base_plugin.FrontendMetadata(es_module_path='/index.js')

  def _read_static_file_impl(self, filename: str) -> bytes:
    """Reads contents from a filename.

    Args:
      filename (str): Name of the file.

    Returns:
      Contents of the file.
    Raises:
      IOError: File could not be read or found.
    """
    filepath = os.path.join(os.path.dirname(__file__), 'static', filename)

    try:
      with open(filepath, 'rb') as infile:
        contents = infile.read()
    except IOError as io_error:
      raise io_error
    return contents

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def static_file_route(self, request: wrappers.Request) -> wrappers.Response:
    """Handles static files."""
    # pytype: enable=wrong-arg-types
    filename = os.path.basename(request.path)
    extension = os.path.splitext(filename)[1]
    if extension == '.html':
      mimetype = 'text/html'
    elif extension == '.css':
      mimetype = 'text/css'
    elif extension == '.js':
      mimetype = 'application/javascript'
    elif extension == '.wasm':
      mimetype = 'application/wasm'
    else:
      mimetype = 'application/octet-stream'
    try:
      contents = self._read_static_file_impl(filename)
    except IOError:
      return respond('Fail to read the files.', 'text/plain', code=404)
    return respond(contents, mimetype)

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def runs_route(self, request: wrappers.Request) -> wrappers.Response:
    # pytype: enable=wrong-arg-types
    runs = self.runs_imp(request)
    return respond(runs, 'application/json')

  def _session_dir_by_run_name_from_request(
      self, request: wrappers.Request | None = None
  ) -> Mapping[str, str] | None:
    """Returns a map of run names to session directories from the request.

    Args:
      request: Optional; werkzeug request used for grabbing session_path and
        run_path arguments.
    """
    session_path_arg = request.args.get('session_path') if request else None
    run_path_arg = (
        request.args.get('run_path')
        if request and not session_path_arg
        else None
    )
    session_dir_by_run_name = None
    if session_path_arg:
      session_dir_by_run_name = {}
      if self._dir_has_xplane_files(session_path_arg):
        run_name = self._epath.Path(session_path_arg).name
        session_dir_by_run_name[run_name] = session_path_arg
    elif run_path_arg:
      fs = profile_io.get_file_system(run_path_arg, self._epath)
      session_dir_by_run_name = fs.get_session_paths(run_path_arg)
    return session_dir_by_run_name

  def _run_dir(
      self, run: str, request: wrappers.Request | None = None
  ) -> str | None:
    """Helper that maps a frontend run name to a profile "run" directory.

    The frontend run name consists of the TensorBoard run name (aka the relative
    path from the logdir root to the directory containing the data) path-joined
    to the Profile plugin's "run" concept (which is a subdirectory of the
    plugins/profile directory representing an individual run of the tool), with
    the special case that TensorBoard run is the logdir root (which is the run
    named '.') then only the Profile plugin "run" name is used, for backwards
    compatibility.

    Args:
      run: the frontend run name, as described above, e.g. train/run1.
      request: Optional; werkzeug request used for grabbing session_path and
        run_path arguments.

    Returns:
      The resolved directory path, e.g. /logdir/train/plugins/profile/run1.

    Raises:
      ValueError: If the run is not found in the run map.
      RuntimeError: If the run directory is not found.
    """
    session_dir_by_run_name = self._session_dir_by_run_name_from_request(
        request
    )
    if session_dir_by_run_name is not None:
      if run in session_dir_by_run_name:
        return session_dir_by_run_name[run]
      else:
        raise ValueError(
            f'Run {run} not found in run map: {session_dir_by_run_name}'
        )

    with self._run_dir_cache_lock:
      if run in self._run_to_profile_run_dir:
        return self._run_to_profile_run_dir[run]

    if not self.logdir:
      raise RuntimeError(
          'No matching run directory for run %s. Logdir is empty.' % run
      )
    tb_run_name, profile_run_name = os.path.split(run.rstrip(os.sep))
    if not tb_run_name:
      tb_run_name = '.'
    tb_run_directory = _tb_run_directory(self.logdir, tb_run_name)
    if not self._epath.Path(tb_run_directory).is_dir():
      raise RuntimeError('No matching run directory for run %s' % run)
    plugin_directory = plugin_asset_util.PluginDirectory(
        tb_run_directory, PLUGIN_NAME
    )
    return os.path.join(plugin_directory, profile_run_name)

  def runs_imp(self, request: wrappers.Request | None = None) -> list[str]:
    """Returns a list all runs for the profile plugin.

    Args:
      request: Optional; werkzeug request used for grabbing ctx and experiment
        id for other host implementations
    """
    session_dir_by_run_name = self._session_dir_by_run_name_from_request(
        request
    )
    if session_dir_by_run_name is not None:
      runs = session_dir_by_run_name.keys()
    else:
      runs = self.generate_runs()
    return sorted(runs, reverse=True)

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def run_tools_route(self, request: wrappers.Request) -> wrappers.Response:
    # pytype: enable=wrong-arg-types
    run = request.args.get('run')
    run_tools = self.run_tools_imp(run, request)
    return respond(run_tools, 'application/json')

  def run_tools_imp(
      self, run, request: wrappers.Request | None = None
  ) -> list[str]:
    """Returns a list of tools given a single run.

    Args:
      run: the frontend run name, item is list returned by runs_imp
      request: Optional; werkzeug request used for grabbing ctx and experiment
        id for other host implementations
    """
    run_dir = self._run_dir(run, request)
    return list(self.generate_tools_of_run(run, run_dir))

  def _run_host_impl(
      self, run: str, run_dir: str, tool: str
  ) -> list[HostMetadata]:
    if not run_dir:
      logger.warning('Cannot find asset directory for: %s', run)
      return []
    return [
        {'hostname': host}
        for host in hosts_from_xplane_filenames(
            self._get_xplane_basenames(run_dir), tool
        )
    ]

  def host_impl(
      self, run: str, tool: str, request: wrappers.Request | None = None
  ) -> list[HostMetadata]:
    """Returns hosts metadata for the run and tool in log directory.

    In the plugin log directory, each directory contains profile data for a
    single run (identified by the directory name), and files in the run
    directory contains data for different tools and hosts. The file that
    contains profile for a specific tool "x" will have extension TOOLS["x"].

    Example:
      log/
        run1/
          plugins/
            profile/
              host1.trace
              host2.trace
              module1.hlo_proto.pb
              module2.hlo_proto.pb
        run2/
          plugins/
            profile/
              host1.trace
              host2.trace

    Args:
      run: the frontend run name, e.g., 'run1' or 'run2' for the example above.
      tool: the requested tool, e.g., 'trace_viewer' for the example above.
      request: Optional; werkzeug request used for grabbing ctx and experiment
        id for other host implementations

    Returns:
      A list of host names, e.g.:
        host_impl(run1, trace_viewer) --> [{"hostname": "host1"}, {"hostname":
        "host2"}]
        host_impl(run1, memory_viewer) --> [{"hostname": "module1"},
        {"hostname":
        "module2"}]
    """
    run_dir = self._run_dir(run, request)
    return self._run_host_impl(run, run_dir, tool)

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def hosts_route(self, request: wrappers.Request) -> wrappers.Response:
    # pytype: enable=wrong-arg-types
    run = request.args.get('run')
    tool = request.args.get('tag')
    hosts = self.host_impl(run, tool, request)
    return respond(hosts, 'application/json')

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def hlo_module_list_route(
      self, request: wrappers.Request
  ) -> wrappers.Response:
    module_names_str = self.hlo_module_list_impl(request)
    return respond(module_names_str, 'text/plain')

  def _get_valid_hosts(
      self, run_dir: str, run: str, tool: str, hosts_param: str, host: str
  ) -> tuple[Sequence[str], Sequence[epath.Path]]:
    """Retrieves and validates the hosts and asset paths for a run and tool.

    Args:
      run_dir: The run directory.
      run: The frontend run name.
      tool: The requested tool.
      hosts_param: Comma-separated list of selected hosts.
      host: The single host parameter.

    Returns:
      A tuple containing (selected_hosts, asset_paths).

    Raises:
      FileNotFoundError: If a required xplane file for the specified host(s)
        is not found.
      IOError: If there is an error reading asset directories.
    """
    asset_paths = []
    selected_hosts = []
    all_xplane_files = {}  # Map host to path

    # Find all available xplane files for the run and map them by host.
    try:
      xplane_basenames = self._get_xplane_basenames(run_dir)
      for basename in xplane_basenames:
        host_name, _ = _parse_filename(basename)
        if host_name:
          all_xplane_files[host_name] = self._epath.Path(
              os.path.join(run_dir, basename)
          )
    except OSError as e:
      logger.warning('Cannot read asset directory: %s, OpError %r', run_dir, e)
      raise IOError(
          'Cannot read asset directory: %s, OpError %r' % (run_dir, e)
      ) from e
    if not all_xplane_files:
      logger.warning('no xplane files found for run: %s, tool: %s', run, tool)
      raise FileNotFoundError(
          'No xplane file found for run: %s, tool: %s' % (run, tool)
      )

    if hosts_param and self._does_tool_support_multi_hosts_processing(tool):
      selected_hosts = hosts_param.split(',')
      for selected_host in selected_hosts:
        if selected_host in all_xplane_files:
          asset_paths.append(all_xplane_files[selected_host])
        else:
          raise FileNotFoundError(
              'No xplane file found for host: %s in run: %s'
              % (selected_host, run)
          )
    elif host == ALL_HOSTS:
      asset_paths = list(all_xplane_files.values())
      selected_hosts = list(all_xplane_files.keys())
    elif host and host in all_xplane_files:
      selected_hosts = [host]
      asset_paths = [all_xplane_files[host]]
    elif host:
      logger.warning('No xplane file found for host: %s in run: %s', host, run)
      if host not in XPLANE_TOOLS_ALL_HOSTS_ONLY:
        raise FileNotFoundError(
            'No xplane file found for host: %s in run: %s' % (host, run)
        )
    # for request that does not specify host or hosts param, use all hosts.
    # would also be no-op for tools that is host-agnostic.
    elif not host and not hosts_param:
      selected_hosts = list(all_xplane_files.keys())
      asset_paths = list(all_xplane_files.values())

    if not asset_paths:
      logger.warning(
          'No matching asset paths found for run %s, tool %s, host(s) %s / %s',
          run,
          tool,
          hosts_param,
          host,
      )
      if not host and tool not in XPLANE_TOOLS_ALL_HOSTS_ONLY:
        raise FileNotFoundError(
            'Host must be specified for tool %s in run %s' % (tool, run)
        )

    return selected_hosts, asset_paths

  def _write_cache_version_file(self, run_dir: str) -> None:
    """Writes the current version to the cache version file."""
    try:
      with self._epath.Path(os.path.join(run_dir, CACHE_VERSION_FILE)).open(
          'w'
      ) as f:
        f.write(self._version.__version__)
    except OSError:
      logger.warning(
          'Cannot write cache version file to %s', run_dir, exc_info=True
      )

  def data_impl(
      self, request: wrappers.Request
  ) -> tuple[bytes | str | None, str, str | None]:
    """Retrieves and processes the tool data for a run and a host.

    Args:
      request: XMLHttpRequest

    Returns:
      A string that can be served to the frontend tool or None if tool,
        run or host is invalid.

    Raises:
      FileNotFoundError: If a required xplane file for the specified host(s)
        is not found.
      IOError: If there is an error reading asset directories.
      AttributeError: If there is an error during xplane to tool data conversion
      ValueError: If xplane conversion fails due to invalid data.
    """
    run = request.args.get('run')
    tool = request.args.get('tag')

    if tool == 'perf_counters' and request.args.get('names_only') == '1':
      device_type = request.args.get('device_type')
      if not device_type:
        raise ValueError(
            'device_type is required for perf_counters with names_only'
        )
      try:
        names = counter_extractor.get_all_counters(device_type)
      except FileNotFoundError:
        logger.warning(
            'Failed to get counter names for device type: %s', device_type
        )
        names = []
      return json.dumps(names), 'application/json', None

    hosts_param = request.args.get('hosts')
    host = request.args.get('host')
    module_name = request.args.get('module_name')
    program_id = request.args.get('program_id')
    tqx = request.args.get('tqx')
    perfetto = _get_bool_arg(request.args, 'perfetto', False)
    use_saved_result = _get_bool_arg(request.args, 'use_saved_result', True)
    full_dma = _get_bool_arg(request.args, 'full_dma', False)
    enable_legacy_dcn = _get_bool_arg(request.args, 'enable_legacy_dcn', False)
    run_dir = self._run_dir(run, request)

    # Check if the cache file exists and if the cache file version is less
    # than the current plugin version, clear the cache.
    try:
      with self._epath.Path(os.path.join(run_dir, CACHE_VERSION_FILE)).open(
          'r'
      ) as f:
        cache_version = f.read().strip()
        if cache_version < self._version.__version__:
          use_saved_result = False
    except FileNotFoundError:
      logger.info('Cache version file not found, invalidating cache.')
      use_saved_result = False
    except OSError:
      logger.warning('Cannot read cache version file', exc_info=True)
      use_saved_result = False

    graph_viewer_options = self._get_graph_viewer_options(request)
    # Host param is used by HLO tools to identify the module.
    params = {
        'graph_viewer_options': graph_viewer_options,
        'tqx': tqx,
        'perfetto': perfetto,
        'host': host,
        'module_name': module_name,
        'program_id': program_id,
        'use_saved_result': use_saved_result,
    }
    if request.args.get('group_by'):
      params['group_by'] = request.args.get('group_by')
    if request.args.get('refresh_suggestion'):
      params['refresh_suggestion'] = request.args.get('refresh_suggestion')
    content_type = 'application/json'

    if tool not in TOOLS and not use_xplane(tool):
      return None, content_type, None
    if tool == 'memory_viewer' and request.args.get(
        'view_memory_allocation_timeline'
    ):
      params['view_memory_allocation_timeline'] = True

    params['memory_space'] = request.args.get('memory_space', '0')

    if tool in ('trace_viewer', 'trace_viewer@'):
      options = {}
      options['resolution'] = request.args.get('resolution', 8000)
      options['full_dma'] = full_dma
      options['enable_legacy_dcn'] = enable_legacy_dcn
      if request.args.get('start_time_ms') is not None:
        options['start_time_ms'] = request.args.get('start_time_ms')
      if request.args.get('end_time_ms') is not None:
        options['end_time_ms'] = request.args.get('end_time_ms')
      event_name = request.args.get('event_name')
      format_arg = request.args.get('format')
      if event_name is not None:
        options['event_name'] = event_name
        # For event selection (fetching event details), do not use the
        # compressed protobuf path; use the regular JSON path instead.
        options['format'] = 'json'
      elif format_arg is not None:
        options['format'] = format_arg
      if request.args.get('duration_ms') is not None:
        options['duration_ms'] = request.args.get('duration_ms')
      if request.args.get('unique_id') is not None:
        options['unique_id'] = request.args.get('unique_id')
      if request.args.get('search_prefix') is not None:
        options['search_prefix'] = request.args.get('search_prefix')
      if request.args.get('search_metadata') is not None:
        # Retrigger presubmits (second attempt).
        options['search_metadata'] = _get_bool_arg(
            request.args, 'search_metadata', False
        )
      params['trace_viewer_options'] = options

    _, content_encoding = None, None
    if use_xplane(tool):
      selected_hosts, asset_paths = self._get_valid_hosts(
          run_dir, run, tool, hosts_param, host
      )
      if not asset_paths:
        return None, content_type, None

      params['hosts'] = selected_hosts
      try:
        data, content_type = self._xspace_to_tool_data(
            asset_paths,
            tool,
            params,
        )
      except AttributeError as e:
        logger.warning('Error generating analysis results due to %r', e)
        raise AttributeError(
            'Error generating analysis results due to %r' % e
        ) from e
      except ValueError as e:
        logger.warning('XPlane convert to tool data failed as %r', e)
        raise
      except FileNotFoundError as e:
        logger.warning('XPlane convert to tool data failed as %r', e)
        raise

      # Write cache version file if use_saved_result is False.
      if not use_saved_result:
        self._write_cache_version_file(run_dir)

      return data, content_type, content_encoding

    logger.info('%s does not use xplane', tool)
    return None, content_type, None

  def hlo_module_list_impl(self, request: wrappers.Request) -> str:
    """Returns a string of HLO module names concatenated by comma for the given run."""
    run = request.args.get('run')
    run_dir = self._run_dir(run, request)
    if not run_dir:
      logger.warning('Cannot find asset directory for: %s', run)
      return ''
    try:
      all_basenames = self._get_all_basenames(run_dir)
      module_list = [
          name
          for f in all_basenames
          if f.endswith('.hlo_proto.pb') and (name := _parse_filename(f)[0])
      ]

      if not module_list:
        xplane_basenames = self._get_xplane_basenames(run_dir)
        xplane_filenames = [os.path.join(run_dir, f) for f in xplane_basenames]
        if xplane_filenames:
          try:
            # This triggers ConvertMultiXSpaceToHloProto in the C++ backend.
            convert.xspace_to_tool_names(xplane_filenames)
          except AttributeError:
            logger.warning(
                'XPlane converters are available after Tensorflow 2.4'
            )

          all_basenames = self._get_all_basenames(run_dir)
          module_list = [
              name
              for f in all_basenames
              if f.endswith('.hlo_proto.pb') and (name := _parse_filename(f)[0])
          ]

      return ','.join(module_list)
    except OSError as e:
      logger.warning('Cannot read asset directory: %s, OpError %r', run_dir, e)
      return ''

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def data_route(self, request: wrappers.Request) -> wrappers.Response:
    """Handlers for data."""
    # pytype: enable=wrong-arg-types
    # params
    #   request: XMLHTTPRequest.
    try:
      data, content_type, content_encoding = self.data_impl(request)
      if data is None:
        return respond('No Data', 'text/plain', code=404)
      return respond(data, content_type, content_encoding=content_encoding)
    # Data fetch error handler
    except TimeoutError as e:
      return respond(str(e), 'text/plain', code=500)
    except AttributeError as e:
      return respond(str(e), 'text/plain', code=500)
    except ValueError as e:
      return respond(str(e), 'text/plain', code=500)
    except FileNotFoundError as e:
      return respond(str(e), 'text/plain', code=500)
    except IOError as e:
      return respond(str(e), 'text/plain', code=500)

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  # pytype: enable=wrong-arg-types
  def data_csv_route(self, request: wrappers.Request) -> wrappers.Response:
    """Retrieves tool data and converts it to CSV before responding."""
    try:
      data, content_type, _ = self.data_impl(request)

      if data is None:
        return respond('No Data Found', 'text/plain', code=404)

      if content_type != 'application/json':
        return respond(
            'CSV format not supported for this tool type',
            'text/plain',
            code=400,
        )

      csv_data = convert.json_to_csv_string(data)
      filename = _generate_csv_filename(request)

      return respond(
          csv_data,
          'text/csv',
          content_encoding=None,
          extra_headers={
              'Content-Disposition': f'attachment; filename="{filename}"'
          },
      )

    except (
        TimeoutError,
        AttributeError,
        ValueError,
        FileNotFoundError,
        IOError,
        TypeError,
    ) as e:
      logger.exception('CSV conversion error')
      return respond(str(e), 'text/plain', code=500)

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def capture_route(self, request: wrappers.Request) -> wrappers.Response:
    # pytype: enable=wrong-arg-types
    return self.capture_route_impl(request)

  def capture_route_impl(self, request: wrappers.Request) -> wrappers.Response:
    """Runs the client trace for capturing profiling information."""
    service_addr = request.args.get('service_addr')
    duration = int(request.args.get('duration', '1000'))
    is_tpu_name = request.args.get('is_tpu_name') == 'true'
    worker_list = request.args.get('worker_list')
    num_tracing_attempts = int(request.args.get('num_retry', '0')) + 1
    options = {
        'host_tracer_level': int(request.args.get('host_tracer_level', '2')),
        'device_tracer_level': int(
            request.args.get('device_tracer_level', '1')
        ),
        'python_tracer_level': int(
            request.args.get('python_tracer_level', '0')
        ),
        'delay_ms': int(request.args.get('delay', '0')),
    }

    if is_tpu_name:
      if self._tf_profiler is None:
        return respond(
            {
                'error': (
                    'TensorFlow is not installed, but is required to use TPU'
                    ' names.'
                )
            },
            'application/json',
            code=500,
        )
      try:
        # Delegate to the helper class for all TF-related logic.
        service_addr, worker_list, master_ip = (
            self._tf_profiler.resolve_tpu_name(service_addr, worker_list or '')
        )
        self.master_tpu_unsecure_channel = master_ip
      except (RuntimeError, ValueError) as err:
        return respond({'error': str(err)}, 'application/json', code=500)

    if not self.logdir:
      return respond(
          {'error': 'logdir is not set, abort capturing.'},
          'application/json',
          code=500,
      )
    try:
      # The core trace call remains, now with cleanly resolved parameters.
      _pywrap_profiler_plugin.trace(
          service_addr.removeprefix('grpc://'),
          str(self.logdir),
          worker_list,
          True,
          duration,
          num_tracing_attempts,
          options,
      )
      return respond(
          {'result': 'Capture profile successfully. Please refresh.'},
          'application/json',
      )
    except Exception as e:  # pylint: disable=broad-except
      return respond({'error': str(e)}, 'application/json', code=500)

  def _get_graph_viewer_options(
      self, request: wrappers.Request
  ) -> dict[str, Any]:
    node_name = request.args.get('node_name')
    module_name = request.args.get('module_name')
    graph_width_str = request.args.get('graph_width') or ''
    graph_width = int(graph_width_str) if graph_width_str.isdigit() else 3
    show_metadata = int(request.args.get('show_metadata') == 'true')
    merge_fusion = int(request.args.get('merge_fusion') == 'true')
    program_id = request.args.get('program_id')
    return {
        'node_name': node_name,
        'module_name': module_name,
        'program_id': program_id,
        'graph_width': graph_width,
        'show_metadata': show_metadata,
        'merge_fusion': merge_fusion,
        'format': request.args.get('format'),
        'type': request.args.get('type'),
    }

  def generate_runs(self) -> Iterator[str]:
    """Generator for a list of runs.

    The "run name" here is a "frontend run name" - see _tb_run_directory() for
    the definition of a "frontend run name" and how it maps to a directory of
    profile data for a specific profile "run". The profile plugin concept of
    "run" is different from the normal TensorBoard run; each run in this case
    represents a single instance of profile data collection, more similar to a
    "step" of data in typical TensorBoard semantics. These runs reside in
    subdirectories of the plugins/profile directory within any regular
    TensorBoard run directory or within the session_dir root directory
    itself (even if it contains no tfevents file and would thus not be
    considered a normal TensorBoard run, for backwards compatibility).

    `generate_runs` will get all runs first, and get tools list from
    `generate_tools_of_run` for a single run due to expensive processing for
    xspace data to parse the tools.
    Example:
      logs/
        plugins/
          profile/
            run1/
              hostA.trace
        train/
          events.out.tfevents.foo
          plugins/
            profile/
              run1/
                hostA.trace
                hostB.trace
              run2/
                hostA.trace
        validation/
          events.out.tfevents.foo
          plugins/
            profile/
              run1/
                hostA.trace
        new_job/
          tensorboard/
            plugins/
              profile/
                run1/
                  hostA.xplane.pb
    Yields:
    A sequence of string that are "frontend run names".
    For the above example, this would be:
        "run1", "train/run1", "train/run2", "validation/run1",
        "new_job/tensorboard/run1"
    """
    if not self.logdir:
      return

    # Ensure that we check the root logdir and all subdirectories.
    # Note that we check if logdir is a directory to handle case where
    # it's actually a multipart directory spec, which this plugin does not
    # support.
    #
    # This change still enforce the requirement that the subdirectories must
    # end with plugins/profile directory, as enforced by TensorBoard.
    logdir_path = self._epath.Path(self.logdir)
    schemeless_logdir = str(logdir_path)
    if '://' in schemeless_logdir:
      schemeless_logdir = schemeless_logdir.split('://', 1)[1]
    tb_runs = {'.'}

    if logdir_path.is_dir():
      try:
        fs = etils.epath.backend.fsspec_backend.fs(self.logdir)
        for path_str in fs.glob(os.path.join(self.logdir, '**', PLUGIN_NAME)):
          path = self._epath.Path(path_str)
          if fs.isdir(path) and path.parent.name == TB_NAME:
            tb_run_dir = path.parent.parent
            tb_run = tb_run_dir.relative_to(schemeless_logdir)
            tb_runs.add(str(tb_run))
      except ValueError:
        # gcsfs not available, fall back to legacy path walk.
        for cur_dir, _, _ in logdir_path.walk():
          if cur_dir.name == PLUGIN_NAME and cur_dir.parent.name == TB_NAME:
            tb_run_dir = cur_dir.parent.parent
            tb_run = tb_run_dir.relative_to(logdir_path)
            tb_runs.add(str(tb_run))
    tb_run_names_to_dirs = {
        run: _tb_run_directory(self.logdir, run) for run in tb_runs
    }
    plugin_assets = _plugin_assets(
        self.logdir, list(tb_run_names_to_dirs), PLUGIN_NAME
    )
    visited_runs = set()
    for tb_run_name, profile_runs in six.iteritems(plugin_assets):
      tb_run_dir = tb_run_names_to_dirs[tb_run_name]
      tb_plugin_dir = plugin_asset_util.PluginDirectory(tb_run_dir, PLUGIN_NAME)

      for profile_run in profile_runs:
        # Remove trailing separator; some filesystem implementations emit this.
        profile_run = profile_run.rstrip(os.sep)
        if tb_run_name == '.':
          frontend_run = profile_run
        else:
          frontend_run = str(self._epath.Path(tb_run_name) / profile_run)
        profile_run_dir = str(self._epath.Path(tb_plugin_dir) / profile_run)
        if self._epath.Path(profile_run_dir).is_dir():
          with self._run_dir_cache_lock:
            self._run_to_profile_run_dir[frontend_run] = profile_run_dir
          if frontend_run not in visited_runs:
            visited_runs.add(frontend_run)
            yield frontend_run

  def generate_tools_of_run(self, run: str, run_dir: str) -> Iterator[str]:
    """Generate a list of tools given a certain run."""
    if not run_dir:
      logger.warning('Cannot find asset directory for: %s', run)
      return
    profile_run_dir = self._epath.Path(run_dir)
    fs = profile_io.get_file_system(run_dir, self._epath)
    cache = ToolsCache(profile_run_dir, fs)

    cached_tools = cache.load()

    if cached_tools is not None:
      for tool in cached_tools:
        yield tool
      return

    # Cache is invalid or doesn't exist, regenerate
    tools = []
    all_filenames = self._get_all_basenames(run_dir)

    if all_filenames:
      tools = self._get_active_tools(all_filenames, str(profile_run_dir))
      cache.save(tools)

    for tool in tools:
      yield tool

  def _get_active_tools(self, filenames, profile_run_dir=''):
    """Get a list of tools available given the filenames created by profiler.

    Args:
      filenames: List of strings that represent filenames
      profile_run_dir: The run directory of the profile.

    Returns:
      A list of strings representing the available tools
    """
    tool_sort_order = [
        'overview_page',
        'trace_viewer',
        'trace_viewer@',
        'graph_viewer',
        'op_profile',
        'hlo_op_profile',
        'input_pipeline_analyzer',
        'input_pipeline',
        'kernel_stats',
        'memory_profile',
        'memory_viewer',
        'roofline_model',
        'perf_counters',
        'pod_viewer',
        'framework_op_stats',
        'tensorflow_stats',  # Legacy name for framework_op_stats
        'hlo_op_stats',
        'hlo_stats',  # Legacy name for hlo_op_stats
        'inference_profile',
        'megascale_stats',
    ]
    tools = _get_tools(filenames, profile_run_dir)
    if 'trace_viewer@' in tools:
      # streaming trace viewer always override normal trace viewer.
      # the trailing '@' is to inform tf-profile-dashboard.html and
      # tf-trace-viewer.html that stream trace viewer should be used.
      tools.discard('trace_viewer')

    sorted_tools = [t for t in tool_sort_order if t in tools]
    remaining_tools = tools.difference(sorted_tools)
    sorted_tools.extend(sorted(remaining_tools))

    return sorted_tools

  # pytype: disable=wrong-arg-types
  @wrappers.Request.application
  def generate_cache_route(
      self, request: wrappers.Request
  ) -> wrappers.Response:
    # pytype: enable=wrong-arg-types
    """Generates tool data cache in the background."""
    return self._generate_cache_impl(request)

  def _generate_cache_impl(
      self, request: wrappers.Request
  ) -> wrappers.Response:
    """Generates tool data cache in the background.

    Args:
      request: The Werkzeug request object. `request.args` may contain the
        following parameters: - session_path: The path to the session directory
        containing XPlane files. (Required) - tools: An optional comma-separated
        list of tool names to generate cache for. If not provided, defaults to
        `DEFAULT_CACHE_TOOLS`.

    Returns:
      A JSON response indicating whether the task was accepted.
    """
    logger.info('Received generate_cache request.')
    if request.method != 'POST':
      return respond('Method Not Allowed', 'text/plain', code=405)

    params = request.args
    session_path = params.get('session_path')
    logger.info('generate_cache called with params: %s', params)

    if not session_path:
      return respond('Missing "session_path" parameter', 'text/plain', code=400)

    try:
      path = self._epath.Path(session_path)
      xplane_basenames = self._get_xplane_basenames(session_path)
      asset_paths = sorted(
          str(path / basename) for basename in xplane_basenames
      )
      if not asset_paths:
        return respond(
            'No XPlane files found in session_path', 'text/plain', code=404
        )
      logger.info(
          'Found %d *.xplane.pb files in %s.', len(asset_paths), session_path
      )
    except OSError as e:
      logger.exception('Error listing files in session_path: %s', session_path)
      return respond(
          f'Error listing files in session_path: {e!r}', 'text/plain', code=500
      )

    runs = self.runs_imp(request)
    if len(runs) != 1:
      # When 'session_path' is provided, runs_imp should return exactly one run
      # corresponding to the session_path's name.
      return respond(
          'Expected exactly one run for the provided session_path, but found'
          ' %d: %s. Please ensure session_path points to a valid profile run'
          ' directory.' % (len(runs), runs),
          'text/plain',
          code=400,
      )
    run_name = runs[0]
    logger.info(
        'Querying available tools for run %s via run_tools_imp.',
        run_name,
    )
    available_run_tools = set(self.run_tools_imp(run_name, request))
    logger.info(
        'Discovered tools for cache generation: %s for run %s',
        available_run_tools,
        run_name,
    )

    tools_str = params.get('tools')
    requested_tools = (
        set(t.strip() for t in tools_str.split(',') if t.strip())
        if tools_str
        else set(DEFAULT_CACHE_TOOLS)
    )
    if tools_str:
      logger.info('Request tools for cache generation: %s', requested_tools)
    else:
      logger.info(
          'No tools specified in request, using default tools: %s',
          DEFAULT_CACHE_TOOLS,
      )

    available_xplane_tools = available_run_tools.intersection(XPLANE_TOOLS_SET)

    filtered_tools = requested_tools.intersection(available_xplane_tools)

    skipped_tools = requested_tools.difference(filtered_tools)
    for tool in skipped_tools:
      if tool not in available_run_tools:
        logger.info(
            'Tool %s was requested for caching but is not available for run %s,'
            ' skipping.',
            tool,
            run_name,
        )
      else:
        logger.warning(
            'Tool %s is available for run %s but not in XPLANE_TOOLS_SET,'
            ' skipping cache generation.',
            tool,
            run_name,
        )

    if not filtered_tools:
      return respond(
          'No valid XPlane tools found or specified for caching in run %s.'
          % run_name,
          'text/plain',
          code=400,
      )

    logger.info(
        'Filtered tools for cache generation: %s for session %s',
        filtered_tools,
        session_path,
    )

    try:
      logger.info(
          'Submitting cache generation task to thread pool for session %s...',
          session_path,
      )
      self._cache_generation_pool.submit(
          self._generate_cache_task,
          asset_paths=asset_paths,
          tool_list=sorted(list(filtered_tools)),
          params=params,
          session_path=session_path,
      )
    except RuntimeError as e:
      logger.exception(
          'Failed to schedule cache generation task for session_path: %s',
          session_path,
      )
      return respond(f'Failed to schedule task: {e!r}', 'text/plain', code=500)
    else:
      return respond(
          {'status': 'ACCEPTED', 'message': 'Cache generation started'},
          'application/json',
          code=202,
      )

  def _generate_cache_task(
      self,
      *,
      asset_paths: Sequence[str],
      tool_list: Sequence[str],
      params: Mapping[str, Any],
      session_path: str,
  ) -> None:
    """Generates and caches tool data from XPlane files in a background thread.

    Args:
      asset_paths: A list of paths to the XPlane files.
      tool_list: A list of tool names for which to generate cache.
      params: Additional parameters from the request.
      session_path: The path to the session directory.
    """
    logger.info(
        'Background cache generation task started for tools: %s', tool_list
    )
    logger.info('Writing cache version file to %s', session_path)
    self._write_cache_version_file(session_path)

    filenames = [os.path.basename(p) for p in asset_paths]

    base_tool_params = dict(params)

    for tool in tool_list:
      try:
        logger.info('Generating cache for tool %s...', tool)
        tool_params = base_tool_params.copy()
        tool_params['hosts'] = hosts_from_xplane_filenames(filenames, tool)
        self._xspace_to_tool_data(
            [self._epath.Path(p) for p in asset_paths], tool, tool_params
        )
        logger.info(
            'Successfully generated cache for tool %s for %d files.',
            tool,
            len(asset_paths),
        )
      # Catch all exceptions to prevent the background thread from crashing.
      # This ensures that even if one tool fails to generate, other tools
      # can still be processed. The error is logged for debugging.
      except (AttributeError, ValueError, OSError):
        logger.exception(
            'Background cache generation failed for tool %s in session %s',
            tool,
            session_path,
        )
      except Exception:  # pylint: disable=broad-except
        logger.exception(
            'Unexpected error during background cache generation for tool %s'
            ' in session %s',
            tool,
            session_path,
        )
