From 89e510287cf4abee1fadf0ae80676994c985077a Mon Sep 17 00:00:00 2001 From: Shuby Mishra Date: Fri, 23 Oct 2026 17:23:50 +0000 Subject: [PATCH] sampling: a dispatcher that serves one pool stream through its replicas Adds sampling_dispatcher with Dispatcher and HttpSampler. The dispatcher is the pool's only reader. It reads the stream incrementally, holds each pending record until its result is saved, keeps up to `concurrency` sends in flight, and fails a record no replica answers before the deadline. HttpSampler sends through the pool's router, finds its replicas by pod label, and drops an answer from a replica that no longer reports the record's revision. In this commit replicas serve only the weights they start with, so a record for any other revision fails at its deadline. SamplingStore gains pending(after=...) and wait_for_new for the incremental reads, and returns prompts as array('I') since the dispatcher holds every pending one in memory. --- Makefile | 2 +- src/server/sampling_dispatcher.py | 223 ++++++++++++++++++++++++++++++ src/server/sampling_store.py | 22 ++- tests/test_sampling_dispatcher.py | 165 ++++++++++++++++++++++ 4 files changed, 408 insertions(+), 4 deletions(-) create mode 100644 src/server/sampling_dispatcher.py create mode 100644 tests/test_sampling_dispatcher.py diff --git a/Makefile b/Makefile index 6646f7b1d..be38e5581 100644 --- a/Makefile +++ b/Makefile @@ -19,7 +19,7 @@ HOST ?= 127.0.0.1 PORT ?= 9003 # The fully qualified base URL used by local CLI tools and clients BASE_URL ?= http://$(HOST):$(PORT) -UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_trainer_logprob_chunks tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_sampler_batch_limits tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_pool_accept +UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_trainer_logprob_chunks tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_sampler_batch_limits tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_pool_accept tests.test_sampling_dispatcher # Only forward BASE_URL to e2e when the user supplied it. The Makefile default # is for local CLI usage; e2e should start its own backend by default. TRAINING_TEST_BASE_URL ?= $(if $(filter environment command line,$(origin BASE_URL)),$(BASE_URL),) diff --git a/src/server/sampling_dispatcher.py b/src/server/sampling_dispatcher.py new file mode 100644 index 000000000..62b8798d2 --- /dev/null +++ b/src/server/sampling_dispatcher.py @@ -0,0 +1,223 @@ +"""Serves one sampling pool's pending records with a sampler and saves the first result for each. + +The sampler is HttpSampler, which sends through the pool's router, Envoy in front of the llm-d EPP, to its vllm serve +replicas.""" + +import asyncio +import contextlib +import logging +import os +import signal +import time +import uuid +from pathlib import Path +from typing import Any + +import httpx +import redis.asyncio as redis + +from server.sampling_store import SamplingStore + +logger = logging.getLogger(__name__) + +# The dispatcher wakes for each new request, or after this long to retry a failed send and look at its replicas. +RETRY_MS = 1000 +# A request no replica has answered for this long fails, so its client stops waiting. +SAMPLE_DEADLINE_SECONDS = 3600 + + +def tmp_root() -> str: + """OPEN_RL_TMP_DIR, or the documented default. The API server renders the readiness probe's path + from it and hands the same value to the dispatcher and the replicas, so they all agree.""" + return os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl") + + +def revision_file(pool: str) -> Path: + """The readiness probe compares vLLM's weight version with the revision in this file.""" + return Path(tmp_root(), "sampling", pool, "revision") + + +def accepted_at(r: dict[str, Any]) -> float: + """When the record's request was accepted. A stream id starts with that time in milliseconds.""" + return int(r["entry_id"].split("-")[0]) / 1000 + + +def failed(message: str) -> dict[str, str]: + return {"type": "RequestFailedResponse", "error_message": message} + + +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"]] + sequences = [{"tokens": c["token_ids"], "logprobs": lp, "stop_reason": c["finish_reason"]} for c, lp in zip(body["choices"], logprobs)] + result: dict[str, Any] = {"type": "sample", "sequences": sequences} + if record["include_prompt_logprobs"]: + pairs = zip(record["prompt_token_ids"], body["prompt_logprobs"], strict=True) + result["prompt_logprobs"] = [None if e is None else e[str(token)]["logprob"] for token, e in pairs] + return result + + +class Dispatcher: + """Serves `pool` with `sampler`, which loads a revision with ensure_weights and answers a record with generate.""" + + def __init__(self, store: SamplingStore, pool: str, sampler: Any, concurrency: int) -> None: + self.store, self.pool, self.sampler, self.concurrency = store, pool, sampler, concurrency + self.revision = "" + # This dispatcher is the pool's only reader, so it reads each record once and drops it when its result is saved. + self.records: dict[str, dict[str, Any]] = {} + self.newest = "0-0" + self.in_flight: dict[str, asyncio.Task] = {} + + async def run(self) -> None: + try: + while True: + try: + await self.dispatch() + await self.store.wait_for_new(self.pool, self.newest, RETRY_MS) + except Exception: + # A Redis or Kubernetes hiccup must not end the pool's dispatching. + logger.exception("Dispatching pool %s failed; retrying", self.pool) + await asyncio.sleep(RETRY_MS / 1000) + finally: + for task in self.in_flight.values(): + task.cancel() + + async def dispatch(self) -> None: + await self.read() + if await self.sampler.ensure_weights(self.revision): + # 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 read(self) -> None: + """Reads new records, fails those past their deadline, and switches revision once nothing pending names the + current one.""" + for r in await self.store.pending(self.pool, after=self.newest): + self.records[r["entry_id"]] = r + self.newest = r["entry_id"] + for r in list(self.records.values()): + if (waited := time.time() - accepted_at(r)) > SAMPLE_DEADLINE_SECONDS: + logger.warning("Failing %s after %.0fs without a result", r["request_id"], waited) + await self.set_result(r, failed(f"No sampler answered in {waited:.0f}s; pool {self.pool} serves {self.revision or 'base weights'}")) + # A send still waiting on a replica that never answers would keep its slot forever. + if task := self.in_flight.get(r["request_id"]): + task.cancel() + records = list(self.records.values()) + # Nothing pending names the current revision, so nothing is in flight on it. + if records and all(r["revision"] != self.revision for r in records): + revision = records[0]["revision"] + logger.info("Pool %s switching from %r to %r", self.pool, self.revision, revision) + self.revision = revision + + def to_send(self) -> list[dict[str, Any]]: + return [r for r in self.records.values() if r["revision"] == self.revision and r["request_id"] not in self.in_flight] + + async def set_result(self, r: dict[str, Any], result: dict[str, Any]) -> None: + await self.store.set_result(r, result) + self.records.pop(r["entry_id"], None) + + async def send(self, r: dict[str, Any]) -> None: + """Saves the sampler's answer. None, or an error, leaves the record pending to be sent again.""" + try: + # A sampler that never answers must not hold the send past its record's deadline. + timeout = max(0.0, accepted_at(r) + SAMPLE_DEADLINE_SECONDS - time.time()) + result = await asyncio.wait_for(self.sampler.generate(r), timeout) + except Exception as exc: + logger.info("Retrying %s after %r", r["request_id"], exc) + return + if result is not None: + await self.set_result(r, result) + + +class HttpSampler: + """Sends a pool's records through `router` to its vllm serve replicas, and keeps each replica on the pool's revision. + `urls`, when given, lists the replicas instead of finding them by pod label.""" + + def __init__(self, http: httpx.AsyncClient, pool: str, router: str, urls: list[str]) -> None: + self.http, self.pool, self.router, self.urls = http, pool, router, urls + self.revision: str | None = None + self.k8s = None + + async def ensure_weights(self, revision: str) -> bool: + """True once one replica serves `revision`.""" + 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 + + 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" + 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. + replica = f"http://{name}" if (name := response.headers.get("x-openrl-upstream")) else url + if response.status_code in (400, 404, 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: + 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 list_replicas(self) -> list[str]: + if self.urls: + return self.urls + if self.k8s is None: + from kubernetes import client, config + + config.load_incluster_config() + self.k8s = client.CoreV1Api() + selector = f"app.kubernetes.io/component=sampler,app.kubernetes.io/instance={self.pool}" + namespace = os.getenv("OPEN_RL_WORKER_NAMESPACE", "openrl-system") + pods = await asyncio.to_thread(self.k8s.list_namespaced_pod, namespace, label_selector=selector, _request_timeout=10) + 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"] + + def write_revision(self) -> None: + path = revision_file(self.pool) + path.parent.mkdir(parents=True, exist_ok=True) + path.with_suffix(".tmp").write_text(self.revision or "default") + path.with_suffix(".tmp").rename(path) + + +def redis_client(url: str) -> redis.Redis: + """When replicas come up, every in-flight send saves its result at once. Past max_connections a + save waits for a free connection; redis-py's default pool of 100 raises instead.""" + return redis.Redis.from_pool(redis.BlockingConnectionPool.from_url(url, decode_responses=True, max_connections=64, timeout=30)) + + +def http_client() -> httpx.AsyncClient: + """Sends have no connection limit, so a burst goes out at once. httpx scans its whole pool for each idle + connection whenever a request starts or ends, so the hundreds a burst leaves idle stall this loop for + seconds. Few are kept.""" + return httpx.AsyncClient(limits=httpx.Limits(max_connections=None, max_keepalive_connections=32)) + + +async def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + logging.getLogger("httpx").setLevel(logging.WARNING) + router = os.getenv("OPEN_RL_SAMPLER_ROUTER", "http://127.0.0.1:8081") + store, pool = SamplingStore(redis_client(os.environ["REDIS_URL"])), os.environ["OPEN_RL_SAMPLER_POOL"] + async with http_client() as http: + sampler = HttpSampler(http, pool, router, []) + 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 + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/src/server/sampling_store.py b/src/server/sampling_store.py index 29af6c33e..6b044acf9 100644 --- a/src/server/sampling_store.py +++ b/src/server/sampling_store.py @@ -2,6 +2,7 @@ Results land in the `open_rl:future:` list that `retrieve_future` already reads.""" import json +from array import array from typing import Any import redis.asyncio as redis @@ -15,6 +16,17 @@ def stream_key(pool: str) -> str: return f"open_rl:sampling:{pool}" +def compact_prompt(record: dict[str, Any]) -> dict[str, Any]: + """A json.loads object_hook that holds a prompt as array('I'), four bytes a token against about 36 as a list of ints, + and drops each list as soon as its record is read. A prompt array('I') cannot hold stays the list, for vLLM to refuse.""" + if isinstance(tokens := record.get("prompt_token_ids"), list): + try: + record["prompt_token_ids"] = array("I", tokens) + except (OverflowError, TypeError): + pass + return record + + def idempotency_key(pool: str, key: str) -> str: return f"open_rl:sampling:{pool}:idem:{key}" @@ -45,9 +57,13 @@ async def add(pipe) -> str: return await self.redis.transaction(add, key, value_from_callable=True) - async def pending(self, pool: str, limit: int | None = 1000) -> list[dict[str, Any]]: - entries = await self.redis.xrange(stream_key(pool), count=limit) - return [{**json.loads(fields["record"]), "entry_id": entry_id} for entry_id, fields in entries] + async def pending(self, pool: str, limit: int | None = 1000, after: str = "0-0") -> list[dict[str, Any]]: + """Pending records in accept order. The dispatcher holds every one until its result is set, so prompts come back compact.""" + entries = await self.redis.xrange(stream_key(pool), min=f"({after}", count=limit) + return [{**json.loads(fields["record"], object_hook=compact_prompt), "entry_id": entry_id} for entry_id, fields in entries] + + async def wait_for_new(self, pool: str, after: str, timeout_ms: int) -> None: + await self.redis.xread({stream_key(pool): after}, count=1, block=timeout_ms) async def set_result(self, record: dict[str, Any], result: dict[str, Any]) -> bool: """Save the first final result and drop the pending entry. False if it already has one.""" diff --git a/tests/test_sampling_dispatcher.py b/tests/test_sampling_dispatcher.py new file mode 100644 index 000000000..df2b79bb5 --- /dev/null +++ b/tests/test_sampling_dispatcher.py @@ -0,0 +1,165 @@ +import asyncio +import inspect +import json +import tempfile +import uuid +from pathlib import Path +from unittest.mock import patch + +import httpx +import uvicorn +from starlette.applications import Starlette +from starlette.responses import JSONResponse, Response +from starlette.routing import Route + +from server.sampling_dispatcher import Dispatcher, HttpSampler, revision_file +from server.sampling_store import SamplingStore, future_key +from tests.redis_server import RedisTestCase + +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}]}} + + +async def until(condition, timeout: float = 10) -> None: + async with asyncio.timeout(timeout): + while not (await result if inspect.isawaitable(result := condition()) else result): + await asyncio.sleep(0.02) + + +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, + and only a finish that names a version, or /update_weight_version, changes the version.""" + + 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.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}) + posts |= {"/update_weights": self.update, "/finish_weight_update": self.finish, "/inference/v1/generate": self.generate} + self.app = Starlette(routes=gets + [Route(path, handler, methods=["POST"]) for path, handler in posts.items()]) + + async def start(self, request) -> Response: + if self.updating: + return Response("start_weight_update called while a weight update is already active. Call finish_weight_update first.", status_code=500) + self.updating = True + return JSONResponse({}) + + async def update(self, request) -> Response: + if not self.updating: + return Response("start_weight_update must be called before update_weights.", status_code=500) + self.log.append(("update", self.name, (await request.json())["update_info"]["target_weights_path"], self.expected.read_text())) + await until(lambda: not self.held.is_set()) + return JSONResponse({}) + + async def pause(self, request) -> Response: + self.paused = True + return JSONResponse({}) + + async def resume(self, request) -> Response: + self.paused = False + return JSONResponse({}) + + async def finish(self, request) -> Response: + body = await request.body() + if not self.updating: + return Response("finish_weight_update called without a matching start_weight_update.", status_code=500) + self.updating = False + if body: + self.version = json.loads(body)["weight_version"] + return JSONResponse({}) + + async def set_version(self, request) -> Response: + self.version = (await request.json())["new_version"] + return JSONResponse({}) + + async def generate(self, request) -> Response: + self.log.append(("generate", self.name, self.version)) + await until(lambda: not self.held.is_set() and not self.paused) + return JSONResponse({"choices": [CHOICE]}) + + +def routed_dispatcher(store: SamplingStore, pool: str, http: httpx.AsyncClient, router_url: str, urls: list[str], concurrency: int) -> Dispatcher: + """The dispatcher of a pool of vllm serve replicas, as sampling_dispatcher.main builds it.""" + return Dispatcher(store, pool, HttpSampler(http, pool, router_url, urls), concurrency) + + +async def run_routed(dispatcher: Dispatcher) -> None: + """Runs it as sampling_dispatcher.main does.""" + await dispatcher.run() + + +def router(log: list, replicas: list[FakeReplica], http: httpx.AsyncClient) -> Starlette: + # Forwards generate to a ready replica, one whose version is the expected one, like the EPP. + async def route(request) -> Response: + ready = [r for r in replicas if r.version == r.expected.read_text()] + if not ready: + log.append(("unrouted", "", "")) + return Response(status_code=503) + upstream = await http.post(f"{ready[0].url}/inference/v1/generate", content=await request.body(), timeout=None) + return Response(upstream.content, upstream.status_code, headers={"x-openrl-upstream": ready[0].url.removeprefix("http://")}) + + return Starlette(routes=[Route("/inference/v1/generate", route, methods=["POST"])]) + + +class SamplingDispatcherTest(RedisTestCase): + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + self.store, self.pool, self.log, self.http = SamplingStore(self.redis), f"sampler-{uuid.uuid4()}", [], httpx.AsyncClient() + self.addAsyncCleanup(self.http.aclose) + self.tmp = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.enterContext(patch.dict("os.environ", {"OPEN_RL_TMP_DIR": str(self.tmp)})) + + async def serve(self, app: Starlette) -> str: + server = uvicorn.Server(uvicorn.Config(app, port=0, log_level="warning", lifespan="off")) + serving = asyncio.create_task(server.serve()) + self.addAsyncCleanup(lambda: (setattr(server, "should_exit", True), serving)[1]) + await until(lambda: server.started) + return f"http://127.0.0.1:{server.servers[0].sockets[0].getsockname()[1]}" + + async def replicas(self, *versions: str) -> list[FakeReplica]: + replicas = [FakeReplica(f"pod-{i}", self.log, version, revision_file(self.pool)) for i, version in enumerate(versions)] + for replica in replicas: + replica.url = await self.serve(replica.app) + self.addAsyncCleanup(replica.held.clear) + return replicas + + 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} + await self.store.accept(r) + return r + + async def dispatcher(self, replicas: list[FakeReplica], concurrency: int = 4) -> asyncio.Task: + router_url = await self.serve(router(self.log, replicas, self.http)) + dispatcher = routed_dispatcher(self.store, self.pool, self.http, router_url, [r.url for r in replicas], concurrency=concurrency) + run = asyncio.create_task(run_routed(dispatcher)) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + return run + + async def results(self, request_id: str) -> list[dict]: + await until(lambda: self.redis.llen(future_key(request_id))) + await asyncio.sleep(0.3) # a second result would show by now + return [json.loads(raw) for raw in await self.redis.lrange(future_key(request_id), 0, -1)] + + async def test_each_request_is_generated_once(self): + replicas = await self.replicas("default", "default") + records = [await self.accepted() for _ in range(40)] + await self.dispatcher(replicas, concurrency=256) + for r in records: + await until(lambda r=r: self.redis.llen(future_key(r["request_id"]))) + await asyncio.sleep(0.5) # a duplicate would be generated by now + self.assertEqual(sum(entry[0] == "generate" for entry in self.log), len(records)) + + async def test_crash_mid_request_saves_exactly_one_result(self): + (replica,) = replicas = await self.replicas("default") + replica.held.set() + r = await self.accepted() + first = await self.dispatcher(replicas) + await until(lambda: self.log.count(("generate", "pod-0", "default")) == 1) + first.cancel() + await self.dispatcher(replicas) + replica.held.clear() + self.assertEqual((await self.results(r["request_id"]), await self.store.pending(self.pool)), ([SAMPLE], []))