fix: make dynamics hook lifecycle status-aware

[ys-teh]

ALCHEMI Toolkit Pull Request

Description

Make dynamics hook execution respect each dispatch's active graphs, align the hook lifecycle across standalone, fused, and domain-parallel dynamics, and initialize forces before the first integration step.

There will be a separate PR to deprecate FusedStage.fused_hooks, revise the order in which the FusedStage.hooks and substage hooks are applied, and add clear documentation.

Items deferred for future work:

• The way FreezeAtomsHook handles frozen atoms may still not be sufficient for every integrator. For example, thermostats and barostats may need an explicit interface to exclude frozen degrees of freedom from their internal statistics and state update.
• Force priming on DistributedPipeline and potential implication on DomainParallel.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

N/A

Changes Made

• Add active_graph_mask to dynamics hook contexts and propagate it through BaseDynamics, FusedStage, and DomainParallel hook dispatch.
• Restrict mutating and validation hooks to active graphs, and fix FreezeAtomsHook so constrained positions, velocities, and forces are handled at the correct split-update stages.
• Prime model forces before integration and expose AFTER_PRE_UPDATE/BEFORE_POST_UPDATE consistently for fused substages.

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Targeted tests run:

• uv run --extra cu13 pytest test/dynamics/ test/hooks/ test/distributed/test_domain_parallel.py test/distributed/test_domain_parallel_convergence.py (passed: 1246 passed, 8 skipped)

Full-suite attempt:

• make pytest (not completed: the first attempt could not collect without the declared ase extra; the rerun was stopped in favor of the scoped dynamics, hooks, and DomainParallel suites above)

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of the code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

The PR appears safe to merge.

Summary

• Propagates active-graph masks through hook contexts and limits mutating or validating hooks to active systems.
• Initializes model outputs before the first integration update and resets priming when batch composition changes.
• Adds fused-stage transition repriming and exposes split-update hook boundaries consistently.
• Validates resolved step counts before hook resources open, including the distributed path.

_{Reviews (12) · Last reviewed commit: "Merge branch 'main' into fix/fusedstage_..."}

[laserkelvin]

Requesting changes for five concrete lifecycle issues. Each inline comment includes a small reproducer and a focused fix.

[laserkelvin]

Can you add consistent jaxtyping for the active graph mask?

[ys-teh]

Added in 4afabe5

[laserkelvin]

While you're at it, could you add jaxtyping shape annotations for active_graph_mask, and converged_mask?

[ys-teh]

Added in 4afabe5

[laserkelvin]

blocking: This primes only the initial batch. After refill_check() replaces all graphs, the new Batch can lack forces and energy, and the next loop calls step() without allocating or priming them. A minimal repro with six two-atom samples, max_batch_size=3, refill_frequency=1, and immediate convergence calls _prime_forces once, then fails on the replacement with KeyError: expected='forces'. Please allocate required output buffers after refill, prime the replacement before its first pre_update, and add this case to test_inflight.py.

[ys-teh]

Thanks for flagging this. It's now fixed. Relevant test: test_refill_allocates_and_primes_replacement_outputs.

I ended up making a bigger change because I realized graphs can enter new stages at different iterations. Repriming the whole batch every time would add an extra model evaluation to many iterations. So now graphs that need repriming will use one iteration to do that by skipping pre-update and post-update. 2de2e76

In the same commit, I also fixed masked updates so inactive graphs won't change the sub-stage optimizer/integrator state, and added an optional reprime_on_entry for stages that need fresh model outputs before their first update.

One item to note is that batch.velocities may still be incorrect when a graph transitions between stages. Fixing this requires a larger refactor, so I’ve deferred it for now. As a temporary workaround, velocities can be adjusted with a hook.

[laserkelvin]

blocking: The final write uses only active_graph_mask, so it drops the periodic check above. With pbc=[[True, True, True], [False, False, False]] and both graphs active, this hook changes the second graph's cell and positions even though it is nonperiodic. Nonperiodic graphs should remain unchanged. Please add a mixed-PBC unit test.

