[tinker] Long-output results: pysimdjson decode, burst-survival settings, harness fidelity (131072 x 212k tokens)

pyproject.toml

@@ -38,6 +38,9 @@ tpu = [
 tinker = [
     "tinker>=0.3.0,<=0.24.1",
     "fastapi[standard]",
+    # Decodes vLLM completion bodies into numeric buffers without per-token
+    # Python objects (skyrl/tinker/extra/completion_decode.py).
+    "pysimdjson",
     "sqlmodel",
     "sqlalchemy[asyncio]",
     "aiosqlite",

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.
@@ -1872,4 +1879,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,
+    )

skyrl/tinker/extra/completion_decode.py

@@ -0,0 +1,139 @@
+"""Decode vLLM ``/v1/completions`` bodies straight into numpy arrays.
+
+A long-output result is almost entirely two numeric arrays (token ids and
+logprobs). Decoding it the ordinary way materializes one Python object per
+element and then converts the lists to numpy, which for a 262k-token result
+(2.8MB of JSON) costs ~38ms and dominates the API server's CPU. pysimdjson
+exposes homogeneous numeric arrays as raw buffers, so the same body decodes in
+~6ms with no per-token Python objects. orjson is the fallback when pysimdjson
+is unavailable or an array is not purely numeric (a null logprob, say).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import orjson
+
+from skyrl.utils.log import logger
+
+try:
+    import simdjson
+except ImportError:  # pragma: no cover - exercised via the forced-fallback tests
+    simdjson = None
+
+
+@dataclass
+class DecodedChoice:
+    finish_reason: str | None
+    tokens: np.ndarray  # int32
+    logprobs: np.ndarray  # float32, same length as tokens
+    # vLLM's raw prompt_logprobs structure, kept only when the request asked for it.
+    prompt_logprobs: list[Any] | None = None
+
+
+def _floats_from_list(values: list[Any], expected_len: int) -> np.ndarray:
+    """Float32 array from a decoded list, filling absent values with zeros.
+
+    vLLM occasionally returns None for logprobs under load; zero-fill so RL
+    advantage computation doesn't see a ragged shape.
+    """
+    if not values:
+        if expected_len:
+            logger.warning("No logprobs returned from vLLM — filling with zeros")
+        return np.zeros(expected_len, dtype=np.float32)
+    if any(value is None for value in values):
+        logger.warning("vLLM returned null logprobs — filling those positions with zeros")
+        values = [0.0 if value is None else value for value in values]
+    return np.asarray(values, dtype=np.float32)
+
+
+class CompletionDecoder:
+    """Decodes completion bodies; one instance per forwarding client."""
+
+    def __init__(self) -> None:
+        self._parser = simdjson.Parser() if simdjson is not None else None
+
+    @property
+    def backend(self) -> str:
+        return "simdjson" if self._parser is not None else "orjson"
+
+    def decode(self, body: bytes, *, want_prompt_logprobs: bool = False) -> list[DecodedChoice]:
+        """Return one :class:`DecodedChoice` per ``choices`` entry.
+
+        Raises ``ValueError`` for a body that is not JSON.
+        """
+        if self._parser is not None:
+            return self._decode_simdjson(body, want_prompt_logprobs)
+        return self._decode_orjson(body, want_prompt_logprobs)
+
+    # -- pysimdjson --------------------------------------------------------
+
+    def _decode_simdjson(self, body: bytes, want_prompt_logprobs: bool) -> list[DecodedChoice]:
+        # The parser owns one document at a time, so everything is copied out
+        # into numpy / Python objects before this method returns.
+        doc = self._parser.parse(body)
+        choices = []
+        for choice in doc.get("choices") or ():
+            tokens = _simd_int32(choice.get("token_ids"))
+            logprobs_obj = choice.get("logprobs")
+            raw_logprobs = logprobs_obj.get("token_logprobs") if logprobs_obj is not None else None
+            logprobs = _simd_float32(raw_logprobs, expected_len=len(tokens))
+            prompt_logprobs = None
+            if want_prompt_logprobs:
+                raw = choice.get("prompt_logprobs")
+                prompt_logprobs = raw.as_list() if raw is not None else None
+            choices.append(
+                DecodedChoice(
+                    finish_reason=choice.get("finish_reason"),
+                    tokens=tokens,
+                    logprobs=logprobs,
+                    prompt_logprobs=prompt_logprobs,
+                )
+            )
+        return choices
+
+    # -- orjson fallback ---------------------------------------------------
+
+    @staticmethod
+    def _decode_orjson(body: bytes, want_prompt_logprobs: bool) -> list[DecodedChoice]:
+        result = orjson.loads(body)
+        choices = []
+        for choice in result.get("choices") or ():
+            tokens = np.asarray(choice.get("token_ids") or (), dtype=np.int32)
+            raw_logprobs = (choice.get("logprobs") or {}).get("token_logprobs") or []
+            choices.append(
+                DecodedChoice(
+                    finish_reason=choice.get("finish_reason"),
+                    tokens=tokens,
+                    logprobs=_floats_from_list(raw_logprobs, expected_len=len(tokens)),
+                    prompt_logprobs=choice.get("prompt_logprobs") if want_prompt_logprobs else None,
+                )
+            )
+        return choices
+
+
+def _simd_int32(array) -> np.ndarray:
+    if array is None:
+        return np.zeros(0, dtype=np.int32)
+    try:
+        return np.frombuffer(array.as_buffer(of_type="i"), dtype=np.int64).astype(np.int32)
+    except TypeError:
+        # Not a homogeneous integer array; take the slow path for this one.
+        return np.asarray(array.as_list(), dtype=np.int32)
+
+
+def _simd_float32(array, *, expected_len: int) -> np.ndarray:
+    if array is None:
+        return _floats_from_list([], expected_len)
+    try:
+        values = np.frombuffer(array.as_buffer(of_type="d"), dtype=np.float64).astype(np.float32)
+    except TypeError:
+        # Nulls or mixed ints/floats: fall back to a list for this array only.
+        return _floats_from_list(array.as_list(), expected_len)
+    if len(values) == 0 and expected_len:
+        # vLLM returned an empty logprob list for a non-empty sequence.
+        return _floats_from_list([], expected_len)
+    return values

skyrl/tinker/extra/skyrl_train_inference_forwarding.py

@@ -8,7 +8,6 @@
 from datetime import datetime, timezone
 
 import aiohttp
-import orjson
 from sqlmodel.ext.asyncio.session import AsyncSession
 
 from skyrl.backends.renderer import render_model_input
@@ -17,12 +16,15 @@
 from skyrl.tinker.config import EngineConfig
 from skyrl.tinker.db_models import EngineStateDB, FutureDB, RequestStatus
 from skyrl.tinker.external_future_store import ExternalFutureStore, PreparedResult
+from skyrl.tinker.extra.completion_decode import CompletionDecoder
 from skyrl.tinker.proto_serialization import (
     sample_output_json_from_proto,
     serialize_sample_output,
 )
 from skyrl.utils.log import logger
 
+_ROUTER_CONNECT_TIMEOUT_SECONDS = 60.0
+
 
 class TransientInferenceError(RuntimeError):
     """A 5xx from vllm-router/vLLM: the request was rejected, not executed, so it is safe to retry."""
@@ -47,6 +49,7 @@ def __init__(
         self._cache_lock = asyncio.Lock()
         # Created on first use so it binds to the serving event loop.
         self._session: aiohttp.ClientSession | None = None
+        self._decoder = CompletionDecoder()
 
     def _get_session(self) -> aiohttp.ClientSession:
         """Return the shared aiohttp session, creating it on first use.
@@ -66,12 +69,18 @@ def _get_session(self) -> aiohttp.ClientSession:
             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)
