# Copyright 2020 The JAX Authors.
#
# 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
#
#     https://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.

from __future__ import annotations

from collections.abc import Callable, Iterator
from contextlib import contextmanager
from functools import wraps
import gzip
import http.server
import json
import logging
import os
import pathlib
import socketserver
import threading
from typing import Any
import warnings

from jax._src import traceback_util
traceback_util.register_exclusion(__file__)

from jax._src import xla_bridge
from jax._src.lib import _profiler
from jax._src.lib import _profile_data
from jax._src.lib import jaxlib_extension_version
from jax import version as jax_version_module
from jax._src.lib import version as version_lib

ProfileData = _profile_data.ProfileData
ProfileEvent = _profile_data.ProfileEvent
ProfileLine = _profile_data.ProfileLine
ProfilePlane = _profile_data.ProfilePlane

_profiler_server: _profiler.ProfilerServer | None = None

logger = logging.getLogger(__name__)


class ProfileOptions(_profiler.ProfileOptions):
  """Profiler Options to configure the collectors for the profiler."""
  # Reject unknown attributes (e.g. typos) rather than silently ignoring them.
  __slots__ = ()


def start_server(
    port: int, requires_backend: bool = True
) -> _profiler.ProfilerServer:
  """Starts the profiler server on port `port`.

  Using the "TensorFlow profiler" feature in `TensorBoard
  <https://www.tensorflow.org/tensorboard>`_ 2.2 or newer, you can
  connect to the profiler server and sample execution traces that show CPU,
  GPU, and/or TPU device activity.

  Args:
    port: The port to start the profiler server on.
    requires_backend: If False, the profiler server will not wait for backends
      to be initialized before starting. Default is True.
  """
  global _profiler_server
  if _profiler_server is not None:
    raise ValueError("Only one profiler server can be active at a time.")

  # Make sure backends are initialized before creating a profiler
  # session. Otherwise on Cloud TPU, libtpu may not be initialized before
  # creating the tracer, which will cause the TPU tracer initialization to
  # fail and no TPU operations will be included in the profile.
  # NOTE(skyewm): I'm not sure this is necessary for start_server (is definitely
  # is for start_trace), but I'm putting it here to be safe.
  if requires_backend:
    xla_bridge.get_backend()

  _profiler_server = _profiler.start_server(port)
  return _profiler_server


def stop_server():
  """Stops the running profiler server."""
  global _profiler_server
  if _profiler_server is None:
    raise ValueError("No active profiler server.")
  _profiler_server = None # Should destroy the profiler server


def register_subprocess(pid: int, port: int) -> Callable[[], None]:
  """Registers a subprocess's profiler server to be profiled alongside the current process.

  When the current process collects a profile (either programmatically or via
  its profiler server), it will propagate the request to all registered
  subprocesses' profiler servers and subsequently, aggregate all their responses
  into the main response returned by this process's profiler server.

  This is helpful when running workers in separate processes that may affect the
  performance of the main process (e.g. PyGrain).

  NOTE: Currently, only CPU profiling of subprocesses is supported.

  Args:
    pid: The process ID of the subprocess.
    port: The port of the profiler server in the subprocess.

  Returns:
    A function that when called or garbage collected will unregister the
    subprocess from the main process's profiler.

  Raises:
    RuntimeError: If the subprocess fails to be registered (e.g. already
    registered, unable to connect, etc.).
  """
  return _profiler.register_subprocess(pid, port)


class _ProfileState:
  def __init__(self):
    self.session: ProfileSession | None = None
    self.lock = threading.Lock()

  def reset(self):
    self.session = None


_profile_state = _ProfileState()


def set_metadata(key: str, value: str) -> None:
  """Sets metadata for the current profiling session."""
  if hasattr(_profiler, "set_metadata"):
    return _profiler.set_metadata(key, value)


def clear_metadata() -> None:
  """Clears metadata for the current profiling session."""
  if hasattr(_profiler, "clear_metadata"):
    return _profiler.clear_metadata()


