Hooks — Core Framework

The nvalchemi.hooks package provides the general-purpose hook system used across all nvalchemi workflows (dynamics, training, custom pipelines). It defines the protocol, context dataclasses, registry, and a set of hooks that are useful regardless of the specific engine type.

See also

• User guide: Hooks — conceptual overview and usage patterns.

• Dynamics hooks: Dynamics Hooks — Stages & Usage — hooks and stages specific to dynamics simulations.

• Training update hooks: Training update hooks — update-stage ownership, veto semantics, and constraints for training hooks.

The Hook protocol

Hook is a runtime_checkable Protocol. Any object that exposes the three required members — stage, frequency, and __call__ — is a valid hook, with no subclassing required:

from enum import Enum
from nvalchemi.hooks import Hook, HookContext

class MyHook:
    """A minimal custom hook — no inheritance required."""

    stage: Enum
    frequency: int = 1

    def __call__(self, ctx: HookContext, stage: Enum) -> None:
        print(f"graphs={ctx.batch.num_graphs}, stage={stage.name}")

Because Hook is a runtime_checkable Protocol, you can also use it as a type hint and check membership with isinstance:

assert isinstance(MyHook(), Hook)  # True ✓

CheckpointableHook

CheckpointableHook is an optional second protocol for hooks that carry restart-critical runtime state. Hook is required by every hook; CheckpointableHook is opt-in — only hooks that implement both state_dict() and load_state_dict() participate in checkpoint save and restore. Developers should meet the checkpointable protocol if the hook has state that needs to be persisted and restartable.

The two required methods:

• state_dict() -> dict — return a serializable snapshot of all runtime state that must survive a restart: accumulated counters, learned parameters, and runtime tensors. Do not include configuration already captured by the constructor; those are restored at construction time.

• load_state_dict(state: Mapping) -> None — restore state from a state_dict() snapshot. Validate critical configuration fields (such as decay rate or step frequency) before restoring runtime tensors to catch checkpoint/config mismatches early.

The training checkpoint loader calls state_dict() on every hook that satisfies this protocol and stores the results alongside model and optimizer state. On resume via TrainingStrategy.load_checkpoint(path, hooks=[...]), load_state_dict() is called on each matching hook by class name. Hooks that do not implement CheckpointableHook are silently skipped; they restart from their initial state, which is correct for purely stateless hooks.

isinstance(hook, CheckpointableHook) is True for any hook that provides both methods, with no subclassing required.

See also

Training update hooks — a concrete CheckpointableHook pattern with Pydantic fields and private runtime tensors.

Context dataclasses

Every hook receives a HookContext or a workflow-specific subclass. The base dataclass contains only fields shared by all hook-enabled engines; specialized contexts add fields that are meaningful only for one workflow category.

HookContext (base, all engines)

Field	Type	Description
batch	Batch | None	Current batch being processed. None is used for lifecycle stages that run before the first batch is available.
model	BaseModelMixin | None	Model being used (if applicable).
global_rank	int	Distributed rank of this process.
workflow	Any	Back-reference to the engine running the hooks. None when the workflow does not inject itself.

DynamicsContext (dynamics workflows)

Field	Type	Description
step_count	int	Current dynamics step number.
converged_mask	torch.Tensor | None	Boolean mask of samples that converged at the current hook stage. None when convergence has not fired for this dispatch.

TrainContext (training workflows)

Field	Type	Description
step_count	int	Current optimizer step number on this worker.
global_step_count	int	Current optimizer step number across all data-parallel workers.
batch_count	int	Number of training batches consumed, including batches whose optimizer step was skipped by update hooks.
epoch_step_count	int	Number of batches consumed within the current training epoch.
epoch	int	Current training epoch.
loss	torch.Tensor | None	Aggregate loss for the current step.
losses	dict[str, torch.Tensor] | None	Named loss components for the current step.
models	dict[str, BaseModelMixin] | ModuleDict | None	Models participating in the training step; this differs from the model attribute which is intended to represent a ‘main’ model in multi-model workflows. The key/model mapping should be semantic, e.g. ‘student’ and ‘teacher’ in distillation workflows, with ‘student’ being the intended ‘main’ model.
optimizers	list[torch.optim.Optimizer]	Optimizers participating in the training step. Empty when no optimizer is attached (e.g. eval-only or manually-driven hook contexts); TrainingUpdateOrchestrator and similar consumers treat an empty list as a no-op.
lr_schedulers	list[LRScheduler | None]	Learning rate schedulers participating in the training step. Aligned positionally with optimizers when populated; entries may be None when an optimizer has no scheduler. Empty when no scheduler is attached.
gradients	dict[str, torch.Tensor] | None	Parameter gradients for the current step.
grad_scaler	torch.amp.GradScaler | None	AMP gradient scaler for mixed-precision training; None when AMP is not in use.
validation	dict[str, Any] | None	Latest validation summary produced by the training strategy’s validation checkpoint (TrainingStrategy.validate()). None until validation has run or after the latest summary is consumed by metric-driven schedulers. In distributed runs, the reduced summary is available on every rank.

