From f7959e3978fa5fbdda8cc519b051c745d9f1dbbb Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Fri, 4 Sep 2026 00:28:53 +0000 Subject: [PATCH 1/2] [tinker] Forward samples with aiohttp instead of httpx httpcore's connection pool rescans every pooled connection and queued request on each event, so the forwarding client's per-request CPU grew with the number of in-flight samples: 14ms at 512, 27ms at 2048, 38ms at 4096 (a standalone benchmark against an instant fake router; profile shows 27.5M is_idle calls for 512 requests). At 2048 in flight the API server needed 119s to forward 2048 instant requests. aiohttp's connector is flat at 0.34ms per request. - SkyRLTrainInferenceForwardingClient uses one aiohttp session: connector limit = forwarding_inference_max_connections (0 = unlimited), sock_read = forwarding_inference_timeout_sec, no total deadline so requests queued behind the engine never hit the old 300s pool timeout, 60s connect timeout (a saturated router takes tens of seconds to accept), Happy Eyeballs off (a wave of cancelled connects left uvloop "File descriptor N is used by transport" errors from aiohappyeyeballs). - Connect-phase errors and 5xx rejections from the router (TransientInferenceError) are retried once after refreshing the proxy URL; read failures stay final since vLLM may still be executing the request. - forwarding_inference_timeout_sec default 300s -> 2048s: with unlimited connections a large rollout burst waits inside vLLM's queue and 128x128 bursts exceed 300s there. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Avi Basnet --- skyrl/tinker/config.py | 6 +- .../extra/skyrl_train_inference_forwarding.py | 109 ++++++++++++------ .../test_inference_forwarding_config.py | 72 +++++++++--- 3 files changed, 135 insertions(+), 52 deletions(-) diff --git a/skyrl/tinker/config.py b/skyrl/tinker/config.py index cd8d5c4a23..a88c005f30 100644 --- a/skyrl/tinker/config.py +++ b/skyrl/tinker/config.py @@ -59,12 +59,14 @@ 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, diff --git a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py index d617d9a7a6..88e63a6776 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 @@ -19,6 +20,13 @@ 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.""" + + +_ROUTER_CONNECT_TIMEOUT_SECONDS = 60.0 + + class SkyRLTrainInferenceForwardingClient: """Forwards EXTERNAL sample requests to the SkyRL-Train-managed vLLM.""" @@ -36,27 +44,48 @@ 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. + # Happy Eyeballs is off: a burst of connect timeouts cancels its + # sock_connect calls mid-flight, and under uvloop the closed sockets' + # descriptors get reused before the loop forgets them ("File + # descriptor N is used by transport"), failing unrelated forwards. + connector = aiohttp.TCPConnector(limit=max_conn or 0, keepalive_timeout=2, happy_eyeballs_delay=None) + self._session = aiohttp.ClientSession( + connector=connector, + timeout=aiohttp.ClientTimeout( + total=None, + # A saturated router can take tens of seconds to accept; + # that is queueing, not failure. + sock_connect=_ROUTER_CONNECT_TIMEOUT_SECONDS, + 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: @@ -119,30 +148,31 @@ async def call_and_store_result( 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 + # 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 " @@ -195,17 +225,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 diff --git a/tests/tinker/test_inference_forwarding_config.py b/tests/tinker/test_inference_forwarding_config.py index 091a49efb3..1987fb1deb 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 == 60.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() From 3bdbb3d88fdc67e45b1f256e5fc251913f21960c Mon Sep 17 00:00:00 2001 From: avigyabb <98926738+avigyabb@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:33:31 -0700 Subject: [PATCH 2/2] [tinker] Keep an undelivered sample result alive for the SDK's retry (#2162) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stack 3/7. Fixes the 128x128 `404 Future not found` seen against j316chuck/SkyRL#18. **Chain.** The SDK polls `retrieve_future` with a 45 s client timeout and gives up; the result lands afterwards; the abandoned handler wakes, builds a response nobody receives (uvicorn drops the send to a dead client silently) and starts the short retrieved-TTL clock; the sweeper evicts the result 120 s later; the SDK's retry of the same request_id gets 404, which the SDK treats as fatal. **Fix.** Start the retrieved clock only if `request.is_disconnected()` is false, and raise the retrieved TTL to 300 s so it outlasts the SDK's worst-case re-poll gap (45 s timeout + up to 30 s backoff, twice). `tests/tinker/test_retrieve_future_lost_response.py` reproduces the chain under a real uvicorn socket with shortened TTLs; it fails on `main` and passes here. A second test checks a delivered result still expires on the short clock, so memory stays bounded. Alternative considered: j316chuck/SkyRL#19 drops the retrieved clock and keeps every result for 2048 s. That also fixes the 404 but retains ~35 minutes of results regardless of delivery; with long-output rollouts (hundreds of KB per result) that is tens of GB. Verified at scale: 131072 requests with 5 s engine queueing and 224k SDK-style abandoned polls completed with zero 404s. **Stack** (each PR retargets to `main` as the one below merges) 1. #2160 2. #2161 3. #2162 4. #2163 5. #2164 6. #2165 7. #2166 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- > [!NOTE] > **Medium Risk** > Changes async polling/delivery semantics and HTTP server tuning on the hot `retrieve_future` path; behavior is covered by new integration tests but affects SDK retry reliability under load. > > **Overview** > Fixes fatal **`404 Future not found`** when the SDK abandons a long `retrieve_future` poll (45s client timeout) and retries the same `request_id` after the result is ready. > > **`retrieve_future`** now calls **`mark_retrieved`** (starting the post-delivery eviction clock) only when the client is still connected (`not await req.is_disconnected()`). If the handler finishes building a response after the client disconnected, the short retrieved TTL no longer starts, so the in-memory store keeps the result for a real retry. > > **`ExternalFutureStore`** raises **`_RETRIEVED_TTL_SECONDS`** from 120s to 300s so delivered results still get a grace window that covers worst-case SDK re-poll gaps (timeout + backoff, twice). > > **Uvicorn** startup sets **`timeout_keep_alive=75`** (vs 5s default) and **`backlog=SKYRL_HTTP_CONNECTION_LIMIT`** to reduce idle disconnects and accept-queue overflows during completion bursts. > > Adds **`test_retrieve_future_lost_response.py`** (real socket on Linux) plus test stubs for **`is_disconnected`** on existing API tests. > > Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 36188da4a70b9bb4fcd8902183695b1668880ef4. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot). --------- Signed-off-by: Avi Basnet Co-authored-by: Claude Fable 5.1 --- skyrl/tinker/api.py | 27 ++- skyrl/tinker/external_future_store.py | 6 +- tests/tinker/test_external_future_store.py | 8 + .../test_retrieve_future_lost_response.py | 159 ++++++++++++++++++ 4 files changed, 195 insertions(+), 5 deletions(-) create mode 100644 tests/tinker/test_retrieve_future_lost_response.py diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index 5432555af5..a7e4e1f450 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -32,6 +32,7 @@ from sqlmodel import SQLModel, func, select from sqlmodel.ext.asyncio.session import AsyncSession +from skyrl.env_vars import SKYRL_HTTP_CONNECTION_LIMIT from skyrl.tinker import types from skyrl.tinker.config import EngineConfig, add_model, config_to_argv from skyrl.tinker.db_models import ( @@ -71,6 +72,12 @@ # How long retrieve_future waits for a result before returning 408 RETRIEVE_FUTURE_TIMEOUT_SECONDS = 300 +# Idle keep-alive for client connections. Under a burst of completions the +# event loop can be busy for many seconds; with uvicorn's 5s default every +# idle SDK connection is closed during such a burst and all clients reconnect +# at once, overflowing the accept backlog. Hold connections across bursts. +HTTP_KEEP_ALIVE_TIMEOUT_SECONDS = 75 + # How often poll_futures looks for newly finished requests. A single query # covers every waiter, so this can stay tight without the load scaling up with # the number of in-flight requests. @@ -1515,8 +1522,13 @@ async def retrieve_future(request: RetrieveFutureRequest, req: Request): else: 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 @@ -1849,4 +1861,13 @@ async def root(): # Store config in app.state so lifespan can access it app.state.engine_config = engine_config - uvicorn.run(app, host=args.host, port=args.port, log_config=get_uvicorn_log_config()) + uvicorn.run( + app, + host=args.host, + port=args.port, + log_config=get_uvicorn_log_config(), + # Pending connections queue in the kernel while the loop is busy instead + # of being refused (effective value is capped by net.core.somaxconn). + backlog=SKYRL_HTTP_CONNECTION_LIMIT, + timeout_keep_alive=HTTP_KEEP_ALIVE_TIMEOUT_SECONDS, + ) diff --git a/skyrl/tinker/external_future_store.py b/skyrl/tinker/external_future_store.py index e69e6999d4..1e58b4b1bc 100644 --- a/skyrl/tinker/external_future_store.py +++ b/skyrl/tinker/external_future_store.py @@ -44,8 +44,10 @@ 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 diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py index 4178dc54dc..c85f9dcac9 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: @@ -294,6 +300,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 +650,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_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