Move TrainerRank micro-batch planning into trainer_rank/_micro_batch_planner

[bradhilton]

Summary

Second of three stacked TrainerRank extractions (base: #999). Moves the 34 micro-batch planning / split / layout-selection / admission methods (~2,000 lines) from src/art/trainer_rank/_impl.py into src/art/trainer_rank/_micro_batch_planner.py, using the same pattern as #999 and the existing _gdn_memory.py/_planner_evidence.py: module-level functions def _name(self: TrainerRank, ...) with verbatim bodies, bound in the class body by assignment. No collaborator object, no delegators, no __init__ change; sibling calls stay self._x(...) so per-instance monkeypatches and TrainerRank.__new__-built ranks keep working, and co_name is preserved.

Moved: _forward_micro_batches, _planning_status (@contextmanager applied to the module function), _recover_admission (@_backward_region applied to the module function — region takes the rank positionally and reads _backward_work() at call time, so decoration site is irrelevant), _plan_structure, _flat_forward_plan, _split_forward_plan, _find_admissible_forward, _admit_plan, _admit_split_rung, _search_next_micro_batch, _select_next_micro_batch, _plan_flat_forward, _compute_group_layout, _select_group_layout, _submit_speculative_wave_planning, _fill_planner_snapshot, report_planner_oom, finish_planner_observation, and the admission/candidate/layout/wave/planner-report helpers. Full list: rg -n '= _micro_batch_planner\.' src/art/trainer_rank/_impl.py.

Deliberately left in _impl.py: forward_micro_batches (its four @overload stubs must be immediately followed by a real def; it dispatches to the moved implementation via self.), _begin_planner_observation (tests/unit/test_trainer_rank_backward_work.py::actual_rank parses _impl.py by name), the _planner_observation property/setter, planner_observation_scope, _telemetry_plan_signature, _execute_admitted_plan, _execute_flat_plan.

Names tests patch on _impl (torch, dist, time, _telemetry_phase, materialize_prefix_tree_layout, estimate_prefix_tree_packed_tokens) are read as _impl.X at call time. Four names that only the planner used and nothing patches (COEFFICIENT_VERSION_FALLBACK, build_canonical_prefix_tree, prefix_tree_layout_candidates, select_prefix_tree_layout) are now imported directly from their origin modules and dropped from _impl.py, along with the dead Generator, asdict, nullcontext, MethodType. Three _impl.py imports that look unused but are read as _impl.X by _memory.py/this module carry # noqa: F401 comments. Def-time = _impl.Unset defaults mean _impl.py imports the module immediately before class TrainerRank: (shared import line with #999).

_micro_batch_planner.py is added to _planner_misses._SOURCE_NAMES so the planner-miss replay drift guard fingerprints the planner code it reruns (total 551 KB < 1 MiB); the existing test asserts it.

Equivalence and stacking

• All 34 function bodies are AST-identical to main's current method bodies (modulo _impl. qualification, the self annotation, docstring dedent). The merged class has 123 methods = main's 183 − (Move TrainerRank memory accounting into trainer_rank/_memory #999's 26 ∪ these 34), with 60 bindings; every remaining method is AST-identical to main's (no stitched methods from the rebase); module level equals main's plus the pre-class import.
• Price warm waves by the fit later plans measured, not the first plan's #991 landed on main after authoring and modified _forward_micro_batches (the interval parameter, self.__dict__.get("_peak_reading") after the tracked flat-plan run, and caller_phase=True, interval=interval on _update_peak_memory_profile); those lines are ported into the moved function and verified against main's body. No other main hunk since the fork touched a moved method.
• All four import orders work.

Validation

• tests/unit/test_trainer_rank_*.py (minus gpu_ci, plus planner_reports and main's new profile_warm) at -n 4: 1,257 passed vs 1,258 on the Move TrainerRank memory accounting into trainer_rank/_memory #999 baseline; the one branch-only failure is a known distributed subprocess-timeout flake that passes serially (6/6). test_trainer_rank_backward_work.py + test_trainer_rank_planning_status.py serially: 32 passed.
• ruff check, ruff format --check, ruff --select F401 clean; ty check shows the same 6 pre-existing _checkpoint.py diagnostics as main.
• Independent review (thermo-nuclear rubric): both decorators analyzed against _backward_work.region's implementation; _impl.X table for all 36 names cross-checked against every _impl name patched in tests; three-way merge with Move TrainerRank memory accounting into trainer_rank/_memory #999 and the slots/optimizer branch dry-run. Bugbot found no bugs. Its two pre-merge items (stale _forward_micro_batches vs Price warm waves by the fit later plans measured, not the first plan's #991; missing _SOURCE_NAMES entry) and two Lows (time via _impl; import tidy) are addressed here.

Observable non-behavioral change: _planner_evidence.failure frame records for frames inside moved functions report module: art.trainer_rank._micro_batch_planner.

_impl.py: 9,843 → 6,706 lines with #999 and this PR. The slots/optimizer extraction follows on top and brings it to ~5,640.

[bradhilton]

Schulman — no concrete blocker found in the reviewed scope.

Reviewed head 35db95460bad370c1fd2222356e71a5327f96436 against stacked base f0a2d8863c91076a6b2529920e26a8b815aec928, including cumulative behavior from 3ad9a5f8d29ed4ea7f8e033e7e1f6cda3fe55607 and its relationship to #1001. Exact head/base were rechecked before posting. An independent delegate performed the focused review and I examined its report. This lane did not author the extraction; shared GitHub account attribution is not proof of agent authorship, and familiarity with earlier trainer fixes is separate.

All 34 moved planner methods, plus the 26 prior memory methods, match main after narrowly normalizing module qualification, self annotations and docstring indentation. Retained method bodies and class state match. Warm-profile pricing, optimizer-counter validation, distributed planning-error coordination, generator cleanup and graph/slot lifetime behavior survive. Dynamic runtime globals remain accessed through _impl, and sibling calls retain instance dispatch. Packaging covers both helper modules; the replay source fingerprint includes their estimator code and remains below its size cap.

Bounded framework-free controls passed for actual extracted bindings, sentinel identity, cancellation/error preservation, late distributed-module replacement, slot dispatch and optimizer-construction rollback. Independent AST/source comparisons passed. No full framework import, model, real optimizer numerics or GPU/distributed execution was performed by this review; reported CI success was observed, not independently rerun.

Limits: function/module/frame metadata and private patch locations change. Frozen code-identity diagnostics need explicit adoption. Source preservation supports this refactor but is not a claim of identical measured gradients, memory peaks or compilation. A reviewer-only metadata-collection failure under an address-space cap is retained; no product source change or product-test failure resulted.

[bradhilton]

Dalinar: memory/planner owner review. No blocking findings.

Reviewed head 35db95460bad370c1fd2222356e71a5327f96436 against its stacked base #999 (f0a2d8863). I also checked the cumulative change against main 0e0c31b32 / c0b1296dc, whose _impl.py is byte-identical to #999's base. I didn't author this PR; its commits carry the shared git identity.

What I checked

• Stacked delta.
  • All 34 moved functions are AST-identical to Move TrainerRank memory accounting into trainer_rank/_memory #999's methods, after normalizing _impl.X, the self annotation and docstrings.
  • Decorators match: @contextmanager on _planning_status and @_backward_region on _recover_admission.
  • Every method left in the class is unchanged, and _memory.py is byte-identical to Move TrainerRank memory accounting into trainer_rank/_memory #999's.
• Cumulative change. Because Move TrainerRank memory accounting into trainer_rank/_memory #999's moved bodies match main, all 60 extracted bodies equal main's current methods. That includes _forward_micro_batches with Price warm waves by the fit later plans measured, not the first plan's #991's interval, caller-phase and peak-reading changes.
• Decoration at module level. _backward_work.region returns a closure that takes the rank positionally and reads rank._backward_work() at call time. contextmanager builds a fresh context manager per call. Decorating the module function and binding it in the class therefore behaves the same. Only __qualname__/__module__ change, and nothing in src/ or the tests reads them for TrainerRank.
• Globals.
  • _micro_batch_planner.py has no unresolved runtime names.
  • time, _telemetry_phase, materialize_prefix_tree_layout, estimate_prefix_tree_packed_tokens, torch and dist are all read through _impl. at call time. Bare torch appears only in annotations, which are unevaluated.
  • So patches on _impl, including the patch.object(_impl, name, …) loop in test_trainer_rank_cache_recovery.py, still intercept.
  • Observers that replace _impl._telemetry_phase and rank._select_next_micro_batch keep working, because _forward_micro_batches dispatches self._select_next_micro_batch.
• Dropped _impl imports. Nothing in ART's src/ or tests reads the eight dropped names through _impl, whether by attribute or by a from … _impl import line. Caladan's source imports only TrainerRankSlotStateError and _CheckpointSlot from _impl; both stay.
• Import order. Importing _micro_batch_planner, _memory, _impl or art.trainer_rank first all work, and the class bindings are the module functions.
• Source fingerprints. _micro_batch_planner.py is in _SOURCE_NAMES. The fingerprinted total is 542 KiB of the 1 MiB limit (main: 545). Replay still builds TrainerRank.__new__ and reaches the moved estimator through self. dispatch.
• Packaging. The wheel packages src/art whole, and the GPU workflow's path filter covers trainer_rank/.

Validation

• Full trainer-rank CPU suite on this head merged with main c0b1296dc: 1482 passed, 8 skipped, 0 failed. That's the same count as Move TrainerRank memory accounting into trainer_rank/_memory #999 merged with main and as main 0e0c31b32.
• ty and ruff are clean on the changed files. Repo-wide F401 flags only a pre-existing unused Literal in _prefix_tree_planner.py.
• A private Fable review with source access found no blocking issues. Its lows are the three below.

Low (none blocking)

1. Planner-miss reports record frame.f_globals["__name__"] per traceback frame. Frames inside moved bodies now report art.trainer_rank._memory or ._micro_batch_planner instead of ._impl. No in-tree consumer or Caladan code filters on the module string, and the PR description calls this out.
2. __qualname__/__module__ change for the moved methods, including the public report_planner_oom, finish_planner_observation and discard_planner_observation. This is cosmetic (repr/autodoc); nothing reads them.
3. test_trainer_rank_backward_work.py::actual_rank AST-extracts eight methods from _impl.py by name. All eight stay defs here, but a later extraction that moves any of them will break that test at load. Worth remembering for the slots/optimizer PR.

Held planner-stack integration (my held drafts #963, #978, #981 and #986)

• The stack edits three methods this PR moves (_plan_cost, _split_chunk_lower_cost, _fill_planner_snapshot), on top of the four Move TrainerRank memory accounting into trainer_rank/_memory #999 moves.
• When it rebases, those edits need a hand port into _micro_batch_planner.py. A conflict resolution that keeps this PR's side would drop them silently.
• This review lifts no hold.

Limits: CPU tests only; no GPU smoke run. Behavioral equivalence comes from AST comparison and static reading.