[ys-teh]

Thanks for highlighting this. The old version (before active_graph_mask was added) also did not take into account the periodic mask, likely because it is rare to encounter cases with mixed boundary condition. I have updated it to account for the periodicity mask here regardless: 6891fd2

[laserkelvin]

blocking: Making active_graph_mask required leaves _test_prime_forces calling dd._prime_forces(local_batch) with the old signature. The two-GPU test will raise TypeError before it checks force priming. Please pass dd._active_graph_mask(local_batch) at that call site and run the multigpu test.

[ys-teh]

Updated in 021aa6d

[laserkelvin]

blocking: run() now initializes forces, but public step() does not, so the two entry points produce different first steps. Minimal repro: clone one batch with zero initial forces, call DemoDynamics.step() on one copy and DemoDynamics.run(..., n_steps=1) on the other; the positions and velocities differ. Since run() is documented as repeatedly calling step(), the first step should match. Please move force readiness into a shared first-step path and add this equality test.

[ys-teh]

I have now standardized the force priming behavior across BaseDynamics, FusedStage, and DomainParallel:

• self._forces_primed=False in .run before any .step.
• run self._prime_forces() and set self._forces_primed=True in the first step.
  I also removed the extra force priming in .run in DomainParallel since it is going to be called anyway in .step.

[laserkelvin]

issue: This override never calls _validate_n_steps(resolved). A small unit repro can force this branch with _dist_model=object() and _forces_primed=True; run(batch, n_steps=-1) then returns the unchanged batch with step_count == 0 instead of raising. DomainParallel should follow the same input contract as BaseDynamics and FusedStage. Please validate resolved before opening hooks and add negative/bool tests.

[ys-teh]

Added in 021aa6d

[laserkelvin]

Additionally want to probably update docs:

• docs/modules/dynamics/fused_stage.rst:73 still shows compute() followed by each stage’s combined masked_update(). PR 181 now performs masked pre_update, one shared compute, then masked post_update.
• docs/modules/dynamics/hooks.rst:319 says fused substages do not fire AFTER_PRE_UPDATE or BEFORE_POST_UPDATE. PR 181 adds both. The diagram below it also shows the old ordering.
• docs/modules/dynamics/implementing_dynamics.rst:223 teaches implementers that FusedStage calls one combined masked_update(). That is no longer how fused execution works.
• docs/modules/dynamics/fused_stage.rst:119 says n_steps is unused and fused termination is purely convergence-driven. PR 181 validates and uses n_steps as a maximum.
• docs/userguide/hooks.md:93 omits the new DynamicsContext.active_graph_mask field. It also still documents _build_context(batch) as the custom-context interface.
• docs/modules/dynamics/hooks.rst:242 describes FreezeAtomsHook as a two-stage hook. PR 181 makes it fire across five stages and changes when zero_forces=False exposes forces.
• The rendered API documentation will also inherit a stale source docstring: nvalchemi/dynamics/base.py:2890 still says fused execution computes first and then applies masked_update().

[ys-teh]

Additionally want to probably update docs:

• docs/modules/dynamics/fused_stage.rst:73 still shows compute() followed by each stage’s combined masked_update(). PR 181 now performs masked pre_update, one shared compute, then masked post_update.
• docs/modules/dynamics/hooks.rst:319 says fused substages do not fire AFTER_PRE_UPDATE or BEFORE_POST_UPDATE. PR 181 adds both. The diagram below it also shows the old ordering.
• docs/modules/dynamics/implementing_dynamics.rst:223 teaches implementers that FusedStage calls one combined masked_update(). That is no longer how fused execution works.
• docs/modules/dynamics/fused_stage.rst:119 says n_steps is unused and fused termination is purely convergence-driven. PR 181 validates and uses n_steps as a maximum.
• docs/userguide/hooks.md:93 omits the new DynamicsContext.active_graph_mask field. It also still documents _build_context(batch) as the custom-context interface.
• docs/modules/dynamics/hooks.rst:242 describes FreezeAtomsHook as a two-stage hook. PR 181 makes it fire across five stages and changes when zero_forces=False exposes forces.
• The rendered API documentation will also inherit a stale source docstring: nvalchemi/dynamics/base.py:2890 still says fused execution computes first and then applies masked_update().