class ProfileSession:
  """Represents a JAX profiling session.

  Can be used as a context manager (via :func:`trace` or direct instantiation)
  or programmatically via :meth:`start` and :meth:`stop`.
  Upon stopping, the captured :class:`ProfileData` and its
  :class:`ProfilePlane`s
  become accessible in memory, and can be explicitly written to disk using
  :meth:`export_profile_data` (or automatically if ``log_dir`` is provided).
  """

  def __init__(
      self,
      log_dir: os.PathLike | str | None = None,
      create_perfetto_link: bool = False,
      create_perfetto_trace: bool = False,
      profiler_options: ProfileOptions | None = None,
  ):
    self.log_dir: str | None = str(log_dir) if log_dir is not None else None
    self.create_perfetto_link: bool = create_perfetto_link
    self.create_perfetto_trace: bool = (
        create_perfetto_trace or create_perfetto_link
    )
    self.profiler_options: ProfileOptions | None = profiler_options
    self._raw_session: _profiler.ProfilerSession | None = None
    self._profile_data: ProfileData | None = None
    self._is_running: bool = False

  def start(self) -> ProfileSession:
    """Starts the profiling session."""
    with _profile_state.lock:
      if _profile_state.session is not None or self._is_running:
        raise RuntimeError(
            "Profile has already been started. Only one profile may be run at"
            " a time."
        )
      clear_metadata()
      # Make sure backends are initialized before creating a profiler
      # session. Otherwise on Cloud TPU, libtpu may not be initialized before
      # creating the tracer, which will cause the TPU tracer initialization to
      # fail and no TPU operations will be included in the profile.
      xla_bridge.get_backend()

      options = self.profiler_options
      if options is None:
        options = ProfileOptions()
      set_metadata("jax_version", jax_version_module.__version__)
      jaxlib_version_str = ".".join(map(str, version_lib))
      set_metadata("jaxlib_version", jaxlib_version_str)
      for backend_name in xla_bridge.backends():
        try:
          backend = xla_bridge.get_backend(backend_name)
          set_metadata(f"{backend.platform}_version", backend.platform_version)
        except RuntimeError:
          pass

      self._raw_session = _profiler.ProfilerSession(options)
      self._is_running = True
      self._profile_data = None
      _profile_state.session = self
    return self

  def stop(self) -> ProfileData | None:
    """Stops the profiling session and returns the captured ProfileData."""
    with _profile_state.lock:
      if not self._is_running or self._raw_session is None:
        raise RuntimeError("No profile started")

      if jaxlib_extension_version >= 501:
        self._profile_data = self._raw_session.stop_and_get_profile_data()
      elif self.log_dir is not None:
        self._raw_session.stop_and_export(self.log_dir)
      else:
        self._profile_data = self._raw_session.stop_and_get_profile_data()
      self._is_running = False
      self._raw_session = None
      _profile_state.reset()
      clear_metadata()

    if self.log_dir is not None:
      if jaxlib_extension_version >= 501:
        self.export_profile_data(
            log_dir=self.log_dir,
            create_perfetto_link=self.create_perfetto_link,
            create_perfetto_trace=self.create_perfetto_trace,
        )
      elif self.create_perfetto_trace or self.create_perfetto_link:
        abs_filename = _write_perfetto_trace_file(self.log_dir)
        if self.create_perfetto_link:
          _host_perfetto_trace_file(abs_filename)

    return self._profile_data

  def export_profile_data(
      self,
      log_dir: os.PathLike | str,
      create_perfetto_link: bool = False,
      create_perfetto_trace: bool = False,
      session_id: str | None = None,
  ) -> None:
    """Exports the captured profile data to the specified directory."""
    export_profile_data(
        self,
        log_dir=log_dir,
        create_perfetto_link=create_perfetto_link,
        create_perfetto_trace=create_perfetto_trace,
        session_id=session_id,
    )

  @property
  def is_running(self) -> bool:
    return self._is_running

  @property
  def profile_data(self) -> ProfileData | None:
    """Returns the captured ProfileData after the session has stopped."""
    return self._profile_data

  def __enter__(self) -> ProfileSession:
    self.start()
    return self

  def __exit__(self, exc_type, exc_val, exc_tb) -> None:
    # Only stop if still running so early session.stop() inside a
    # with-block does not raise RuntimeError("No profile started") on exit.
    if self.is_running:
      self.stop()


