dew.telemetry.profile

Explicit process-local JAX capture and unmodified native XProf reports.

Only start/stop drain the default backend’s live arrays and effects. This is not a distributed barrier. Normal execution installs no instrumentation hooks.

Name	Summary
require_profile_support	Resolve the optional native report converter before starting work.
capture_options	The options a Dew capture starts with, for any JAX trace that should read like one.
active_profile	The explicitly enabled process-local profiler, or None.
Profiler	One reusable explicitly enabled profiler for context or manual use.
region	A named span in the active capture, or a no-op when nothing captures.
profile	Configure native profiling; capture starts only on enter or start().

require_profile_support

function source

def require_profile_support() -> _Converter

Resolve the optional native report converter before starting work.

capture_options

function source

def capture_options() -> jax.profiler.ProfileOptions

The options a Dew capture starts with, for any JAX trace that should read like one.

Host events at level 2 and the compiled HLO, and no Python tracer: it records every Python and C call, which slows Python-heavy host work several times over (a server’s 512-row settle loop took 0.85 ms a step traced against 0.18 ms untraced), so the trace would show host time, and the device idle behind it, that the run never spends. On TPU the trace also carries compute and sync events.

active_profile

function source

def active_profile() -> Profiler | None

The explicitly enabled process-local profiler, or None.

Profiler

class source

class Profiler(
    directory: str | os.PathLike[str] | None = None,
    *,
    options: jax.profiler.ProfileOptions | None = None,
)

One reusable explicitly enabled profiler for context or manual use.

Each start creates a fresh capture child below directory. Without a supplied directory, the first start allocates a persistent temporary root. Collection covers this process and drains only JAX’s default backend, without copying array values to the host. Other processes must enter their own captures.

Profiler.region

def region(name: str) -> AbstractContextManager[None]

A named span in the trace while this profiler is capturing.

Wrap the work whose device and host time should show under name. Outside a capture the span is a no-op, so the same code runs unprofiled without a branch at every site.

Profiler.start

def start() -> Self

Profiler.stop

def stop() -> None

region

function source

def region(name: str) -> AbstractContextManager[None]

A named span in the active capture, or a no-op when nothing captures.

The site does not need to know whether a profiler is running or which one; with region("input.wait"): reads the same either way.

profile

function source

def profile(
    directory: str | os.PathLike[str] | None = None,
    *,
    options: jax.profiler.ProfileOptions | None = None,
) -> Profiler

Configure native profiling; capture starts only on enter or start().