From 3af93786e391b4d57f73af82546dd2fa73330d10 Mon Sep 17 00:00:00 2001 From: Shuby Mishra Date: Fri, 23 Oct 2026 19:33:28 +0000 Subject: [PATCH] sampling: load each pool revision onto its replicas HttpSampler starts a load on every replica not at the pool's revision. A load pauses the replica, runs vLLM's weight update, resets the prefix cache, finishes the update with the revision as its version, and resumes it. The reset comes before the finish, so a replica never reports the revision while its cache holds the old weights' entries. vLLM refuses a reset while a request holds cache blocks, so a refused reset is tried up to three times, and one that keeps failing forgets the weights. An answer from a replica that paused to load weights after the request went out is dropped and sent again. The version check after an answer tries three times on a connection error or timeout, so a slow check does not throw away a finished answer. A revision that retention deleted fails its records, before the switch or while the pool serves it, so the pool moves on. A 404 sends the request again until its deadline, since a replica's view of the shared volume may not show a just-frozen LoRA snapshot yet. A 400 or 422 still fails it. The first time the dispatcher sees a replica, and after one of its loads breaks off, it closes any weight update an earlier load left open and resumes a replica left paused. Such a replica may hold part of a revision, so it reports "unknown" until a full load replaces its weights. The check runs as a task of its own for each replica, like a load, so a wedged replica holds up no send. Until its check finishes the replica gets no load and counts for nothing, and a check that fails runs again on a later pass. The FFT trainer writes `bases` in each delta's metadata, the sampler saves that hold the weights the delta starts from or reaches. A replica loads a delta only when it is at one of its bases. When no save holds the starting weights, as after load_from_state, the sampler save is written in full. The pool still switches revision once nothing pending names the current one, as before, so a sparse delta serves every pending record for it before the pool moves past it. When every replica answers, none holds a delta revision or weights it applies on, and none is loading, as for an older delta after a newer revision loaded, HttpSampler raises RevisionUnloadable. The dispatcher fails that revision's records at once, saying full weights must be saved to sample it, instead of letting them wait out the deadline. A replica that does not answer may hold the revision, so the decision waits for it. --- src/server/sampling_dispatcher.py | 166 +++++++++++++++++++++++++++-- src/training/fft_trainer_worker.py | 16 ++- tests/test_sampling_dispatcher.py | 98 ++++++++++++++++- 3 files changed, 265 insertions(+), 15 deletions(-) diff --git a/src/server/sampling_dispatcher.py b/src/server/sampling_dispatcher.py index 62b8798d2..1cb84fd68 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 c547ab093..6e69e150b 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 df2b79bb5..e3458f414 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)))