Docs are deferred to #182 to prevent merge conflict.

[dallasfoster]

This test is not sufficient to cover new code changes. Please consider adding a hook and testing, like so:

seen = []

class _Probe:
    stage, frequency = DynamicsStage.AFTER_STEP, 1
    def __call__(self, ctx, stage):
        seen.append(ctx.active_graph_mask)

dd.register_hook(_Probe())
local_batch, _ = dd.step(local_batch)

mask = seen[-1]
assert mask is not None
assert mask.shape == (local_batch.num_graphs,)
assert bool(mask.all())            # nothing has graduated yet

[dallasfoster]

Worth adding a HookScope.GLOBAL probe too, since that branch builds its
context from self._gather_all(batch) rather than the local batch — a different
object, and the one place a shape mismatch could appear.

[ys-teh]

Updated the test in e224792.

Note that active_graph_mask is mostly redundant for the current domain parallel use case, where each rank owns a shard of the same system rather than a subset of multiple graphs. So active_graph_mask should be None. If domain parallel later supports ranks carrying shards from multiple graphs with mixed active statuses, that behavior should be covered already but it may still need additional testing.

[dallasfoster]

The DomainParallel I believe has active_graph_mask: torch.Tensor | None. I would suggest aligning on a single definition/contract. I even believe that HookRegistryMixin uses kwargs for capturing something like this.

[ys-teh]

Updated to initialize active_graph_mask = None. e224792

[laserkelvin]

Two in-line comments: one on compute hooks firing twice at step 0 during force priming (with a small reproducer), one lint nit that will fail pre-commit CI.

[laserkelvin]

suggestion: Priming runs the model once here, then the first step runs it again — and both fire the BEFORE_COMPUTE/AFTER_COMPUTE hooks at step_count == 0. A hook that watches compute (e.g. an energy logger) sees step 0 twice. Greptile flagged the same thing above; confirming it with a runnable check:

import torch
from nvalchemi.data import AtomicData, Batch
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.dynamics.demo import DemoDynamics
from nvalchemi.models.demo import DemoModel, DemoModelWrapper

class Counter:
    stage = DynamicsStage.AFTER_COMPUTE
    frequency = 1
    def __init__(self):
        self.calls = []
    def __call__(self, ctx, stage):
        self.calls.append(ctx.step_count)

batch = Batch.from_data_list(
    [AtomicData(atomic_numbers=torch.tensor([6]), positions=torch.zeros(1, 3))]
)
batch.forces = torch.zeros(1, 3)
batch.energies = torch.zeros(1, 1)

dyn = DemoDynamics(model=DemoModelWrapper(DemoModel()), n_steps=2, dt=1.0)
hook = Counter()
dyn.register_hook(hook)
dyn.run(batch)

print(hook.calls)  # expect [0, 1] for two steps; got [0, 0, 1]

I assume the contract is one hook call per model compute, stamped with the step that compute belongs to. If the extra step-0 call is intended, one line in the step() docstring saying so would be enough; otherwise skip hook dispatch during priming.

[ys-teh]

Yes, I just updated the docstring to explain it. Previously I considered handling initial force priming as a dedicated iteration by masking the update phases and using that iteration’s normal compute, similar to transition repriming in the updated FusedStage. That would prevent the step count issue. But, that would change the existing first-step behavior in FusedStage and require a broader change, so I didn't do it.

[laserkelvin]

nit: This file currently fails ruff format --check — trailing whitespace here, plus a few blank-line spots — and docs/modules/dynamics/fused_stage.rst has a trailing space in the new section. make lint (pre-commit ruff-format + trailing-whitespace) will go red on CI. One make format pass clears all of it.

[ys-teh]

/ok to test 3a286fb