feat: load FFT revisions onto sampling pool replicas

src/server/sampling_dispatcher.py

@@ -5,6 +5,7 @@
 
 import asyncio
 import contextlib
+import json
 import logging
 import os
 import signal
@@ -24,6 +25,10 @@
 RETRY_MS = 1000
 # A request no replica has answered for this long fails, so its client stops waiting.
 SAMPLE_DEADLINE_SECONDS = 3600
+# The version a replica reports once its weights may not match any revision. Never a path, so no pool revision equals it.
+UNKNOWN_WEIGHTS = "unknown"
+# How long one /weight_info call may take.
+WEIGHT_INFO_TIMEOUT_SECONDS = 5
 
 
 def tmp_root() -> str:
@@ -46,6 +51,22 @@ def failed(message: str) -> dict[str, str]:
   return {"type": "RequestFailedResponse", "error_message": message}
 
 
+def delta_applies(revision: str, current: str) -> bool:
+  """Whether weights at `current` can load `revision`. A delta applies on any save at the weights it starts from or
+  reaches. Deltas written before the trainer listed them name one parent."""
+  metadata = json.loads(Path(revision, "metadata.json").read_text())
+  bases = [base or "default" for base in metadata.get("bases", [metadata.get("parent")])]
+  return metadata.get("format") != "sparse_delta" or current in bases
+
+
+class RevisionUnloadable(Exception):
+  """No replica holds `revision`, a sparse delta, or weights it applies on. `held` is the weights they hold."""
+
+  def __init__(self, revision: str, held: list[str]) -> None:
+    super().__init__(f"{revision} does not apply on {held}")
+    self.revision, self.held = revision, held
+
+
 def to_result(body: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
   """The old sampler's result shape, from a vLLM /inference/v1/generate response."""
   logprobs = [[e["logprob"] for e in c["logprobs"]["content"]] for c in body["choices"]]
@@ -84,12 +105,32 @@ async def run(self) -> None:
 
   async def dispatch(self) -> None:
     await self.read()
-    if await self.sampler.ensure_weights(self.revision):
+    if await self.ready():
       # Read after the await above, so a request whose result was saved meanwhile is not sent again.
       for r in self.to_send()[: self.concurrency - len(self.in_flight)]:
         self.in_flight[r["request_id"]] = task = asyncio.create_task(self.send(r))
         task.add_done_callback(lambda _, request_id=r["request_id"]: self.in_flight.pop(request_id))
 
+  async def ready(self) -> bool:
+    """Whether the sampler serves the current revision, loading it if it can. A revision it can never load fails its
+    records, and the next read moves the pool on."""
+    revision = self.revision
+    if revision and not Path(revision, "metadata.json").is_file():
+      # Retention can delete the revision before or after the switch, and no sampler that lacks it can load it again.
+      await self.fail_revision(revision, f"Sampler weights {revision} no longer exist")
+      return False
+    try:
+      return await self.sampler.ensure_weights(revision)
+    except RevisionUnloadable as exc:
+      held = ", ".join(exc.held)
+      message = f"Sampler weights {exc.revision} cannot be loaded on top of the weights the pool holds ({held}). Save full weights to sample them."
+      await self.fail_revision(exc.revision, message)
+      return False
+
+  async def fail_revision(self, revision: str, message: str) -> None:
+    for r in [r for r in self.records.values() if r["revision"] == revision]:
+      await self.set_result(r, failed(message))
+
   async def read(self) -> None:
     """Reads new records, fails those past their deadline, and switches revision once nothing pending names the
     current one."""
@@ -138,37 +179,137 @@ def __init__(self, http: httpx.AsyncClient, pool: str, router: str, urls: list[s
     self.http, self.pool, self.router, self.urls = http, pool, router, urls
     self.revision: str | None = None
     self.k8s = None
+    self.loading: dict[str, asyncio.Task] = {}
+    self.checking: dict[str, asyncio.Task] = {}
+    self.paused: dict[str, float] = {}  # when each replica last paused to load weights
+    self.checked: set[str] = set()  # replicas checked for a pause or an open update, checked again after a load of ours breaks off
 
   async def ensure_weights(self, revision: str) -> bool:
-    """True once one replica serves `revision`."""
+    """Starts a load on every checked replica not at `revision`. True once one replica serves it. Raises
+    RevisionUnloadable when every replica answers, none holds weights `revision` applies on, and none is loading."""
     if revision != self.revision:
       self.revision = revision
       self.write_revision()
     target = revision or "default"
     urls = await self.list_replicas()
-    versions = await asyncio.gather(*(self.weight_version(url) for url in urls), return_exceptions=True)
-    return target in versions
+    for url in urls:
+      if url not in self.checked and url not in self.loading and url not in self.checking:
+        self.checking[url] = asyncio.create_task(self.check(url))
+    # A replica under its first-sight check is left out until the check finishes.
+    seen = [url for url in urls if url not in self.checking]
+    versions = await asyncio.gather(*(self.weight_version(url) for url in seen), return_exceptions=True)
+    held = [version for version in versions if isinstance(version, str)]
+    # A replica that did not answer may hold the revision, and one mid-load is about to hold other weights.
+    stuck = revision and held and len(held) == len(urls) and target not in held and not self.loading
+    if stuck and Path(revision, "metadata.json").is_file() and not any(delta_applies(revision, version) for version in held):
+      raise RevisionUnloadable(revision, held)
+    for url, version in zip(seen, versions):
+      # A replica still starting does not answer yet, and a pool on base weights never loads any.
+      if revision and isinstance(version, str) and version != target and url not in self.loading:
+        self.loading[url] = asyncio.create_task(self.load_weights(url, revision))
+    return target in held
 
   async def generate(self, r: dict[str, Any]) -> dict[str, Any] | None:
     """The result, or None when the answer did not come from the record's revision."""
-    target = r["revision"] or "default"
+    target, sent = r["revision"] or "default", time.monotonic()
     params = {**r["sampling_params"], "logprobs": 0, "prompt_logprobs": 0 if r["include_prompt_logprobs"] else None}
     attempt_id = f"{r['request_id']}-{uuid.uuid4().hex[:8]}"
     # The store hands prompts back as array('I'), which JSON cannot encode.
     body = {"request_id": attempt_id, "token_ids": list(r["prompt_token_ids"]), "model": r["model"], "sampling_params": params}
     url = self.router
     response = await self.http.post(f"{url}/inference/v1/generate", json=body, timeout=None)
-    # Envoy names the replica that answered. Its weights must still be the target.
+    # Envoy names the replica that answered. Its weights must still be the target, with no pause to load others since this was sent.
     replica = f"http://{name}" if (name := response.headers.get("x-openrl-upstream")) else url
-    if response.status_code in (400, 404, 422):
+    # A 404 can be a just-frozen LoRA snapshot that the replica's view of the shared volume does not show yet.
+    if response.status_code in (400, 422):
       return failed(f"Sampler rejected the request ({response.status_code}): {response.text[:500]}")
-    if response.status_code != 200 or await self.weight_version(replica) != target:
+    if response.status_code != 200 or await self.version_after_answer(replica) != target or self.paused.get(replica, 0.0) >= sent:
       return None
     try:
       return to_result(response.json(), r)
     except (KeyError, TypeError, ValueError) as exc:
       return failed(f"Sampler response broke the contract: {exc!r}")
 
+  async def version_after_answer(self, replica: str) -> str | None:
+    """The weight version of the replica that answered, or None when three tries cannot reach it. A slow check
+    must not throw away a finished answer."""
+    for _ in range(3):
+      try:
+        return await self.weight_version(replica)
+      except httpx.TransportError as exc:
+        logger.info("Checking the weight version of %s failed: %r", replica, exc)
+    return None
+
+  async def close(self) -> None:
+    """A load cut off midway leaves its replica paused, so loads get a moment to finish and resume before the HTTP client closes."""
+    if tasks := [*self.loading.values(), *self.checking.values()]:
+      _, pending = await asyncio.wait(tasks, timeout=5)
+      for task in pending:
+        task.cancel()
+
+  async def load_weights(self, url: str, revision: str) -> None:
+    """Pauses the replica, loads `revision` and resumes it. A load that fails or is cut off leaves the replica unchecked,
+    so the next dispatch's check closes any update it left open, forgets partial weights and resumes the replica."""
+    try:
+      if not Path(revision, "metadata.json").is_file():
+        return  # the dispatcher fails the revision's records
+      if not delta_applies(revision, current := await self.weight_version(url)):
+        raise RuntimeError(f"{url} holds {current!r}, not a revision delta {revision!r} applies on")
+      self.paused[url] = time.monotonic()
+      self.checked.discard(url)
+      await self.post(url, "/pause?mode=wait&clear_cache=false")
+      await self.post(url, "/start_weight_update")
+      await self.post(url, "/update_weights", {"update_info": {"target_weights_path": revision}})
+      # Before the finish names the version, so a replica never reports the revision with the old weights' cache.
+      if not await self.reset_prefix_cache(url):
+        raise RuntimeError("vLLM did not reset the prefix cache")
+      await self.post(url, "/finish_weight_update", {"weight_version": revision})
+      await self.post(url, "/resume")
+      self.checked.add(url)
+    except Exception as exc:
+      logger.warning("Loading %r on %s failed: %r", revision, url, exc)  # the next dispatch checks the replica and tries again
+    finally:
+      self.loading.pop(url)
+
+  async def reset_prefix_cache(self, url: str) -> bool:
+    """vLLM refuses while any request holds cache blocks, so a refusal is tried again a moment later."""
+    for attempt in range(3):
+      if attempt:
+        await asyncio.sleep(0.5)
+      if (await self.post(url, "/reset_prefix_cache")).json()["success"]:
+        return True
+    return False
+
+  async def check(self, url: str) -> None:
+    """A replica's check on first sight and after a load fails. Any weight update a load left open is closed and its
+    weights forgotten, paused or not, since a dispatcher stopped mid-update used to resume the replica. A failed or
+    cut-off load leaves a replica paused, where the router would queue requests forever, so one found paused is resumed.
+    It runs beside the dispatch, since a wedged replica can take minutes to answer."""
+    try:
+      await self.forget_weights(url)
+      if (await self.http.get(f"{url}/is_paused", timeout=5)).raise_for_status().json()["is_paused"]:
+        logger.info("Resuming %s, left paused by an earlier load", url)
+        await self.post(url, "/resume")
+      self.checked.add(url)
+    except Exception as exc:
+      logger.info("Checking %s failed: %r", url, exc)  # a replica still starting does not answer, and a later dispatch checks again
+    finally:
+      self.checking.pop(url)
+
+  async def forget_weights(self, url: str) -> None:
+    """Closes an update a cut-off load left active, which refuses every later start, and has the replica report
+    UNKNOWN_WEIGHTS until a full load, since it may hold part of a revision. vLLM has no abort, but a bare finish closes
+    an active update and is refused when none is, so a replica that was not updating keeps its version."""
+    # It waits behind an update a dead dispatcher left running, but a wedged replica must not hold up the pool.
+    response = await self.http.post(f"{url}/finish_weight_update", timeout=300)
+    if response.status_code != 200:
+      return
+    logger.warning("Weights on %s may be partial, so it reports %r until a full load replaces them", url, UNKNOWN_WEIGHTS)
+    await self.post(url, "/update_weight_version", {"new_version": UNKNOWN_WEIGHTS})
+
+  async def post(self, url: str, path: str, body: dict[str, Any] | None = None) -> httpx.Response:
+    return (await self.http.post(url + path, json=body, timeout=None)).raise_for_status()
+
   async def list_replicas(self) -> list[str]:
     if self.urls:
       return self.urls
@@ -183,7 +324,7 @@ async def list_replicas(self) -> list[str]:
     return [f"http://{p.status.pod_ip}:8000" for p in pods.items if p.status.pod_ip and p.status.phase == "Running"]
 
   async def weight_version(self, url: str) -> str:
-    return (await self.http.get(f"{url}/weight_info", timeout=5)).raise_for_status().json()["weight_version"]
+    return (await self.http.get(f"{url}/weight_info", timeout=WEIGHT_INFO_TIMEOUT_SECONDS)).raise_for_status().json()["weight_version"]
 
   def write_revision(self) -> None:
     path = revision_file(self.pool)
@@ -215,8 +356,11 @@ async def main() -> None:
     task = asyncio.create_task(Dispatcher(store, pool, sampler, int(os.getenv("OPEN_RL_DISPATCH_CONCURRENCY", "256"))).run())
     for sig in (signal.SIGTERM, signal.SIGINT):
       asyncio.get_running_loop().add_signal_handler(sig, task.cancel)
-    with contextlib.suppress(asyncio.CancelledError):
-      await task
+    try:
+      with contextlib.suppress(asyncio.CancelledError):
+        await task
+    finally:
+      await sampler.close()
 
 
 if __name__ == "__main__":

src/training/fft_trainer_worker.py

@@ -50,6 +50,11 @@ def __init__(self):
     self._prev_weights_shadow: dict[str, torch.Tensor] = {}
     self.model_layer_shapes: dict[str, tuple[int, ...]] = {}
     self.total_model_elements: int = 0
+    # Sampler saves at the current weights, oldest first. "" names the base weights a replica starts with.
+    self.sampler_saves: list[str] = [""]
+    # Sampler saves at the weights the last optimizer step started from, which its sparse delta applies on.
+    # Before any step the delta is empty, and the list is empty when no save holds those weights.
+    self.delta_from: list[str] = [""]
 
   def set_weight_sync_strategy(self, strategy: str) -> None:
     if strategy not in ("full", "delta"):
@@ -178,7 +183,8 @@ def save_state(self, model_id: str, state_path: str, include_optimizer: bool = F
 
     # Under the delta strategy save_state writes the sparse delta the sampler
     # consumes. load_from_state cannot open it, so FFT is not resumable yet.
-    if self.weight_sync_cfg.strategy == "delta":
+    # A sampler save that no replica could apply the delta on is written in full.
+    if self.weight_sync_cfg.strategy == "delta" and (kind != "sampler" or self.delta_from):
       if kind != "sampler":
         logger.warning("save_state for %s under the delta strategy writes a delta, not a resumable checkpoint", model_id)
       return self.save_state_delta(model_id=model_id, state_path=state_path, kind=kind)
@@ -271,6 +277,9 @@ def save_state_delta(
       "total_elements": total_elements,
       "layer_names": layer_names_list,
       "density_pct": round(100.0 * total_changed / max(1, total_elements), 3),
+      # The last sampler save before the last optimizer step. A replica holding any of the bases can apply this delta.
+      "parent": self.delta_from[-1] if self.delta_from else None,
+      "bases": list(dict.fromkeys(self.delta_from + self.sampler_saves)),
       "timestamp": time.time(),
     }
     with open(os.path.join(state_path, "metadata.json"), "w") as f:
@@ -286,6 +295,7 @@ def save_for_sampler(self, model_id: str, alias: str | None, ref: str | None) ->
     rel_path = ref[len("tinker://") :] if ref.startswith("tinker://") else ref.lstrip("/")
     local_path = os.path.join(os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl"), "sampler_full", rel_path)
     self.save_state(model_id, local_path, False, "sampler")
+    self.sampler_saves.append(local_path)
     return local_path
 
   def load_from_state(self, model_id: str, state_path: str, restore_optimizer: bool = False) -> dict[str, Any]:
@@ -307,6 +317,8 @@ def load_from_state(self, model_id: str, state_path: str, restore_optimizer: boo
     target_device = "auto" if num_gpus > 1 else self.device
     self.model = AutoModelForCausalLM.from_pretrained(state_path, dtype=dtype, device_map=target_device)
     self.prepare_model_for_training()
+    # No sampler save holds the loaded weights.
+    self.sampler_saves, self.delta_from = [], []
 
     if restore_optimizer and metadata.get("has_optimizer"):
       optimizer_path = os.path.join(state_path, "optimizer.pt")
@@ -372,6 +384,8 @@ def optim_step(self, adam_params: dict[str, Any], model_id: str | None = None) -
     self.optimizer.zero_grad()
     t_step_end = time.perf_counter()
     step_time = t_step_end - t_step_start
+    # The weights moved, so the sampler saves made so far hold the weights this step's delta starts from.
+    self.delta_from, self.sampler_saves = self.sampler_saves, []
 
     delta_compute_time = 0.0
     if self.weight_sync_cfg.strategy == "delta" and self.model is not None and hasattr(self.model, "named_parameters"):