fix(mlx): reconcile #325 with main and restore recovered adapters

[erniesg]

Problem and changes

Follow-up to #325 and #134. MLX could recover a scenario step after restart while serving base weights, or expose a runtime ID inconsistent with resident weights.

This PR keeps the default no-op recovery hook on WeightRuntime, checks whether recovery needs bytes before downloading, restores MLX adapters before admission, and refuses metadata read from a changed durable head. Failed adapter loads keep admission closed. Composed artifacts remain lazy.

Refreshed against #325 at 44ca5bf9494388c7e009a63894fbd600ea685400, which includes main at b9d3ffb481e6f3a6be98da07f682d459d48e4727. The merge retains upstream experiment reporting and the no-download guard. Three new upstream fixtures now use class hook overrides recognized by that guard; all original concurrency and failure assertions remain intact. #827 remains superseded.

Compatibility and integration

Final head: 4fe6fc5644f6bf65490fa9df47871dbcbaa2e676. The recovery change adds no configuration, persisted-format, wire-contract, or dependency changes relative to the current MLX base. Existing runtimes inherit a no-op hook; interface docstrings describe the contract. The base requires Python 3.12.

Keep this stacked on worktree-issue-134. Per the base workflow, #325 should land first; then rebase this dependent PR onto resulting main and retarget it. A maintainer may explicitly choose to incorporate this head into #325 instead. This refresh does not modify #325 or change this PR's target. The base-is-main failure is expected while stacked.

Verification

• Python 3.12.8 full source suite: 6,614 passed, 118 skipped; 90.98% coverage. Command: PATH="$PWD/.venv/bin:$PATH" GIT_CONFIG_GLOBAL=/dev/null GIT_CONFIG_SYSTEM=/dev/null HF_HUB_OFFLINE=1 .venv/bin/python -m pytest tests -n4 --dist loadfile --cov=reef --cov-report=term -ra.
• Regression rerun: 134 passed, covering recovery and upstream experiment-reporting integration. Pre-commit passed for all files and again for the resolution/fixture files; mypy passed for 387 source files.
• 58 real MLX tests passed, using Metal, including engine/training tests on cached Qwen2.5-0.5B-Instruct-4bit. Runtime contract tests also use fake engines and are separate from this hardware result.
• Separate-process offline qualification: one real optimizer step, safetensors export, real Git LFS/SQLite persistence, durable-head repair, restored scenario step 1, exact equality of 16 adapter tensors and probe logits, and corrupt-load admission closure.
• Clean CPU wheel: 4,924 passed, 144 skipped, one failure in the unchanged upstream process-guard forced-parent-death test. It timed out with repeated signal-handler entry; the full source run passed it, and an isolated rerun passed all 6 process-guard tests. This is disclosed as an intermittent unresolved package-check failure, not a fully green package run. Wheel import, CLI, and dependency isolation smoke checks passed.

Python 3.10.18 and 3.11.13 package-install dry runs both reject Python>=3.12. Upstream CONTRIBUTING.md still requests these versions, while pyproject.toml, AGENTS.md, and CI require/select only 3.12. The current workflow does not expand its matrix on rerun. Maintainers need to reconcile that policy; no requirements were bypassed.

The 118 source skips include PostgreSQL, opt-in Ray, external harness binaries, and host-specific sandbox checks. Optimizer-state restart equivalence and a full post-refactor GSM8K stream remain unqualified. Adapter restoration does not restore process-local AdamW state or prove identical subsequent updates.

Latest-head CI, Docs, Docker, and harness smoke report action_required and need an authorized repository maintainer/admin to approve and run them. Prior checks on a3c2563 do not qualify this new head.

AI assistance

Codex assisted with integration, test-fixture adaptation, source/package checks, and real MLX qualification. The human contributor remains responsible for review and claims.

Checklist

• Existing PR and history preserved; scoped recovery behavior retained.
• Local source, policy, formatting, type, and real MLX checks reported.
• Latest-head hosted checks and required maintainer review complete.
• Base PR landed and stacked PR integration completed.

[erniesg]

Updated the final patch to 93665dad. Startup restoration is now an explicit MLX runtime capability, so it does not force unsupported checkpoint loading on executor-backed runtimes. Independent source review found no actionable defects; 355 focused tests passed. The full run was 4,075 passed, 58 skipped, and two previously reproduced base failures, detailed in the description.

Could a repository admin approve the final CI, Docker, and Docs runs? All report action_required with zero jobs; my approval attempt returned 403 requiring admin rights. I would appreciate the affected codeowner review and merge into worktree-issue-134 once required checks pass.

The description also links the separately reviewed current-main integration, including current objective/scheduling contracts and a passing real Metal lifecycle pilot. That integration has not been merged upstream.

[BobbyZhouZijian]

Hi @erniesg thanks for the PR. Could we use a default no-op startup recovery hook on WeightRuntime instead of adding another capability subclass? Also, check whether recovery is needed before materializing the checkpoint; currently, opted-out runtimes still incur the download and can fail startup.

[erniesg]

Thanks Bobby, both changes are implemented in a3c2563. WeightRuntime now provides a default no-op restore_recovered_checkpoint hook, with MLX opting in. Recovery checks happen before checkpoint materialization, so opted-out runtimes avoid the download.

Regression tests cover default/no-download behavior, lazy component recovery, and failed adapter loads keeping admission closed. I also reconciled the branch with current main and preserved the updated MLX interfaces.

The combined local suite passed: 6,553 tests, 118 skipped, with 90.94% coverage. Formatting and type checks passed, and a fresh independent review found no actionable code defects. Hosted CI still needs approval.