feat(vllm-model): pin sampling params for on-policy training

responses_api_models/vllm_model/app.py

@@ -167,6 +167,23 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig):
 
     chat_template_kwargs: Optional[Dict[str, Any]] = None
 
+    # Sampling params this server puts on every request it sends to the engine, replacing what the caller sent.
+    # On-policy training requires generation to use the sampling distribution the policy is optimized under,
+    # and a caller outside the training loop has no way to know it.
+    #
+    # The common case is an absent parameter rather than a conflicting one.
+    # A caller need not send sampling params at all.
+    # Converters forward a field only when it was set, so the outbound body can carry no temperature or top_p,
+    # and the engine applies a default of its own that has no relation to the configured one.
+    # Replacing rather than filling in covers the other case, a caller that sends values it chose itself.
+    #
+    # Read from config only, never from a request, so the server and not the caller decides them.
+    # Applied at every site that builds a request for the engine,
+    # since a pin that covers some endpoints and not others is off-policy while reporting that sampling is pinned.
+    #
+    # Unset means no pin.
+    sampling_overrides: Optional[Dict[str, Any]] = None
+
     # Corresponds to the extra_body of OpenAI Client.
     extra_body: Optional[Dict[str, Any]] = None
 
@@ -226,6 +243,13 @@ def model_post_init(self, context):
         return super().model_post_init(context)
 
     def _post_init(self) -> None:
