diff --git a/skyrl/benchmarks/load_test_tinker_sampling.py b/skyrl/benchmarks/load_test_tinker_sampling.py new file mode 100644 index 0000000000..a00da01b8a --- /dev/null +++ b/skyrl/benchmarks/load_test_tinker_sampling.py @@ -0,0 +1,865 @@ +#!/usr/bin/env python3 +"""Load test: is the Tinker API server the bottleneck at N concurrent sample requests? + +Exercises the API server layers on the non-colocated SkyRL-Train path +(``backend=megatron``, ``trainer.placement.colocate_all=false``): FastAPI + +uvicorn, the SQLite-backed session/model tables, the in-memory +``ExternalFutureStore``, and ``SkyRLTrainInferenceForwardingClient`` which +forwards each sample to the engine-managed vLLM router. Everything *below* the +server is stubbed so no GPU is needed and the numbers isolate the server: + + fake router ``--role vllm`` serves ``/v1/completions`` in place of the + vllm-router + vLLM. ``--vllm-mode barrier`` holds requests until + ``--barrier-requests`` are in flight and releases them at once + (worst-case completion burst); ``--vllm-mode latency`` admits + ``--max-num-seqs`` at a time for ``--gen-seconds`` each (a real + engine's queueing). + API server ``--role server`` the real ``skyrl.tinker.api`` app under + uvicorn, configured exactly as the non-colocated megatron server + is, with the engine subprocess replaced by a sleeper and the + router URL the engine would publish to ``EngineStateDB`` seeded + to the fake router. Pass ``--server-url`` to test a real server + instead (engine + vLLM must be up). + load client ``--role load``. ``--client raw`` speaks the HTTP API directly from + ``--workers`` processes, each bound to its own loopback source IP so + the client is not capped by one IP's ~28k ephemeral ports. ``--client + sdk`` uses the public Tinker SDK (``sample_async``); the SDK caps + in-flight samples per SamplingClient at the server-advertised + ``sample_max_concurrent_requests`` (default 2000). + +Usage (everything in one command, CPU only): + + uv run --extra tinker python skyrl/benchmarks/load_test_tinker_sampling.py \\ + --num-requests 131072 --forwarding-max-connections 2048 + + # Realistic engine queueing: 2048 concurrent generations of 5s each + uv run --extra tinker python skyrl/benchmarks/load_test_tinker_sampling.py \\ + --num-requests 131072 --forwarding-max-connections 2048 \\ + --vllm-mode latency --max-num-seqs 2048 --gen-seconds 5 + + # Against a server you started yourself: + uv run --extra tinker python skyrl/benchmarks/load_test_tinker_sampling.py \\ + --role load --url http://127.0.0.1:8000 --num-requests 4096 + +The summary reports completed/failed counts by error class, submit and +end-to-end latency percentiles, the peak number of requests the fake router saw +in flight at once (the concurrency the server actually sustained), and the API +server's peak RSS / file descriptors and ``/healthz`` latency under load. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import os +import signal +import statistics +import subprocess +import sys +import tempfile +import time +from collections import Counter +from pathlib import Path +from typing import Any + +import aiohttp +import psutil + +TINY_MODEL = "trl-internal-testing/tiny-Qwen3ForCausalLM" +DEFAULT_API_PORT = 18779 +DEFAULT_VLLM_PORT = 18879 +# Per-IP ephemeral port budget with the Linux default range (32768-60999). +DEFAULT_REQUESTS_PER_WORKER = 16384 +# The Tinker SDK gives up on each retrieve_future poll after 45s and re-polls +# the same request_id (the server would hold it for up to 300s). +SDK_RETRIEVE_POLL_TIMEOUT_SECONDS = 45.0 + + +# --------------------------------------------------------------------------- # +# Fake vLLM +# --------------------------------------------------------------------------- # + + +def run_fake_vllm(args: argparse.Namespace) -> None: + from aiohttp import web + + state = {"received": 0, "completed": 0, "in_flight": 0, "peak_in_flight": 0} + barrier_requests = args.barrier_requests or args.num_requests + release = asyncio.Event() + admit = asyncio.Semaphore(args.max_num_seqs) + body_cache: dict[tuple[int, int], bytes] = {} + + def body_for(n: int, max_tokens: int) -> bytes: + key = (n, max_tokens) + if key not in body_cache: + choice = { + "token_ids": [1] * max_tokens, + "logprobs": {"token_logprobs": [-0.5] * max_tokens}, + "finish_reason": "length", + } + body_cache[key] = json.dumps({"choices": [choice] * n}, separators=(",", ":")).encode() + return body_cache[key] + + async def barrier_timeout_watch() -> None: + await asyncio.sleep(args.barrier_timeout) + if not release.is_set(): + print( + f"[fake-vllm] barrier timeout after {args.barrier_timeout}s with " + f"{state['received']}/{barrier_requests} received -- releasing", + flush=True, + ) + release.set() + + async def completions(request: web.Request) -> web.Response: + payload = await request.json() + state["received"] += 1 + state["in_flight"] += 1 + state["peak_in_flight"] = max(state["peak_in_flight"], state["in_flight"]) + try: + if args.vllm_mode == "barrier": + if state["received"] >= barrier_requests: + release.set() + await release.wait() + else: + async with admit: + await asyncio.sleep(args.gen_seconds) + state["completed"] += 1 + return web.Response( + body=body_for(payload.get("n", 1), payload["max_tokens"]), content_type="application/json" + ) + finally: + state["in_flight"] -= 1 + + async def stats(_: web.Request) -> web.Response: + return web.json_response(state) + + async def serve() -> None: + app = web.Application(client_max_size=64 * 1024 * 1024) + app.router.add_post("/v1/completions", completions) + app.router.add_get("/healthz", stats) + app.router.add_get("/stats", stats) + runner = web.AppRunner(app, access_log=None) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", args.vllm_port, backlog=65535) + await site.start() + if args.vllm_mode == "barrier": + asyncio.create_task(barrier_timeout_watch()) + print(f"[fake-vllm] mode={args.vllm_mode} listening on 127.0.0.1:{args.vllm_port}", flush=True) + await asyncio.Event().wait() + + asyncio.run(serve()) + + +# --------------------------------------------------------------------------- # +# Raw HTTP load client (one worker process) +# --------------------------------------------------------------------------- # + + +def worker_source_ip(worker_index: int) -> str: + # 127.0.0.0/8 is entirely bound to lo on Linux, so any 127.x.y.z is a free + # source address with its own ephemeral-port budget. + return f"127.1.{worker_index // 250}.{worker_index % 250 + 1}" + + +async def raw_worker(args: argparse.Namespace, worker_index: int, count: int) -> dict[str, Any]: + base = args.url.rstrip("/") + "/api/v1" + payload = { + "num_samples": args.num_samples, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [1000] * args.prompt_tokens}]}, + "sampling_params": {"max_tokens": args.max_tokens, "temperature": 1.0, "seed": 0}, + "base_model": args.base_model, + } + retrieve_headers = {"Accept": "application/x-protobuf, application/json"} if args.proto else {} + + submit_latency: list[float] = [] + e2e_latency: list[float] = [] + errors: Counter[str] = Counter() + error_samples: dict[str, str] = {} + retries_408 = 0 + reconnects = 0 + asample_retries = 0 + poll_timeouts = 0 + completed = 0 + + def record_error(kind: str, detail: str) -> None: + errors[kind] += 1 + error_samples.setdefault(kind, detail[:300]) + + local_addr = (worker_source_ip(worker_index), 0) if args.source_ips else None + connector = aiohttp.TCPConnector(limit=0, local_addr=local_addr, force_close=False) + timeout = aiohttp.ClientTimeout(total=None, sock_connect=60, sock_read=args.poll_timeout) + + async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session: + + # The Tinker SDK holds at most 400 asample POSTs in flight per client + # (its sample dispatch semaphore); results are then awaited with one + # long-poll each. Mirror that so the submit burst matches production. + submit_gate = asyncio.Semaphore(max(1, args.submit_concurrency)) + + async def one(i: int) -> None: + nonlocal retries_408, reconnects, asample_retries, poll_timeouts, completed + t0 = time.monotonic() + phase = "asample" + try: + async with submit_gate: + while True: + try: + async with session.post(f"{base}/asample", json=payload) as resp: + if resp.status != 200: + record_error(f"asample_http_{resp.status}", await resp.text()) + return + request_id = (await resp.json())["request_id"] + break + except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as e: + # The SDK re-sends an asample whose connection was + # reset (the server may or may not have accepted + # the first copy). Mirror that, once per request. + if asample_retries >= count: + record_error("asample_retry_budget_exhausted", str(e)) + return + asample_retries += 1 + await asyncio.sleep(0.05) + submit_latency.append(time.monotonic() - t0) + phase = "retrieve" + while True: + if time.monotonic() - t0 > args.request_timeout: + record_error("request_timeout", f"no result after {args.request_timeout}s") + return + try: + async with session.post( + f"{base}/retrieve_future", json={"request_id": request_id}, headers=retrieve_headers + ) as resp: + if resp.status == 408: + retries_408 += 1 + continue + if resp.status != 200: + record_error(f"retrieve_http_{resp.status}", await resp.text()) + return + await resp.read() + break + except asyncio.TimeoutError: + # Client-side poll timeout: the SDK abandons the poll + # and immediately re-polls the same request_id. + poll_timeouts += 1 + continue + except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as e: + # The server dropped the connection (e.g. uvicorn's + # keep-alive timer firing during an event-loop stall). + # retrieve_future is idempotent and the Tinker SDK + # retries it, so do the same. + reconnects += 1 + if reconnects > args.max_reconnects * count: + record_error("retrieve_reconnect_budget_exhausted", str(e)) + return + await asyncio.sleep(0.05) + e2e_latency.append(time.monotonic() - t0) + completed += 1 + except aiohttp.ClientConnectorError as e: + record_error(f"{phase}_connect_error", str(e)) + except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError, aiohttp.ClientPayloadError) as e: + record_error(f"{phase}_{type(e).__name__}", str(e)) + except asyncio.TimeoutError: + record_error(f"{phase}_client_timeout", "") + except Exception as e: # noqa: BLE001 - a load test wants every failure class counted + record_error(f"{phase}_{type(e).__name__}", str(e)) + + started = time.monotonic() + tasks = [] + interval = args.ramp_seconds / count if args.ramp_seconds > 0 and count else 0.0 + for i in range(count): + tasks.append(asyncio.ensure_future(one(i))) + if interval: + await asyncio.sleep(interval) + await asyncio.gather(*tasks) + wall = time.monotonic() - started + + return { + "worker": worker_index, + "requested": count, + "completed": completed, + "retries_408": retries_408, + "reconnects": reconnects, + "asample_retries": asample_retries, + "poll_timeouts": poll_timeouts, + "errors": dict(errors), + "error_samples": error_samples, + "wall_seconds": wall, + "submit_latency": submit_latency, + "e2e_latency": e2e_latency, + } + + +def run_raw_worker(args: argparse.Namespace, worker_index: int, count: int, out_path: Path) -> None: + result = asyncio.run(raw_worker(args, worker_index, count)) + out_path.write_text(json.dumps(result)) + + +# --------------------------------------------------------------------------- # +# Tinker SDK load client +# --------------------------------------------------------------------------- # + + +async def sdk_load(args: argparse.Namespace) -> dict[str, Any]: + os.environ.setdefault("TINKER_API_KEY", "tml-dummy") + import tinker + from tinker import types as ttypes + + service_client = tinker.ServiceClient(base_url=args.url, timeout=SDK_RETRIEVE_POLL_TIMEOUT_SECONDS) + sampling_client = await service_client.create_sampling_client_async(base_model=args.base_model) + prompt = ttypes.ModelInput.from_ints([1000] * args.prompt_tokens) + params = ttypes.SamplingParams(max_tokens=args.max_tokens, temperature=1.0, seed=0) + + e2e_latency: list[float] = [] + errors: Counter[str] = Counter() + error_samples: dict[str, str] = {} + + async def one() -> None: + t0 = time.monotonic() + try: + await sampling_client.sample_async(prompt=prompt, num_samples=args.num_samples, sampling_params=params) + e2e_latency.append(time.monotonic() - t0) + except Exception as e: # noqa: BLE001 + kind = type(e).__name__ + errors[kind] += 1 + error_samples.setdefault(kind, str(e)[:300]) + + started = time.monotonic() + await asyncio.gather(*(one() for _ in range(args.num_requests))) + return { + "worker": 0, + "requested": args.num_requests, + "completed": len(e2e_latency), + "retries_408": 0, + "reconnects": 0, + "asample_retries": 0, + "poll_timeouts": 0, + "errors": dict(errors), + "error_samples": error_samples, + "wall_seconds": time.monotonic() - started, + "submit_latency": [], + "e2e_latency": e2e_latency, + } + + +# --------------------------------------------------------------------------- # +# Load role: fan out to worker processes and merge +# --------------------------------------------------------------------------- # + + +def percentiles(values: list[float]) -> dict[str, float]: + if not values: + return {} + ordered = sorted(values) + + def pct(p: float) -> float: + return ordered[min(len(ordered) - 1, int(math.ceil(p / 100 * len(ordered))) - 1)] + + return { + "p50": round(statistics.median(ordered), 3), + "p90": round(pct(90), 3), + "p99": round(pct(99), 3), + "max": round(ordered[-1], 3), + } + + +def merge_worker_results(results: list[dict[str, Any]]) -> dict[str, Any]: + errors: Counter[str] = Counter() + error_samples: dict[str, str] = {} + submit: list[float] = [] + e2e: list[float] = [] + for r in results: + errors.update(r["errors"]) + for kind, sample in r["error_samples"].items(): + error_samples.setdefault(kind, sample) + submit.extend(r["submit_latency"]) + e2e.extend(r["e2e_latency"]) + completed = sum(r["completed"] for r in results) + wall = max(r["wall_seconds"] for r in results) + return { + "requested": sum(r["requested"] for r in results), + "completed": completed, + "failed": sum(errors.values()), + "retries_408": sum(r["retries_408"] for r in results), + "reconnects": sum(r.get("reconnects", 0) for r in results), + "asample_retries": sum(r.get("asample_retries", 0) for r in results), + "poll_timeouts": sum(r.get("poll_timeouts", 0) for r in results), + "errors": dict(errors.most_common()), + "error_samples": error_samples, + "wall_seconds": round(wall, 2), + "throughput_per_s": round(completed / wall, 1) if wall else None, + "submit_latency_s": percentiles(submit), + "e2e_latency_s": percentiles(e2e), + } + + +def run_load(args: argparse.Namespace) -> dict[str, Any]: + if args.client == "sdk": + return merge_worker_results([asyncio.run(sdk_load(args))]) + + workers = args.workers or max(1, math.ceil(args.num_requests / DEFAULT_REQUESTS_PER_WORKER)) + per_worker = [args.num_requests // workers + (1 if i < args.num_requests % workers else 0) for i in range(workers)] + with tempfile.TemporaryDirectory(prefix="tinker_load_") as tmp: + procs = [] + for i, count in enumerate(per_worker): + out = Path(tmp) / f"worker_{i}.json" + cmd = [ + sys.executable, + __file__, + "--role", + "raw-worker", + "--url", + args.url, + "--base-model", + args.base_model, + "--num-requests", + str(count), + "--max-tokens", + str(args.max_tokens), + "--prompt-tokens", + str(args.prompt_tokens), + "--num-samples", + str(args.num_samples), + "--ramp-seconds", + str(args.ramp_seconds), + "--request-timeout", + str(args.request_timeout), + "--max-reconnects", + str(args.max_reconnects), + "--poll-timeout", + str(args.poll_timeout), + "--submit-concurrency", + str(max(1, args.submit_concurrency // workers)), + "--worker-index", + str(i), + "--worker-out", + str(out), + ] + if args.proto: + cmd.append("--proto") + if not args.source_ips: + cmd.append("--no-source-ips") + procs.append((subprocess.Popen(cmd), out)) + results = [] + for proc, out in procs: + proc.wait() + if proc.returncode != 0 or not out.exists(): + raise RuntimeError(f"load worker {out.stem} exited with {proc.returncode}") + results.append(json.loads(out.read_text())) + merged = merge_worker_results(results) + merged["workers"] = workers + return merged + + +# --------------------------------------------------------------------------- # +# API server role: the real app, configured like the non-colocated megatron server +# --------------------------------------------------------------------------- # + +NON_COLOCATED_MEGATRON_BACKEND_CONFIG = { + "strategy": "megatron", + "trainer.placement.policy_num_gpus_per_node": 2, + "trainer.placement.policy_num_nodes": 1, + "trainer.placement.colocate_all": False, + "trainer.policy.megatron_config.tensor_model_parallel_size": 1, + "trainer.policy.megatron_config.pipeline_model_parallel_size": 1, + "trainer.policy.megatron_config.lora_config.merge_lora": False, + "trainer.policy.model.lora.max_loras": 4, + "trainer.policy.model.lora.max_cpu_loras": 4, +} + + +def run_server(args: argparse.Namespace) -> None: + """Serve ``skyrl.tinker.api`` with the engine stubbed out and the router URL pre-published. + + The lifespan runs unmodified: with ``backend=megatron`` and ``colocate_all`` + false it installs ``SkyRLTrainInferenceForwardingClient``, which resolves + the vLLM router URL from ``EngineStateDB``. The engine (which would need + GPUs and a real vLLM) is the only thing replaced -- by a sleeping process -- + since it is not on the async sample path being measured. + """ + import uvicorn + from sqlmodel import Session, SQLModel, create_engine + + from skyrl.tinker import api as tinker_api + from skyrl.tinker.config import EngineConfig + from skyrl.tinker.db_models import EngineStateDB + from skyrl.utils.log import get_uvicorn_log_config + + workdir = Path(args.server_workdir) + workdir.mkdir(parents=True, exist_ok=True) + database_url = f"sqlite:///{workdir / 'tinker.db'}" + sync_engine = create_engine(database_url) + SQLModel.metadata.create_all(sync_engine) + with Session(sync_engine) as session: + session.merge(EngineStateDB(singleton_id=1, inference_proxy_url=args.vllm_url)) + session.commit() + sync_engine.dispose() + + tinker_api._build_uv_run_cmd_engine = lambda parent_cmd, engine_config: [ + sys.executable, + "-c", + "import time; time.sleep(10**9)", + ] + tinker_api.app.state.engine_config = EngineConfig( + base_model=args.base_model, + backend="megatron", + backend_config=NON_COLOCATED_MEGATRON_BACKEND_CONFIG, + database_url=database_url, + checkpoints_base=str(workdir / "checkpoints"), + forwarding_inference_max_connections=args.forwarding_max_connections, + forwarding_inference_timeout_sec=args.forwarding_timeout, + ) + profile_path = os.environ.get("TINKER_LOADTEST_PROFILE") + if profile_path: + # Profile the whole server process; the orchestrator sends SIGUSR1 + # after the load finishes and the stats are dumped from the handler. + import cProfile + + profiler = cProfile.Profile() + + def dump_profile(signum, frame): + profiler.disable() + profiler.dump_stats(profile_path) + + signal.signal(signal.SIGUSR1, dump_profile) + profiler.enable() + uvicorn.run(tinker_api.app, host="127.0.0.1", port=args.api_port, log_config=get_uvicorn_log_config()) + + +# --------------------------------------------------------------------------- # +# Orchestrator: fake vLLM + real API server + monitor + load +# --------------------------------------------------------------------------- # + + +async def http_get_json(session: aiohttp.ClientSession, url: str) -> dict | None: + try: + async with session.get(url, timeout=aiohttp.ClientTimeout(total=10)) as resp: + return await resp.json() + except Exception: # noqa: BLE001 + return None + + +def wait_for_http(url: str, timeout: float) -> None: + async def poll() -> None: + deadline = time.monotonic() + timeout + async with aiohttp.ClientSession() as session: + while time.monotonic() < deadline: + if await http_get_json(session, url) is not None: + return + await asyncio.sleep(0.5) + raise TimeoutError(f"{url} did not come up within {timeout}s") + + asyncio.run(poll()) + + +class ServerMonitor: + """Samples the API server's RSS / fds / CPU and /healthz latency while load runs.""" + + def __init__(self, pid: int | None, health_url: str): + self._proc = subprocess.Popen( + [sys.executable, __file__, "--role", "monitor", "--url", health_url, "--monitor-pid", str(pid or 0)], + stdout=subprocess.PIPE, + text=True, + ) + + def stop(self) -> dict[str, Any]: + self._proc.send_signal(signal.SIGINT) + out, _ = self._proc.communicate(timeout=30) + return json.loads(out.strip().splitlines()[-1]) if out.strip() else {} + + +def run_monitor(args: argparse.Namespace) -> None: + """Subprocess body for ServerMonitor: prints one JSON summary on SIGINT.""" + proc = psutil.Process(args.monitor_pid) if args.monitor_pid else None + peak_rss = peak_fds = 0 + cpu: list[float] = [] + health: list[float] = [] + failures = 0 + + async def loop() -> None: + nonlocal peak_rss, peak_fds, failures + async with aiohttp.ClientSession() as session: + while True: + if proc is not None: + try: + # The API process only; the engine child is idle on this path. + peak_rss = max(peak_rss, proc.memory_info().rss) + peak_fds = max(peak_fds, proc.num_fds()) + cpu.append(proc.cpu_percent(interval=None)) + except psutil.Error: + pass + t0 = time.monotonic() + if await http_get_json(session, args.url) is None: + failures += 1 + else: + health.append(time.monotonic() - t0) + await asyncio.sleep(1.0) + + try: + asyncio.run(loop()) + except KeyboardInterrupt: + pass + print( + json.dumps( + { + "peak_rss_gb": round(peak_rss / 1e9, 2), + "peak_fds": peak_fds, + "cpu_percent_p50": round(statistics.median(cpu), 1) if cpu else None, + "cpu_percent_max": round(max(cpu), 1) if cpu else None, + "healthz_latency_s": percentiles(health), + "healthz_failures": failures, + } + ), + flush=True, + ) + + +def start_api_server(args: argparse.Namespace, workdir: Path, vllm_url: str) -> tuple[subprocess.Popen, Path]: + log_path = workdir / "server.log" + cmd = [ + sys.executable, + __file__, + "--role", + "server", + "--api-port", + str(args.api_port), + "--base-model", + args.base_model, + "--vllm-url", + vllm_url, + "--server-workdir", + str(workdir), + "--forwarding-timeout", + str(args.forwarding_timeout), + ] + if args.forwarding_max_connections is not None: + cmd += ["--forwarding-max-connections", str(args.forwarding_max_connections)] + print(f"[orchestrator] starting API server: {' '.join(cmd)}") + print(f"[orchestrator] server log: {log_path}") + proc = subprocess.Popen(cmd, stdout=open(log_path, "w"), stderr=subprocess.STDOUT, start_new_session=True) + return proc, log_path + + +def print_sysctl_hints(args: argparse.Namespace) -> None: + try: + lo, hi = map(int, Path("/proc/sys/net/ipv4/ip_local_port_range").read_text().split()) + somaxconn = int(Path("/proc/sys/net/core/somaxconn").read_text()) + except OSError: + return + per_ip = hi - lo + 1 + workers = args.workers or max(1, math.ceil(args.num_requests / DEFAULT_REQUESTS_PER_WORKER)) + print( + f"[orchestrator] ephemeral ports per source IP: {per_ip}, somaxconn: {somaxconn}, nofile: {os.sysconf('SC_OPEN_MAX')}" + ) + if args.client == "raw" and not args.source_ips and args.num_requests > per_ip * 0.9: + print("[orchestrator] WARNING: a single source IP cannot hold this many connections; drop --no-source-ips") + if args.client == "raw" and args.source_ips and args.num_requests / workers > per_ip * 0.9: + print( + f"[orchestrator] WARNING: {args.num_requests / workers:.0f} connections per worker exceeds one IP's ports; raise --workers" + ) + if args.num_requests > per_ip * 0.9: + print( + "[orchestrator] NOTE: the API server forwards every in-flight sample over its own outbound connection " + f"from one IP, so more than ~{per_ip} simultaneously forwarded requests will fail with EADDRNOTAVAIL" + ) + + +def _raise_keyboard_interrupt(signum, frame): + raise KeyboardInterrupt + + +def run_all(args: argparse.Namespace) -> dict[str, Any]: + signal.signal(signal.SIGTERM, _raise_keyboard_interrupt) + print_sysctl_hints(args) + workdir = Path(tempfile.mkdtemp(prefix="tinker_sampling_load_")) + vllm_url = f"http://127.0.0.1:{args.vllm_port}" + api_url = args.server_url or f"http://127.0.0.1:{args.api_port}" + children: list[subprocess.Popen] = [] + server_proc: subprocess.Popen | None = None + log_path: Path | None = None + try: + vllm_cmd = [ + sys.executable, + __file__, + "--role", + "vllm", + "--vllm-port", + str(args.vllm_port), + "--vllm-mode", + args.vllm_mode, + "--barrier-requests", + str(args.barrier_requests or args.num_requests), + "--barrier-timeout", + str(args.barrier_timeout), + "--max-num-seqs", + str(args.max_num_seqs), + "--gen-seconds", + str(args.gen_seconds), + ] + children.append(subprocess.Popen(vllm_cmd)) + wait_for_http(f"{vllm_url}/healthz", 30) + + if args.server_url is None: + server_proc, log_path = start_api_server(args, workdir, vllm_url) + children.append(server_proc) + wait_for_http(f"{api_url}/api/v1/healthz", args.server_startup_timeout) + api_pid = server_proc.pid if server_proc is not None else None + print( + f"[orchestrator] API server up (pid {api_pid}); starting load: {args.num_requests} requests via {args.client}" + ) + + monitor = ServerMonitor(api_pid, f"{api_url}/api/v1/healthz") + args.url = api_url + load = run_load(args) + server_stats = monitor.stop() + if server_proc is not None and os.environ.get("TINKER_LOADTEST_PROFILE"): + server_proc.send_signal(signal.SIGUSR1) + time.sleep(3) + + async def vllm_stats() -> dict | None: + async with aiohttp.ClientSession() as session: + return await http_get_json(session, f"{vllm_url}/stats") + + report = { + "config": { + "num_requests": args.num_requests, + "client": args.client, + "vllm_mode": args.vllm_mode, + "max_tokens": args.max_tokens, + "prompt_tokens": args.prompt_tokens, + "proto": args.proto, + "forwarding_max_connections": args.forwarding_max_connections, + "forwarding_timeout": args.forwarding_timeout, + }, + "load": load, + "fake_vllm": asyncio.run(vllm_stats()), + "api_server": server_stats, + "server_log": str(log_path) if log_path else None, + } + return report + finally: + for proc in reversed(children): + try: + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) if proc is server_proc else proc.terminate() + proc.wait(timeout=20) + except Exception: # noqa: BLE001 + proc.kill() + + +def print_report(report: dict[str, Any]) -> None: + load = report["load"] + print("\n=== Tinker sampling load test ===") + print(f"config: {json.dumps(report['config'])}") + print( + f"requested={load['requested']} completed={load['completed']} failed={load['failed']} " + f"408_retries={load['retries_408']} reconnects={load.get('reconnects', 0)} " + f"asample_retries={load.get('asample_retries', 0)} poll_timeouts={load.get('poll_timeouts', 0)}" + ) + print(f"wall={load['wall_seconds']}s throughput={load['throughput_per_s']} req/s") + print(f"submit latency (s): {load['submit_latency_s']}") + print(f"e2e latency (s): {load['e2e_latency_s']}") + if load["errors"]: + print("errors by class:") + for kind, n in load["errors"].items(): + print(f" {kind}: {n} e.g. {load['error_samples'].get(kind, '')!r}") + if report.get("fake_vllm"): + print( + f"fake vLLM: {json.dumps(report['fake_vllm'])} <- peak_in_flight is the concurrency the server sustained" + ) + if report.get("api_server"): + print(f"API server: {json.dumps(report['api_server'])}") + if report.get("server_log"): + print(f"server log: {report['server_log']}") + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--role", choices=["all", "load", "vllm", "server", "raw-worker", "monitor"], default="all") + p.add_argument("--num-requests", type=int, default=131072) + p.add_argument("--client", choices=["raw", "sdk"], default="raw") + p.add_argument("--workers", type=int, default=0, help="raw client processes (default: ceil(n / 16384))") + p.add_argument("--no-source-ips", dest="source_ips", action="store_false", help="bind all workers to 127.0.0.1") + p.add_argument("--ramp-seconds", type=float, default=0.0, help="spread request launch over this many seconds") + p.add_argument("--request-timeout", type=float, default=900.0, help="give up on a request after this many seconds") + p.add_argument( + "--poll-timeout", + type=float, + default=SDK_RETRIEVE_POLL_TIMEOUT_SECONDS, + help="client-side timeout per retrieve_future poll before re-polling (Tinker SDK: 45s)", + ) + p.add_argument( + "--submit-concurrency", + type=int, + default=400, + help="max asample POSTs in flight across all workers (Tinker SDK default: 400 per client)", + ) + p.add_argument( + "--max-reconnects", + type=float, + default=2.0, + help="retrieve_future reconnect budget per worker, as a multiple of its request count (SDK-style retries)", + ) + p.add_argument("--proto", action="store_true", help="retrieve results as protobuf (Accept: application/x-protobuf)") + p.add_argument("--max-tokens", type=int, default=128) + p.add_argument("--prompt-tokens", type=int, default=64) + p.add_argument("--num-samples", type=int, default=1) + p.add_argument("--base-model", default=TINY_MODEL) + p.add_argument("--url", default=f"http://127.0.0.1:{DEFAULT_API_PORT}", help="API server URL for --role load") + p.add_argument( + "--server-url", default=None, help="--role all: use this running API server instead of launching one" + ) + p.add_argument( + "--forwarding-max-connections", + type=int, + default=None, + help="EngineConfig.forwarding_inference_max_connections for the launched server (default: unlimited)", + ) + p.add_argument( + "--forwarding-timeout", type=float, default=300.0, help="EngineConfig.forwarding_inference_timeout_sec" + ) + p.add_argument("--server-startup-timeout", type=float, default=120.0) + p.add_argument("--api-port", type=int, default=DEFAULT_API_PORT) + p.add_argument("--vllm-port", type=int, default=DEFAULT_VLLM_PORT) + p.add_argument("--vllm-mode", choices=["barrier", "latency"], default="barrier") + p.add_argument("--barrier-requests", type=int, default=0, help="barrier size (default: --num-requests)") + p.add_argument("--barrier-timeout", type=float, default=600.0, help="release the barrier anyway after this long") + p.add_argument("--max-num-seqs", type=int, default=1024, help="latency mode: concurrent generations") + p.add_argument("--gen-seconds", type=float, default=2.0, help="latency mode: seconds per generation") + p.add_argument("--json-out", type=Path, default=None) + # internal + p.add_argument("--vllm-url", default=f"http://127.0.0.1:{DEFAULT_VLLM_PORT}") + p.add_argument("--server-workdir", default=None) + p.add_argument("--worker-index", type=int, default=0) + p.add_argument("--worker-out", type=Path, default=None) + p.add_argument("--monitor-pid", type=int, default=0) + return p.parse_args() + + +def main() -> int: + sys.stdout.reconfigure(line_buffering=True) + args = parse_args() + if args.role == "vllm": + run_fake_vllm(args) + return 0 + if args.role == "server": + run_server(args) + return 0 + if args.role == "raw-worker": + run_raw_worker(args, args.worker_index, args.num_requests, args.worker_out) + return 0 + if args.role == "monitor": + run_monitor(args) + return 0 + if args.role == "load": + report = {"config": {"num_requests": args.num_requests, "client": args.client}, "load": run_load(args)} + else: + report = run_all(args) + print_report(report) + if args.json_out: + args.json_out.write_text(json.dumps(report, indent=2) + "\n") + return 0 if report["load"]["failed"] == 0 and report["load"]["completed"] == report["load"]["requested"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index 5432555af5..2328a850cc 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -308,15 +308,19 @@ async def lifespan(app: FastAPI): # SkyRL-Train default is colocate_all=True; only opt into forwarding # when the operator explicitly sets it to False. is_colocated = bool(backend_cfg.get("trainer.placement.colocate_all", True)) + store_ttls = dict( + retrieved_ttl_sec=app.state.engine_config.external_future_retrieved_ttl_sec, + completed_ttl_sec=app.state.engine_config.external_future_completed_ttl_sec, + ) if app.state.engine_config.external_inference_url: - app.state.external_future_store = ExternalFutureStore() + app.state.external_future_store = ExternalFutureStore(**store_ttls) await app.state.external_future_store.start() app.state.external_inference_client = ExternalInferenceClient( app.state.engine_config, app.state.db_engine, app.state.external_future_store ) logger.info(f"External engine configured: {app.state.engine_config.external_inference_url}") elif backend_name in ("megatron", "fsdp") and not is_colocated: - app.state.external_future_store = ExternalFutureStore() + app.state.external_future_store = ExternalFutureStore(**store_ttls) await app.state.external_future_store.start() app.state.external_inference_client = SkyRLTrainInferenceForwardingClient( app.state.engine_config, app.state.db_engine, app.state.external_future_store @@ -929,12 +933,17 @@ class WeightsInfoResponse(BaseModel): class ClientConfigResponse(BaseModel): pjwt_auth_enabled: bool = False + # Cap on in-flight samples per SamplingClient; the SDK applies it as its + # sampling RetryConfig.max_connections. + sample_max_concurrent_requests: int = 2000 @app.post("/api/v1/client/config", response_model=ClientConfigResponse) -async def client_config(): - """Stub for tinker SDK client_config handshake.""" - return ClientConfigResponse() +async def client_config(req: Request): + """Tinker SDK client_config handshake: server-side flags for the client.""" + return ClientConfigResponse( + sample_max_concurrent_requests=req.app.state.engine_config.sample_max_concurrent_requests, + ) @app.get("/api/v1/healthz", response_model=HealthResponse) @@ -1505,18 +1514,32 @@ async def retrieve_future(request: RetrieveFutureRequest, req: Request): types.RequestType(request_type) in PROTO_SERIALIZABLE_REQUEST_TYPES and PROTO_CONTENT_TYPE in req.headers.get("accept", "").lower() ): - async with req.app.state.proto_serialization_lock: - content = await asyncio.to_thread( - _serialize_proto_result, - types.RequestType(request_type), - result_data, - ) + # Forwarded samples are stored as proto already and go out as-is; + # anything stored as JSON is encoded once here and cached. + content = external_future_store.proto_result(request_id) if found_in_memory else None + if content is None: + async with req.app.state.proto_serialization_lock: + content = await asyncio.to_thread( + _serialize_proto_result, + types.RequestType(request_type), + result_data, + ) + if found_in_memory: + external_future_store.cache_proto(request_id, content) response: Response = Response(content=content, media_type=PROTO_CONTENT_TYPE) else: + if result_data is None and found_in_memory: + # Stored as proto only; a pre-proto client wants JSON. + result_data = external_future_store.json_result(request_id) response = raw_json_response(result_data) # Start the retry-grace clock now that the response is built and about to - # be sent, so a large result is never evicted mid-delivery. - if found_in_memory: + # be sent, so a large result is never evicted mid-delivery -- but only if + # this client is still there to receive it. The SDK abandons a poll after + # 45s and retries the same request_id; if the result lands after that, + # this handler wakes on a dead connection (uvicorn drops the send + # silently) and starting the short clock here would let the sweeper + # evict a result nobody received, turning the retry into a 404. + if found_in_memory and not await req.is_disconnected(): external_future_store.mark_retrieved(request_id) return response diff --git a/skyrl/tinker/config.py b/skyrl/tinker/config.py index cd8d5c4a23..b81628f2c0 100644 --- a/skyrl/tinker/config.py +++ b/skyrl/tinker/config.py @@ -59,18 +59,46 @@ class EngineConfig(BaseModel): json_schema_extra={"argparse_type": lambda v: None if v == "None" else int(v)}, ) forwarding_inference_timeout_sec: float = Field( - default=300.0, + default=2048.0, gt=0, description=( "Read timeout in seconds for API-side requests forwarded to the " "SkyRL-Train-managed inference engine. This must cover time spent " - "queued behind other requests as well as generation time." + "queued behind other requests as well as generation time: with the " + "default unlimited connection count a large rollout burst waits inside " + "vLLM's queue, and 128x128 bursts routinely exceed 300s there." ), json_schema_extra={ "argparse_type": float, "env_var": "SKYRL_FORWARDING_INFERENCE_TIMEOUT_SEC", }, ) + external_future_retrieved_ttl_sec: float = Field( + default=300.0, + gt=0, + description=( + "How long a forwarded sample result stays in memory after it was delivered, so an " + "SDK retry after a lost HTTP response still finds it. Must outlast the SDK's worst-case " + "re-poll gap (45s poll timeout + up to 30s backoff, twice). Memory for long-output " + "rollouts is roughly completion rate x result size x this window." + ), + ) + external_future_completed_ttl_sec: float = Field( + default=600.0, + gt=0, + description="How long a completed but never-delivered forwarded sample result stays in memory.", + ) + sample_max_concurrent_requests: int = Field( + default=2000, + gt=0, + description=( + "Advertised to the Tinker SDK via /api/v1/client/config as the maximum " + "number of sample requests one SamplingClient keeps in flight; the SDK " + "queues the rest client-side. Each in-flight sample costs the API server " + "one open long-poll connection and one retrieve_future re-poll every 45s. " + "2000 is the SDK's own default." + ), + ) session_cleanup_interval_sec: int = Field( default=60, description="How often to check for stale sessions (seconds). Set to -1 to disable cleanup.", diff --git a/skyrl/tinker/external_future_store.py b/skyrl/tinker/external_future_store.py index e69e6999d4..14e1a7b9c6 100644 --- a/skyrl/tinker/external_future_store.py +++ b/skyrl/tinker/external_future_store.py @@ -9,16 +9,31 @@ from skyrl.tinker import types from skyrl.tinker.db_models import RequestStatus +from skyrl.tinker.proto_serialization import sample_output_json_from_proto from skyrl.utils.log import logger +@dataclass +class PreparedResult: + """A completed sample result already in wire form. + + Forwarded samples are encoded to proto once, straight from the decoded vLLM + body; JSON text is produced only if a pre-proto client asks for it. + """ + + proto: bytes | None = None + json: str | None = None + + @dataclass class ExternalFuture: request_id: int model_id: str | None - request_data: dict status: RequestStatus = RequestStatus.PENDING + # At most one of these is populated at completion; the other is derived + # lazily on the first request for it and then cached for retries. result_data: str | None = None + result_proto: bytes | None = None created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) completed_at: datetime | None = None retrieved_at: datetime | None = None @@ -44,15 +59,23 @@ class ExternalFutureStore: # retry following a lost HTTP response still finds it. Measured from # delivery (mark_retrieved), never from the in-store read: a large result # can spend minutes being serialized and sent, and starting the clock at - # read would evict it mid-delivery. - _RETRIEVED_TTL_SECONDS = 120.0 + # read would evict it mid-delivery. The SDK re-polls after a 45s client + # timeout plus up to 30s of backoff, so two consecutive misses span 150s; + # the grace has to outlast that. + _RETRIEVED_TTL_SECONDS = 300.0 # Completed but not yet delivered — governs the read/serialize/send window # and clients that never come back. _COMPLETED_TTL_SECONDS = 600.0 # Pending entries whose forwarding task died without completing them. _PENDING_TTL_SECONDS = 3600.0 - def __init__(self): + def __init__(self, *, retrieved_ttl_sec: float | None = None, completed_ttl_sec: float | None = None): + # Retention after delivery is the dominant memory term for long-output + # rollouts (results x retrieved TTL), so operators can tune it. + if retrieved_ttl_sec is not None: + self._RETRIEVED_TTL_SECONDS = retrieved_ttl_sec + if completed_ttl_sec is not None: + self._COMPLETED_TTL_SECONDS = completed_ttl_sec self._entries: dict[int, ExternalFuture] = {} # Boot-epoch id space: each server process starts below every id an # earlier process could plausibly have handed out (2^20 ids per @@ -66,15 +89,33 @@ async def start(self) -> None: self._sweeper = asyncio.create_task(self._sweep_loop()) def create(self, model_id: str | None, request_data: BaseModel) -> int: + # The request body is not retained: nothing reads it back on this + # path, and a long prompt would cost ~100KB per pending entry. request_id = self._next_request_id self._next_request_id -= 1 - self._entries[request_id] = ExternalFuture( - request_id=request_id, - model_id=model_id, - request_data=request_data.model_dump(mode="json"), - ) + self._entries[request_id] = ExternalFuture(request_id=request_id, model_id=model_id) return request_id + def proto_result(self, request_id: int) -> bytes | None: + """Proto wire bytes for a completed result, if it has them.""" + entry = self._entries.get(request_id) + return entry.result_proto if entry is not None else None + + def cache_proto(self, request_id: int, proto: bytes) -> None: + """Keep a proto encoding produced at retrieval so retries skip re-encoding.""" + entry = self._entries.get(request_id) + if entry is not None: + entry.result_proto = proto + + def json_result(self, request_id: int) -> str | None: + """JSON text for a completed result, deriving it from proto on first use.""" + entry = self._entries.get(request_id) + if entry is None: + return None + if entry.result_data is None and entry.result_proto is not None: + entry.result_data = sample_output_json_from_proto(entry.result_proto) + return entry.result_data + async def wait(self, request_id: int, timeout: float) -> tuple[RequestStatus, types.RequestType, str | None] | None: entry = self._entries.get(request_id) if entry is None: @@ -97,13 +138,17 @@ def mark_retrieved(self, request_id: int) -> None: if entry is not None: entry.retrieved_at = datetime.now(timezone.utc) - async def complete(self, request_id: int, result_data: BaseModel, status: RequestStatus) -> None: + async def complete(self, request_id: int, result_data: BaseModel | PreparedResult, status: RequestStatus) -> None: entry = self._entries.get(request_id) if entry is None: # Swept as abandoned before the forwarding task finished. logger.warning("External future %s was evicted before its result arrived — dropping", request_id) return - entry.result_data = result_data.model_dump_json() + if isinstance(result_data, PreparedResult): + entry.result_data = result_data.json + entry.result_proto = result_data.proto + else: + entry.result_data = result_data.model_dump_json() entry.status = status entry.completed_at = datetime.now(timezone.utc) entry.event.set() diff --git a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py index d617d9a7a6..681768a5d5 100644 --- a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py +++ b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py @@ -7,7 +7,8 @@ import asyncio from datetime import datetime, timezone -import httpx +import aiohttp +import orjson from sqlmodel.ext.asyncio.session import AsyncSession from skyrl.backends.renderer import render_model_input @@ -15,10 +16,18 @@ from skyrl.tinker import types from skyrl.tinker.config import EngineConfig from skyrl.tinker.db_models import EngineStateDB, FutureDB, RequestStatus -from skyrl.tinker.external_future_store import ExternalFutureStore +from skyrl.tinker.external_future_store import ExternalFutureStore, PreparedResult +from skyrl.tinker.proto_serialization import ( + sample_output_json_from_proto, + serialize_sample_output, +) from skyrl.utils.log import logger +class TransientInferenceError(RuntimeError): + """A 5xx from vllm-router/vLLM: the request was rejected, not executed, so it is safe to retry.""" + + class SkyRLTrainInferenceForwardingClient: """Forwards EXTERNAL sample requests to the SkyRL-Train-managed vLLM.""" @@ -36,27 +45,42 @@ def __init__( self.external_future_store = external_future_store self._cached_proxy_url: str | None = None self._cache_lock = asyncio.Lock() - # Backpressure layered: httpx pool -> vllm-router -> vLLM max_num_seqs. - # Default `forwarding_inference_max_connections=None` is unlimited; - # the only cost is file descriptors (raise `ulimit -n` accordingly). - max_conn = engine_config.forwarding_inference_max_connections - max_keepalive = max(max_conn // 4, 32) if max_conn is not None else None - self._http_client: httpx.AsyncClient = httpx.AsyncClient( - timeout=httpx.Timeout( - connect=10.0, - read=engine_config.forwarding_inference_timeout_sec, - write=300.0, - pool=300.0, - ), - limits=httpx.Limits( - max_connections=max_conn, - max_keepalive_connections=max_keepalive, - ), - ) + # Created on first use so it binds to the serving event loop. + self._session: aiohttp.ClientSession | None = None + + def _get_session(self) -> aiohttp.ClientSession: + """Return the shared aiohttp session, creating it on first use. + + Backpressure is layered: connector limit -> vllm-router -> vLLM + max_num_seqs. Default `forwarding_inference_max_connections=None` is + unlimited; the only cost is file descriptors (raise `ulimit -n` + accordingly). Requests beyond the limit wait in the connector's FIFO + queue with no deadline, so a backlog of many thousands of samples + drains at the engine's pace instead of failing. + + aiohttp rather than httpx: httpcore's pool rescans every connection + for every request, so its per-request CPU grows with the number of + in-flight samples (~28ms each at 512 in flight); aiohttp stays flat. + """ + if self._session is None or self._session.closed: + max_conn = self.engine_config.forwarding_inference_max_connections + # keepalive_timeout must stay under the router's idle timeout so a + # pooled connection is never reused after the server closed it. + connector = aiohttp.TCPConnector(limit=max_conn or 0, keepalive_timeout=2) + self._session = aiohttp.ClientSession( + connector=connector, + timeout=aiohttp.ClientTimeout( + total=None, + sock_connect=10.0, + sock_read=self.engine_config.forwarding_inference_timeout_sec, + ), + ) + return self._session async def aclose(self) -> None: - """Close the persistent httpx client. Called from api.py lifespan shutdown.""" - await self._http_client.aclose() + """Close the shared aiohttp session. Called from api.py lifespan shutdown.""" + if self._session is not None and not self._session.closed: + await self._session.close() async def _read_proxy_url_from_db(self) -> str | None: async with AsyncSession(self.db_engine) as session: @@ -113,45 +137,47 @@ async def call_and_store_result( logger.warning("FutureDB row %s missing on completion write — skipping", request_id) return # `result_data` is a text column holding pre-serialized JSON. - future.result_data = result.model_dump_json() + if isinstance(result, PreparedResult): + future.result_data = result.json or sample_output_json_from_proto(result.proto) + else: + future.result_data = result.model_dump_json() future.status = status future.completed_at = datetime.now(timezone.utc) await session.commit() - async def _forward_with_retry(self, sample_req, model_id: str, *, base_model: str | None) -> types.SampleOutput: - # Retry only failures that occur before a request can reach vLLM. Read - # and write failures are ambiguous: vLLM may still be executing the + async def _forward_with_retry(self, sample_req, model_id: str, *, base_model: str | None) -> PreparedResult: + # Retry only failures where the request demonstrably did not execute: + # connect-phase errors and 5xx rejections from the router. Read and + # write failures are ambiguous: vLLM may still be executing the # request, so retrying would duplicate generation load. try: try: proxy_url = await self._resolve_proxy_url() return await self._forward(proxy_url, sample_req, model_id, base_model=base_model) - except (httpx.ConnectError, httpx.ConnectTimeout) as e: + except (aiohttp.ClientConnectorError, aiohttp.ConnectionTimeoutError, TransientInferenceError) as e: logger.warning( - "Connection error talking to %s (%s: %s) — refreshing proxy URL and retrying once", + "Transient error talking to %s (%s: %s) — refreshing proxy URL and retrying once", self._cached_proxy_url, type(e).__name__, e, ) proxy_url = await self._resolve_proxy_url(force_refresh=True) return await self._forward(proxy_url, sample_req, model_id, base_model=base_model) - except httpx.ReadTimeout as e: + except aiohttp.SocketTimeoutError as e: # Not retried (see above). Long-context requests routinely exceed the # default read deadline, so tell the caller how to raise it. The # message is stored in the FutureDB ErrorResponse and shown to clients. timeout_sec = self.engine_config.forwarding_inference_timeout_sec raise RuntimeError( f"Inference request to {self._cached_proxy_url} timed out after {timeout_sec:g}s waiting for " - "a response (httpx.ReadTimeout). The request was not retried because vLLM may still be " + "a response (read timeout). The request was not retried because vLLM may still be " "executing it. If requests are expected to take this long (long prompts, large max_tokens, " "or queueing behind other requests), increase the deadline with " "`--forwarding-inference-timeout-sec` (EngineConfig.forwarding_inference_timeout_sec) or " "the SKYRL_FORWARDING_INFERENCE_TIMEOUT_SEC environment variable." ) from e - async def _forward( - self, proxy_url: str, sample_req, model_id: str, *, base_model: str | None - ) -> types.SampleOutput: + async def _forward(self, proxy_url: str, sample_req, model_id: str, *, base_model: str | None) -> PreparedResult: # model_id matches the LoRA name registered with vLLM during # save_weights_for_sampler; base_model is used for non-LoRA sampling. model_name = base_model if base_model else model_id @@ -195,17 +221,22 @@ async def _forward( headers["X-Session-ID"] = session_id url = f"{proxy_url}/v1/completions" - response = await self._http_client.post(url, json=payload, headers=headers) - if response.status_code >= 400: - raise RuntimeError(f"vLLM /v1/completions returned {response.status_code}: {response.text}") - try: - result = response.json() - except ValueError as e: - # vllm-router can return HTML on transient errors even with 2xx status. - raise RuntimeError( - f"vLLM /v1/completions returned non-JSON ({response.status_code}, " - f"content-type={response.headers.get('content-type')!r}): {response.text[:512]}" - ) from e + async with self._get_session().post(url, json=payload, headers=headers) as response: + body = await response.read() + if response.status >= 500: + raise TransientInferenceError( + f"vLLM /v1/completions returned {response.status}: {body.decode(errors='replace')}" + ) + if response.status >= 400: + raise RuntimeError(f"vLLM /v1/completions returned {response.status}: {body.decode(errors='replace')}") + try: + result = orjson.loads(body) + except orjson.JSONDecodeError as e: + # vllm-router can return HTML on transient errors even with 2xx status. + raise RuntimeError( + f"vLLM /v1/completions returned non-JSON ({response.status}, " + f"content-type={response.headers.get('content-type')!r}): {body[:512].decode(errors='replace')}" + ) from e prompt_logprobs = None topk = None @@ -231,16 +262,8 @@ async def _forward( # Tinker's stop_reason is Literal["stop", "length"]; vLLM emits a wider set. finish_reason = choice.get("finish_reason") stop_reason = "stop" if finish_reason in ("stop", "stop_token") else "length" - sequences.append( - types.GeneratedSequence( - tokens=tokens, - logprobs=logprobs, - stop_reason=stop_reason, - ) - ) + sequences.append((stop_reason, tokens, logprobs)) - return types.SampleOutput( - sequences=sequences, - prompt_logprobs=prompt_logprobs, - topk_prompt_logprobs=topk, - ) + # Encode straight to the proto wire form the SDK retrieves; no pydantic + # model or JSON text is built for the result (see PreparedResult). + return PreparedResult(proto=serialize_sample_output(sequences, prompt_logprobs, topk)) diff --git a/skyrl/tinker/proto_serialization.py b/skyrl/tinker/proto_serialization.py index 03ec7f4af3..474be5bcf0 100644 --- a/skyrl/tinker/proto_serialization.py +++ b/skyrl/tinker/proto_serialization.py @@ -19,8 +19,10 @@ """ import base64 +from collections.abc import Iterable, Sequence import numpy as np +import orjson from tinker.proto import tinker_public_pb2 as pb from skyrl.tinker import types @@ -124,24 +126,44 @@ def serialize_result(request_type: types.RequestType, result_data: dict) -> byte def _serialize_sample_output(result_data: dict) -> bytes: output = types.SampleOutput.model_validate(result_data) + return serialize_sample_output( + [(seq.stop_reason, seq.tokens, seq.logprobs) for seq in output.sequences], + output.prompt_logprobs, + output.topk_prompt_logprobs, + ) + + +def serialize_sample_output( + sequences: Iterable[tuple[str, Sequence[int], Sequence[float]]], + prompt_logprobs: Sequence[float | None] | None, + topk_prompt_logprobs: Sequence[Sequence[tuple[int, float]] | None] | None, +) -> bytes: + """Build ``SampleResponse`` wire bytes from plain Python data. + + ``sequences`` holds ``(stop_reason, tokens, logprobs)`` per sequence. This is + the hot path for forwarded samples: the vLLM body is decoded once and + encoded straight to proto, with no pydantic model and no JSON text in + between (each of which costs about as much as this whole function for a + 32k-token result). + """ proto = pb.SampleResponse() - for seq in output.sequences: + for stop_reason, tokens, logprobs in sequences: proto.sequences.append( pb.SampledSequence( - stop_reason=_STOP_REASON_TO_PROTO[seq.stop_reason], - tokens=np.asarray(seq.tokens, dtype=np.int32).tobytes(), - logprobs=np.asarray(seq.logprobs, dtype=np.float32).tobytes(), + stop_reason=_STOP_REASON_TO_PROTO[stop_reason], + tokens=np.asarray(tokens, dtype=np.int32).tobytes(), + logprobs=np.asarray(logprobs, dtype=np.float32).tobytes(), ) ) - if output.prompt_logprobs is not None: + if prompt_logprobs is not None: proto.prompt_logprobs = np.array( - [np.nan if lp is None else lp for lp in output.prompt_logprobs], dtype=np.float32 + [np.nan if lp is None else lp for lp in prompt_logprobs], dtype=np.float32 ).tobytes() - if output.topk_prompt_logprobs is not None: - rows = output.topk_prompt_logprobs + if topk_prompt_logprobs is not None: + rows = topk_prompt_logprobs # k is not recorded in the result, so recover it from the widest row. # With every row undefined, use k=1 so prompt_length stays encoded # (the client maps fully-masked rows back to None). @@ -164,6 +186,49 @@ def _serialize_sample_output(result_data: dict) -> bytes: return proto.SerializeToString() +_PROTO_TO_STOP_REASON = {value: key for key, value in _STOP_REASON_TO_PROTO.items()} + + +def sample_output_json_from_proto(proto_bytes: bytes) -> str: + """Inverse of :func:`serialize_sample_output`, as ``SampleOutput`` JSON text. + + Serves clients that predate proto results (SDK < 0.25) from a result that + was stored as proto. Logprobs come back at float32 precision, the same + values proto clients receive. + """ + proto = pb.SampleResponse.FromString(proto_bytes) + sequences = [ + { + "stop_reason": _PROTO_TO_STOP_REASON[seq.stop_reason], + "tokens": np.frombuffer(seq.tokens, dtype=np.int32).tolist(), + "logprobs": np.frombuffer(seq.logprobs, dtype=np.float32).tolist(), + } + for seq in proto.sequences + ] + + prompt_logprobs = None + if proto.prompt_logprobs: + values = np.frombuffer(proto.prompt_logprobs, dtype=np.float32) + prompt_logprobs = [None if np.isnan(value) else float(value) for value in values] + + topk = None + if proto.HasField("topk_prompt_logprobs"): + block = proto.topk_prompt_logprobs + shape = (block.prompt_length, block.k) + token_ids = np.frombuffer(block.token_ids, dtype=np.int32).reshape(shape) + logprobs = np.frombuffer(block.logprobs, dtype=np.float32).reshape(shape) + masked = (token_ids == _TOPK_MASK_TOKEN_ID) & (logprobs == np.float32(_TOPK_MASK_LOGPROB)) + topk = [] + for row_ids, row_lps, row_mask in zip(token_ids, logprobs, masked): + row = [[int(t), float(lp)] for t, lp, m in zip(row_ids, row_lps, row_mask) if not m] + # A fully masked row encodes an undefined position (None). + topk.append(row or None) + + return orjson.dumps( + {"sequences": sequences, "prompt_logprobs": prompt_logprobs, "topk_prompt_logprobs": topk} + ).decode() + + def _serialize_forward_backward_output(result_data: dict) -> bytes: output = types.ForwardBackwardOutput.model_validate(result_data) proto = pb.ForwardBackwardOutput() diff --git a/skyrl/utils/log.py b/skyrl/utils/log.py index cd5be129d5..dad5cde16e 100644 --- a/skyrl/utils/log.py +++ b/skyrl/utils/log.py @@ -53,6 +53,15 @@ def get_uvicorn_log_config() -> dict: **RICH_HANDLER_KWARGS, "formatter": "default", }, + # Plain handler for the per-request access log: RichHandler renders + # each record through a rich Table (~1.5ms of event-loop CPU per + # record), which at thousands of requests per second is the API + # server's single largest CPU cost. + "access": { + "class": "logging.StreamHandler", + "formatter": "default", + "stream": "ext://sys.stderr", + }, }, "loggers": { # Main uvicorn logger (general server messages) @@ -60,7 +69,7 @@ def get_uvicorn_log_config() -> dict: # Uvicorn error logger (startup, shutdown, exceptions) "uvicorn.error": {"handlers": ["default"], "level": "INFO", "propagate": False}, # HTTP access logs (request/response logging) - "uvicorn.access": {"handlers": ["default"], "level": "INFO", "propagate": False}, + "uvicorn.access": {"handlers": ["access"], "level": "INFO", "propagate": False}, }, } diff --git a/tests/tinker/test_api_validation.py b/tests/tinker/test_api_validation.py index e644a0eb34..2e5294164b 100644 --- a/tests/tinker/test_api_validation.py +++ b/tests/tinker/test_api_validation.py @@ -18,6 +18,19 @@ def _make_datum() -> api.Datum: ) +@pytest.mark.asyncio +async def test_client_config_advertises_configured_sample_concurrency_cap(): + from types import SimpleNamespace + + from skyrl.tinker.config import EngineConfig + + req = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(engine_config=EngineConfig(base_model="m")))) + assert (await api.client_config(req)).sample_max_concurrent_requests == 2000 + + req.app.state.engine_config = EngineConfig(base_model="m", sample_max_concurrent_requests=16384) + assert (await api.client_config(req)).sample_max_concurrent_requests == 16384 + + def test_forward_backward_input_accepts_ppo_threshold_keys(): req = api.ForwardBackwardInput( data=[_make_datum()], diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py index 4178dc54dc..d7b75c9bb6 100644 --- a/tests/tinker/test_external_future_store.py +++ b/tests/tinker/test_external_future_store.py @@ -42,6 +42,11 @@ def _sample_input(seq_id: int) -> types.SampleInput: ) +async def _still_connected() -> bool: + """Stand-in for ``Request.is_disconnected`` on a client that is still waiting.""" + return False + + class _CompletingForwarder: def __init__(self, store: ExternalFutureStore): self.store = store @@ -140,6 +145,7 @@ async def test_sustained_model_path_rollouts_training_futures_and_heartbeats(fut ) ), headers={}, + is_disconnected=_still_connected, ) async with AsyncSession(engine) as session: @@ -271,6 +277,12 @@ async def wait(self, request_id, timeout): def mark_retrieved(self, request_id): pass + def proto_result(self, request_id): + return None + + def cache_proto(self, request_id, proto): + pass + def serialize_result_in_thread(request_type, result_data): nonlocal active_serializations, max_active_serializations serialization_thread_ids.append(threading.get_ident()) @@ -294,6 +306,7 @@ def serialize_result_in_thread(request_type, result_data): ) ), headers={"accept": api.PROTO_CONTENT_TYPE}, + is_disconnected=_still_connected, ) responses = await asyncio.gather( @@ -643,6 +656,7 @@ async def test_retrieve_future_serializes_in_memory_result_as_proto(future_store ) ), headers={"accept": "application/x-protobuf, application/json"}, + is_disconnected=_still_connected, ) response = await api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_id)), request) diff --git a/tests/tinker/test_inference_forwarding_config.py b/tests/tinker/test_inference_forwarding_config.py index 091a49efb3..9b8a8bdcf1 100644 --- a/tests/tinker/test_inference_forwarding_config.py +++ b/tests/tinker/test_inference_forwarding_config.py @@ -1,12 +1,14 @@ import argparse -from unittest.mock import AsyncMock, call, patch +from types import SimpleNamespace +from unittest.mock import AsyncMock, call -import httpx +import aiohttp import pytest from skyrl.tinker.config import EngineConfig, add_model from skyrl.tinker.extra.skyrl_train_inference_forwarding import ( SkyRLTrainInferenceForwardingClient, + TransientInferenceError, ) @@ -21,20 +23,37 @@ def test_forwarding_timeout_reads_environment(monkeypatch) -> None: assert config.forwarding_inference_timeout_sec == 1800.0 -def test_forwarding_client_uses_configured_timeout() -> None: +@pytest.mark.asyncio +async def test_forwarding_client_uses_configured_timeout_and_connection_limit() -> None: config = EngineConfig( base_model="test-model", forwarding_inference_timeout_sec=1800.0, + forwarding_inference_max_connections=64, ) + client = SkyRLTrainInferenceForwardingClient(config, db_engine=None) + try: + session = client._get_session() + assert session.timeout.sock_connect == 10.0 + assert session.timeout.sock_read == 1800.0 + # No overall deadline: a request may wait in the connector queue for + # as long as the engine takes to get to it. + assert session.timeout.total is None + assert session.connector.limit == 64 + finally: + await client.aclose() + + +@pytest.mark.asyncio +async def test_forwarding_client_default_connection_limit_is_unlimited() -> None: + client = SkyRLTrainInferenceForwardingClient(EngineConfig(base_model="test-model"), db_engine=None) + try: + assert client._get_session().connector.limit == 0 + finally: + await client.aclose() - with patch("skyrl.tinker.extra.skyrl_train_inference_forwarding.httpx.AsyncClient") as async_client: - SkyRLTrainInferenceForwardingClient(config, db_engine=None) - timeout = async_client.call_args.kwargs["timeout"] - assert timeout.connect == 10.0 - assert timeout.read == 1800.0 - assert timeout.write == 300.0 - assert timeout.pool == 300.0 +def _connect_error(message: str) -> aiohttp.ClientConnectorError: + return aiohttp.ClientConnectorError(SimpleNamespace(ssl=None, host="inference", port=8000), OSError(message)) @pytest.mark.asyncio @@ -43,7 +62,7 @@ async def test_forwarding_retries_connection_failure() -> None: client._cached_proxy_url = "http://old" client._resolve_proxy_url = AsyncMock(side_effect=["http://old", "http://new"]) expected = object() - client._forward = AsyncMock(side_effect=[httpx.ConnectError("unreachable"), expected]) + client._forward = AsyncMock(side_effect=[_connect_error("unreachable"), expected]) result = await client._forward_with_retry(object(), "model", base_model=None) @@ -52,19 +71,46 @@ async def test_forwarding_retries_connection_failure() -> None: assert client._forward.await_count == 2 +@pytest.mark.asyncio +async def test_forwarding_retries_transient_5xx_once() -> None: + client = object.__new__(SkyRLTrainInferenceForwardingClient) + client._cached_proxy_url = "http://old" + client._resolve_proxy_url = AsyncMock(side_effect=["http://old", "http://new"]) + expected = object() + client._forward = AsyncMock(side_effect=[TransientInferenceError("503 from router"), expected]) + + result = await client._forward_with_retry(object(), "model", base_model=None) + + assert result is expected + assert client._forward.await_count == 2 + + +@pytest.mark.asyncio +async def test_forwarding_does_not_retry_4xx() -> None: + client = object.__new__(SkyRLTrainInferenceForwardingClient) + client._cached_proxy_url = "http://inference" + client._resolve_proxy_url = AsyncMock(return_value="http://inference") + client._forward = AsyncMock(side_effect=RuntimeError("vLLM /v1/completions returned 400: bad request")) + + with pytest.raises(RuntimeError, match="returned 400"): + await client._forward_with_retry(object(), "model", base_model=None) + + client._forward.assert_awaited_once() + + @pytest.mark.asyncio async def test_forwarding_does_not_retry_read_timeout() -> None: client = object.__new__(SkyRLTrainInferenceForwardingClient) client.engine_config = EngineConfig(base_model="test-model", forwarding_inference_timeout_sec=123.0) client._cached_proxy_url = "http://inference" client._resolve_proxy_url = AsyncMock(return_value="http://inference") - client._forward = AsyncMock(side_effect=httpx.ReadTimeout("slow response")) + client._forward = AsyncMock(side_effect=aiohttp.SocketTimeoutError("slow response")) with pytest.raises(RuntimeError) as exc_info: await client._forward_with_retry(object(), "model", base_model=None) message = str(exc_info.value) - assert isinstance(exc_info.value.__cause__, httpx.ReadTimeout) + assert isinstance(exc_info.value.__cause__, aiohttp.SocketTimeoutError) assert "http://inference" in message assert "timed out after 123s" in message client._resolve_proxy_url.assert_awaited_once_with() diff --git a/tests/tinker/test_retrieve_future_lost_response.py b/tests/tinker/test_retrieve_future_lost_response.py new file mode 100644 index 0000000000..a8ce73eb1a --- /dev/null +++ b/tests/tinker/test_retrieve_future_lost_response.py @@ -0,0 +1,159 @@ +"""A result whose delivery the client never received must survive for the SDK's retry. + +Reproduces the 128x128 failure Chuck hit against PR j316chuck/SkyRL#18 (same +code as main): the SDK polls ``retrieve_future`` with a 45s client timeout and +gives up; the result lands afterwards; the abandoned handler still builds a +response and starts the short *retrieved* TTL clock even though nobody got the +bytes; the sweeper evicts the entry; the SDK's retry of the same request_id +gets ``404 Future not found``, which the SDK treats as fatal. + +The server runs under a real uvicorn socket so the abandoned poll is a genuine +TCP disconnect, exactly as with the SDK. TTLs are shortened so the whole chain +takes a few seconds. +""" + +import asyncio +import sys +from contextlib import suppress +from types import SimpleNamespace + +import aiohttp +import pytest +import pytest_asyncio +import uvicorn +from sqlalchemy.ext.asyncio import create_async_engine +from sqlmodel import SQLModel + +from skyrl.tinker import api, types +from skyrl.tinker.config import EngineConfig +from skyrl.tinker.db_models import ( + RequestStatus, + enable_sqlite_wal, + get_async_database_url, +) +from skyrl.tinker.external_future_store import ExternalFutureStore + +BASE_MODEL = "test-model" +RETRIEVED_TTL_SECONDS = 1.0 +SWEEP_INTERVAL_SECONDS = 0.2 + + +class _GatedForwarder: + """Completes each forwarded sample only once the test releases it.""" + + def __init__(self, store: ExternalFutureStore): + self.store = store + self.release = asyncio.Event() + + async def call_and_store_result(self, request_id, sample_req, model_id, checkpoint_id, *, base_model=None): + await self.release.wait() + result = types.SampleOutput( + sequences=[types.GeneratedSequence(stop_reason="length", tokens=[1, 2, 3], logprobs=[-0.1, -0.2, -0.3])] + ) + await self.store.complete(request_id, result, RequestStatus.COMPLETED) + + +@pytest_asyncio.fixture() +async def served_app(tmp_path, monkeypatch): + """The real API app on a real uvicorn socket, with app.state wired the way the lifespan does.""" + monkeypatch.setattr(ExternalFutureStore, "_RETRIEVED_TTL_SECONDS", RETRIEVED_TTL_SECONDS) + monkeypatch.setattr(ExternalFutureStore, "_SWEEP_INTERVAL_SECONDS", SWEEP_INTERVAL_SECONDS) + + engine = create_async_engine(get_async_database_url(f"sqlite:///{tmp_path / 'tinker.db'}")) + enable_sqlite_wal(engine.sync_engine) + async with engine.begin() as connection: + await connection.run_sync(SQLModel.metadata.create_all) + + store = ExternalFutureStore() + await store.start() + forwarder = _GatedForwarder(store) + + state = api.app.state + state.engine_config = EngineConfig(base_model=BASE_MODEL) + state.db_engine = engine + state.future_waiters = {} + state.future_poller = asyncio.create_task(api.poll_futures(engine, state.future_waiters, poll_interval_sec=0.01)) + state.proto_serialization_lock = asyncio.Lock() + state.db_write_lock = asyncio.Lock() + state.sampling_model_cache = {} + state.sampling_model_cache_lock = asyncio.Lock() + state.validated_sampler_checkpoints = set() + state.sampler_checkpoint_validation_lock = asyncio.Lock() + state.external_future_store = store + state.external_inference_client = forwarder + + config = uvicorn.Config(api.app, host="127.0.0.1", port=0, log_level="warning", lifespan="off") + server = uvicorn.Server(config) + serve_task = asyncio.create_task(server.serve()) + while not server.started: + await asyncio.sleep(0.01) + port = server.servers[0].sockets[0].getsockname()[1] + + yield SimpleNamespace(url=f"http://127.0.0.1:{port}/api/v1", store=store, forwarder=forwarder) + + server.should_exit = True + await serve_task + state.future_poller.cancel() + with suppress(asyncio.CancelledError): + await state.future_poller + await store.close() + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.skipif(sys.platform != "linux", reason="relies on uvicorn disconnect handling over a real socket") +async def test_retry_after_client_abandoned_poll_is_served(served_app): + payload = { + "num_samples": 1, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [1, 2, 3]}]}, + "sampling_params": {"max_tokens": 3, "temperature": 1.0, "seed": 0}, + "base_model": BASE_MODEL, + } + async with aiohttp.ClientSession() as client: + async with client.post(f"{served_app.url}/asample", json=payload) as resp: + assert resp.status == 200 + request_id = (await resp.json())["request_id"] + + # The SDK's retrieve_future poll times out client-side (45s in the SDK) + # while the result is still pending, and the connection is closed. + with pytest.raises(asyncio.TimeoutError): + await client.post( + f"{served_app.url}/retrieve_future", + json={"request_id": request_id}, + timeout=aiohttp.ClientTimeout(total=0.3), + ) + await asyncio.sleep(0.2) # let the server observe the disconnect + + # The result arrives after the client gave up. The abandoned handler + # wakes, builds a response nobody will receive, and must NOT start the + # short retrieved-TTL clock. + served_app.forwarder.release.set() + await asyncio.sleep(RETRIEVED_TTL_SECONDS + 3 * SWEEP_INTERVAL_SECONDS) + + # The SDK retries the same request_id once its backoff elapses. + async with client.post(f"{served_app.url}/retrieve_future", json={"request_id": request_id}) as resp: + body = await resp.text() + assert resp.status == 200, f"retry of an undelivered result got {resp.status}: {body}" + assert types.SampleOutput.model_validate_json(body).sequences[0].tokens == [1, 2, 3] + + +@pytest.mark.asyncio +@pytest.mark.skipif(sys.platform != "linux", reason="relies on uvicorn disconnect handling over a real socket") +async def test_delivered_result_still_expires_on_retrieved_ttl(served_app): + """A result the client actually received is reclaimed on the short clock as before.""" + payload = { + "num_samples": 1, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [1, 2, 3]}]}, + "sampling_params": {"max_tokens": 3, "temperature": 1.0, "seed": 0}, + "base_model": BASE_MODEL, + } + served_app.forwarder.release.set() + async with aiohttp.ClientSession() as client: + async with client.post(f"{served_app.url}/asample", json=payload) as resp: + request_id = (await resp.json())["request_id"] + async with client.post(f"{served_app.url}/retrieve_future", json={"request_id": request_id}) as resp: + assert resp.status == 200 + await resp.read() + assert int(request_id) in served_app.store._entries + await asyncio.sleep(RETRIEVED_TTL_SECONDS + 3 * SWEEP_INTERVAL_SECONDS) + assert int(request_id) not in served_app.store._entries diff --git a/tests/tinker/test_sample_result_fast_path.py b/tests/tinker/test_sample_result_fast_path.py new file mode 100644 index 0000000000..0cff804059 --- /dev/null +++ b/tests/tinker/test_sample_result_fast_path.py @@ -0,0 +1,178 @@ +"""Forwarded sample results are encoded to proto once and served as-is. + +The forwarding client decodes the vLLM body and encodes straight to the +``SampleResponse`` wire form; no pydantic model or JSON text is built unless a +pre-proto client asks for JSON, in which case it is derived from the proto and +cached. These tests pin the fast path to the validated path byte for byte. +""" + +import asyncio +import json +from types import SimpleNamespace + +import numpy as np +import pytest + +from skyrl.tinker import api, types +from skyrl.tinker.db_models import RequestStatus +from skyrl.tinker.external_future_store import ExternalFutureStore, PreparedResult +from skyrl.tinker.proto_serialization import ( + PROTO_CONTENT_TYPE, + sample_output_json_from_proto, + serialize_result, + serialize_sample_output, +) + +SEQUENCES = [ + ("length", [1, 2, 3, 40000], [-0.5, -1.25, -0.03125, -7.0]), + ("stop", [7], [-2.0]), +] +PROMPT_LOGPROBS = [None, -0.75, -3.5] +TOPK = [None, [(11, -0.5), (12, -1.5)], [(13, -0.25)]] + + +def _validated_bytes(prompt_logprobs=None, topk=None) -> bytes: + output = types.SampleOutput( + sequences=[types.GeneratedSequence(stop_reason=s, tokens=t, logprobs=lp) for s, t, lp in SEQUENCES], + prompt_logprobs=prompt_logprobs, + topk_prompt_logprobs=topk, + ) + return serialize_result(types.RequestType.SAMPLE, output.model_dump()) + + +def test_fast_path_matches_validated_serialization_bytes(): + assert serialize_sample_output(SEQUENCES, None, None) == _validated_bytes() + assert serialize_sample_output(SEQUENCES, PROMPT_LOGPROBS, TOPK) == _validated_bytes(PROMPT_LOGPROBS, TOPK) + + +def test_json_from_proto_round_trips_sample_output(): + text = sample_output_json_from_proto(serialize_sample_output(SEQUENCES, PROMPT_LOGPROBS, TOPK)) + output = types.SampleOutput.model_validate_json(text) + + assert [(s.stop_reason, s.tokens) for s in output.sequences] == [(s, t) for s, t, _ in SEQUENCES] + for seq, (_, _, logprobs) in zip(output.sequences, SEQUENCES): + # Logprobs travel as float32 on the wire. + assert seq.logprobs == np.asarray(logprobs, dtype=np.float32).tolist() + assert output.prompt_logprobs[0] is None + assert output.prompt_logprobs[1:] == np.asarray(PROMPT_LOGPROBS[1:], dtype=np.float32).tolist() + assert output.topk_prompt_logprobs[0] is None + assert output.topk_prompt_logprobs[1] == [(11, -0.5), (12, -1.5)] + assert output.topk_prompt_logprobs[2] == [(13, -0.25)] + # Same key layout as pydantic's own dump, so JSON clients see nothing new. + assert list(json.loads(text)) == ["sequences", "prompt_logprobs", "topk_prompt_logprobs"] + + +def test_json_from_proto_without_optional_fields(): + output = types.SampleOutput.model_validate_json( + sample_output_json_from_proto(serialize_sample_output(SEQUENCES, None, None)) + ) + assert output.prompt_logprobs is None + assert output.topk_prompt_logprobs is None + + +@pytest.mark.asyncio +async def test_store_serves_proto_directly_and_derives_json_lazily(): + store = ExternalFutureStore() + request_id = store.create("model_a", SimpleNamespace()) + proto = serialize_sample_output(SEQUENCES, None, None) + + await store.complete(request_id, PreparedResult(proto=proto), RequestStatus.COMPLETED) + + status, request_type, result_data = await store.wait(request_id, timeout=1) + assert (status, request_type, result_data) == (RequestStatus.COMPLETED, types.RequestType.EXTERNAL, None) + assert store.proto_result(request_id) is proto + text = store.json_result(request_id) + assert types.SampleOutput.model_validate_json(text).sequences[1].tokens == [7] + # Derived once, then cached for retries. + assert store.json_result(request_id) is text + + +@pytest.mark.asyncio +async def test_store_still_accepts_pydantic_results(): + store = ExternalFutureStore() + request_id = store.create("model_a", SimpleNamespace()) + output = types.SampleOutput(sequences=[]) + + await store.complete(request_id, output, RequestStatus.COMPLETED) + + assert (await store.wait(request_id, timeout=1))[2] == output.model_dump_json() + assert store.proto_result(request_id) is None + + +def test_store_ttl_overrides_apply_per_instance(): + store = ExternalFutureStore(retrieved_ttl_sec=5.0, completed_ttl_sec=7.0) + assert (store._RETRIEVED_TTL_SECONDS, store._COMPLETED_TTL_SECONDS) == (5.0, 7.0) + assert ExternalFutureStore._RETRIEVED_TTL_SECONDS == 300.0 + assert ExternalFutureStore()._RETRIEVED_TTL_SECONDS == 300.0 + + +async def _connected() -> bool: + return False + + +def _request(store: ExternalFutureStore, accept: str, serialize_calls: list) -> SimpleNamespace: + return SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + external_future_store=store, + future_waiters={}, + proto_serialization_lock=asyncio.Lock(), + ) + ), + headers={"accept": accept}, + is_disconnected=_connected, + ) + + +@pytest.mark.asyncio +async def test_retrieve_future_passes_stored_proto_through_without_reencoding(monkeypatch): + store = ExternalFutureStore() + request_id = store.create("model_a", SimpleNamespace()) + proto = serialize_sample_output(SEQUENCES, None, None) + await store.complete(request_id, PreparedResult(proto=proto), RequestStatus.COMPLETED) + serialize_calls: list = [] + monkeypatch.setattr(api, "_serialize_proto_result", lambda *a: serialize_calls.append(a) or b"unexpected") + + response = await api.retrieve_future( + api.RetrieveFutureRequest(request_id=str(request_id)), _request(store, PROTO_CONTENT_TYPE, serialize_calls) + ) + + assert response.media_type == PROTO_CONTENT_TYPE + assert response.body == proto + assert serialize_calls == [] + + +@pytest.mark.asyncio +async def test_retrieve_future_serves_json_client_from_stored_proto(): + store = ExternalFutureStore() + request_id = store.create("model_a", SimpleNamespace()) + await store.complete( + request_id, PreparedResult(proto=serialize_sample_output(SEQUENCES, None, None)), RequestStatus.COMPLETED + ) + + response = await api.retrieve_future( + api.RetrieveFutureRequest(request_id=str(request_id)), _request(store, "application/json", []) + ) + + assert response.media_type == "application/json" + assert types.SampleOutput.model_validate_json(response.body).sequences[0].tokens == [1, 2, 3, 40000] + + +@pytest.mark.asyncio +async def test_retrieve_future_encodes_json_stored_result_once_for_proto_clients(monkeypatch): + store = ExternalFutureStore() + request_id = store.create("model_a", SimpleNamespace()) + output = types.SampleOutput( + sequences=[types.GeneratedSequence(stop_reason=s, tokens=t, logprobs=lp) for s, t, lp in SEQUENCES] + ) + await store.complete(request_id, output, RequestStatus.COMPLETED) + calls: list = [] + real = api._serialize_proto_result + monkeypatch.setattr(api, "_serialize_proto_result", lambda *a: calls.append(a) or real(*a)) + request = _request(store, PROTO_CONTENT_TYPE, calls) + + first = await api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_id)), request) + second = await api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_id)), request) + + assert first.body == second.body == serialize_sample_output(SEQUENCES, None, None) + assert len(calls) == 1