From d6227f35e0f300dc4da2e5cc1e7098df943d9f20 Mon Sep 17 00:00:00 2001 From: sergiopaniego Date: Fri, 31 Jul 2026 13:26:21 +0200 Subject: [PATCH] Fix interception proxy capturing 0 token ids on current vLLM vLLM with --return-tokens-as-token-ids encodes the token id in the logprobs "token" field as "token_id:" rather than a separate "token_id" key, so the capture read zero token ids. Every env that trains on the captured trace (opencode, pi, ...) then loses the token ids GRPO needs. Parse the id from the token field. Adds a unit test. --- envs/opencode_env/sandbox/interception.py | 29 +++++++------ tests/envs/test_interception_capture.py | 51 +++++++++++++++++++++++ 2 files changed, 65 insertions(+), 15 deletions(-) create mode 100644 tests/envs/test_interception_capture.py diff --git a/envs/opencode_env/sandbox/interception.py b/envs/opencode_env/sandbox/interception.py index 131d41024..65d3a9d57 100644 --- a/envs/opencode_env/sandbox/interception.py +++ b/envs/opencode_env/sandbox/interception.py @@ -130,9 +130,7 @@ async def chat_completions(request: Request) -> Response: try: body = json.loads(raw_body) except json.JSONDecodeError: - return JSONResponse( - status_code=400, content={"error": "invalid json body"} - ) + return JSONResponse(status_code=400, content={"error": "invalid json body"}) forwarded_body = _prepare_forwarded_body(body, cfg) headers = { @@ -338,7 +336,7 @@ async def _stream() -> Any: yield line + "\n" if not line.startswith("data:"): continue - data = line[len("data:"):].strip() + data = line[len("data:") :].strip() if data == "[DONE]": continue try: @@ -381,7 +379,11 @@ def _accumulate_stream_chunk(chunk: dict[str, Any], acc: dict[str, Any]) -> None tc_idx = tc.get("index", 0) bucket = acc["tool_calls_by_idx"].setdefault( (idx, tc_idx), - {"id": None, "type": "function", "function": {"name": "", "arguments": ""}}, + { + "id": None, + "type": "function", + "function": {"name": "", "arguments": ""}, + }, ) if tc.get("id"): bucket["id"] = tc["id"] @@ -458,11 +460,11 @@ def _build_turn_record( token_ids: list[int] = [] per_token_logps: list[float] = [] for entry in content_lp: - tokens.append(entry.get("token", "")) - # OpenAI returns no raw token ids; vLLM returns them as ``token_id``. - token_id = entry.get("token_id") - if token_id is not None: - token_ids.append(int(token_id)) + token = entry.get("token", "") + tokens.append(token) + # vLLM (--return-tokens-as-token-ids) returns the id in the token field as "token_id:". + if token.startswith("token_id:"): + token_ids.append(int(token.split(":", 1)[1])) lp = entry.get("logprob") if lp is not None: per_token_logps.append(float(lp)) @@ -487,8 +489,7 @@ def _strip_logprobs(response_json: dict[str, Any]) -> dict[str, Any]: choices = out.get("choices") if isinstance(choices, list): out["choices"] = [ - {k: v for k, v in (ch or {}).items() if k != "logprobs"} - for ch in choices + {k: v for k, v in (ch or {}).items() if k != "logprobs"} for ch in choices ] return out @@ -537,9 +538,7 @@ def start(self) -> None: lifespan="on", ) self._server = uvicorn.Server(config) - self._thread = threading.Thread( - target=self._run_server, daemon=True - ) + self._thread = threading.Thread(target=self._run_server, daemon=True) self._thread.start() # Wait for the server to accept connections. deadline = time.time() + 10 diff --git a/tests/envs/test_interception_capture.py b/tests/envs/test_interception_capture.py new file mode 100644 index 000000000..2acc15f5b --- /dev/null +++ b/tests/envs/test_interception_capture.py @@ -0,0 +1,51 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""The interception proxy must capture per-token ids from a vLLM logprobs response. + +vLLM with ``--return-tokens-as-token-ids`` encodes the id in the ``token`` field as +``"token_id:"`` (no separate ``token_id`` key), which the capture must parse so +GRPO trains on real token ids.""" + +import os +import sys + +# Make ``envs/`` importable when running from the repository root. +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) +_ENVS_DIR = os.path.join(_REPO_ROOT, "envs") +if _ENVS_DIR not in sys.path: + sys.path.insert(0, _ENVS_DIR) + +from opencode_env.sandbox.interception import _build_turn_record + + +def _response(content): + return {"choices": [{"finish_reason": "stop", "logprobs": {"content": content}}]} + + +def test_captures_ids_from_vllm_token_id_prefix(): + # vLLM --return-tokens-as-token-ids: id lives in the ``token`` field as "token_id:N". + content = [ + {"token": "token_id:15339", "logprob": -0.1}, + {"token": "token_id:1917", "logprob": -0.2}, + ] + rec = _build_turn_record( + turn_idx=0, request_body={}, response_json=_response(content), latency_s=0.0 + ) + assert rec.completion_token_ids == [15339, 1917] + assert rec.per_token_logps == [-0.1, -0.2] + + +def test_plain_openai_has_logprobs_but_no_ids(): + # A plain OpenAI response (no ids) still yields logprobs, just no token ids. + content = [{"token": "hi", "logprob": -0.3}] + rec = _build_turn_record( + turn_idx=0, request_body={}, response_json=_response(content), latency_s=0.0 + ) + assert rec.completion_token_ids == [] + assert rec.per_token_logps == [-0.3]