DiLoCo: persist work-unit dispatch state across server restart (#105)

[jdinalt]

Closes #105.

Problem

The DiLoCo server is the authority for which dataset rows each worker has consumed — the worker keeps no dataset-progress state of its own, by design, because the server was supposed to track and persist it. But the server's work-queue state was never written to server_state.pt, so a server restart rebuilt empty queues and re-issued already-trained units within the epoch.

The "intentionally not persisted (#46)" comments were doubly wrong: issue #46 is unrelated (torchao quantization), and the queue was never actually persisted in the committed work-unit-dispatch feature (#80) — it was an unfinished/untested piece with a fabricated rationale citation, not a sanctioned decision.

Change

• Persist + restore the per-(dataset_id, shuffle_seed) issued/completed bitmaps + counters and the _dataset_lengths integrity snapshot in server_state.pt (keyed on disk by a "dataset_id|seed" string; bitmaps as bytes). A worker re-registering its dataset reuses the restored queue (_handle_register_dataset already reuses any existing queue for the key), so issuance resumes at the next un-issued unit, not unit 0.
• Cross-experiment safety is intrinsic (the original ghost-queue worry): a changed dataset hashes to a new dataset_id → fresh key → stale queues are never matched and sit inert; the 409 length-mismatch guard handles same-id/different-length; a hard reset is "restart from model weights + purge output_dir". No expiry machinery needed.
• Flush on graceful shutdown: stop() now saves, and run() handles SIGTERM (how the scheduler stops a server job) in addition to SIGINT — so a clean stop doesn't lose units issued since the last autosave.
• Removes the incorrect #46 citations from server.py and the two docs; rewrites the design-doc restart section and the diloco.md crash-recovery note to describe the persisted behavior as-is.

Audit note

This came out of a full audit of what the server does/doesn't persist (requested on #105). Conclusion: the optimizer/model side (weights + outer-optimizer momentum + _sync_round + param ordering + known_workers) was already covered; the work-unit dispatch subsystem was the one correctness gap — everything else unpersisted is transient in-flight (re-synced by workers) or pure stats.

Testing

• Rewrote the (now-inverted) persistence tests: round-trip of bitmaps/counters/dataset_lengths + persisted-file keys, resume-at-next-unit after restart, and malformed-key resilience.
• Full tests/unit/ml/diloco/ green (317).

🤖 Generated with Claude Code