def start_trace(
    log_dir: os.PathLike | str,
    create_perfetto_link: bool = False,
    create_perfetto_trace: bool = False,
    profiler_options: ProfileOptions | None = None,
) -> ProfileSession | None:
  """Starts a profiler trace.

  The trace will capture CPU, GPU, and/or TPU activity, including Python
  functions and JAX on-device operations. Use :func:`stop_trace` to end the
  trace and save the results to ``log_dir``.

  The resulting trace can be viewed with Xprof. Note that Xprof
  doesn't need to be running when collecting the trace.

  Only one trace may be collected at a time. A RuntimeError will be raised if
  :func:`start_trace` is called while another trace is running.

  Args:
    log_dir: The directory to save the profiler trace to (usually the Xprof log
      directory).
    create_perfetto_link: A boolean which, if true, creates and prints link to
      the Perfetto trace viewer UI (https://ui.perfetto.dev). The program will
      block until the link is opened and Perfetto loads the trace.
    create_perfetto_trace: A boolean which, if true, additionally dumps a
      ``perfetto_trace.json.gz`` file that is compatible for upload with the
      Perfetto trace viewer UI (https://ui.perfetto.dev). The file will also be
      generated if ``create_perfetto_link`` is true. This could be useful if you
      want to generate a Perfetto-compatible trace without blocking the process.
    profiler_options: Profiler options to configure the profiler for collection.

  Returns:
    The active :class:`ProfileSession`.
  """
  session = ProfileSession(
      log_dir=log_dir,
      create_perfetto_link=create_perfetto_link,
      create_perfetto_trace=create_perfetto_trace,
      profiler_options=profiler_options,
  )
  session.start()
  return session


def _write_perfetto_trace_file(log_dir: os.PathLike | str):
  # Navigate to folder with the latest trace dump to find `trace.json.jz`
  trace_folders = (pathlib.Path(log_dir).absolute() / "plugins" / "profile").iterdir()
  latest_trace_folder = max(trace_folders, key=os.path.getmtime)
  trace_jsons = latest_trace_folder.glob("*.trace.json.gz")
  try:
    trace_json, = trace_jsons
  except ValueError as value_error:
    raise ValueError(f"Invalid trace folder: {latest_trace_folder}") from value_error

  logger.info("Loading trace.json.gz and removing its metadata...")
  # Perfetto doesn't like the `metadata` field in `trace.json` so we remove
  # it.
  # TODO(sharadmv): speed this up by updating the generated `trace.json`
  # to not include metadata if possible.
  with gzip.open(trace_json, "rb") as fp:
    trace = json.load(fp)
    del trace["metadata"]
  perfetto_trace = latest_trace_folder / "perfetto_trace.json.gz"
  logger.info("Writing perfetto_trace.json.gz...")
  with gzip.open(perfetto_trace, "w") as fp:
    fp.write(json.dumps(trace).encode("utf-8"))
  return perfetto_trace

class _PerfettoServer(http.server.SimpleHTTPRequestHandler):
  """Handles requests from `ui.perfetto.dev` for the `trace.json`"""

  def end_headers(self):
    self.send_header('Access-Control-Allow-Origin', '*')
    return super().end_headers()

  def do_GET(self):
    self.server.last_request = self.path  # pyrefly: ignore[missing-attribute]
    return super().do_GET()

  def do_POST(self):
    self.send_error(404, "File not found")

