From 94b5591377bc622623ac411c7236a8a7ee7ff96b Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Thu, 3 Sep 2026 20:47:10 +0000 Subject: [PATCH 1/2] [tinker] Encode forwarded sample results to proto once and serve them as-is Long-output rollouts made the per-result payload work the API server's main cost. For a 32k-token result (356KB JSON) the forwarding path spent 4ms (orjson decode, pydantic validate, pydantic JSON dump) and the proto path the SDK >= 0.25 uses on retrieve_future spent another 7ms (stdlib json.loads plus proto build) inside a single global lock, capping proto delivery near 150 results/s regardless of concurrency; the proto build holds the GIL, so the thread hop bought nothing. - The forwarding client decodes the vLLM body once and encodes straight to SampleResponse wire bytes (serialize_sample_output, shared with the validated path and pinned to it byte for byte by tests). No pydantic model or JSON text is built for the result. - ExternalFutureStore keeps the proto bytes (8 bytes/token, 26% smaller than the JSON text); retrieve_future passes them through for proto clients and derives JSON lazily, cached, for pre-proto clients (sample_output_json_from_proto). Results stored as JSON (DB path, errors) keep the existing encode-in-thread path, now cached per entry. - Pending entries no longer retain the request body (never read back on this path; ~100KB per entry for long prompts). - Store TTLs are EngineConfig fields (external_future_retrieved_ttl_sec, external_future_completed_ttl_sec): retention after delivery is the dominant memory term, roughly completion rate x result size x window. 512 concurrent 32k-token results, proto client: server CPU 7.5s -> 3.5s, 73 -> 167 results/s. 131072 requests with 8k-token results, 2048-way engine queueing and SDK-style 45s re-polls: 131072/131072, 0 failures, 267/s, peak RSS 8.9GB. 32768 requests with 32k-token results: 32768/32768, 163/s, 9.6GB. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Avi Basnet (cherry picked from commit ac488182a397b6bef9737e1fd4abdb3ef85d4d93) Signed-off-by: Avi Basnet --- skyrl/tinker/api.py | 29 ++- skyrl/tinker/config.py | 15 ++ skyrl/tinker/external_future_store.py | 61 +++++- .../extra/skyrl_train_inference_forwarding.py | 33 ++-- skyrl/tinker/proto_serialization.py | 81 +++++++- tests/tinker/test_external_future_store.py | 6 + tests/tinker/test_sample_result_fast_path.py | 178 ++++++++++++++++++ 7 files changed, 360 insertions(+), 43 deletions(-) create mode 100644 tests/tinker/test_sample_result_fast_path.py diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index a7e4e1f450..60529013f7 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -315,15 +315,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 @@ -1512,14 +1516,23 @@ 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 -- but only if diff --git a/skyrl/tinker/config.py b/skyrl/tinker/config.py index a88c005f30..0b224ca2e3 100644 --- a/skyrl/tinker/config.py +++ b/skyrl/tinker/config.py @@ -73,6 +73,21 @@ class EngineConfig(BaseModel): "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.", + ) 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 1e58b4b1bc..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 @@ -54,7 +69,13 @@ class ExternalFutureStore: # 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 @@ -68,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: @@ -99,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 88e63a6776..5b8ff252cc 100644 --- a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py +++ b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py @@ -16,7 +16,11 @@ 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 @@ -142,12 +146,15 @@ 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: + 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 @@ -179,9 +186,7 @@ async def _forward_with_retry(self, sample_req, model_id: str, *, base_model: st "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 @@ -266,16 +271,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/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py index c85f9dcac9..d7b75c9bb6 100644 --- a/tests/tinker/test_external_future_store.py +++ b/tests/tinker/test_external_future_store.py @@ -277,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()) 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 From f38200ac6cc636448548008eeeff75d77fad1205 Mon Sep 17 00:00:00 2001 From: avigyabb <98926738+avigyabb@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:01:44 -0700 Subject: [PATCH 2/2] Clean up comments in ExternalFuture class Removed comments explaining retention and population of result data in ExternalFuture class. --- skyrl/tinker/external_future_store.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/skyrl/tinker/external_future_store.py b/skyrl/tinker/external_future_store.py index 14e1a7b9c6..e8a405cfd7 100644 --- a/skyrl/tinker/external_future_store.py +++ b/skyrl/tinker/external_future_store.py @@ -30,8 +30,6 @@ class ExternalFuture: request_id: int model_id: str | None 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)) @@ -70,8 +68,6 @@ class ExternalFutureStore: _PENDING_TTL_SECONDS = 3600.0 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: @@ -89,8 +85,6 @@ 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)