feat(session): fill omitted sampling fields from defaults set at session creation

docs/user-guide/agentic-rollout.md

@@ -133,11 +133,11 @@ sequence, trims model-specific boundary tokens, and builds the training sample.
 
 ### Choose the session behavior
 
-Session purpose is fixed at creation. `POST /sessions` accepts `{"evaluation": true}` for evaluation; omitting the body or field defaults to training. The field must be a boolean. The agentic generator passes `GenerateFnInput.evaluation` automatically, and the purpose survives v1 rollback and v2 branching.
+Use `{"evaluation": true}` in `POST /sessions` for evaluation; omitting it defaults to training. The agentic generator sets this automatically, and the purpose stays fixed across turns, retries, and branches.
 
-Evaluation sessions force `return_sampling_mask`, `return_routed_experts`, and `return_indexer_topk` to `false`, including when a request or model rule asks to enable them. They ignore `routed_experts_start_len`. Sampling defaults and override order stay unchanged, and temperature can vary between turns. TITO rendering, token IDs, logprobs, and template compatibility checks still apply; agentic evaluation still collects samples.
+`temperature`, `top_p`, and `top_k` in `POST /sessions` provide defaults for omitted or `null` chat fields. Training requests must match the registered temperature or receive HTTP 400; eval temperature and `top_p`/`top_k` remain overridable. The agentic generator registers the sample's resolved values automatically.
 
-Disabling replay outputs does not disable all internal capture work in a shared SGLang engine with capture enabled globally.
+Evaluation forces `return_sampling_mask`, `return_routed_experts`, and `return_indexer_topk` off and ignores `routed_experts_start_len`. TITO, logprobs, and sample collection still apply; engine-internal capture may remain enabled.
 
 History handling depends on the selected server version:
 

docs/user-guide/cli-reference.md

@@ -201,6 +201,7 @@ Sections mirror the launch-script argument groups.
 | `--eval-max-response-len` | int | – | Max eval response length. Inherits from rollout if unset. |
 | `--eval-temperature` | float | – | Eval temperature. Inherits from rollout if unset. |
 | `--eval-top-p` | float | – | Eval top-p. Inherits from rollout if unset. |
+| `--eval-top-k` | int | – | Eval top-k. Inherits from rollout if unset. |
 | `--eval-num-gpus` | int | `0` | Dedicated eval fleet size. `0` = shared-engine eval. Requires `train_async.py`. |
 | `--eval-num-gpus-per-engine` | int | `1` | Eval engine TP, independent of rollout TP. |
 | `--eval-hf-dir` | str | – | Staging dir for per-eval HF snapshots (tmpfs recommended). Unset + `--save-hf` = reuse mode. |

miles/rollout/generate_hub/agentic_tool_call.py

@@ -50,7 +50,9 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     )
     use_v2 = getattr(input.args, "use_session_server", None) == "v2"
     collect_spec_metrics = use_v2 and input.args.sglang_speculative_algorithm is not None
-    tracer = await OpenAIEndpointTracer.create(input.args, evaluation=input.evaluation)
+    tracer = await OpenAIEndpointTracer.create(
+        input.args, evaluation=input.evaluation, sampling_params=input.sampling_params
+    )
 
     custom_agent_function: Callable = load_function(input.args.custom_agent_function_path)
     assert (

miles/rollout/generate_utils/openai_endpoint_utils.py

@@ -13,6 +13,7 @@
     SamplesReply,
     decode_samples_and_merge_input_sample,
 )
+from miles.rollout.session.types import CreateSessionRequest
 from miles.utils.http_utils import post, post_bytes_no_retry
 from miles.utils.types import Sample
 
@@ -44,7 +45,7 @@ def session_server_id(self) -> str:
         return self.router_url.removeprefix("http://")
 
     @staticmethod