def _host_perfetto_trace_file(path: os.PathLike | str):
  # ui.perfetto.dev looks for files hosted on `127.0.0.1:9001`. We set up a
  # TCP server that is hosting the `perfetto_trace.json.gz` file.
  port = 9001
  orig_directory = pathlib.Path.cwd()
  directory, filename = os.path.split(path)
  try:
    os.chdir(directory)
    socketserver.TCPServer.allow_reuse_address = True
    with socketserver.TCPServer(('127.0.0.1', port), _PerfettoServer) as httpd:
      url = f"https://ui.perfetto.dev/#!/?url=http://127.0.0.1:{port}/{filename}"
      print(f"Open URL in browser: {url}")

      # Once ui.perfetto.dev acquires trace.json from this server we can close
      # it down.
      while httpd.__dict__.get('last_request') != '/' + filename:
        httpd.handle_request()
  finally:
    os.chdir(orig_directory)


def export_profile_data(
    profile_data: ProfileData | ProfileSession,
    log_dir: os.PathLike | str,
    create_perfetto_link: bool = False,
    create_perfetto_trace: bool = False,
    session_id: str | None = None,
) -> None:
  """Exports captured profile data to the specified directory.

  Args:
    profile_data: The :class:`ProfileData` or completed :class:`ProfileSession`
      to export.
    log_dir: The directory to save the profiler trace to (usually the Xprof log
      directory).
    create_perfetto_link: A boolean which, if true, creates and prints link to
      the Perfetto trace viewer UI (https://ui.perfetto.dev). The program will
      block until the link is opened and Perfetto loads the trace.
    create_perfetto_trace: A boolean which, if true, additionally dumps a
      ``perfetto_trace.json.gz`` file that is compatible for upload with the
      Perfetto trace viewer UI (https://ui.perfetto.dev).
    session_id: Optional session ID string for the export folder.
  """
  if isinstance(profile_data, ProfileSession):
    if profile_data.is_running:
      raise RuntimeError(
          "Cannot export a ProfileSession while it is running. Call stop()"
          " first."
      )
    if profile_data.profile_data is None:
      raise RuntimeError(
          "No profile data to export. Ensure the session has completed."
      )
    data = profile_data.profile_data
    if session_id is None and profile_data.profiler_options is not None:
      session_id = profile_data.profiler_options.session_id
  elif isinstance(profile_data, ProfileData):
    data = profile_data
  else:
    raise TypeError(
        "Expected ProfileData or ProfileSession, got"
        f" {type(profile_data).__name__}"
    )

  if jaxlib_extension_version >= 501:
    _profiler.export_to_xprof(data, str(log_dir), session_id=session_id or "")
  else:
    raise RuntimeError("export_profile_data requires JAX 0.10.3 or newer.")
  if create_perfetto_trace or create_perfetto_link:
    abs_filename = _write_perfetto_trace_file(str(log_dir))
    if create_perfetto_link:
      _host_perfetto_trace_file(abs_filename)


def stop_trace() -> ProfileData | None:
  """Stops the currently-running profiler trace.

  The trace will be saved to the ``log_dir`` passed to the corresponding
  :func:`start_trace` call (if provided). Raises a RuntimeError if a trace
  hasn't been started.

  Returns:
    The captured :class:`ProfileData`.
  """
  with _profile_state.lock:
    session = _profile_state.session
    if session is None:
      raise RuntimeError("No profile started")

  return session.stop()


def stop_and_get_fdo_profile() -> bytes | str:
  """Stops the currently-running profiler trace and export fdo_profile.

  Currently, this is only supported for GPU.
  Raises a RuntimeError if a trace hasn't been started.
  """
  with _profile_state.lock:
    session = _profile_state.session
    if session is None or session._raw_session is None:
      raise RuntimeError("No profile started")
    xspace = session._raw_session.stop()
    fdo_profile = _profiler.get_fdo_profile(xspace)
    session._is_running = False
    session._raw_session = None
    _profile_state.reset()
    clear_metadata()
    return fdo_profile