+        if self.config.sampling_overrides:
+            LOG.info(
+                "`%s` pins sampling on every request to the engine: %s",
+                self.config.name,
+                self.config.sampling_overrides,
+            )
+
         self._clients = [
             NeMoGymAsyncOpenAI(
                 base_url=base_url,
@@ -301,6 +325,19 @@ async def responses(
             responses_create_params=body, chat_completion=chat_completion_response
         )
 
+    def _apply_sampling_overrides(self, body_dict: Dict[str, Any]) -> Dict[str, Any]:
+        """Force ``config.sampling_overrides`` onto an outbound body, in place.
+
+        Applied last at every site that builds a request for the engine, so the pinned values win
+        over both what the client sent and anything ``extra_body`` merged in, and are present when
+        the client sent nothing. Every path has to call this: a harness picks its own endpoint, and
+        a pin that covers only one of them yields off-policy generation while reporting that
+        sampling is pinned.
+        """
+        if self.config.sampling_overrides:
+            body_dict.update(self.config.sampling_overrides)
+        return body_dict
+
     async def _responses_native(
         self, request: Request, body: NeMoGymResponseCreateParamsNonStreaming
     ) -> NeMoGymResponse:
@@ -322,6 +359,7 @@ async def _responses_native(
             body_dict["chat_template_kwargs"] = deepcopy(self.config.chat_template_kwargs)
         if self.config.extra_body:
             body_dict = self.config.extra_body | body_dict
+        self._apply_sampling_overrides(body_dict)
 
         client = self._resolve_client(request)
         response_dict = await client.create_response(**body_dict)
@@ -552,7 +590,7 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict:
                 # No user message found — create one with just the audio blocks.
                 body_dict.setdefault("messages", []).append({"role": "user", "content": list(audio_blocks)})
 
-        return body_dict
+        return self._apply_sampling_overrides(body_dict)
 
     async def chat_completions(
         self, request: Request, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body()
@@ -1000,7 +1038,11 @@ def _build_completion_body_from_chat_body(self, chat_body_dict: Dict[str, Any],
             out["return_token_ids"] = True
             out["return_tokens_as_token_ids"] = True
 
-        return out
+        # This path never runs _preprocess_chat_completion_create_params;
+        # chat_completions() branches here before preprocessing, so the pin has to be applied again.
+        # vLLM accepts the same sampling field names on /v1/completions, and the body is forwarded as raw JSON,
+        # so params without a first-class OpenAI completion field (top_k, min_p) pass through.
+        return self._apply_sampling_overrides(out)
 
     def _completion_dict_to_chat_completion(self, completion_dict: Dict[str, Any]) -> NeMoGymChatCompletion:
         """Wrap a /v1/completions response as a NeMoGymChatCompletion.

responses_api_models/vllm_model/tests/test_app.py

@@ -4797,3 +4797,84 @@ async def mock_create_chat_completion(**kwargs):
             )
         # The tokenize endpoint must not be reached once the contract check fails.
         mock_client.create_tokenize.assert_not_called()
+
+
+class TestSamplingOverrides:
+    """Forcing the sampling params on every request.
+
+    An external harness picks its own temperature and top_p. On-policy RL requires
+    generation to match the distribution the policy is optimized under, so the
+    server overrides whatever the client sent rather than trusting it.
+    """
+
+    @staticmethod
+    def _server(overrides: dict[str, object] | None, **kwargs: object) -> VLLMModel:
+        config = VLLMModelConfig(
+            host="0.0.0.0",
+            port=8081,
+            base_url="http://api.openai.com/v1",
+            api_key="dummy_key",  # pragma: allowlist secret
+            model="dummy_model",
+            entrypoint="",
+            name="",
+            return_token_id_information=False,
+            uses_reasoning_parser=False,
+            sampling_overrides=overrides,
+            **kwargs,
+        )
+        return VLLMModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={}))
+
+    def test_overrides_replace_what_the_client_sent(self) -> None:
+        server = self._server({"temperature": 1.0, "top_p": 1.0})
+        out = server._preprocess_chat_completion_create_params(
+            MagicMock(), {"messages": [{"role": "user", "content": "hi"}], "temperature": 0.2, "top_p": 0.5}
+        )
+        assert out["temperature"] == 1.0
+        assert out["top_p"] == 1.0
+
+    def test_overrides_apply_even_when_the_client_sent_nothing(self) -> None:
+        server = self._server({"temperature": 1.0})
+        out = server._preprocess_chat_completion_create_params(
+            MagicMock(), {"messages": [{"role": "user", "content": "hi"}]}
+        )
+        assert out["temperature"] == 1.0
+
+    def test_unset_leaves_the_request_alone(self) -> None:
+        server = self._server(None)
+        out = server._preprocess_chat_completion_create_params(
+            MagicMock(), {"messages": [{"role": "user", "content": "hi"}], "temperature": 0.2}
+        )
+        assert out["temperature"] == 0.2
+
+    def test_overrides_reach_the_completions_api_path(self) -> None:
+        """``use_completions_api`` skips chat preprocessing entirely.
+
+        ``chat_completions`` branches into ``_chat_completions_via_completions_api`` before
+        ``_preprocess_chat_completion_create_params`` runs, so a pin applied only there would be
+        silently inert on the path base-model training uses.
+        """
+        server = self._server({"temperature": 1.0, "top_p": 1.0, "top_k": -1}, use_completions_api=True)
+        out = server._build_completion_body_from_chat_body(
+            {"messages": [{"role": "user", "content": "hi"}], "temperature": 0.2, "top_p": 0.5, "top_k": 50},
+            "hi",
+        )
+        assert out["temperature"] == 1.0
+        assert out["top_p"] == 1.0
+        # No first-class /v1/completions field; the body is forwarded as raw JSON so it still lands.
+        assert out["top_k"] == -1
+
+    def test_overrides_win_over_extra_body_on_the_completions_path(self) -> None:
+        """``extra_body`` merges under the request body; the pin is applied after both."""
+        server = self._server(
+            {"temperature": 1.0},
+            use_completions_api=True,
+            extra_body={"temperature": 0.7, "min_p": 0.05},
+        )
+        out = server._build_completion_body_from_chat_body({"messages": [], "temperature": 0.2}, "hi")
+        assert out["temperature"] == 1.0
+        assert out["min_p"] == 0.05
+
+    def test_overrides_reach_the_responses_native_path(self) -> None:
+        server = self._server({"temperature": 1.0}, is_responses_native=True)
+        body = {"model": "dummy_model", "temperature": 0.2}
+        assert server._apply_sampling_overrides(body)["temperature"] == 1.0