Explicit pool claim/reset: one trainer call owns one pool epoch

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 —
     # `<run_id>/weight_vN` — while the trainer wrote the version dir under the run
-    # partition (update_weight_disk_dir = <root>/<run_id>), so a new run is a
-    # forward move of the pointer, never a colliding rewind.
+    # partition (update_weight_disk_dir = <root>/<run_id>). `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 ``<run_id>/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:

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 <run_id>/weight_v000000 (empty)
+trainer.publish()    -> board.advance(run_id, N) # write <run_id>/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 `<run_id>/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.

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`.
+"""

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 (``<root>/<run_id>/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",
+    )