(2/3) refactor(session): assemble training samples on the session server; records never leave it

docs/user-guide/rollout-endpoints.md

@@ -217,8 +217,9 @@ is a thin wrapper around the custom agent. It:
 1. Creates a session on MilesRouter and builds a session-scoped `base_url`.
 2. Calls the custom agent (from `--custom-agent-function-path`) to send one or more
    chat requests.
-3. Collects session records via `OpenAIEndpointTracer`.
-4. Converts records into `Sample` objects via `compute_samples_from_openai_records`.
+3. Collects server-assembled `Sample` objects via `OpenAIEndpointTracer.collect_samples`
+   (the session server converts records into samples, truncates and merges on the
+   owning instance; records never leave the server).
 
 For broader customization beyond the OpenAI wrapper, see the `/generate` path above.
 

miles/rollout/generate_hub/agentic_tool_call.py

@@ -4,8 +4,9 @@
 The agent logic is fully encapsulated in a user-provided async function
 (--custom-agent-function-path). This generate function only handles:
   1. TITO session tracing (OpenAIEndpointTracer)
-  2. Converting session records to training samples
-  3. Multi-turn merge
+  2. Collecting the worker-assembled training samples (the session server
+     converts records to samples, truncates and merges in the owning worker)
+  3. Driver-side metadata application (agent_metadata, session_metadata)
 
 Agent function contract:
   async def my_agent(
@@ -33,8 +34,6 @@ async def my_agent(
 
 from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
 from miles.rollout.generate_utils.openai_endpoint_utils import OpenAIEndpointTracer
-from miles.rollout.generate_utils.sample_utils import merge_samples
-from miles.rollout.session.samples.merge import compute_samples_from_openai_records, truncate_samples_by_total_tokens
 from miles.utils.misc import load_function
 from miles.utils.types import Sample
 
@@ -81,55 +80,48 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         logger.warning(f"{log_prefix} Agent function failed: {e}", exc_info=True)
 
     finally:
-        logger.debug(f"{log_prefix} Calling collect_records...")
-        records, session_metadata = await tracer.collect_records()
-        logger.debug(f"{log_prefix} collect_records done: {len(records)} records")
+        # The session server assembles the samples on the owning instance; records
+        # never leave it. Runs even when the agent function failed, like the old
+        # collect_records; a collect failure (422/5xx/timeout) raises loudly.
+        logger.debug(f"{log_prefix} Calling collect_samples...")
+        result = await tracer.collect_samples(
+            input.sample,
+            multi_samples=input.args.generate_multi_samples,
+            max_seq_len=max_seq_len,
+        )
+        logger.debug(
+            f"{log_prefix} collect_samples done: {len(result.samples)} samples, "
+            f"total_time={time.monotonic()-t_start:.1f}s"
+        )
 
-    if not records:
-        logger.warning("No model calls recorded for sample")
+    if not result.samples:
+        if result.empty_reason == "all_truncated":
+            logger.warning("All samples truncated (prompt already exceeds max_seq_len)")
+        else:
+            logger.warning("No model calls recorded for sample")
         sample = deepcopy(input.sample)
         sample.status = Sample.Status.ABORTED
         return GenerateFnOutput(samples=sample)
 
-    logger.debug(f"{log_prefix} Computing samples from {len(records)} records...")
-    samples = compute_samples_from_openai_records(
-        input.args,
-        input.sample,
-        records,
-        input.state.tokenizer,
-        accumulated_token_ids=session_metadata.get("accumulated_token_ids"),
-        max_trim_tokens=session_metadata.get("max_trim_tokens", 0),
-    )
-
-    logger.debug(
-        f"{log_prefix} compute_samples done: {len(samples)} samples, total_time={time.monotonic()-t_start:.1f}s"
-    )
+    samples = result.samples
     for s in samples:
         s.metadata.update(agent_metadata or {})
 
     # If the agent function reports wall-clock time spent outside policy generation
     # (env/tool steps), surface it on Sample.non_generation_time so throughput
-    # accounting subtracts it. Must be equal across all turn-samples: merge_samples
-    # collapses them with _merge_equal_value, which asserts the values match.
+    # accounting subtracts it (applied to every returned sample, merged or per-turn).
     ngt = ((agent_metadata or {}).get("agent_metrics") or {}).get("total_tool_time")
     if ngt is not None:
         for s in samples:
             s.non_generation_time = ngt
 
-    if max_seq_len is not None:
-        samples = truncate_samples_by_total_tokens(samples, max_seq_len, input.state.tokenizer)
-
-    if not samples:
-        logger.warning("All samples truncated (prompt already exceeds max_seq_len)")
-        sample = deepcopy(input.sample)
-        sample.status = Sample.Status.ABORTED
-        return GenerateFnOutput(samples=sample)
-
     if not input.args.generate_multi_samples:
-        samples = merge_samples(samples, input.state.tokenizer)
-        samples.metadata.update(session_metadata)
-    else:
-        samples[-1].metadata.update(session_metadata)
+        # The server merged already; unwrap to a scalar Sample — downstream
+        # forks on isinstance(sample, list) to pick the multi-samples branch.
+        (merged,) = samples
+        merged.metadata.update(result.session_metadata)
+        return GenerateFnOutput(samples=merged)
+    samples[-1].metadata.update(result.session_metadata)
     return GenerateFnOutput(samples=samples)
 
 
@@ -142,8 +134,8 @@ def _add_arguments(parser: argparse.ArgumentParser):
         default=None,
         dest="max_seq_len",
         help="Max sequence length in tokens (prompt + completion, including env responses) "
-        "per session. Truncates samples on the Miles side and is forwarded to the "
-        "Harbor agent server (as max_seq_len) to abort the trial early.",
+        "per session. Truncation happens inside the session server during sample assembly; "
+        "also forwarded to the Harbor agent server (as max_seq_len) to abort the trial early.",
     )
 
 

miles/rollout/generate_utils/openai_endpoint_utils.py

@@ -7,8 +7,9 @@
 import random
 from argparse import Namespace
 
-from miles.rollout.session.types import GetSessionResponse, SessionRecord
-from miles.utils.http_utils import post
+from miles.rollout.session.samples.codec import SamplesReply, decode_samples_reply
+from miles.utils.http_utils import post, post_bytes_no_retry
+from miles.utils.types import Sample
 
 logger = logging.getLogger(__name__)
 
@@ -50,39 +51,33 @@ async def create(args: Namespace):
             session_server_instance_id=session_server_instance_id,
         )
 
