nvalchemi-dynamics-hooks

nvalchemi Hooks

Overview

Hooks are callbacks that fire at specific points during each workflow step. They observe or modify batch state without changing the engine itself. The hook system is framework-wide: the same Hook protocol works for dynamics and custom pipelines. Dynamics engines pass DynamicsContext; custom engines can pass HookContext or their own context subclass.

from nvalchemi.hooks import (
    BiasedPotentialHook,
    DynamicsContext,
    Hook,
    HookContext,
    HookRegistryMixin,
    NeighborListHook,
    WrapPeriodicHook,
)
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.dynamics.hooks import (
    EnergyDriftMonitorHook,
    LoggingHook,
    MaxForceClampHook,
    NaNDetectorHook,
    SnapshotHook,
    StageTimingHook,
    TorchProfilerHook,
)

---

Hook protocol

Any object with these attributes satisfies the Hook protocol (runtime-checkable):

class Hook(Protocol):
    frequency: int        # execute every N steps (1 = every step)
    stage: Enum | None    # stage enum value (None for stage-agnostic hooks)

    def __call__(self, ctx: HookContext, stage: Enum) -> None:
        """Called with a context snapshot and the current stage."""
        ...

A hook fires when step_count % hook.frequency == 0 (so all hooks fire at step 0), except ON_ADMISSION, which fires once per admission regardless of frequency.

HookContext — base snapshot shared by hook-enabled workflows:

@dataclass(kw_only=True)
class HookContext:
    batch: Batch              # current batch (all engines)
    model: BaseModelMixin | None = None
    global_rank: int = 0      # distributed rank
    workflow: Any = None      # back-reference to the engine

DynamicsContext — context passed by dynamics engines:

@dataclass(kw_only=True)
class DynamicsContext(HookContext):
    step_count: int = 0
    converged_mask: torch.Tensor | None = None
    active_graph_mask: torch.Tensor | None = None
    graduated_mask: torch.Tensor | None = None   # ON_GRADUATE only