+            # 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,
-                    sock_connect=10.0,
+                    # 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,
                 ),
             )
@@ -230,8 +239,10 @@ async def _forward(self, proxy_url: str, sample_req, model_id: str, *, base_mode
             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:
+                # Token ids and logprobs land directly in int32/float32 arrays;
+                # no per-token Python objects are built (see completion_decode).
+                choices = self._decoder.decode(body, want_prompt_logprobs=want_prompt_logprobs)
+            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}, "
@@ -243,26 +254,16 @@ async def _forward(self, proxy_url: str, sample_req, model_id: str, *, base_mode
         if want_prompt_logprobs:
             # All `n` choices share one prompt, so vLLM repeats the same prompt
             # logprobs on each choice; read them off the first.
-            choices = result.get("choices") or []
-            raw = choices[0].get("prompt_logprobs") if choices else None
+            raw = choices[0].prompt_logprobs if choices else None
             if raw is None:
                 logger.warning("Requested prompt logprobs but vLLM /v1/completions returned none")
             prompt_logprobs, topk = convert_vllm_prompt_logprobs(prompt_tokens, raw, topk=topk_prompt_logprobs)
 
-        sequences = []
-        for choice in result.get("choices", []):
-            tokens = choice.get("token_ids", [])
-            lp = choice.get("logprobs") or {}
-            logprobs = lp.get("token_logprobs") or []
-            # vLLM occasionally returns None for logprobs under load; zero-fill so
-            # RL advantage computation doesn't see a ragged shape.
-            if not logprobs and tokens:
-                logger.warning("No logprobs returned from vLLM — filling with zeros")
-                logprobs = [0.0] * len(tokens)
-            # 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((stop_reason, tokens, logprobs))
+        # Tinker's stop_reason is Literal["stop", "length"]; vLLM emits a wider set.
+        sequences = [
+            ("stop" if choice.finish_reason in ("stop", "stop_token") else "length", choice.tokens, choice.logprobs)
+            for choice in choices
+        ]
 
         # Encode straight to the proto wire form the SDK retrieves; no pydantic
         # model or JSON text is built for the result (see PreparedResult).

tests/tinker/test_completion_decode.py

@@ -0,0 +1,108 @@
+"""vLLM completion bodies decode to the same arrays with pysimdjson and with the orjson fallback."""
+
+import numpy as np
+import orjson
+import pytest
+
+from skyrl.tinker.extra import completion_decode
+from skyrl.tinker.extra.completion_decode import CompletionDecoder
+
+
+def _body(choices: list[dict]) -> bytes:
+    return orjson.dumps({"id": "cmpl", "object": "text_completion", "choices": choices, "usage": {}})
+
+
+def _decoders() -> list[CompletionDecoder]:
+    fast = CompletionDecoder()
+    slow = CompletionDecoder()
+    slow._parser = None
+    return [fast, slow]
+
+
+@pytest.mark.parametrize("decoder", _decoders(), ids=lambda d: d.backend)
+def test_decodes_tokens_and_logprobs_to_typed_arrays(decoder):
+    choices = decoder.decode(
+        _body(
+            [
+                {
+                    "token_ids": [5, 6, 70000],
+                    "logprobs": {"token_logprobs": [-0.5, -1.25, -3.0]},
+                    "finish_reason": "stop",
+                },
+                {"token_ids": [1], "logprobs": {"token_logprobs": [0]}, "finish_reason": "length"},
+            ]
+        )
+    )
+
+    assert [c.finish_reason for c in choices] == ["stop", "length"]
+    assert choices[0].tokens.dtype == np.int32 and choices[0].logprobs.dtype == np.float32
+    assert choices[0].tokens.tolist() == [5, 6, 70000]
+    assert choices[0].logprobs.tolist() == [-0.5, -1.25, -3.0]
+    # An integer-valued logprob still comes out as float32.
+    assert choices[1].logprobs.tolist() == [0.0]
+    assert choices[0].prompt_logprobs is None
+
+
+@pytest.mark.parametrize("decoder", _decoders(), ids=lambda d: d.backend)
+def test_missing_or_null_logprobs_are_zero_filled(decoder):
+    choices = decoder.decode(
+        _body(
+            [
+                {"token_ids": [1, 2, 3], "logprobs": None, "finish_reason": "length"},
+                {"token_ids": [1, 2], "logprobs": {"token_logprobs": []}, "finish_reason": "length"},
+                {"token_ids": [1, 2, 3], "logprobs": {"token_logprobs": [None, -0.5, -1.0]}, "finish_reason": "stop"},
+                {"token_ids": [], "logprobs": {"token_logprobs": []}, "finish_reason": "stop"},
+            ]
+        )
+    )
+
+    assert choices[0].logprobs.tolist() == [0.0, 0.0, 0.0]
+    assert choices[1].logprobs.tolist() == [0.0, 0.0]
+    assert choices[2].logprobs.tolist() == [0.0, -0.5, -1.0]
+    assert choices[3].tokens.tolist() == [] and choices[3].logprobs.tolist() == []
+    assert all(c.logprobs.dtype == np.float32 for c in choices)
+
+
+@pytest.mark.parametrize("decoder", _decoders(), ids=lambda d: d.backend)
+def test_prompt_logprobs_kept_only_when_requested(decoder):
+    raw = [None, {"5": {"logprob": -0.1, "rank": 1, "decoded_token": "a"}}]
+    body = _body(
+        [{"token_ids": [1], "logprobs": {"token_logprobs": [-1.0]}, "finish_reason": "stop", "prompt_logprobs": raw}]
+    )
+
+    assert decoder.decode(body)[0].prompt_logprobs is None
+    assert decoder.decode(body, want_prompt_logprobs=True)[0].prompt_logprobs == raw
+
+
+@pytest.mark.parametrize("decoder", _decoders(), ids=lambda d: d.backend)
+def test_non_json_body_raises_value_error(decoder):
+    with pytest.raises(ValueError):
+        decoder.decode(b"<html>502 Bad Gateway</html>")
+
+
+@pytest.mark.skipif(completion_decode.simdjson is None, reason="pysimdjson not installed")
+def test_fast_and_fallback_agree_on_large_arrays():
+    rng = np.random.default_rng(0)
+    tokens = rng.integers(0, 150_000, size=200_000).tolist()
+    logprobs = (-rng.random(200_000) * 20).tolist()
+    body = _body([{"token_ids": tokens, "logprobs": {"token_logprobs": logprobs}, "finish_reason": "length"}])
+    fast, slow = _decoders()
+
+    a, b = fast.decode(body)[0], slow.decode(body)[0]
+
+    assert np.array_equal(a.tokens, b.tokens)
+    assert np.array_equal(a.logprobs, b.logprobs)
+    assert fast.backend == "simdjson" and slow.backend == "orjson"
+
+
+def test_decoder_can_be_reused_across_bodies():
+    decoder = CompletionDecoder()
+    first = decoder.decode(
+        _body([{"token_ids": [1, 2], "logprobs": {"token_logprobs": [-1.0, -2.0]}, "finish_reason": "stop"}])
+    )
+    second = decoder.decode(
+        _body([{"token_ids": [9], "logprobs": {"token_logprobs": [-9.0]}, "finish_reason": "length"}])
+    )
+    # Arrays from the first body survive the second parse (they were copied out).
+    assert first[0].tokens.tolist() == [1, 2] and first[0].logprobs.tolist() == [-1.0, -2.0]
+    assert second[0].tokens.tolist() == [9]