fix: keep OPD teacher scoring on native vLLM route

Author: tim0120

skills/configs/SKILL.md

@@ -70,7 +70,7 @@ In `opd`, rollouts are generated by the student. The orchestrator scores the stu
 
 `[inference]` is required for the usual online path because it starts the student inference server and auto-configures `orchestrator.student.client.base_url`. The student pool is used for online evals and policy weight sync. For externally started student inference, set `orchestrator.student.client.base_url` explicitly instead.
 
-Teacher logprob scoring supports both self-hosted vLLM and Prime API teacher clients: `/inference/v1/generate` for vLLM server roots, `/api/v1/generate` when the teacher client base URL ends in `/api/v1`.
+Teacher logprob scoring uses PrimeRL's vLLM-native `/inference/v1/generate` route. The request field is `token_ids`, meaning the prompt plus completion tokens to score; response `choices[].token_ids` remains generated completion tokens and is not used for OPD scoring.
 
 ## RL trainer token exports
 

src/prime_rl/orchestrator/utils.py

@@ -107,20 +107,15 @@ async def compute_teacher_logprobs(
     """Compute teacher model logprobs for a batch of training samples via prefill."""
     import httpx
 
-    def _teacher_generate_request(base_url: str, model_name: str, token_ids: list[int]) -> tuple[str, dict[str, Any]]:
+    def _teacher_generate_request(
+        base_url: str,
+        model_name: str,
+        scored_token_ids: list[int],
+    ) -> tuple[str, dict[str, Any]]:
         base = base_url.rstrip("/")
-        if base.endswith("/api/v1"):
-            return f"{base}/generate", {
-                "model": model_name,
-                "prompt_token_ids": token_ids,
-                "max_tokens": 1,
-                "temperature": 1.0,
-                "top_p": 1.0,
-                "prompt_logprobs": 1,
-            }
         return f"{base.removesuffix('/v1')}/inference/v1/generate", {
             "model": model_name,
-            "token_ids": token_ids,
+            "token_ids": scored_token_ids,
             "sampling_params": {
                 "max_tokens": 1,
                 "temperature": 1.0,
@@ -157,27 +152,26 @@ def _flatten_prompt_logprobs(response: dict[str, Any], token_ids: list[int]) ->
 
     async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample) -> list[float]:
         client = setup_openai_client(client_config)
-        token_ids = list(sample.prompt_ids) + list(sample.completion_ids)
+        scored_token_ids = list(sample.prompt_ids) + list(sample.completion_ids)
 
         # Two escape hatches from ``AsyncOpenAI.post``:
         #   1. URL — vLLM mounts ``/inference/v1/generate`` at server root,
-        #      while Prime Inference exposes ``/api/v1/generate``. Pass an
-        #      absolute URL so the SDK's ``_prepare_url`` skips base-url merge.
+        #      so pass an absolute URL and skip the SDK's base-url merge.
         #   2. Parse — vLLM's ``GenerateResponse`` is a plain
         #      ``pydantic.BaseModel`` and the SDK's parse layer rejects any
         #      ``cast_to`` that doesn't subclass ``openai.BaseModel``. Use
         #      ``cast_to=httpx.Response`` so the SDK still builds the request
         #      (preserving ``auth_headers``, retries, timeouts, idempotency
         #      keys) and just hands us the raw response to validate ourselves.
-        url, body = _teacher_generate_request(str(client.base_url), model_name, token_ids)
+        url, body = _teacher_generate_request(str(client.base_url), model_name, scored_token_ids)
         http_response = await client.post(
             url,
             cast_to=httpx.Response,
             body=body,
         )
         http_response.raise_for_status()
         response = http_response.json()
-        return _flatten_prompt_logprobs(response, token_ids)
+        return _flatten_prompt_logprobs(response, scored_token_ids)
 
     return await asyncio.gather(*[_compute_single(client, sample) for client, sample in zip(cycle(clients), samples)])
 

tests/unit/orchestrator/test_teacher_logprobs.py

@@ -85,53 +85,6 @@ async def _run():
     asyncio.run(_run())
 
 
-def test_compute_teacher_logprobs_uses_prime_generate_for_api_base_url(monkeypatch):
-    async def _run():
-        fake_client = _FakeOpenAIClient(
-            {
-                "request_id": "gen-test",
-                "choices": [],
-                "prompt_logprobs": [
-                    None,
-                    {"13": {"logprob": -0.1}, "2": {"logprob": -0.7}},
-                    {"198": {"logprob": -0.2}, "3": {"logprob": -0.3}},
-                ],
-                "kv_transfer_params": None,
-            },
-            base_url="https://api.primeintellect.ai/api/v1",
-        )
-        monkeypatch.setattr(orchestrator_utils, "setup_openai_client", lambda _: fake_client)
-
-        sample = SimpleNamespace(
-            prompt_ids=[1],
-            prompt_mask=[True],
-            completion_ids=[2, 3],
-            completion_mask=[True, True],
-            completion_logprobs=[-0.1, -0.2],
-            completion_temperatures=[1.0, 1.0],
-            env_name="test-env",
-        )
-
-        result = await orchestrator_utils.compute_teacher_logprobs(
-            clients=[vf.ClientConfig()],
-            model_name="teacher-model",
-            samples=[sample],
-        )
-
-        assert result == [[0.0, -0.7, -0.3]]
-        assert fake_client.calls[0]["url"] == "https://api.primeintellect.ai/api/v1/generate"
-        assert fake_client.calls[0]["body"] == {
-            "model": "teacher-model",
-            "prompt_token_ids": [1, 2, 3],
-            "max_tokens": 1,
-            "temperature": 1.0,
-            "top_p": 1.0,
-            "prompt_logprobs": 1,
-        }
-
-    asyncio.run(_run())
-
-
 def test_compute_teacher_logprobs_rejects_wrong_length(monkeypatch):
     async def _run():
         fake_client = _FakeOpenAIClient(