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.
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.

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().