diff --git a/src/server/sampling_dispatcher.py b/src/server/sampling_dispatcher.py index 62b8798d..1cb84fd6 100644 --- a/src/server/sampling_dispatcher.py +++ b/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__": diff --git a/src/training/fft_trainer_worker.py b/src/training/fft_trainer_worker.py index c547ab09..6e69e150 100644 --- a/src/training/fft_trainer_worker.py +++ b/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"): diff --git a/tests/test_sampling_dispatcher.py b/tests/test_sampling_dispatcher.py index df2b79bb..e3458f41 100644 --- a/tests/test_sampling_dispatcher.py +++ b/tests/test_sampling_dispatcher.py @@ -7,6 +7,8 @@ from unittest.mock import patch import httpx +import safetensors.torch +import torch import uvicorn from starlette.applications import Starlette from starlette.responses import JSONResponse, Response @@ -15,6 +17,9 @@ from server.sampling_dispatcher import Dispatcher, HttpSampler, revision_file from server.sampling_store import SamplingStore, future_key from tests.redis_server import RedisTestCase +from tests.test_exclusive_trainer import tiny_worker +from training.fft_trainer_worker import FFTTrainingWorker +from training.types import Datum, TensorData SAMPLE = {"type": "sample", "sequences": [{"tokens": [100, 101], "logprobs": [-0.5, -0.25], "stop_reason": "length"}]} CHOICE = {"token_ids": [100, 101], "finish_reason": "length", "logprobs": {"content": [{"logprob": -0.5}, {"logprob": -0.25}]}} @@ -26,6 +31,45 @@ async def until(condition, timeout: float = 10) -> None: await asyncio.sleep(0.02) +def fft_worker() -> FFTTrainingWorker: + """A real FFT trainer on a tiny Llama, resident on the CPU, publishing sparse deltas.""" + worker = tiny_worker() + worker.cpu_offload = False + worker.set_weight_sync_strategy("delta") + worker.prepare_model_for_training() + return worker + + +def train_step(worker: FFTTrainingWorker) -> None: + row = [3, 5, 8, 13, 21, 34] + datum = Datum(model_input=row[:-1], loss_fn_inputs={"target_tokens": TensorData(data=row[1:]), "weights": TensorData(data=[1.0] * 5)}) + worker.forward_backward([datum], "cross_entropy") + worker.optim_step({"learning_rate": 1e-2}) + + +def weights(worker: FFTTrainingWorker) -> dict[str, torch.Tensor]: + return {name: param.detach().clone() for name, param in worker.model.named_parameters()} + + +def replayed(start: dict[str, torch.Tensor], revisions: list[str]) -> dict[str, torch.Tensor]: + """A replica's weights after it loads `revisions` in order, written the way the delta_snapshot engine writes them.""" + out = {name: tensor.clone() for name, tensor in start.items()} + for revision in revisions: + metadata = json.loads(Path(revision, "metadata.json").read_text()) + if metadata.get("format") == "sparse_delta": + tensors = safetensors.torch.load_file(str(Path(revision, "delta.safetensors"))) + for i, name in enumerate(metadata["layer_names"]): + out[name].view(-1)[tensors[f"{i}.indices"].long()] = tensors[f"{i}.values"] + else: + for file in Path(revision).glob("*.safetensors"): + out.update((name, tensor) for name, tensor in safetensors.torch.load_file(str(file)).items() if name in out) + return out + + +def same(a: dict[str, torch.Tensor], b: dict[str, torch.Tensor]) -> bool: + return a.keys() == b.keys() and all(torch.equal(a[name], b[name]) for name in a) + + class FakeReplica: """vLLM's weight endpoints and generate with a real weight version. `held` pauses both; an update logs what the readiness probe expects. Like vLLM, a paused replica queues generate until it is resumed, one weight update is active from start to finish, @@ -34,7 +78,7 @@ class FakeReplica: def __init__(self, name: str, log: list, version: str, expected: Path): self.name, self.log, self.version, self.expected, self.held, self.url = name, log, version, expected, asyncio.Event(), "" self.paused, self.updating = False, False - gets = [Route("/weight_info", lambda r: JSONResponse({"weight_version": self.version}))] + gets = [Route("/weight_info", self.weight_info)] gets.append(Route("/is_paused", lambda r: JSONResponse({"is_paused": self.paused}))) posts = {"/pause": self.pause, "/resume": self.resume, "/start_weight_update": self.start, "/update_weight_version": self.set_version} posts["/reset_prefix_cache"] = lambda r: JSONResponse({"success": True}) @@ -62,6 +106,9 @@ async def resume(self, request) -> Response: self.paused = False return JSONResponse({}) + async def weight_info(self, request) -> Response: + return JSONResponse({"weight_version": self.version}) + async def finish(self, request) -> Response: body = await request.body() if not self.updating: @@ -87,8 +134,11 @@ def routed_dispatcher(store: SamplingStore, pool: str, http: httpx.AsyncClient, async def run_routed(dispatcher: Dispatcher) -> None: - """Runs it as sampling_dispatcher.main does.""" - await dispatcher.run() + """Runs it as sampling_dispatcher.main does, closing the sampler when it stops.""" + try: + await dispatcher.run() + finally: + await dispatcher.sampler.close() def router(log: list, replicas: list[FakeReplica], http: httpx.AsyncClient) -> Starlette: @@ -126,6 +176,14 @@ async def replicas(self, *versions: str) -> list[FakeReplica]: self.addAsyncCleanup(replica.held.clear) return replicas + def revision(self, name: str, **metadata) -> str: + (self.tmp / name).mkdir() + (self.tmp / name / "metadata.json").write_text(json.dumps(metadata)) + return str(self.tmp / name) + + def loads(self, *names: str) -> dict[str, list[str]]: + return {name: [entry[2] for entry in self.log if entry[:2] == ("update", name)] for name in names} + async def accepted(self, revision: str = "", prompt: list | None = None) -> dict: identity = {"request_id": str(uuid.uuid4()), "pool": self.pool, "revision": revision, "model": "base", "idempotency_key": None} r = {**identity, "prompt_token_ids": prompt or [1, 2, 3], "sampling_params": {"n": 1, "max_tokens": 2}, "include_prompt_logprobs": False} @@ -163,3 +221,37 @@ async def test_crash_mid_request_saves_exactly_one_result(self): await self.dispatcher(replicas) replica.held.clear() self.assertEqual((await self.results(r["request_id"]), await self.store.pending(self.pool)), ([SAMPLE], [])) + + async def test_an_answer_from_the_wrong_weights_is_not_saved_and_the_pod_is_prepared_again(self): + (replica,) = replicas = await self.replicas(w1 := self.revision("w1")) + replica.held.set() + r = await self.accepted(w1) + await self.dispatcher(replicas) + await until(lambda: ("generate", "pod-0", w1) in self.log) + replica.version = "default" # restarted in place while the request ran + replica.held.clear() + self.assertEqual(await self.results(r["request_id"]), [SAMPLE]) + self.assertEqual(self.log, [("generate", "pod-0", w1), ("update", "pod-0", w1, w1), ("generate", "pod-0", w1)]) + + async def test_a_replica_on_any_save_at_the_weights_a_delta_starts_from_or_reaches_applies_it(self): + worker = fft_worker() + train_step(worker) + # An RL step with nothing to learn from takes no optimizer step, so its rollouts sample a second save at the same weights. + first, at_first = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-0"), weights(worker) + second = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-1") + # The late replica started on the first save and never loads the second. + early, late = await self.replicas(first, first) + r = await self.accepted(second) + run = await self.dispatcher([early]) + self.assertEqual(await self.results(r["request_id"]), [SAMPLE]) + run.cancel() + await asyncio.gather(run, return_exceptions=True) + train_step(worker) + third = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-2") + r = await self.accepted(third) + await self.dispatcher([early, late]) + self.assertEqual(await self.results(r["request_id"]), [SAMPLE]) + await until(lambda: self.loads(late.name) == {late.name: [third]}) + self.assertEqual(self.loads(early.name), {early.name: [second, third]}) + self.assertTrue(same(replayed(at_first, [second, third]), weights(worker))) + self.assertTrue(same(replayed(at_first, [third]), weights(worker)))