nvalchemi.dynamics.hooks.ProfilerHook

class nvalchemi.dynamics.hooks.ProfilerHook(profiled_stages='all', *, frequency=1, enable_nvtx=True, timer_backend='auto', log_path=None, show_console=False, console_frequency=1, stage=DynamicsStage.BEFORE_STEP)

Per-stage timing hook for dynamics workflows.

A single ProfilerHook instance registers itself at every requested stage. On each call it records a timestamp; when the last profiled stage in a step fires, it computes the elapsed time between consecutive stages and (optionally) writes to CSV / console.

The hook uses stages (plural) so that register_hook() registers it at all listed stages in one call.

The hook supports DynamicsStage and custom enum types via plum dispatch, automatically annotating NVTX ranges with the appropriate domain (dynamics or custom).

Parameters:

• profiled_stages (set[Enum] | {"all", "step", "detailed"}) –

  Which stages to instrument.

  • "all" (default): every DynamicsStage except ON_CONVERGE.

  • "step": BEFORE_STEP and AFTER_STEP only.

  • "detailed": all stages from BEFORE_STEP through AFTER_STEP (excluding ON_CONVERGE).

  • A custom set[Enum] for fine-grained control.

• frequency (int, optional) – Profile every frequency steps. Default 1.

• enable_nvtx (bool, optional) – Emit NVTX push/pop ranges for Nsight Systems. Default True.

• timer_backend ({"cuda_event", "perf_counter", "auto"}, optional) – Timing backend. "auto" selects cuda_event on GPU devices and perf_counter on CPU. Default "auto".

• log_path (str | Path | None, optional) – Path to a CSV file for persistent timing logs. Each row records the rank, step, stage transition, wall-clock offset, and delta. Default None (no file).

• show_console (bool, optional) – Print a formatted timing table via loguru at each profiled step. Default False.

• console_frequency (int, optional) – When show_console is True, print every console_frequency profiled steps. Default 1.

• stage (Enum)

_profiled_stages

Profiled stages in execution order (private).

Type:

list[Enum]

frequency

Execution frequency in steps.

Type:

int

timings

Accumulated per-transition timing data (seconds).

Type:

dict[Enum, list[float]]

Examples

>>> from nvalchemi.dynamics.hooks import ProfilerHook
>>> profiler = ProfilerHook()
>>> dynamics = DemoDynamics(model=model, n_steps=100, dt=0.5, hooks=[profiler])
>>> dynamics.run(batch)
>>> print(profiler.summary())

With CSV logging and console output:

>>> profiler = ProfilerHook(
...     "detailed",
...     log_path="profiler.csv",
...     show_console=True,
...     console_frequency=10,
... )
>>> dynamics = DemoDynamics(model=model, n_steps=1000, dt=0.5, hooks=[profiler])
>>> dynamics.run(batch)

__init__(profiled_stages='all', *, frequency=1, enable_nvtx=True, timer_backend='auto', log_path=None, show_console=False, console_frequency=1, stage=DynamicsStage.BEFORE_STEP)

Parameters:

• profiled_stages (set[Enum] | Literal['all', 'step', 'detailed'])

• frequency (int)

• enable_nvtx (bool)

• timer_backend (Literal['cuda_event', 'perf_counter', 'auto'])

• log_path (str | Path | None)

• show_console (bool)

• console_frequency (int)

• stage (Enum)

Return type:

None

Methods

__init__([profiled_stages, frequency, ...])
close()	Flush and close the CSV log file, if open.
reset()	Clear all accumulated timing data.
summary()	Return per-transition timing statistics.