-    async def collect_records(self) -> tuple[list[SessionRecord], dict]:
+    async def collect_samples(
+        self, input_sample: Sample, *, multi_samples: bool, max_seq_len: int | None
+    ) -> SamplesReply:
+        """Fetch the server-assembled training samples for this session.
+
+        Single direct POST, no retries: a 5xx means the owning instance died and
+        the session's records died with it, and a 422 is a deterministic
+        assembly failure whose assertion text is
+        the body — both must raise loudly, immediately. A timeout raises too
+        (assembly is seconds server-side; the old records path silently ABORTed
+        the sample on timeout and lost data). The session DELETE is attempted
+        on every path, success or failure, matching the old cleanup semantics;
+        a DELETE failure is only a warning.
+        """
         try:
-            response = await asyncio.wait_for(
-                post(self.base_url, {}, action="get"),
+            payload = await post_bytes_no_retry(
+                f"{self.base_url}/samples",
+                {"multi_samples": multi_samples, "max_seq_len": max_seq_len},
                 timeout=_SESSION_REQUEST_TIMEOUT,
             )
-        except asyncio.TimeoutError:
-            logger.error(
-                f"Timed out waiting for session {self.session_id} records after {_SESSION_REQUEST_TIMEOUT}s "
-                f"(likely stale HTTP keepalive connection). Returning empty records."
-            )
-            # Still attempt to clean up the session.
+        finally:
             try:
                 await asyncio.wait_for(
                     post(self.base_url, {}, action="delete"),
                     timeout=_SESSION_REQUEST_TIMEOUT,
                 )
-            except Exception:
-                logger.warning(f"Failed to delete session {self.session_id} after timeout")
-            return [], {}
-        except Exception as e:
-            logger.warning(f"Failed to get session {self.session_id} records: {e}")
-            raise
-        response = GetSessionResponse.model_validate(response)
-        records = response.records
-        metadata = response.metadata
-
-        try:
-            await asyncio.wait_for(
-                post(self.base_url, {}, action="delete"),
-                timeout=_SESSION_REQUEST_TIMEOUT,
-            )
-        except Exception as e:
-            logger.warning(f"Failed to delete session {self.session_id} after collecting records: {e}")
+            except Exception as e:
+                logger.warning(f"Failed to delete session {self.session_id} after collecting samples: {e}")
 
-        return (records or []), metadata
+        return decode_samples_reply(payload, input_sample)

miles/rollout/session/core.py

@@ -5,6 +5,7 @@
 - ``chat_completions`` strips the R3 replay payloads (``routed_experts`` / ``indexer_topk``) from the client reply copy-on-write; the ``SessionRecord`` keeps the full response for the training path (``GET /sessions/{id}``).
 - ``chat_completions`` holds the per-session lock for prep and state update but not across the proxy call; ``closing`` re-checks and the ``num_assistant`` check gate concurrent DELETE/chat.
 - ``stream: true`` is served as fake streaming: the backend call stays non-streaming (TITO needs the complete message + meta_info) and the full response is re-rendered as a single SSE chunk plus ``data: [DONE]``. Errors all happen before the SSE body is built, so they keep their real status codes as JSON.
+- ``collect_samples`` assembles training Samples from the session's records on the server (compute -> truncate -> merge, synchronously on the loop like the lock-free ``get_session``); deterministic assembly failures return 422 with the assertion text.
 """
 
 import json
@@ -14,13 +15,16 @@
 
 from starlette.responses import Response
 
+from miles.rollout.generate_utils.sample_utils import merge_samples
 from miles.rollout.session.errors import (
     MessageValidationError,
     SessionNotFoundError,
     TokenizationError,
     UpstreamResponseError,
 )
 from miles.rollout.session.linear_trajectory import SessionRegistry
+from miles.rollout.session.samples.codec import encode_samples_reply
+from miles.rollout.session.samples.merge import compute_samples_from_openai_records, truncate_samples_by_total_tokens
 from miles.rollout.session.types import GetSessionResponse, SessionRecord
 
 logger = logging.getLogger(__name__)
@@ -45,6 +49,11 @@ def _render_json(payload) -> bytes:
     return json.dumps(payload, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode("utf-8")
 
 
+def _samples_response(payload: bytes) -> Response:
+    """The samples-op reply: one safetensors binary payload."""
+    return Response(content=payload, status_code=200, media_type="application/octet-stream")
+
+
 _CLIENT_STRIPPED_META_KEYS = ("routed_experts", "indexer_topk")
 
 
@@ -148,8 +157,10 @@ async def create_session(self) -> Response:
         session_id = self.registry.create_session()
         return Response(content=_render_json({"session_id": session_id}), status_code=200, media_type=JSON_MEDIA_TYPE)
 
-    async def get_session(self, session_id: str) -> Response:
-        session = self.registry.get_session(session_id)
+    def _session_metadata(self, session_id: str, session) -> dict:
+        """The per-session assembly/inspection metadata dict, shared by
+        `get_session` (records debug dump) and `collect_samples` (samples op)
+        so the two can never drift."""
         metadata: dict = {}
         try:
             mismatch = self.registry.compute_session_mismatch(session)
@@ -160,11 +171,55 @@ async def get_session(self, session_id: str) -> Response:
             metadata["tito_session_mismatch"] = mismatch
         metadata["accumulated_token_ids"] = session.token_ids
         metadata["max_trim_tokens"] = self.registry.tito_tokenizer.max_trim_tokens
+        return metadata
+
+    async def get_session(self, session_id: str) -> Response:
+        session = self.registry.get_session(session_id)
+        metadata = self._session_metadata(session_id, session)
         payload = GetSessionResponse(session_id=session_id, records=session.records, metadata=metadata)
         return Response(
             content=_render_json(payload.model_dump(mode="json")), status_code=200, media_type=JSON_MEDIA_TYPE
         )
 
+    async def collect_samples(self, session_id: str, *, multi_samples: bool, max_seq_len: int | None) -> Response:
+        """Assemble training Samples from this session's records, on the server.
+
+        Runs synchronously on the server loop — no await between reading the
+        session state and finishing the reply — the same invariant that makes
+        the lock-free `get_session` safe against concurrent chat updates. Do
+        not offload the assembly to an executor without snapshotting records
+        or holding the session lock.
+
+        Deterministic assembly failures map to 422 with the assertion text as
+        the body. They are caught HERE so they never escape
+        as an unhandled 500; the ValueError catch also
+        covers corrupt stored R3 payloads (binascii/reshape errors) — equally
+        deterministic record damage. Unknown exceptions still propagate (a real
+        bug must not masquerade as 422).
+        """
+        session = self.registry.get_session(session_id)
+        metadata = self._session_metadata(session_id, session)
+        tokenizer = self.registry.tokenizer
+        if not session.records:
+            return _samples_response(encode_samples_reply([], metadata, empty_reason="no_records"))
+        try:
+            samples = compute_samples_from_openai_records(
+                self.args,
+                session.records,
+                tokenizer,
+                accumulated_token_ids=metadata.get("accumulated_token_ids"),
+                max_trim_tokens=metadata.get("max_trim_tokens", 0),
+            )
+            if max_seq_len is not None:
+                samples = truncate_samples_by_total_tokens(samples, max_seq_len, tokenizer)
+            if not samples:
+                return _samples_response(encode_samples_reply([], metadata, empty_reason="all_truncated"))
+            if not multi_samples:
+                samples = [merge_samples(samples, tokenizer)]
+        except (AssertionError, ValueError) as exc:
+            return Response(content=str(exc).encode(), status_code=422, media_type="text/plain")
+        return _samples_response(encode_samples_reply(samples, metadata))
+
     async def delete_session(self, session_id: str) -> Response:
         session = self.registry.get_session(session_id)
         if session.closing: