feat(vllm_model): add extra_body_override_keys for server-authoritative sampling

responses_api_models/vllm_model/app.py

@@ -67,6 +67,15 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig):
     # Corresponds to the extra_body of OpenAI Client.
     extra_body: Optional[Dict[str, Any]] = None
 
+    # Keys of ``extra_body`` whose configured value takes precedence over a value set explicitly on
+    # the request. By default ``extra_body`` only supplies defaults and the request always wins.
+    #
+    # This exists for RL training, where generation must be on-policy: the sampling parameters used
+    # to generate have to match the training config exactly. An agent that hardcodes e.g.
+    # ``temperature`` on an internal/auxiliary call would otherwise silently produce off-policy
+    # rollouts. Set to e.g. ``["temperature", "top_p"]`` to make the server authoritative for those.
+    extra_body_override_keys: Optional[List[str]] = None
+
     default_headers: Dict[str, str] = Field(default_factory=dict)
     # Optional prefix for resolving relative ``metadata.audio_path`` (or
     # entries in ``metadata.audio_paths``) against. Absolute paths are used
@@ -180,6 +189,19 @@ def _load_chat_template_tokenizer(self):
 
         return tokenizer
 
+    def _merge_extra_body(self, extra_body: Dict[str, Any], body_dict: Dict[str, Any]) -> Dict[str, Any]:
+        """Merge ``extra_body`` under ``body_dict`` so the request wins, then re-apply override keys.
+
+        ``extra_body`` normally supplies defaults only: anything the request set explicitly takes
+        precedence. Keys listed in ``config.extra_body_override_keys`` invert that, so the server
+        value wins even against an explicit request value. See ``VLLMModelConfig`` for why.
+        """
+        body_dict = extra_body | body_dict
+        for key in self.config.extra_body_override_keys or ():
+            if key in extra_body:
+                body_dict[key] = extra_body[key]
+        return body_dict
+
     async def responses(
         self, request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body()
     ) -> NeMoGymResponse:
@@ -217,7 +239,7 @@ async def _responses_native(
         if self.config.chat_template_kwargs:
             body_dict["chat_template_kwargs"] = deepcopy(self.config.chat_template_kwargs)
         if self.config.extra_body:
-            body_dict = self.config.extra_body | body_dict
+            body_dict = self._merge_extra_body(self.config.extra_body, body_dict)
 
         client = self._resolve_client(request)
         response_dict = await client.create_response(**body_dict)
@@ -379,7 +401,7 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict:
             body_dict.pop("top_logprobs", None)
 
         if extra_body:
-            body_dict = extra_body | body_dict
+            body_dict = self._merge_extra_body(extra_body, body_dict)
 
         # Audio sidechannel: rows can carry audio on
         # ``responses_create_params.metadata`` via three mutually exclusive
@@ -813,10 +835,10 @@ def _build_completion_body_from_chat_body(self, chat_body_dict: Dict[str, Any],
 
         # Operator-level extra_body merges in (e.g. return_tokens_as_token_ids).
         # Same precedence as the chat path: extra_body fields do NOT override
-        # request-level fields.
+        # request-level fields, except for keys listed in extra_body_override_keys.
         if self.config.extra_body:
             extra_body = deepcopy(self.config.extra_body)
-            out = extra_body | out
+            out = self._merge_extra_body(extra_body, out)
 
         if self.config.return_token_id_information:
             # Prefer vLLM's inline prompt and generation token IDs. Keep the

responses_api_models/vllm_model/configs/vllm_model.yaml

@@ -10,4 +10,5 @@ policy_model:
       uses_interleaved_reasoning: true
       chat_template_kwargs: null
       extra_body: null
+      extra_body_override_keys: null
       default_headers: {}

responses_api_models/vllm_model/tests/test_app.py

@@ -4684,3 +4684,76 @@ 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()
+
+
+def _make_extra_body_model(extra_body_override_keys: list[str] | None = None) -> VLLMModel:
+    """A VLLMModel instance with the minimum config needed to exercise _merge_extra_body."""
+    config = VLLMModelConfig(
+        host="0.0.0.0",
+        port=8080,
+        entrypoint="",
+        name="vllm_model",
+        base_url="http://localhost:9999/v1",
+        api_key="dummy_key",  # pragma: allowlist secret
+        model="dummy-model",
+        return_token_id_information=False,
+        uses_reasoning_parser=False,
+        uses_interleaved_reasoning=False,
+        extra_body_override_keys=extra_body_override_keys,
+    )
+    return VLLMModel(config=config, server_client=MagicMock(spec=ServerClient))
+
+
+class TestMergeExtraBody:
+    def test_request_wins_by_default(self) -> None:
+        """Without override keys, an explicit request value beats the configured default."""
+        model = _make_extra_body_model()
+        merged = model._merge_extra_body({"temperature": 1.0}, {"temperature": 0.3})
+        assert merged["temperature"] == 0.3
+
+    def test_extra_body_supplies_defaults(self) -> None:
+        """Keys absent from the request still come from extra_body."""
+        model = _make_extra_body_model()
+        merged = model._merge_extra_body({"temperature": 1.0, "top_p": 1.0}, {"model": "m"})
+        assert merged == {"temperature": 1.0, "top_p": 1.0, "model": "m"}
+
+    def test_override_key_wins_over_explicit_request_value(self) -> None:
+        """The RL case: a hardcoded request temperature must not override the training config."""
+        model = _make_extra_body_model(extra_body_override_keys=["temperature", "top_p"])
+        merged = model._merge_extra_body({"temperature": 1.0, "top_p": 1.0}, {"temperature": 0.3})
+        assert merged["temperature"] == 1.0
+        assert merged["top_p"] == 1.0
+
+    def test_override_only_applies_to_listed_keys(self) -> None:
+        """Keys not listed keep the normal request-wins behaviour."""
+        model = _make_extra_body_model(extra_body_override_keys=["temperature"])
+        merged = model._merge_extra_body({"temperature": 1.0, "top_p": 1.0}, {"temperature": 0.3, "top_p": 0.5})
+        assert merged["temperature"] == 1.0
+        assert merged["top_p"] == 0.5
+
+    def test_override_key_absent_from_extra_body_is_ignored(self) -> None:
+        """An override key the server did not configure must not inject a value."""
+        model = _make_extra_body_model(extra_body_override_keys=["temperature", "top_p"])
+        merged = model._merge_extra_body({"temperature": 1.0}, {"top_p": 0.5})
+        assert merged["temperature"] == 1.0
+        assert merged["top_p"] == 0.5
+
+    def test_does_not_mutate_configured_extra_body(self) -> None:
+        """The configured dict is shared across requests and must not be written through."""
+        model = _make_extra_body_model(extra_body_override_keys=["temperature"])
+        extra_body = {"temperature": 1.0}
+        model._merge_extra_body(extra_body, {"temperature": 0.3, "seed": 7})
+        assert extra_body == {"temperature": 1.0}
+
+    def test_completions_path_honours_override_keys(self) -> None:
+        """The /v1/completions body builder must apply overrides too.
+
+        It merges extra_body separately from the chat and Responses paths. If it kept the
+        plain request-wins merge, an operator who set extra_body_override_keys would get the
+        on-policy guarantee on two of the three paths and silently lose it on the third.
+        """
+        model = _make_extra_body_model(extra_body_override_keys=["temperature"])
+        model.config.extra_body = {"temperature": 1.0, "top_p": 0.95}
+        out = model._build_completion_body_from_chat_body({"temperature": 0.3, "top_p": 0.5}, prompt="x")
+        assert out["temperature"] == 1.0
+        assert out["top_p"] == 0.5