nvalchemi.hooks.StageTimingHook

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

Per-stage timing hook for hook-enabled workflows.

A single StageTimingHook 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 implements _runs_on_stage so hook registries can dispatch it at every selected stage.

The hook supports dynamics presets and custom enum stage sets. Contexts should provide batch, global_rank, and step_count attributes.

Parameters:

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

  Which stages to instrument.

  • "all" (default): every dynamics 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[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 | None, optional) – Single stage to bind the hook to when it is not dispatched across the full profiled_stages set. Default None, letting the hook fire at every profiled stage via _runs_on_stage.

_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.hooks import StageTimingHook
>>> profiler = StageTimingHook()
>>> 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 = StageTimingHook(
...     "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)

close()

Flush and close the CSV log file, if open.

Return type:

None

reset()

Clear all accumulated timing data.

Return type:

None

summary()

Return per-transition timing statistics.

Returns:

Mapping from "PREV_STAGE->STAGE" label to a stats dict with keys mean_s, std_s, min_s, max_s, total_s, n_samples.

Return type:

dict[str, dict[str, float]]