diff --git a/cookbook/bulletin_hooks.py b/cookbook/bulletin_hooks.py index d2c0e744..1e742f6d 100644 --- a/cookbook/bulletin_hooks.py +++ b/cookbook/bulletin_hooks.py @@ -23,7 +23,7 @@ from typing import Any from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import parse_weight_identity +from stitch.protocol import BASE_VERSION, PointerRewind, parse_weight_identity from stitch.providers.modal import commit_volume, discover_flash_targets, volume_reloader, wake_targets @@ -56,23 +56,56 @@ def commit_and_wake( version = parse_weight_identity(Path(version_dir).name) rank = distributed_rank() - # Rank 0 owns the `latest` pointer and writes it before committing, so the + # Rank 0 owns the `latest` pointer and advances it before committing, so the # committed bulletin is self-consistent for the poll/startup path. The pointer # lives at the transport root (the Volume mount) and is self-identifying — # `/weight_vN` — while the trainer wrote the version dir under the run - # partition (update_weight_disk_dir = /), so a new run is a - # forward move of the pointer, never a colliding rewind. + # partition (update_weight_disk_dir = /). `advance` enforces the + # monotonic-within-run rule (the new run's first publish forks at base); a + # same-run rewind (e.g. a republish) is dropped rather than serving stale + # weights — never silently overwritten. if version is not None and rank in (None, 0): - FilesystemBulletinBoard(_transport_root(args), layout="slime").write_latest( - _run_id(args), version - ) + board = FilesystemBulletinBoard(_transport_root(args), layout="slime") + try: + board.advance(_run_id(args), version) + except PointerRewind: + logger.warning( + "publish of version %s would rewind latest; dropping (run %r)", + version, + _run_id(args), + exc_info=True, + ) + return commit_volume(_volume_name(args)) if version is None or rank not in (None, 0): return - # Waking warm containers is a best-effort latency optimization: a transient - # Modal control-plane error must not kill the training step — `latest` is - # already committed and sidecars self-sync on their next poll. + _best_effort_wake(args, version, app_name_env=app_name_env, cls_name_env=cls_name_env) + + +def claim_pool(args: Any, *, app_name_env: str, cls_name_env: str) -> None: + """Trainer launch hook (rank 0): claim the rollout pool for this run. + + Write the empty pointer ``/weight_v000000``, commit the Volume, and + wake the pool — so every replica (cold or already-warm on a finished run) + resets to base *before* the first delta publishes, instead of inferring the + reset from the first publish's run mismatch. ``run_id`` must be fresh per + launch (the run's epoch/fence token); claiming a run already at the pointer + raises :class:`PointerRewind`, which fails the launch loudly rather than + leaving the pool pinned to a dead incarnation's high-water mark. + """ + if distributed_rank() not in (None, 0): + return + board = FilesystemBulletinBoard(_transport_root(args), layout="slime") + board.claim(_run_id(args)) + commit_volume(_volume_name(args)) + _best_effort_wake(args, BASE_VERSION, app_name_env=app_name_env, cls_name_env=cls_name_env) + + +def _best_effort_wake(args: Any, version: int, *, app_name_env: str, cls_name_env: str) -> None: + """Nudge warm Flash containers to reconcile now. Best-effort: a transient + Modal control-plane error must not kill the training step — `latest` is + already committed and sidecars self-sync on their next poll/startup.""" try: app_name = getattr(args, "rollout_modal_flash_app_name", None) or os.environ[app_name_env] cls_name = getattr(args, "rollout_modal_flash_server_cls_name", None) or os.getenv( @@ -195,9 +228,21 @@ def _transport_root(args: Any) -> str: def _run_id(args: Any) -> str: - """The run partition (chain identity). Passed explicitly via custom_config, - falling back to the basename of the per-run write dir.""" - return str(getattr(args, "run_id", None) or Path(bulletin_root(args)).name) + """The run partition (chain identity), passed explicitly via custom_config. + + Required — never derived from the write-dir basename. A fresh run_id per + launch is the epoch/fence token that makes restart safe (a restart is just a + new epoch that claims and resets the pool); deriving it from a per-run dir + that a crash-restart could reuse is exactly the reuse hazard this design + removes, so a missing run_id is a launch misconfiguration, not a fallback. + """ + run_id = getattr(args, "run_id", None) + if not run_id: + raise ValueError( + "run_id is required (pass it via custom_config_path); the bulletin " + "hooks no longer derive it from the write-dir basename" + ) + return str(run_id) def _gate_board(args: Any) -> FilesystemBulletinBoard: diff --git a/cookbook/local_disagg/README.md b/cookbook/local_disagg/README.md new file mode 100644 index 00000000..9948802c --- /dev/null +++ b/cookbook/local_disagg/README.md @@ -0,0 +1,50 @@ +# local_disagg — minimal pool-claim harness + +The smallest possible disaggregated-rollout setup: **no Modal, no slime/miles, +no GPUs.** A local-filesystem bulletin board, an in-memory rollout pool, and a +trainer-as-writer, wired with the *same* claim / advance / reconcile primitives +the production cookbooks use. It exists so the pool-ownership invariants can be +exercised and iterated on in milliseconds. + +## The model + +One **trainer** call owns one **run**, which owns one **pool epoch**: + +``` +trainer.claim() -> board.claim(run_id) # write /weight_v000000 (empty) +trainer.publish() -> board.advance(run_id, N) # write /weight_vN, monotonic +replica.reconcile() -> WeightSyncManager.sync_to() # converge to latest (reset on run switch) +``` + +- The **bulletin board** is the single source of truth: a self-identifying + `latest` pointer `/weight_v{N}`. +- The **trainer** is the single writer. `claim` resets the pool to base for a + fresh run; `publish` advances monotonically within the run. Both go through + the board's guarded writers (`stitch.protocol.decide_pointer_move`), so a + reused `run_id` or a non-monotonic publish raises `PointerRewind` instead of + serving stale weights. +- Each **replica** is a pure reconciler — it reads `latest` and converges + (replay the chain forward, or reset-to-base then replay on a run switch). + +`run_id` is a per-launch epoch/fence token (default: a fresh `uuid4`). A restart +is just a new epoch that claims and resets the pool — there is no special restart +path, and a crash-restart can never reuse a `run_id` to resurrect a dead pointer. + +## Invariants (see `harness_test.py`) + +- A claim resets every replica to base (v0) under the new run's id. +- Within a run the pool converges to each published version, in order. +- A new run forks at base: the pool resets even from a higher prior version. +- Re-claiming a run already at the pointer (a reused `run_id`) is rejected as a + rewind; the correct restart mints a fresh `run_id`. +- A non-monotonic publish within a run is rejected. +- A late-joining (cold / scaled-up) replica reconciles to the *current* run. + +## Run it + +```bash +uv run python -m pytest cookbook/local_disagg/harness_test.py -q +``` + +To iterate by hand, build a board + trainer + pool and step through claim / +publish / reconcile; see `harness.py` for the (tiny) surface. diff --git a/cookbook/local_disagg/__init__.py b/cookbook/local_disagg/__init__.py new file mode 100644 index 00000000..f587cbb9 --- /dev/null +++ b/cookbook/local_disagg/__init__.py @@ -0,0 +1,9 @@ +"""Minimal, dependency-free disaggregated-rollout harness. + +No Modal, no slime/miles, no GPUs: a filesystem bulletin board, an in-memory +rollout pool, and a trainer-as-writer, wired with the *same* claim/advance/ +reconcile primitives the real cookbooks use. It exists to exercise — and pin +down with tests — the pool-claim invariants quickly and locally. + +See :mod:`cookbook.local_disagg.harness`. +""" diff --git a/cookbook/local_disagg/harness.py b/cookbook/local_disagg/harness.py new file mode 100644 index 00000000..cc814f61 --- /dev/null +++ b/cookbook/local_disagg/harness.py @@ -0,0 +1,150 @@ +"""A minimal in-memory disaggregated-rollout harness. + +Three pieces mirror the real cookbooks with none of the infrastructure: + +- :class:`MemoryEngine` — an in-memory rollout server. It "serves" exactly one + weight version; ``apply_manifest`` steps it forward one delta and ``reset`` + drops it back to base (v0). Stands in for the SGLang adapter. +- :class:`LocalReplica` — one ``WeightSyncManager`` driving one ``MemoryEngine`` + against a shared bulletin board. The pool is a list of these; each is a pure + reconciler that converges to ``latest`` on ``reconcile()``. +- :class:`LocalTrainer` — the single writer for one run. It ``claim``s the pool + (resets every replica to base) at launch, then ``publish``es monotonic delta + versions. One trainer ↔ one run ↔ one pool epoch. + +The trainer writes the *same* slime-layout chain the real disk-delta publisher +writes (``//weight_v{N}/model.safetensors.index.json``), so the +replicas reconcile through the production ``WeightSyncManager`` path — only the +engine and the transport (a local dir instead of a Modal Volume / S3) are fakes. +""" + +from __future__ import annotations + +import json +import uuid +from pathlib import Path + +from stitch.bulletin import FilesystemBulletinBoard +from stitch.protocol import BASE_VERSION, PointerMove, VersionManifest + + +class MemoryEngine: + """In-memory rollout engine: tracks the single weight version it serves. + + base is v0; each ``apply_manifest`` advances exactly one delta and ``reset`` + re-materializes base. ``applied`` / ``resets`` are recorded so tests can + assert the reconcile path (replay the chain, reset on a run switch). + """ + + backend = "memory" + + def __init__(self) -> None: + self.version = 0 + self.applied: list[int] = [] + self.resets = 0 + + async def prepare(self) -> None: + pass + + async def flush_cache(self) -> None: + pass + + async def apply_manifest(self, manifest: VersionManifest, version_path: str) -> None: + self.version = manifest.version + self.applied.append(manifest.version) + + async def reset(self) -> None: + self.version = BASE_VERSION + self.resets += 1 + + async def pause_generation(self) -> None: + pass + + async def continue_generation(self) -> None: + pass + + +class LocalReplica: + """One rollout-pool replica: a ``WeightSyncManager`` + its ``MemoryEngine``. + + Constructed lazily (the sync manager needs a running event loop), so the + pool can be sized before any reconcile. ``reconcile`` converges the replica + to the board's current ``(run_id, version)`` — following a run switch + (reset → replay) exactly like the production sidecar. + """ + + def __init__(self, board: FilesystemBulletinBoard) -> None: + from stitch.sync import WeightSyncManager + + self.engine = MemoryEngine() + self.manager = WeightSyncManager(board=board, engine=self.engine, commit_mode="in_place") + + async def reconcile(self) -> None: + await self.manager.sync_to() + + @property + def served_version(self) -> int: + return self.engine.version + + @property + def served_run_id(self) -> str | None: + return self.manager.current_run_id + + +class LocalTrainer: + """The single writer for one run: claim the pool, then publish deltas. + + ``run_id`` defaults to a fresh per-launch token (the epoch/fence id that + makes restart a clean new-run reset, never a colliding rewind). ``claim`` + writes the empty pointer; ``publish`` writes the next version dir and + advances the pointer. Both go through the board's guarded writers, so a + reused run_id or a non-monotonic publish raises rather than serving stale + weights. + """ + + def __init__(self, board: FilesystemBulletinBoard, run_id: str | None = None) -> None: + self.board = board + self.run_id = run_id or uuid.uuid4().hex[:12] + self.version = BASE_VERSION + + def claim(self) -> PointerMove: + move = self.board.claim(self.run_id) + self.version = BASE_VERSION + return move + + def publish(self) -> PointerMove: + """Materialize the next delta version and advance ``latest`` to it.""" + nxt = self.version + 1 + _write_version_dir(self.board.version_dir(self.run_id, nxt), version=nxt, base=self.version) + move = self.board.advance(self.run_id, nxt) + self.version = nxt + return move + + +def make_pool(board: FilesystemBulletinBoard, size: int) -> list[LocalReplica]: + return [LocalReplica(board) for _ in range(size)] + + +async def reconcile_pool(pool: list[LocalReplica]) -> None: + for replica in pool: + await replica.reconcile() + + +def open_board(root: str | Path) -> FilesystemBulletinBoard: + """The slime-layout board both trainer and pool share (run-scoped chains).""" + return FilesystemBulletinBoard(str(root), layout="slime") + + +def _write_version_dir(version_dir: Path, *, version: int, base: int) -> None: + """Write the slime disk-delta version dir the real publisher would write: + a canonical HF index whose metadata carries the version lineage.""" + version_dir.mkdir(parents=True, exist_ok=True) + (version_dir / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {"version": f"{version:06d}", "base_version": f"{base:06d}"}, + "weight_map": {"w": "model-00001-of-00001.safetensors"}, + } + ), + encoding="utf-8", + ) diff --git a/cookbook/local_disagg/harness_test.py b/cookbook/local_disagg/harness_test.py new file mode 100644 index 00000000..9617d97b --- /dev/null +++ b/cookbook/local_disagg/harness_test.py @@ -0,0 +1,183 @@ +"""Pool-claim invariants, exercised on the minimal in-memory harness. + +These are the guarantees the explicit claim/advance design is meant to provide, +phrased against the same primitives the real cookbooks use (board.claim / +board.advance + a WeightSyncManager pool). Each test reads as one invariant. +""" + +from __future__ import annotations + +import asyncio +import tempfile +import unittest + +from cookbook.local_disagg.harness import ( + LocalReplica, + LocalTrainer, + make_pool, + open_board, + reconcile_pool, +) +from stitch.protocol import PointerRewind + + +class PoolClaimInvariantTest(unittest.TestCase): + def test_claim_resets_pool_to_base(self) -> None: + """A claim drops every replica to base (v0) under the new run's id — + the explicit 'empty' starting state, written before any delta.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + trainer = LocalTrainer(board) + pool = make_pool(board, size=3) + + move = trainer.claim() + self.assertTrue(move.reset) + self.assertEqual(move.version, 0) + await reconcile_pool(pool) + + self.assertTrue(all(r.served_version == 0 for r in pool)) + self.assertTrue(all(r.served_run_id == trainer.run_id for r in pool)) + + asyncio.run(run()) + + def test_publish_advances_pool_monotonically(self) -> None: + """Within a run the pool converges to each published version in order.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + trainer = LocalTrainer(board) + pool = make_pool(board, size=2) + trainer.claim() + + for expected in (1, 2, 3): + self.assertEqual(trainer.publish().version, expected) + await reconcile_pool(pool) + self.assertTrue(all(r.served_version == expected for r in pool)) + + # The chain was replayed delta-by-delta, never skipped. + self.assertEqual(pool[0].engine.applied, [1, 2, 3]) + + asyncio.run(run()) + + def test_new_run_resets_pool_even_from_higher_version(self) -> None: + """A fresh run forks at base: the pool re-materializes base and replays + the new chain, even though the finished run reached a higher version.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + pool = make_pool(board, size=2) + + old = LocalTrainer(board) + old.claim() + for _ in range(5): + old.publish() + await reconcile_pool(pool) + self.assertTrue(all(r.served_version == 5 for r in pool)) + + new = LocalTrainer(board) + self.assertNotEqual(new.run_id, old.run_id) + new.claim() + new.publish() # the new run is only at v1 + await reconcile_pool(pool) + + self.assertTrue(all(r.served_version == 1 for r in pool)) + self.assertTrue(all(r.served_run_id == new.run_id for r in pool)) + # The drop from v5 to the new run's base went through an engine reset. + self.assertTrue(all(r.engine.resets >= 1 for r in pool)) + + asyncio.run(run()) + + def test_restart_with_reused_run_id_is_rejected_as_rewind(self) -> None: + """The restart hazard, made impossible: re-claiming a run already at the + pointer is a rewind, not a silent stale-pointer reuse. A restart must + mint a fresh run_id (a new epoch), which claims cleanly.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + crashed = LocalTrainer(board, run_id="run-fixed") + crashed.claim() + for _ in range(3): + crashed.publish() + + # Crash-restart that reused the same run_id would rewind latest + # (v3 -> v0) onto stale weights — rejected. + restarted_same = LocalTrainer(board, run_id="run-fixed") + with self.assertRaises(PointerRewind): + restarted_same.claim() + + # The correct restart is a new epoch: a fresh run_id claims clean. + pool = make_pool(board, size=2) + restarted_fresh = LocalTrainer(board) + restarted_fresh.claim() + await reconcile_pool(pool) + self.assertTrue(all(r.served_version == 0 for r in pool)) + self.assertTrue(all(r.served_run_id == restarted_fresh.run_id for r in pool)) + + asyncio.run(run()) + + def test_non_monotonic_publish_within_run_is_rejected(self) -> None: + """Within a run the pointer only moves forward; re-advancing to an + already-published version is a rewind.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + trainer = LocalTrainer(board) + trainer.claim() + trainer.publish() + trainer.publish() # at v2 + + with self.assertRaises(PointerRewind): + board.advance(trainer.run_id, 2) + with self.assertRaises(PointerRewind): + board.advance(trainer.run_id, 1) + + asyncio.run(run()) + + def test_late_joining_replica_reconciles_to_current_run(self) -> None: + """A replica that joins after the claim (a scaled-up / cold container) + converges to the *current* run's chain, never a finished run's pointer.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + + old = LocalTrainer(board) + old.claim() + for _ in range(4): + old.publish() + + new = LocalTrainer(board) + new.claim() + new.publish() + new.publish() # current run at v2 + + # Replica created only now, with no prior state. + latecomer = LocalReplica(board) + await latecomer.reconcile() + + self.assertEqual(latecomer.served_version, 2) + self.assertEqual(latecomer.served_run_id, new.run_id) + # The cold catch-up over the v1..v2 tail composes into a single + # engine apply at the target version (not one apply per + # intermediate version) — see WeightSyncManager._sync_once. + self.assertEqual(latecomer.engine.applied, [2]) + + asyncio.run(run()) + + def test_claim_requires_run_id(self) -> None: + """A claim must name its run (the per-launch epoch token); an empty run + id is a launch misconfiguration, not a usable claim.""" + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + with self.assertRaises(ValueError): + board.claim("") + + +if __name__ == "__main__": + unittest.main() diff --git a/cookbook/miles_disagg/hooks.py b/cookbook/miles_disagg/hooks.py index 67b8f3e4..9f75160c 100644 --- a/cookbook/miles_disagg/hooks.py +++ b/cookbook/miles_disagg/hooks.py @@ -9,21 +9,31 @@ from typing import Any from cookbook.bulletin_hooks import ( + claim_pool as _claim_pool, commit_and_wake as _commit_and_wake, gated_rollout_request_hook, ) +_APP_NAME_ENV = "MILES_DELTA_APP_NAME" +_CLS_NAME_ENV = "MILES_DELTA_SERVER_CLS_NAME" + + +def claim_pool(args: Any) -> None: + """Claim the rollout pool for this run at launch (resets every replica to base).""" + _claim_pool(args, app_name_env=_APP_NAME_ENV, cls_name_env=_CLS_NAME_ENV) + + def commit_and_wake(args: Any, version_dir: str, rollout_engines: list[Any]) -> None: """miles ``custom_delta_pre_push_path`` hook (publish-only, bulletin board).""" _commit_and_wake( args, version_dir, rollout_engines, - app_name_env="MILES_DELTA_APP_NAME", - cls_name_env="MILES_DELTA_SERVER_CLS_NAME", + app_name_env=_APP_NAME_ENV, + cls_name_env=_CLS_NAME_ENV, ) # Re-export for the trainer's custom_rollout_request_hook_path. -__all__ = ["commit_and_wake", "gated_rollout_request_hook"] +__all__ = ["claim_pool", "commit_and_wake", "gated_rollout_request_hook"] diff --git a/cookbook/miles_disagg/modal_train.py b/cookbook/miles_disagg/modal_train.py index bbed7d4e..c01aa7c6 100644 --- a/cookbook/miles_disagg/modal_train.py +++ b/cookbook/miles_disagg/modal_train.py @@ -21,6 +21,7 @@ import tempfile import uuid from pathlib import Path +from types import SimpleNamespace import modal import modal.experimental @@ -465,6 +466,16 @@ def train(self, experiment: str, payload: dict) -> None: helpers.prepare_miles_config(cfg, tempfile.mkdtemp()) cmd = helpers.build_train_cmd(cfg, MILES_ROOT) + # Claim the pool for this run *before* miles starts publishing: write the + # empty pointer (/weight_v000000) and wake the pool so every + # replica resets to base now, closing the window where a replica could + # reconcile to a finished run's stale high-water version. + from cookbook.miles_disagg import hooks + + hooks.claim_pool( + SimpleNamespace(update_weight_disk_dir=cfg.update_weight_disk_dir, **cfg.custom_config_path) + ) + print(f"Training {experiment}: nodes={N_TRAIN_NODES}, rollout_endpoint={cfg.rollout_endpoint_url}") print(f"Command: {cmd}") # Tee the full training output to a committed Volume file so failures are diff --git a/cookbook/slime_disagg/hooks.py b/cookbook/slime_disagg/hooks.py index 93ea67f9..d4b1814d 100644 --- a/cookbook/slime_disagg/hooks.py +++ b/cookbook/slime_disagg/hooks.py @@ -9,21 +9,31 @@ from typing import Any from cookbook.bulletin_hooks import ( + claim_pool as _claim_pool, commit_and_wake as _commit_and_wake, gated_rollout_request_hook, ) +_APP_NAME_ENV = "SLIME_DELTA_APP_NAME" +_CLS_NAME_ENV = "SLIME_DELTA_SERVER_CLS_NAME" + + +def claim_pool(args: Any) -> None: + """Claim the rollout pool for this run at launch (resets every replica to base).""" + _claim_pool(args, app_name_env=_APP_NAME_ENV, cls_name_env=_CLS_NAME_ENV) + + def commit_and_wake(args: Any, version_dir: str, rollout_engines: list[Any]) -> None: """SLIME ``custom_delta_pre_push_path`` hook (publish-only, bulletin board).""" _commit_and_wake( args, version_dir, rollout_engines, - app_name_env="SLIME_DELTA_APP_NAME", - cls_name_env="SLIME_DELTA_SERVER_CLS_NAME", + app_name_env=_APP_NAME_ENV, + cls_name_env=_CLS_NAME_ENV, ) # Re-export for the trainer's custom_rollout_request_hook_path. -__all__ = ["commit_and_wake", "gated_rollout_request_hook"] +__all__ = ["claim_pool", "commit_and_wake", "gated_rollout_request_hook"] diff --git a/cookbook/slime_disagg/modal_train.py b/cookbook/slime_disagg/modal_train.py index b4db9b5e..6ecd0468 100644 --- a/cookbook/slime_disagg/modal_train.py +++ b/cookbook/slime_disagg/modal_train.py @@ -17,6 +17,7 @@ import tempfile import uuid from pathlib import Path +from types import SimpleNamespace import modal import modal.experimental @@ -369,6 +370,16 @@ def train(self, experiment: str, payload: dict) -> None: helpers.prepare_slime_config(cfg, tempfile.mkdtemp()) cmd = helpers.build_train_cmd(cfg, SLIME_ROOT) + # Claim the pool for this run *before* slime starts publishing: write the + # empty pointer (/weight_v000000) and wake the pool so every + # replica resets to base now, closing the window where a replica could + # reconcile to a finished run's stale high-water version. + from cookbook.slime_disagg import hooks + + hooks.claim_pool( + SimpleNamespace(update_weight_disk_dir=cfg.update_weight_disk_dir, **cfg.custom_config_path) + ) + print( f"Training {experiment}: nodes={N_TRAIN_NODES}, rollout_endpoint={cfg.rollout_endpoint_url}" ) diff --git a/cookbook/standalone_rollouts/frontdoor.py b/cookbook/standalone_rollouts/frontdoor.py index 299180dc..09ac7e85 100644 --- a/cookbook/standalone_rollouts/frontdoor.py +++ b/cookbook/standalone_rollouts/frontdoor.py @@ -27,8 +27,10 @@ from typing import Any from stitch.protocol import ( + PointerRewind, RolloutPoolState, RolloutReplicaState, + decide_pointer_move, format_snapshot_identity, parse_weight_identity, ) @@ -54,6 +56,11 @@ def advance_latest_decision( Returns ``{"run_id": str|None, "version": int, "reset": bool}`` to accept, or ``{"error": {...}}`` to reject (``InvalidIdentity`` / ``WeightRewindRejected``). + The accept/reset/rewind call is :func:`stitch.protocol.decide_pointer_move`, + the same rule the bulletin-board publish path advances through; this wrapper + only parses the wire identity and shapes the error dict. An explicit claim is + just a signal at ``weight_v000000`` with a fresh ``run_id`` (a cross-run move, + so ``reset=True``). """ version = parse_weight_identity(identity) if version is None: @@ -63,22 +70,20 @@ def advance_latest_decision( "message": f"identity {identity!r} is not weight_v", } } - if request_run_id != current_run_id: - # New run: its version space restarts, so accepting it is not a rewind. - return {"run_id": request_run_id, "version": int(version), "reset": True} - if version <= current_version: + try: + move = decide_pointer_move( + current_run_id, current_version, run_id=request_run_id, version=version + ) + except PointerRewind as rewind: return { "error": { "type": "WeightRewindRejected", - "message": ( - f"latest is at version {current_version} (run {current_run_id!r}); " - f"refusing to rewind to {version}" - ), - "current_version": int(current_version), - "requested_version": int(version), + "message": str(rewind), + "current_version": rewind.current_version, + "requested_version": rewind.requested_version, } } - return {"run_id": request_run_id, "version": int(version), "reset": False} + return {"run_id": move.run_id, "version": move.version, "reset": move.reset} def pool_state_from_server_infos(infos: list[dict[str, Any]]) -> RolloutPoolState: diff --git a/cookbook/standalone_rollouts/frontdoor_test.py b/cookbook/standalone_rollouts/frontdoor_test.py index 439199f1..7d7523f4 100644 --- a/cookbook/standalone_rollouts/frontdoor_test.py +++ b/cookbook/standalone_rollouts/frontdoor_test.py @@ -37,6 +37,14 @@ def test_runless_layout_stays_monotonic(self) -> None: rewind = advance_latest_decision(None, 5, "weight_v000005", None) self.assertEqual(rewind["error"]["type"], "WeightRewindRejected") + def test_claim_is_a_base_version_signal_for_a_fresh_run(self) -> None: + # An explicit claim is just weight_v000000 with a fresh run id: a + # cross-run move, so it resets the pool to base before any delta. + self.assertEqual( + advance_latest_decision("run-a", 5, "weight_v000000", "run-b"), + {"run_id": "run-b", "version": 0, "reset": True}, + ) + def test_rejects_unparseable_identity(self) -> None: self.assertEqual( advance_latest_decision(None, 0, "base", None)["error"]["type"], "InvalidIdentity" diff --git a/cookbook/standalone_rollouts/modal_serve.py b/cookbook/standalone_rollouts/modal_serve.py index 9adec368..f06a963a 100644 --- a/cookbook/standalone_rollouts/modal_serve.py +++ b/cookbook/standalone_rollouts/modal_serve.py @@ -24,7 +24,6 @@ from cookbook.slime_disagg import helpers from cookbook.standalone_rollouts import frontdoor as frontdoor_mod from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import format_snapshot_identity from stitch.providers.modal import ( discover_flash_targets, resolve_flash_gateway_url, @@ -482,10 +481,10 @@ async def advance_to(run_id: str | None, version: int) -> None: # Singleton writer: a single small write is one atomic S3 PutObject, # so no rename dance is needed. The pointer is self-identifying # (`/weight_vN`), so a new run is a forward move (not a rewind) - # and there is no separate run pointer to flip. - (S3_TRANSPORT_MOUNT_PATH / "latest").write_text( - format_snapshot_identity(run_id, version), encoding="utf-8" - ) + # and there is no separate run pointer to flip. The monotonic/reset + # decision already ran in advance_latest_decision, so this writes the + # decided move through the same board the pool reconciles against. + board.write_latest(run_id, version) async def list_server_infos() -> list[dict]: targets = await asyncio.to_thread( diff --git a/cookbook/standalone_rollouts/provider_test.py b/cookbook/standalone_rollouts/provider_test.py index d94f0974..4af57fe0 100644 --- a/cookbook/standalone_rollouts/provider_test.py +++ b/cookbook/standalone_rollouts/provider_test.py @@ -89,9 +89,11 @@ async def run() -> None: ) await manager.startup_sync() - # Reconciled to the `latest` pointer, applying the chain in order. + # Reconciled to the `latest` pointer: the v1..v2 tail composes + # into a single engine apply at the target version (not one apply + # per intermediate version) — see WeightSyncManager._sync_once. self.assertEqual(manager.current_version, 2) - self.assertEqual(engine.applies, [1, 2]) + self.assertEqual(engine.applies, [2]) asyncio.run(run()) diff --git a/src/stitch/bulletin.py b/src/stitch/bulletin.py index 473a7ecd..b91a703c 100644 --- a/src/stitch/bulletin.py +++ b/src/stitch/bulletin.py @@ -9,8 +9,11 @@ from typing import Any, Protocol from stitch.protocol import ( + BASE_VERSION, + PointerMove, VersionManifest, atomic_write_text, + decide_pointer_move, format_snapshot_identity, parse_snapshot_identity, read_latest, @@ -29,6 +32,10 @@ def read_latest(self) -> tuple[str | None, int]: ... def write_latest(self, run_id: str | None, version: int) -> None: ... + def advance(self, run_id: str | None, version: int) -> PointerMove: ... + + def claim(self, run_id: str) -> PointerMove: ... + def version_dir(self, run_id: str | None, version: int) -> Path: ... def read_manifest(self, run_id: str | None, version: int) -> VersionManifest: ... @@ -99,11 +106,44 @@ def read_latest(self) -> tuple[str | None, int]: return (None, read_latest(self.root)) def write_latest(self, run_id: str | None, version: int) -> None: + """Overwrite the pointer unconditionally. Prefer :meth:`advance` / + :meth:`claim`, which enforce the monotonic-within-run rule; this raw + write is for callers that have already made the move decision.""" if self.layout == "slime": atomic_write_text(self.root / "latest", format_snapshot_identity(run_id, version)) else: write_latest(self.root, version) + def advance(self, run_id: str | None, version: int) -> PointerMove: + """Move ``latest`` to ``(run_id, version)`` under the shared monotonic + rule (see :func:`decide_pointer_move`), then write it. + + Returns the :class:`PointerMove` (``reset=True`` when the move crosses + runs); raises :class:`PointerRewind` on a same-run rewind. This is the + single guarded writer both cookbook patterns publish through, so neither + can silently rewind the pointer onto stale weights. + """ + current_run_id, current_version = self.read_latest() + move = decide_pointer_move( + current_run_id, current_version, run_id=run_id, version=version + ) + self.write_latest(move.run_id, move.version) + return move + + def claim(self, run_id: str) -> PointerMove: + """Claim the pool for a fresh run: advance to the empty pointer + ``/weight_v000000``, resetting every replica to base before any + delta is published. + + ``run_id`` is the run's epoch/fence token and must be unique per launch; + claiming a run already at the pointer is a rewind (a reused run_id after + a restart), which surfaces as :class:`PointerRewind` rather than leaving + the pool pinned to the dead incarnation's high-water mark. + """ + if not run_id: + raise ValueError("claim requires a run_id (the run's per-launch epoch token)") + return self.advance(run_id, BASE_VERSION) + def version_dir(self, run_id: str | None, version: int) -> Path: if self.layout == "slime": base = self.root / run_id if run_id else self.root diff --git a/src/stitch/bulletin_test.py b/src/stitch/bulletin_test.py index 7380ee67..c011a1b1 100644 --- a/src/stitch/bulletin_test.py +++ b/src/stitch/bulletin_test.py @@ -7,6 +7,7 @@ from stitch.bulletin import FilesystemBulletinBoard from stitch.protocol import ( + PointerRewind, VersionManifest, format_snapshot_identity, parse_snapshot_identity, @@ -106,6 +107,46 @@ def test_publish_manifest_advances_run_pointer(self) -> None: ) self.assertEqual(board.read_latest(), ("run-a", 1)) + def test_advance_is_monotonic_within_run_and_resets_across_runs(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + board = FilesystemBulletinBoard(Path(tmp), layout="slime") + + first = board.advance("run-a", 1) + self.assertEqual((first.run_id, first.version, first.reset), ("run-a", 1, True)) + self.assertFalse(board.advance("run-a", 2).reset) + self.assertEqual(board.read_latest(), ("run-a", 2)) + + with self.assertRaises(PointerRewind): + board.advance("run-a", 2) + # A rejected advance leaves the pointer untouched. + self.assertEqual(board.read_latest(), ("run-a", 2)) + + crossed = board.advance("run-b", 1) + self.assertTrue(crossed.reset) + self.assertEqual(board.read_latest(), ("run-b", 1)) + + def test_claim_writes_empty_pointer_and_rejects_reuse(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + board = FilesystemBulletinBoard(Path(tmp), layout="slime") + + move = board.claim("run-a") + self.assertTrue(move.reset) + self.assertEqual(board.read_latest(), ("run-a", 0)) + + board.advance("run-a", 1) + # Re-claiming the same run (a restart that reused its run_id) rewinds. + with self.assertRaises(PointerRewind): + board.claim("run-a") + # A fresh run id claims cleanly. + self.assertTrue(board.claim("run-b").reset) + self.assertEqual(board.read_latest(), ("run-b", 0)) + + def test_claim_requires_run_id(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + board = FilesystemBulletinBoard(Path(tmp), layout="slime") + with self.assertRaises(ValueError): + board.claim("") + def test_unknown_layout_rejected(self) -> None: with self.assertRaises(ValueError): FilesystemBulletinBoard("/tmp", layout="bogus") diff --git a/src/stitch/protocol.py b/src/stitch/protocol.py index 56440066..195199be 100644 --- a/src/stitch/protocol.py +++ b/src/stitch/protocol.py @@ -14,6 +14,11 @@ PROTOCOL_VERSION = 1 LATEST_FILE = "latest.json" +# The "empty" pointer a launch claims before publishing any delta: version 0 is +# the run's base (every run forks at base, so a chain starts at v1 on top of v0). +# A pool reconciling to ``/weight_v000000`` resets to base weights. +BASE_VERSION = 0 + class SyncState(str, Enum): IDLE = "IDLE" @@ -403,6 +408,73 @@ def parse_snapshot_identity(text: str) -> tuple[str | None, int]: return (run_id, version) +class PointerRewind(Exception): + """A pointer move would rewind ``latest`` within the same run. + + The only writer of ``latest`` advances it monotonically *within a run*; a + move to an equal-or-lower version on the same run (e.g. a restarted trainer + that reused its run_id, or a duplicate claim) is rejected as a rewind rather + than silently serving stale weights. Crossing to a *different* run is not a + rewind — it forks at base (see :func:`decide_pointer_move`). + """ + + def __init__( + self, *, run_id: str | None, current_version: int, requested_version: int + ) -> None: + super().__init__( + f"latest is at version {current_version} (run {run_id!r}); " + f"refusing to rewind to {requested_version}" + ) + self.run_id = run_id + self.current_version = int(current_version) + self.requested_version = int(requested_version) + + +@dataclass(frozen=True) +class PointerMove: + """An accepted move of the ``latest`` pointer. + + ``reset`` is True when the move crosses to a different run, i.e. the pool + must re-materialize base and restart the version space (a claim, or a new + run's first publish); False for an ordinary monotonic advance within a run. + """ + + run_id: str | None + version: int + reset: bool + + +def decide_pointer_move( + current_run_id: str | None, + current_version: int, + *, + run_id: str | None, + version: int, +) -> PointerMove: + """Decide whether the single ``latest`` writer may move to ``(run_id, version)``. + + The one rule both the trainer-as-writer (bulletin board) and the + frontdoor-as-writer (hot-load API) paths share, so they can't bake in + divergent semantics: + + - A *different* run forks at base, so its version space restarts; accepting + it (even at a lower number, including the ``BASE_VERSION`` claim) is a + reset, not a rewind. + - Within the same run (and the run-less layout where both ids are None) the + move must be strictly newer; otherwise :class:`PointerRewind` is raised. + """ + version = int(version) + if run_id != current_run_id: + return PointerMove(run_id=run_id, version=version, reset=True) + if version <= int(current_version): + raise PointerRewind( + run_id=current_run_id, + current_version=current_version, + requested_version=version, + ) + return PointerMove(run_id=run_id, version=version, reset=False) + + def version_dir(root: str | Path, version: int) -> Path: return Path(root) / "versions" / weight_identity(version) diff --git a/src/stitch/protocol_test.py b/src/stitch/protocol_test.py index 37abe901..97acf765 100644 --- a/src/stitch/protocol_test.py +++ b/src/stitch/protocol_test.py @@ -6,11 +6,14 @@ from stitch.bulletin import FilesystemBulletinBoard from stitch.protocol import ( + BASE_VERSION, Artifact, + PointerRewind, RolloutPoolState, RolloutReplicaState, VersionManifest, WeightVersionPolicy, + decide_pointer_move, evaluate_version_policy, parse_weight_identity, read_latest, @@ -18,6 +21,41 @@ ) +class DecidePointerMoveTest(unittest.TestCase): + """The single accept/reset/rewind rule both the bulletin-board publish path + and the frontdoor hot-load path share (so they can't diverge).""" + + def test_forward_within_run_is_a_non_reset_advance(self) -> None: + move = decide_pointer_move("run-a", 4, run_id="run-a", version=5) + self.assertEqual((move.run_id, move.version, move.reset), ("run-a", 5, False)) + + def test_same_or_lower_version_within_run_rewinds(self) -> None: + with self.assertRaises(PointerRewind): + decide_pointer_move("run-a", 5, run_id="run-a", version=5) + with self.assertRaises(PointerRewind) as cm: + decide_pointer_move("run-a", 5, run_id="run-a", version=3) + self.assertEqual(cm.exception.current_version, 5) + self.assertEqual(cm.exception.requested_version, 3) + + def test_different_run_forks_at_base_as_a_reset(self) -> None: + # A new run is accepted even at a lower version (its space restarts) ... + move = decide_pointer_move("run-a", 5, run_id="run-b", version=1) + self.assertEqual((move.run_id, move.version, move.reset), ("run-b", 1, True)) + # ... including the empty BASE_VERSION claim. + claim = decide_pointer_move("run-a", 5, run_id="run-b", version=BASE_VERSION) + self.assertEqual((claim.run_id, claim.version, claim.reset), ("run-b", 0, True)) + + def test_first_claim_against_empty_pointer_is_a_reset(self) -> None: + move = decide_pointer_move(None, 0, run_id="run-a", version=BASE_VERSION) + self.assertEqual((move.run_id, move.version, move.reset), ("run-a", 0, True)) + + def test_runless_layout_keeps_monotonic_cas(self) -> None: + move = decide_pointer_move(None, 2, run_id=None, version=3) + self.assertFalse(move.reset) + with self.assertRaises(PointerRewind): + decide_pointer_move(None, 3, run_id=None, version=3) + + class ProtocolTest(unittest.TestCase): def test_manifest_round_trips_extended_and_legacy_fields(self) -> None: with tempfile.TemporaryDirectory() as tmp: