Fix/dynamics hooks and order

[ys-teh]

ALCHEMI Toolkit Pull Request

Description

Fix dynamics hook dispatch so hooks are no longer silently skipped and their
ordering is consistent across BaseDynamics, FusedStage, and fused
sub-stages. This also adds the one-time ON_ADMISSION stage, provides clear
warnings for unsupported fused-level stages, and deprecates the separate
FusedStage.register_fused_hook() API.

#181 needs to be merged before this.

Deferred for future PRs:

• ON_ADMISSION stage on DistributedPipeline / DomainParallel
• Efficient force priming on DistributedPipeline when using in-flight batching
• Better handling of optimizer / integrator states with hooks, e.g., FreezeAtomsHook, for correct behavior.

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

Changes Made

• Complete the hook lifecycle for FusedStage, including previously missing hook calls. Warn when a hook is registered on FusedStage for a sub-stage-only lifecycle
  boundary instead of silently accepting a hook that will never run.
• Add DynamicsStage.ON_ADMISSION = -1 for validation, shape-dependent
  allocation, and other one-time batch setup before force priming. Admission
  remains outside the compiled fused step.
• Deprecate FusedStage.fused_hooks and FusedStage.register_fused_hook() in favor of the standard hook storage and register_hook().
• Update the dynamics lifecycle documentation, diagrams, API guidance, and changelog.

Testing

• uv run --extra cu13 pytest test/dynamics/test_base.py test/dynamics/test_single_loop.py test/dynamics/test_bias_hook.py test/dynamics/test_cell_align.py test/dynamics/test_freeze_hook.py test/dynamics/test_periodic_hook.py test/dynamics/test_safety_hooks.py test/hooks/test_context.py
  — 351 passed, 1 slow test skipped.
• pytest test/dynamics test/hooks/test_neighbor_list_hook.py test/hooks/test_neighbor_list_nve_conservation.py test/hooks/test_stage_timing_hook.py on 2 GPUs
  — 1109 passed, 2 skipped.
• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

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

Below are the rendered images in the updated documentation.

BaseDynamics workflow (docs/modules/dynamics/hooks.rst line 29):

FusedStage workflow (docs/modules/dynamics/hooks.rst line 341):

Overall workflow (docs/modules/dynamics/overview.rst line 27):

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 your 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.

[dallasfoster]

_ensure_admission_initialized is the only thing that dispatches the new stage,
and it is called from just two places — BaseDynamics.step and the FusedStage
step. DomainParallel overrides step and calls neither; ON_ADMISSION does
not appear anywhere in domain_parallel.py:

I think the fix is to call self._ensure_admission_initialized(batch) at the top
of DomainParallel.step, before the _dist_model is None fallback, so the
wrapper fires admission on its own registry once per system on either branch.

While you are there: DomainParallel.__init__ forwards only the model, so
self.exit_status defaults to 1 and is never taken from the wrapped dynamics,
while _active_graph_mask (domain_parallel.py:681) reads
self._dynamics.exit_status.

[ys-teh]

Thanks, it should now be resolved in f75d7b1. I added self._ensure_admission_initialized(batch) after the _dist_model is None fallback instead since there will be readmission with dynamics.run anyway.

[dallasfoster]

ON_ADMISSION is covered in test_single_loop.py and test_base.py, but the
three test files this PR touches carry zero multigpu marks, so none of the new
stage or reordering runs on the 2-GPU job.

domain_parallel.py never mentions
ON_ADMISSION, but it inherits the dispatch through its own _call_hooks
override — so under DD the new stage fires on a DP-specific path that computes
its mask differently and has no test at any world size.

Could we add one multigpu test that registers a counting hook on a
DomainParallel, runs a couple of steps at world_size=2, and asserts
ON_ADMISSION fired once per system with the expected maskT

This comment is similar to the hole in #181 that opens from the other side, so one test settles both.

[ys-teh]

Updated in f75d7b1.