@contextmanager
def trace(
    log_dir: os.PathLike | str | None = None,
    create_perfetto_link: bool = False,
    create_perfetto_trace: bool = False,
    profiler_options: ProfileOptions | None = None,
) -> Iterator[ProfileSession]:
  """Context manager to take a profiler trace.

  The trace will capture CPU, GPU, and/or TPU activity, including Python
  functions and JAX on-device operations.

  The yielded :class:`ProfileSession` object provides in-memory access to
  the :class:`ProfileData` and :class:`ProfilePlane`s upon exiting the context.
  If ``log_dir`` is specified, the trace is also exported to ``log_dir``
  for viewing in Xprof.

  Only one trace may be collected at a time. A RuntimeError will be raised if a
  trace is started while another trace is running.

  Args:
    log_dir: Optional directory to save the profiler trace to (usually the Xprof
      log directory). If None, the trace is kept in memory only.
    create_perfetto_link: A boolean which, if true, creates and prints link to
      the Perfetto trace viewer UI (https://ui.perfetto.dev). The program will
      block until the link is opened and Perfetto loads the trace.
    create_perfetto_trace: A boolean which, if true, additionally dumps a
      ``perfetto_trace.json.gz`` file that is compatible for upload with the
      Perfetto trace viewer UI (https://ui.perfetto.dev).
    profiler_options: Profiler options to configure the profiler for collection.

  Yields:
    A :class:`ProfileSession` object that gives access to the profile data
    and XPlanes after exiting the context.
  """
  with ProfileSession(
      log_dir=log_dir,
      create_perfetto_link=create_perfetto_link,
      create_perfetto_trace=create_perfetto_trace,
      profiler_options=profiler_options,
  ) as session:
    yield session


class TraceAnnotation(_profiler.TraceMe):
  """Context manager that generates a trace event in the profiler.

  The trace event spans the duration of the code enclosed by the context.

  For example:

  >>> x = jnp.ones((1000, 1000))
  >>> with jax.profiler.TraceAnnotation("my_label"):
  ...   result = jnp.dot(x, x.T).block_until_ready()

  This will cause a "my_label" event to show up on the trace timeline if the
  event occurs while the process is being traced.
  """


class StepTraceAnnotation(TraceAnnotation):
  """Context manager that generates a step trace event in the profiler.

  The step trace event spans the duration of the code enclosed by the context.
  The profiler will provide the performance analysis for each step trace event.

  For example, it can be used to mark training steps and enable the profiler to
  provide the performance analysis per step:

  >>> while global_step < NUM_STEPS:                                           # doctest: +SKIP
  ...   with jax.profiler.StepTraceAnnotation("train", step_num=global_step):  # doctest: +SKIP
  ...     train_step()                                                         # doctest: +SKIP
  ...     global_step += 1                                                     # doctest: +SKIP

  This will cause a "train xx" event to show up on the trace timeline if the
  event occurs while the process is being traced by TensorBoard. In addition,
  if using accelerators, the device trace timeline will also show a "train xx"
  event. Note that "step_num" can be set as a keyword argument to pass the
  global step number to the profiler.

  """

  def __init__(self, name: str, **kwargs):
    super().__init__(name, _r=1, **kwargs)


def annotate_function(func: Callable, name: str | None = None,
                      **decorator_kwargs):
  """Decorator that generates a trace event for the execution of a function.

  For example:

  >>> @jax.profiler.annotate_function
  ... def f(x):
  ...   return jnp.dot(x, x.T).block_until_ready()
  >>>
  >>> result = f(jnp.ones((1000, 1000)))

  This will cause an "f" event to show up on the trace timeline if the
  function execution occurs while the process is being traced by TensorBoard.

  Arguments can be passed to the decorator via :py:func:`functools.partial`.

  >>> from functools import partial

  >>> @partial(jax.profiler.annotate_function, name="event_name")
  ... def f(x):
  ...   return jnp.dot(x, x.T).block_until_ready()

  >>> result = f(jnp.ones((1000, 1000)))
  """

  name = name or getattr(func, '__qualname__', None)
  name = name or func.__name__
  @wraps(func)
  def wrapper(*args, **kwargs):
    with TraceAnnotation(name, **decorator_kwargs):
      return func(*args, **kwargs)
  return wrapper