Registration and dispatch

Hooks are registered either at construction or manually via register_hook(). The HookRegistryMixin provides flat-list storage and dispatch logic for any engine.

# At construction (recommended for most cases)
engine = MyEngine(hooks=[MyHook()])

# Or register later
engine.register_hook(AnotherHook())

At each stage, all registered hooks for that stage fire in registration order, but only if step_count % hook.frequency == 0.

The dispatch logic for each hook is:

1. If the hook defines _runs_on_stage(stage) -> bool, call it.

2. Otherwise, check stage == hook.stage.

3. If matched, call hook(ctx, stage) with a fresh context object.

Note

At step_count == 0 all hooks fire (since 0 % n == 0 for any n).

Stage enums and multi-stage hooks

Each workflow engine fires hooks at named lifecycle points defined by a stage enum. The two built-in enums are:

• DynamicsStage — 9 stages from BEFORE_STEP through ON_CONVERGE. See Dynamics Hooks — Stages & Usage.

• TrainingStage — stages from SETUP through AFTER_TRAINING. See Training update hooks.

Custom pipelines may use any Enum type. For hooks that fire at more than one stage, define _runs_on_stage(stage) -> bool instead of a single stage attribute. Hooks that must support multiple enum types can overload __call__ with plum-dispatch; see Hooks.

General-purpose hooks

These hooks live in nvalchemi.hooks and work with any engine that uses the hook system, not just dynamics.

Hook	Purpose
NeighborListHook	Compute or refresh the neighbor list (MATRIX or COO format) with optional Verlet-skin buffering to skip redundant rebuilds. Fires at BEFORE_COMPUTE.
BiasedPotentialHook	Add an external bias potential (energy + forces) for enhanced sampling: umbrella sampling, metadynamics, steered MD, harmonic restraints, wall potentials.
WrapPeriodicHook	Wrap atomic positions back into the unit cell under PBC. Fires at AFTER_POST_UPDATE, respects per-system batch.pbc flags.
StageTimingHook	Measure elapsed time between hook stages, with optional NVTX ranges, CSV output, and console summaries.
TorchProfilerHook	Capture PyTorch profiler Chrome traces for training and dynamics through PhysicsNeMo’s profiler wrapper, with rank-specific output directories.

Reporting

ReportingOrchestrator is a standard hook that fans reporting events to a list of reporter objects at a configured stage and frequency. The Reporter protocol requires one method:

def report(ctx: HookContext, stage: Enum, state: ReportingState) -> None: ...

Two optional class attributes control distributed behavior:

• rank_zero_only = True — the orchestrator does not call the reporter on nonzero ranks.

• requires_all_ranks = True — all ranks participate in a collective reduction; only rank zero calls report() with the merged snapshot.

collect_scalars() assembles a ScalarSnapshot — a frozen payload of scalar values, counters, and metadata — from the current hook context. Custom reporters call it directly; built-in reporters call it internally.

TensorBoardReporter and RichReporter are provided implementations. RichLayout and BaseRichLayout control the dashboard surface for RichReporter.

See also

Reporting — setup, layout design, and custom reporters.

API Reference

Protocol

Hook	Protocol for hooks that observe or modify workflow state.
CheckpointableHook	Protocol for hooks that own restart-critical runtime state.
HookContext	Common context object passed to hooks.
DynamicsContext	Context object passed to dynamics hooks.
TrainContext	Context object passed to training hooks.
HookRegistryMixin	Mixin providing flat-list hook storage and dispatch.

General-purpose hooks

BiasedPotentialHook	Add an external bias potential to forces and energy after the forward pass.
NeighborListHook	Compute and cache neighbor lists before each model evaluation.
StageTimingHook	Per-stage timing hook for hook-enabled workflows.
TorchProfilerHook	Capture PyTorch profiler traces through PhysicsNeMo's profiler wrapper.
WrapPeriodicHook	Wrap atomic positions back into the simulation cell under PBC.

Reporting

Reporter	Protocol for reporting sinks owned by ReportingOrchestrator.
ReportingOrchestrator	Fan out hook contexts to reporting sinks.
ReportingState	Mutable state shared by a reporting orchestrator and its reporters.
ScalarSnapshot	Scalar reporting payload for one hook event.
collect_scalars	Collect scalar values from a hook context.
TensorBoardReporter	Write scalar reporting snapshots to TensorBoard.
RichReporter	Render scalar reporting snapshots as a live Rich dashboard.
RichLayout	Layout policy used by RichReporter.
BaseRichLayout	Reusable Rich dashboard layout for scalar tables and plot panels.
TrainingRichLayout	Rich dashboard layout for training workflows.
DynamicsRichLayout	Rich dashboard layout for dynamics workflows.