-    async def create(args: Namespace, *, evaluation: bool = False):
+    async def create(args: Namespace, *, evaluation: bool = False, sampling_params: dict | None = None):
         session_addrs = getattr(args, "session_server_addrs", None)
         if not session_addrs:
             raise RuntimeError(
@@ -56,7 +57,12 @@ async def create(args: Namespace, *, evaluation: bool = False):
         session_url = f"http://{session_addr}"
         instance_ids = getattr(args, "session_server_instance_ids", None) or {}
         session_server_instance_id = instance_ids.get(session_addr)
-        response = await post(f"{session_url}/sessions", {"evaluation": evaluation}, action="post")
+        # Drop engine-only sampling fields before validating the session creation body.
+        session_params = {
+            key: value for key, value in (sampling_params or {}).items() if key in CreateSessionRequest.model_fields
+        }
+        body = CreateSessionRequest.model_validate({**session_params, "evaluation": evaluation})
+        response = await post(f"{session_url}/sessions", body.model_dump(exclude_none=True), action="post")
         session_id = response["session_id"]
         use_v2 = getattr(args, "use_session_server", None) == "v2"
         return OpenAIEndpointTracer(

miles/rollout/session/core.py

@@ -231,8 +231,8 @@ async def health(self) -> Response:
             body["session_server_instance_id"] = self.instance_id
         return Response(content=_render_json(body), status_code=200, media_type=JSON_MEDIA_TYPE)
 
-    async def create_session(self, *, evaluation: bool = False) -> Response:
-        session_id = self.registry.create_session(evaluation=evaluation)
+    async def create_session(self, *, evaluation: bool = False, sampling_defaults: dict | None = None) -> Response:
+        session_id = self.registry.create_session(evaluation=evaluation, sampling_defaults=sampling_defaults)
         return Response(content=_render_json({"session_id": session_id}), status_code=200, media_type=JSON_MEDIA_TYPE)
 
     def _session_metadata(self, session_id: str, session) -> dict:

miles/rollout/session/linear_trajectory.py

@@ -89,6 +89,7 @@ class LinearTrajectory:
     num_assistant: int = 0
     turn_args_history: list[dict[str, Any]] = field(default_factory=list)
     evaluation: bool = False
+    sampling_defaults: dict[str, Any] = field(default_factory=dict)
 
     @property
     def turn_args(self) -> dict[str, Any]:
@@ -128,7 +129,12 @@ def prepare_token_ids_and_request_args(
         checkpoint_index = self._find_rollback_checkpoint(request_messages, matcher)
         turn_args = self.turn_args_history[checkpoint_index] if checkpoint_index >= 0 else None
         prepared = prepare_chat_request(
-            client_args, tito_tokenizer, config=config, turn_args=turn_args, evaluation=self.evaluation
+            client_args,
+            tito_tokenizer,
+            config=config,
+            turn_args=turn_args,
+            evaluation=self.evaluation,
+            sampling_defaults=self.sampling_defaults,
         )
         prepared.body["input_ids"] = self._render_token_ids(
             request_messages,
@@ -376,9 +382,11 @@ def __init__(
             message_matcher if message_matcher is not None else strict_message_matches
         )
 
-    def create_session(self, *, evaluation: bool = False) -> str:
+    def create_session(self, *, evaluation: bool = False, sampling_defaults: dict[str, Any] | None = None) -> str:
         session_id = uuid.uuid4().hex
-        self.sessions[session_id] = LinearTrajectory(evaluation=evaluation)
+        self.sessions[session_id] = LinearTrajectory(
+            evaluation=evaluation, sampling_defaults=dict(sampling_defaults or {})
+        )
         return session_id
 
     def get_session(self, session_id: str) -> LinearTrajectory:

miles/rollout/session/request_args.py

@@ -2,7 +2,8 @@
 
 ``prepare_chat_request`` owns a copy of the client's input. Config rules apply
 Session server constraints and retain the client's streaming preference. The
-TITO tokenizer then applies model rules to the same full request.
+session's sampling defaults fill the fields the client omitted. The TITO
+tokenizer then applies model rules to the same full request.
 
 After rendering and a successful generation, the session records the complete
 resolved request as ``turn_args``. A continuation supplies that history to the
@@ -54,13 +55,16 @@ def prepare_chat_request(
     config: SessionServerConfig,
     turn_args: dict[str, Any] | None,
     evaluation: bool = False,
+    sampling_defaults: dict[str, Any] | None = None,
 ) -> PreparedChatRequest:
-    """Resolve an owned request using server rules, model rules, and prior turn args.
+    """Resolve an owned request using server rules, session defaults, model rules, and prior turn args.
 
     ``turn_args`` is the continued turn's full request; ``None`` starts a root.
+    ``sampling_defaults`` are the session's values for sampling fields the client omits.
     Client input and recorded history remain unchanged.
     """
     request_args, client_stream = resolve_request_args_by_config(deepcopy(client_args), config, evaluation=evaluation)
+    apply_session_sampling_defaults(request_args, sampling_defaults or {}, evaluation=evaluation)
     try:
         request_args = tito_tokenizer.resolve_request_args(request_args, turn_args=turn_args)
     except ValueError as e:
@@ -74,6 +78,23 @@ def prepare_chat_request(
     )
 
 
+def apply_session_sampling_defaults(
+    request_args: dict[str, Any], sampling_defaults: dict[str, Any], *, evaluation: bool = False
+) -> None:
+    """Fill in place the sampling fields the client left unset from the session's defaults.
+
+    ``None`` counts as unset. Explicit values are kept, except a training temperature
+    mismatch raises ``MessageValidationError`` to keep rollout and training aligned.
+    """
+    for key, value in sampling_defaults.items():
+        if request_args.get(key) is None:
+            request_args[key] = value
+        elif not evaluation and key == "temperature" and request_args[key] != value:
+            raise MessageValidationError(
+                f"temperature={request_args[key]!r} does not match the training session temperature={value!r}"
+            )
+
+
 def resolve_request_args_by_config(
     request_args: dict[str, Any], config: SessionServerConfig, *, evaluation: bool = False
 ) -> tuple[dict[str, Any], bool]:

miles/rollout/session/sessions.py

@@ -107,7 +107,10 @@ async def create_session(request: Request):
             params = CreateSessionRequest.model_validate_json(await request.body() or b"{}")
         except ValidationError as exc:
             return JSONResponse(status_code=400, content={"error": str(exc)})
-        return await core.create_session(evaluation=params.evaluation)
+        return await core.create_session(
+            evaluation=params.evaluation,
+            sampling_defaults=params.model_dump(exclude={"evaluation"}, exclude_none=True),
+        )
 
     @app.get("/sessions/{session_id}")
     async def get_session(session_id: str):

miles/rollout/session/types.py

@@ -1,10 +1,22 @@
-from pydantic import BaseModel, Field, StrictBool
+from pydantic import BaseModel, Field, StrictBool, StrictFloat, StrictInt, field_validator
 
 from miles.utils.pydantic_utils import StrictBaseModel
 
 
 class CreateSessionRequest(StrictBaseModel):
     evaluation: StrictBool = False
+    # the caller resolves rollout/eval/dataset values; the session only fills fields a request omits
+    temperature: StrictFloat | None = None
+    top_p: StrictFloat | None = None
+    top_k: StrictInt | None = None
+
+    @field_validator("top_k", mode="before")
+    @classmethod
+    def _integral_float_as_int(cls, value):
+        # an eval dataset YAML may spell top_k as 40.0; a fractional top_k stays an error
+        if isinstance(value, float) and value.is_integer():
+            return int(value)
+        return value
 
 
 class SessionRecord(BaseModel):

miles/rollout/session/v2/session_state.py

@@ -50,6 +50,7 @@ class SessionStateV2:
     closing: bool = field(default=False, repr=False, compare=False)
     tree: SessionTree = field(default_factory=SessionTree)
     evaluation: bool = False
+    sampling_defaults: dict[str, Any] = field(default_factory=dict)
 
     def latest(self) -> TrajectoryNode | None:
         """The most recently committed generation (always a leaf), or ``None``
@@ -105,6 +106,7 @@ def prepare_token_ids_and_request_args(
         config=config,
         turn_args=parent.turn_args if parent is not None else None,
         evaluation=state.evaluation,
+        sampling_defaults=state.sampling_defaults,
     )
     prepared.body["input_ids"] = _render_token_ids(
         parent, request_messages, template_args=prepared.template_args, tito_tokenizer=tito_tokenizer
@@ -214,9 +216,11 @@ class SessionRegistryV2(SessionRegistry):
 
     sessions: dict[str, SessionStateV2]
 
-    def create_session(self, *, evaluation: bool = False) -> str:
+    def create_session(self, *, evaluation: bool = False, sampling_defaults: dict[str, Any] | None = None) -> str:
         session_id = uuid.uuid4().hex
-        self.sessions[session_id] = SessionStateV2(evaluation=evaluation)
+        self.sessions[session_id] = SessionStateV2(
+            evaluation=evaluation, sampling_defaults=dict(sampling_defaults or {})
+        )
         return session_id
 
     def compute_mismatch(

miles/utils/arguments.py

@@ -548,6 +548,9 @@ def add_rollout_arguments(parser):
                     "Requires train_async.py."
                 ),
             )
+            # Sampling values reach the engine per request only: the built-in generate path sends them
+            # itself and the session server fills fields an agent omits from its session's defaults.
+            # They are never engine launch arguments: an engine shared by rollout and eval has no single default.
             parser.add_argument(
                 "--rollout-temperature",
                 type=float,

tests/fast/rollout/generate_hub/test_agentic_v2.py

@@ -28,7 +28,7 @@ async def collect_samples(self, input_sample, *, max_seq_len, agent_metadata=Non
         return self.reply
 
 
-def _generate_input(*, evaluation=False, **args_kwargs) -> GenerateFnInput:
+def _generate_input(*, evaluation=False, sampling_params=None, **args_kwargs) -> GenerateFnInput:
     args = SimpleNamespace(
         **{
             "session_server_addrs": ["127.0.0.1:12345"],
@@ -48,7 +48,7 @@ def _generate_input(*, evaluation=False, **args_kwargs) -> GenerateFnInput:
         label="label",
         metadata={"source": "test"},
     )
-    return GenerateFnInput(state=state, sample=sample, sampling_params={}, evaluation=evaluation)
+    return GenerateFnInput(state=state, sample=sample, sampling_params=sampling_params or {}, evaluation=evaluation)
 
 
 async def _fake_agent(**kwargs):
@@ -65,8 +65,9 @@ def _session_metadata(spec_info=None):
 
 
 def _patch_agent(monkeypatch, tracer):
-    async def fake_create(args, *, evaluation=False):
+    async def fake_create(args, *, evaluation=False, sampling_params=None):
         tracer.evaluation = evaluation
+        tracer.sampling_params = sampling_params
         return tracer
 
     monkeypatch.setattr(agentic_tool_call.OpenAIEndpointTracer, "create", fake_create)
@@ -80,10 +81,12 @@ async def test_success_returns_list_and_forwards_agent_metadata(monkeypatch, eva
     tracer = _Tracer(SamplesReply(samples=[sample], session_metadata={}, empty_reason=None))
     _patch_agent(monkeypatch, tracer)
 
-    generate_input = _generate_input(evaluation=evaluation)
+    generate_input = _generate_input(evaluation=evaluation, sampling_params={"temperature": 0.7, "max_new_tokens": 8})
     output = await agentic_tool_call.generate(generate_input)
 
     assert tracer.evaluation is evaluation
+    # The session fills these into requests the agent sends without them.
+    assert tracer.sampling_params == {"temperature": 0.7, "max_new_tokens": 8}
     assert output.samples == [sample]
     assert output.samples[0].rollout_id is None
     assert tracer.agent_metadata == {"agent_result": "done"}
@@ -216,7 +219,7 @@ async def test_empty_session_server_addrs_is_rejected(self, monkeypatch, addrs):
         """generate() raises the documented AssertionError when session_server_addrs is absent, null or empty, without creating a tracer."""
         created_for: list[object] = []
 
-        async def fake_create(args, *, evaluation=False):
+        async def fake_create(args, *, evaluation=False, sampling_params=None):
             created_for.append(args)
             return _Tracer(SamplesReply(samples=[], session_metadata={}, empty_reason="no_records"))
 

tests/fast/rollout/generate_hub/test_multi_turn.py

@@ -173,6 +173,9 @@ def expected_openai_request(messages: list[dict], **extra) -> dict:
         # The R3 replay flags follow the launch flags and are always present.
         "return_routed_experts": False,
         "return_indexer_topk": False,
+        # The mock agent drops request_kwargs; the session fills the temperature the
+        # generator registered at creation.
+        "temperature": DEFAULT_SAMPLING_PARAMS["temperature"],
         "chat_template_kwargs": {"clear_thinking": False},
         **extra,
     }

tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py

@@ -22,13 +22,31 @@
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("create_kwargs", [{}, {"evaluation": False}, {"evaluation": True}])
-async def test_create_reads_session_server_instance_id_from_args(monkeypatch, create_kwargs):
+@pytest.mark.parametrize(
+    ("create_kwargs", "expected_payload"),
+    [
+        ({}, {"evaluation": False}),
+        ({"evaluation": False}, {"evaluation": False}),
+        ({"evaluation": True}, {"evaluation": True}),
+        # Only the sampling fields the session fills travel; None and other keys stay behind.
+        (
+            {"sampling_params": {"temperature": 0.6, "top_p": None, "top_k": 20, "max_new_tokens": 8}},
+            {"evaluation": False, "temperature": 0.6, "top_k": 20},
+        ),
+        # An eval dataset YAML can spell top_k as 40.0; the session body carries the integer.
+        ({"sampling_params": {"top_k": 40.0}}, {"evaluation": False, "top_k": 40}),
+    ],
+)
+async def test_create_reads_session_server_instance_id_from_args(monkeypatch, create_kwargs, expected_payload):
     calls: list[tuple[str, str]] = []
 
     async def fake_post(url: str, payload: dict, action: str = "post"):
         calls.append((action, url))
-        assert payload == {"evaluation": create_kwargs.get("evaluation", False)}
+        assert payload == expected_payload
+        # 40.0 == 40 in Python, so the equality above cannot see a float leaking through
+        assert {key: type(value) for key, value in payload.items()} == {
+            key: type(value) for key, value in expected_payload.items()
+        }
         assert action == "post"
         assert url == "http://127.0.0.1:12345/sessions"
         return {"session_id": "session-123"}