def device_memory_profile(backend: str | None = None) -> bytes:
  """Captures a JAX device memory profile as ``pprof``-format protocol buffer.

  A device memory profile is a snapshot of the state of memory, that describes the JAX
  :class:`~jax.Array` and executable objects present in memory and their
  allocation sites.

  For more information how to use the device memory profiler, see
  :ref:`jax-201-memory-profiling`.

  The profiling system works by instrumenting JAX on-device allocations,
  capturing a Python stack trace for each allocation. The instrumentation is
  always enabled; :func:`device_memory_profile` provides an API to capture it.

  The output of :func:`device_memory_profile` is a binary protocol buffer that
  can be interpreted and visualized by the `pprof tool
  <https://github.com/google/pprof>`_.

  Args:
    backend: optional; the name of the JAX backend for which the device memory
      profile should be collected.

  Returns:
    A byte string containing a binary `pprof`-format protocol buffer.
  """
  client = xla_bridge.get_backend(backend)
  return gzip.compress(client.heap_profile())


def save_device_memory_profile(filename, backend: str | None = None) -> None:
  """Collects a device memory profile and writes it to a file.

  :func:`save_device_memory_profile` is a convenience wrapper around :func:`device_memory_profile`
  that saves its output to a ``filename``. See the
  :func:`device_memory_profile` documentation for more information.

  Args:
    filename: the filename to which the profile should be written.
    backend: optional; the name of the JAX backend for which the device memory
      profile should be collected.
  """
  profile = device_memory_profile(backend)
  with open(filename, "wb") as f:
    f.write(profile)


# Allows to run model with profiler given amount of times. After required amount
# of retries achieved client can collect FDO data.
class PGLEProfiler:

  def __init__(self, retries: int, percentile: int):
    self.retries: int = retries
    self.percentile: int = percentile
    self.collected_fdo: bytes | None = None
    self.called_times: int = 0
    self.fdo_profiles: list[Any] = []
    self.current_session: _profiler.ProfilerSession | None = None

  def consume_fdo_profile(self) -> bytes | None:
    if self.collected_fdo is not None:
      return self.collected_fdo

    if not self.is_enabled() or self.called_times != self.retries:
      return None

    self.collected_fdo = _profiler.aggregate_profiled_instructions(
        self.fdo_profiles, self.percentile
    )
    return self.collected_fdo

  def is_fdo_consumed(self):
    return self.collected_fdo is not None

  def disable(self):
    self.retries = 0

  def is_enabled(self):
    return self.retries > 0

  def is_running(self):
    return self.current_session is not None

  @classmethod
  @contextmanager
  def trace(cls, runner: PGLEProfiler | None):
    if (runner is None or runner.is_running()
        or not runner.is_enabled() or runner.is_fdo_consumed()):
      yield
    else:
      options = _profiler.ProfileOptions()
      options.enable_hlo_proto = True
      options.raise_error_on_start_failure = True
      runner.current_session = _profiler.ProfilerSession(options)

      try:
        yield
      finally:
        xspace = runner.current_session.stop()
        runner.fdo_profiles.append(
            _profiler.get_fdo_profile(xspace)
        )
        runner.current_session = None

        runner.called_times += 1
        if runner.fdo_profiles[-1] == b'':
          warnings.warn(
              "PGLE collected an empty trace, may be due to contention with "
              "another tool that subscribes to CUPTI, such as Nsight Systems - check "
              "for CUPTI_ERROR_MULTIPLE_SUBSCRIBERS_NOT_SUPPORTED from XLA. "
              "Consider populating a persistent compilation cache with PGLE enabled, "
              "and then profiling a second run that has the "
              "JAX_COMPILATION_CACHE_EXPECT_PGLE option enabled.",
              RuntimeWarning)