nvalchemi.dynamics.hooks.ProfilerHook

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

Per-stage timing hook for dynamics simulations.

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.

Parameters:

• stages (set[HookStageEnum] | {"all", "step", "detailed"}) –

  Which stages to instrument.

  • "all" (default): every stage 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[HookStageEnum] 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.

stages

Profiled stages in execution order (used by register_hook).

Type:

list[HookStageEnum]

frequency

Execution frequency in steps.

Type:

int

timings

Accumulated per-transition timing data (seconds).

Type:

dict[HookStageEnum, 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__(stages='all', *, frequency=1, enable_nvtx=True, timer_backend='auto', log_path=None, show_console=False, console_frequency=1)

Parameters:

• stages (set[HookStageEnum] | 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)

Return type:

None

Methods

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