diff --git a/renderers/client.py b/renderers/client.py index fd13e6d61..d663f82cf 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio +import base64 import json import logging import math @@ -14,6 +15,7 @@ from typing import Any, cast import httpx +import numpy as np from openai import AsyncOpenAI from renderers.base import ( @@ -134,9 +136,26 @@ def parse_generate_response(raw: bytes) -> dict[str, Any]: return payload +def _parse_packed_completion_logprobs( + packed: Mapping[str, Any], completion_ids: list[int] +) -> list[float]: + """prime-rl's packed form: base64 float32 ``{data, shape, dtype}``, one per token.""" + values = np.frombuffer(base64.b64decode(packed["data"]), dtype=packed["dtype"]) + valid = np.isfinite(values) & (values > VLLM_LOGPROB_SENTINEL) + if len(values) != len(completion_ids) or not valid.all(): + raise MalformedGenerateResponseError( + f"Engine response choice.completion_logprobs: {len(values)} values for " + f"{len(completion_ids)} completion tokens, or a non-finite or missing (-9999) logprob." + ) + return values.astype(np.float64).tolist() + + def _parse_completion_logprobs( choice: Mapping[str, Any], completion_ids: list[int] ) -> list[float]: + packed = choice.get("completion_logprobs") + if packed is not None: + return _parse_packed_completion_logprobs(packed, completion_ids) raw_logprobs = choice.get("logprobs") if not isinstance(raw_logprobs, Mapping): raise MalformedGenerateResponseError( @@ -376,7 +395,8 @@ def _prepare(): routed_experts = choice.get("routed_experts") # vLLM's native kept-set sampling masks (``--return-sampling-mask``): - # one list of surviving vocab ids per completion token. + # one list of surviving vocab ids per completion token, or (prime-rl's + # server) packed CSR ``{"ids", "counts"}`` base64 int32 arrays. sampling_mask = choice.get("sampling_mask") # /inference/v1/generate returns finish_reason in {"stop","length",...} — diff --git a/tests/test_client.py b/tests/test_client.py index 16fdafef6..d9fc53724 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -336,6 +336,33 @@ def test_generate_preserves_zero_completion_logprob(): assert result["completion_logprobs"] == [0.0, -0.2] +def test_generate_reads_packed_completion_logprobs(): + client = _FakeClient() + client.choice.pop("logprobs") + values = np.array([-0.1, 0.0], dtype=np.float32) + client.choice["completion_logprobs"] = { + "data": base64.b64encode(values.tobytes()).decode(), + "shape": [2], + "dtype": "float32", + } + + result = _run_generate(client) + + assert result["completion_logprobs"] == values.tolist() + + client.choice["completion_logprobs"]["data"] = base64.b64encode( + np.array([-0.1, -9999.0], dtype=np.float32).tobytes() + ).decode() + with pytest.raises(MalformedGenerateResponseError, match="missing"): + _run_generate(client) + + client.choice["completion_logprobs"]["data"] = base64.b64encode( + np.array([-0.1, np.inf], dtype=np.float32).tobytes() + ).decode() + with pytest.raises(MalformedGenerateResponseError, match="non-finite"): + _run_generate(client) + + class _MalformedToolRenderer(_FakeRenderer): """Returns only a malformed tool-call attempt — finish_reason must stay "stop"."""