fix(opd): score teacher over /inference/v1/generate (text + multimodal)

[aoshen02]

Problem (test_qwen2.5_0.5B_opd_vllm, CI #12)

The OPD test launches a stock vllm.entrypoints.openai.api_server teacher and points --rm-url at /v1/completions. But the built-in reward function (vime/rollout/on_policy_distillation.py) — rewritten in #18 when SGLang was removed — POSTs vime's disaggregated /inference/v1/generate contract instead:

• request: top-level token_ids + nested sampling_params
• response: top-level prompt_logprobs

/inference/v1/generate is mounted only on the disagg serve app (vllm/entrypoints/serve/disagg/api_router.py), not on the stock OpenAI server. So the teacher returns HTTP 400.

Worse, reward_func used resp.raise_for_status(): on non-200 aiohttp raises ClientResponseError, whose response headers are a CIMultiDictProxy that can't be pickled. When Ray ships the exception to the driver the real 400 is replaced by an opaque can't pickle CIMultiDictProxy.

The intended contract is already documented by merged PR #62 ("OPD teacher uses /v1/completions with echo=True + prompt_logprobs=1", consuming choices[0].prompt_logprobs) and matched by the test's --rm-url. This PR realigns the code to it.

Note vs #99: #99 flips the test's --rm-url to /inference/v1/generate, but the teacher there is still the stock OpenAI server, which doesn't serve that route → it would 404. Fixing the code to speak /v1/completions (this PR) keeps the standard, portable contract and needs no teacher/endpoint change in the test.

Fix (vime/rollout/on_policy_distillation.py, single file)

• reward_func → OpenAI /v1/completions body: prompt=sample.tokens, max_tokens=1, temperature=0, echo=True, prompt_logprobs=1, skip_special_tokens=False. model only sent when --opd-teacher-model is set (vLLM accepts a missing model → default served model; _is_model_supported(None)→True).
• replace raise_for_status() with an explicit if resp.status != 200: raise RuntimeError(status, body[:1000]) — picklable across Ray, and surfaces the real teacher error instead of the CIMultiDictProxy mask.
• post_process_rewards → read choices[0].prompt_logprobs (was top-level). Per-position shape is identical (dict[token_id→Logprob] | None, pos 0 = None), so _logprob_for_token / length-assert / trim logic is unchanged.

Matches the verl gold standard (extract_prompt_logprobs reads output.prompt_logprobs[1:] from vLLM completions) and slime's parent HTTP-teacher design.

Validation

gb200 single node, 2 student + 1 teacher GPU (shrink of the 4+1 layout), image vime-vllm:sync1916-arm-cpuvalidate, EXIT_RC=0:

• Teacher vLLM server ready (HTTP 200) — no 400, no CIMultiDictProxy.
• teacher_log_probs flow into the OPD KL penalty:
  • rollout 0: rollout_log_probs=-0.1547, teacher_log_probs=-0.1589, opd_reverse_kl=0.0044
  • rollout 1: rollout_log_probs=-0.1546, teacher_log_probs=-0.1596, opd_reverse_kl=0.0041
• Full job (2 rollouts + train + eval) exits cleanly.

Fixes #12.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request migrates the On-Policy Distillation (OPD) teacher logprobs endpoint from the custom /inference/v1/generate to the standard, OpenAI-compatible /v1/completions endpoint. This involves updating the request payload format (e.g., using prompt instead of token_ids and enabling echo) and adjusting the response parsing to extract logprobs from choices[0].prompt_logprobs. Additionally, it replaces resp.raise_for_status() with custom error handling to prevent serialization issues with Ray. The review feedback suggests reusing a persistent aiohttp.ClientSession to avoid performance bottlenecks and using errors="replace" when decoding the error response text to prevent potential decoding exceptions.

[gemini-code-assist]

Instantiating a new aiohttp.ClientSession on every call to reward_func is a performance bottleneck. It prevents connection pooling and keep-alive reuse, resulting in a new TCP connection and SSL handshake for every request. Under high concurrency, this can also lead to socket exhaustion.

Consider reusing a persistent session. For example, you can lazily initialize a module-level session:

_session: aiohttp.ClientSession | None = None

async def get_session() -> aiohttp.ClientSession:
    global _session
    if _session is None or _session.closed:
        _session = aiohttp.ClientSession()
    return _session

And then use it in reward_func without the async with block for the session itself:

    session = await get_session()
    async with session.post(args.rm_url, json=payload) as resp:
        ...

[gemini-code-assist]

Using await resp.text() without specifying an error handling policy can raise a UnicodeDecodeError if the teacher server returns non-UTF-8 binary data or corrupted text on error (e.g., due to a proxy issue or misconfigured gateway). Specifying errors="replace" ensures that the response body is always decoded successfully into a string, preventing the original HTTP error from being masked by a decoding exception.

Suggested change

	body = await resp.text()
	body = await resp.text(errors="replace")

[CalvinXKY]

The PR description says fixing the code to speak /v1/completions avoids changing the teacher, but the actual code uses /inference/v1/generate and the test also flips --rm-url to it — while the teacher is still the stock OpenAI server which doesn't serve that route. The test will 404? Need to either update the teacher launch to use the disagg serve app, or update the PR description to match the actual approach.