(2/2) 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(
@@ -23,6 +24,7 @@ async def my_agent(
 """
 
 import argparse
+import asyncio
 import logging
 import time
 from collections.abc import Callable
@@ -33,8 +35,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
 
@@ -67,6 +67,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     metadata = {**metadata, "session_server_id": tracer.session_server_id}
 
     agent_metadata = None
+    collect_timed_out = False
     t_start = time.monotonic()
     try:
         logger.debug(f"{log_prefix} Starting agent function call")
@@ -81,52 +82,50 @@ 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")
-
-    if not records:
-        logger.warning("No model calls recorded for sample")
+        # Collect even if the agent failed.
+        logger.debug(f"{log_prefix} Calling collect_samples...")
+        try:
+            result = await tracer.collect_samples(
+                input.sample,
+                max_seq_len=max_seq_len,
+            )
+        except asyncio.TimeoutError:
+            collect_timed_out = True
+            logger.warning(f"{log_prefix} Timed out collecting samples", exc_info=True)
+        else:
+            logger.debug(
+                f"{log_prefix} collect_samples done: {len(result.samples)} samples, "
+                f"total_time={time.monotonic()-t_start:.1f}s"
+            )
+
+    if collect_timed_out:
         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),
-    )
+    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} 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.
     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)
-
-    sample = merge_samples(samples, input.state.tokenizer)
-    sample.metadata.update(session_metadata)
+    (sample,) = samples
+    sample.metadata.update(result.session_metadata)
     return GenerateFnOutput(samples=sample)
 
 
@@ -138,8 +137,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_and_merge_input_sample
+from miles.utils.http_utils import post, post_bytes_no_retry
+from miles.utils.types import Sample
 
 logger = logging.getLogger(__name__)
 
@@ -50,39 +51,22 @@ 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, *, max_seq_len: int | None) -> SamplesReply:
+        """Fetch server-assembled training samples for this session."""
         try:
-            response = await asyncio.wait_for(
-                post(self.base_url, {}, action="get"),
+            # `asyncio.TimeoutError` propagates after cleanup is attempted for `agentic_tool_call.generate` to handle.
+            payload = await post_bytes_no_retry(
+                f"{self.base_url}/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_and_merge_input_sample(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
+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__)
@@ -48,6 +52,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")
 
 
@@ -151,8 +160,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)
@@ -163,11 +174,43 @@ 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, *, max_seq_len: int | None) -> Response:
+        """Assemble training Samples from this session's records.
+
+        Validation failures return 422; unexpected errors propagate.
+        """
+        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([], 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([], metadata, empty_reason="all_truncated"))
+            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(samples, metadata))
+
     async def delete_session(self, session_id: str) -> Response:
         session = self.registry.get_session(session_id)
         if session.closing: