# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Per-stage wall-clock timing for hook-enabled workflows."""
from __future__ import annotations
import csv
import io
import statistics
import time
from enum import Enum
from pathlib import Path
from typing import Literal
import torch
from loguru import logger
from nvalchemi.data import Batch
from nvalchemi.hooks._context import HookContext
try:
import nvtx
except ImportError:
nvtx = None
__all__ = ["StageTimingHook"]
def _get_dynamics_stage_type() -> type[Enum]:
"""Import and return the dynamics stage enum on demand."""
from nvalchemi.dynamics.base import DynamicsStage
return DynamicsStage
def _stage_domain(stage: Enum) -> str:
"""Return a stable domain label for a stage enum."""
stage_type = type(stage)
if (
stage_type.__name__ == "DynamicsStage"
and stage_type.__module__ == "nvalchemi.dynamics.base"
):
return "dynamics"
return "custom"
def _sort_stages(stages: set[Enum]) -> list[Enum]:
"""Sort stage enum members by their integer value."""
return sorted(stages, key=lambda s: s.value)
[docs]
class StageTimingHook:
"""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``.
Attributes
----------
_profiled_stages : list[Enum]
Profiled stages in execution order (private).
frequency : int
Execution frequency in steps.
timings : dict[Enum, list[float]]
Accumulated per-transition timing data (seconds).
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)
"""
[docs]
def __init__(
self,
profiled_stages: set[Enum] | Literal["all", "step", "detailed"] = "all",
*,
frequency: int = 1,
enable_nvtx: bool = True,
timer_backend: Literal["cuda_event", "perf_counter", "auto"] = "auto",
log_path: str | Path | None = None,
show_console: bool = False,
console_frequency: int = 1,
stage: Enum | None = None,
) -> None:
# Init file handle early so __del__ is safe on validation errors.
self._csv_file: io.TextIOWrapper | None = None
self._csv_writer: csv.DictWriter | None = None
if isinstance(profiled_stages, str):
dynamics_stage = _get_dynamics_stage_type()
if profiled_stages == "all":
resolved = {
s for s in dynamics_stage if s != dynamics_stage.ON_CONVERGE
}
elif profiled_stages == "step":
resolved = {dynamics_stage.BEFORE_STEP, dynamics_stage.AFTER_STEP}
elif profiled_stages == "detailed":
resolved = {
dynamics_stage.BEFORE_STEP,
dynamics_stage.BEFORE_PRE_UPDATE,
dynamics_stage.AFTER_PRE_UPDATE,
dynamics_stage.BEFORE_COMPUTE,
dynamics_stage.AFTER_COMPUTE,
dynamics_stage.BEFORE_POST_UPDATE,
dynamics_stage.AFTER_POST_UPDATE,
dynamics_stage.AFTER_STEP,
}
else:
raise ValueError(
f"Unknown stages preset {profiled_stages!r}. "
f"Use 'all', 'step', 'detailed', or a set of Enum."
)
else:
resolved = set(profiled_stages)
if len(resolved) < 2:
raise ValueError(
"At least two stages are required to measure timing deltas."
)
# Sorted by execution order — private profiled stages list.
self._profiled_stages: list[Enum] = _sort_stages(resolved)
# Primary stage for protocol compliance.
self.stage = stage or self._profiled_stages[0]
self.frequency = frequency
self.enable_nvtx = enable_nvtx
self.timer_backend = timer_backend
self.log_path = Path(log_path) if log_path is not None else None
self.show_console = show_console
self.console_frequency = console_frequency
# Per-step scratch — separate dicts for type safety.
self._current_step: int = -1
self._step_cuda_events: dict[Enum, torch.cuda.Event] = {}
self._step_cpu_timestamps: dict[Enum, int] = {}
# Device the CUDA events live on; all record/sync/measure pin to it so
# multi-GPU runs don't hit an invalid-handle across device contexts.
self._cuda_device: torch.device | None = None
# Accumulated timing: transition endpoint -> list of delta_s.
self.timings: dict[Enum, list[float]] = {s: [] for s in self._profiled_stages}
self._t0_ns: int = time.perf_counter_ns()
self._backend_resolved: str | None = None
self._steps_recorded: int = 0
# ------------------------------------------------------------------
# Hook entry point
# ------------------------------------------------------------------
def _runs_on_stage(self, stage: Enum) -> bool:
"""Check if this hook should run on the given stage.
Parameters
----------
stage : Enum
The stage to check.
Returns
-------
bool
True if this hook runs on the given stage.
"""
return stage in set(self._profiled_stages)
@torch.compiler.disable
def _record(
self,
batch: Batch,
current_stage: Enum,
step_count: int,
global_rank: int,
domain: str = "dynamics",
) -> None:
"""Record a timestamp for the current stage.
Parameters
----------
batch : Batch
The current batch of atomic data.
current_stage : Enum
The current dynamics stage being executed.
step_count : int
The current step number.
global_rank : int
The distributed rank of this process.
domain : str, optional
The domain for NVTX annotation (e.g., "dynamics", "custom").
Default ``"dynamics"``.
"""
# New step: flush the previous one, then reset scratch.
if step_count != self._current_step:
if self._current_step >= 0:
self._flush_step(global_rank)
self._current_step = step_count
self._step_cuda_events.clear()
self._step_cpu_timestamps.clear()
# NVTX annotation.
if self.enable_nvtx and nvtx is not None:
idx = self._profiled_stages.index(current_stage)
if idx > 0:
nvtx.pop_range()
nvtx.push_range(f"{domain}/{current_stage.name}/{step_count}")
# Timestamp.
dev = batch.device
if isinstance(dev, str):
dev = torch.device(dev)
if self._backend_resolved is None:
self._backend_resolved = self._resolve_backend(dev)
if self._backend_resolved == "cuda_event":
# Create + record on the batch's device so the event handle is valid
# when measured, regardless of the ambient current device.
with torch.cuda.device(dev):
event = torch.cuda.Event(enable_timing=True)
event.record()
self._cuda_device = dev
self._step_cuda_events[current_stage] = event
else:
self._step_cpu_timestamps[current_stage] = time.perf_counter_ns()
# If this is the last profiled stage in the step, flush now.
if current_stage == self._profiled_stages[-1]:
self._flush_step(global_rank)
self._current_step = -1
self._step_cuda_events.clear()
self._step_cpu_timestamps.clear()
def __call__(self, ctx: HookContext, stage: Enum) -> None:
"""Record timing for a profiled stage.
Parameters
----------
ctx : HookContext
Workflow context. Dynamics and training contexts provide
``step_count``; base contexts default to step ``0``.
stage : Enum
Current workflow stage.
"""
batch = ctx.batch
if batch is None:
raise ValueError("StageTimingHook requires ctx.batch to record timing.")
step_count = int(getattr(ctx, "step_count", 0))
domain = _stage_domain(stage)
self._record(batch, stage, step_count, ctx.global_rank or 0, domain=domain)
# ------------------------------------------------------------------
# Backend resolution
# ------------------------------------------------------------------
def _resolve_backend(self, device: torch.device) -> str:
"""Resolve the timing backend based on configuration and device."""
if self.timer_backend != "auto":
return self.timer_backend
if device.type == "cuda":
return "cuda_event"
return "perf_counter"
# ------------------------------------------------------------------
# Step flush — compute deltas, log
# ------------------------------------------------------------------
def _flush_step(self, rank: int) -> None:
"""Compute per-transition deltas for the current step and log."""
use_cuda = self._backend_resolved == "cuda_event"
if use_cuda:
ordered = [s for s in self._profiled_stages if s in self._step_cuda_events]
else:
ordered = [
s for s in self._profiled_stages if s in self._step_cpu_timestamps
]
if len(ordered) < 2:
return
# Measure on the events' own device; a stray CUDA error degrades to a
# skipped sample rather than aborting the run (timing is best-effort).
try:
if use_cuda and self._cuda_device is not None:
with torch.cuda.device(self._cuda_device):
torch.cuda.synchronize(self._cuda_device)
deltas = self._cuda_deltas(ordered)
elif use_cuda:
torch.cuda.synchronize()
deltas = self._cuda_deltas(ordered)
else:
deltas = self._cpu_deltas(ordered)
except RuntimeError as exc:
logger.debug("StageTimingHook: skipped step timing ({})", exc)
if self.enable_nvtx and nvtx is not None:
nvtx.pop_range()
return
for stage, delta_s in deltas.items():
self.timings[stage].append(delta_s)
t_since_init_s = (time.perf_counter_ns() - self._t0_ns) / 1e9
self._steps_recorded += 1
if self.log_path is not None:
self._write_csv(rank, self._current_step, t_since_init_s, ordered, deltas)
if self.show_console and (self._steps_recorded % self.console_frequency == 0):
self._print_console(
rank, self._current_step, t_since_init_s, ordered, deltas
)
# Close NVTX range for the last stage in this step.
if self.enable_nvtx and nvtx is not None:
nvtx.pop_range()
def _cuda_deltas(self, ordered: list[Enum]) -> dict[Enum, float]:
"""Per-transition deltas (s) from recorded CUDA events."""
deltas: dict[Enum, float] = {}
for i in range(1, len(ordered)):
prev_ev = self._step_cuda_events[ordered[i - 1]]
curr_ev = self._step_cuda_events[ordered[i]]
delta_s = prev_ev.elapsed_time(curr_ev) / 1000.0
deltas[ordered[i]] = delta_s
return deltas
def _cpu_deltas(self, ordered: list[Enum]) -> dict[Enum, float]:
"""Per-transition deltas (s) from perf-counter timestamps."""
deltas: dict[Enum, float] = {}
for i in range(1, len(ordered)):
prev_ts = self._step_cpu_timestamps[ordered[i - 1]]
curr_ts = self._step_cpu_timestamps[ordered[i]]
delta_s = (curr_ts - prev_ts) / 1e9
deltas[ordered[i]] = delta_s
return deltas
# ------------------------------------------------------------------
# CSV output
# ------------------------------------------------------------------
def _write_csv(
self,
rank: int,
step: int,
t_since_init: float,
ordered: list[Enum],
deltas: dict[Enum, float],
) -> None:
"""Append one row per transition to the CSV log."""
rows = []
for i, stage in enumerate(ordered[1:], start=1):
rows.append(
{
"rank": rank,
"step": step,
"stage": f"{ordered[i - 1].name}->{stage.name}",
"t_since_init_s": f"{t_since_init:.6f}",
"delta_s": f"{deltas[stage]:.6f}",
}
)
if self._csv_writer is None:
log_path = self.log_path
if log_path is None:
return
fh = open(log_path, "w", newline="") # noqa: SIM115
self._csv_file = fh
self._csv_writer = csv.DictWriter(
fh,
fieldnames=["rank", "step", "stage", "t_since_init_s", "delta_s"],
)
self._csv_writer.writeheader()
self._csv_writer.writerows(rows)
if self._csv_file is not None:
self._csv_file.flush()
# ------------------------------------------------------------------
# Console output
# ------------------------------------------------------------------
def _print_console(
self,
rank: int,
step: int,
t_since_init: float,
ordered: list[Enum],
deltas: dict[Enum, float],
) -> None:
"""Print a formatted timing table for the current step."""
lines = [f"[Profiler] rank={rank} step={step} t={t_since_init:.3f}s"]
for i, stage in enumerate(ordered[1:], start=1):
prev_name = ordered[i - 1].name
lines.append(
f" {prev_name} -> {stage.name}: {deltas[stage] * 1000:.3f} ms"
)
logger.info("\n".join(lines))
# ------------------------------------------------------------------
# Summary / reset / close
# ------------------------------------------------------------------
[docs]
def summary(self) -> dict[str, dict[str, float]]:
"""Return per-transition timing statistics.
Returns
-------
dict[str, dict[str, float]]
Mapping from ``"PREV_STAGE->STAGE"`` label to a stats dict
with keys ``mean_s``, ``std_s``, ``min_s``, ``max_s``,
``total_s``, ``n_samples``.
"""
result: dict[str, dict[str, float]] = {}
for idx, stage in enumerate(self._profiled_stages):
samples = self.timings[stage]
if not samples:
continue
prev_name = self._profiled_stages[idx - 1].name
label = f"{prev_name}->{stage.name}"
n = len(samples)
result[label] = {
"mean_s": statistics.mean(samples),
"std_s": statistics.stdev(samples) if n > 1 else 0.0,
"min_s": min(samples),
"max_s": max(samples),
"total_s": sum(samples),
"n_samples": float(n),
}
return result
[docs]
def reset(self) -> None:
"""Clear all accumulated timing data."""
for stage in self.timings:
self.timings[stage].clear()
self._step_cuda_events.clear()
self._step_cpu_timestamps.clear()
self._cuda_device = None
self._current_step = -1
self._backend_resolved = None
self._t0_ns = time.perf_counter_ns()
self._steps_recorded = 0
[docs]
def close(self) -> None:
"""Flush and close the CSV log file, if open."""
if self._csv_file is not None:
self._csv_file.close()
self._csv_file = None
self._csv_writer = None
def __del__(self) -> None:
self.close()
