From c95ee80420371658c882d74a9906926af7656475 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Fri, 24 Jul 2026 20:54:35 +0000 Subject: [PATCH] feat(session): assemble training samples on the session server; records never leave it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit POST /sessions/{session_id}/samples (registered before the catch-all proxy, which would otherwise forward it to the inference backend): the owning instance runs compute -> truncate -> merge synchronously on its event loop (the lock-free get_session invariant) and replies with one safetensors payload — per-sample tensors named {field}.{index} for the array fields (loss_mask crosses as uint8, one byte per token), with the json-codec fields riding as the rank-one uint8 tensor _samples_meta (safetensors.numpy.load exposes no header metadata and __metadata__ is reserved by the format), so malformed payloads fail inside safetensors' validated parser instead of hand-rolled framing checks. The wire contract is one declarative table, SAMPLES_VALUE_SPEC: every computed field maps to a ValueSpec naming its codec ("tensor" | "tensor_list" | "json") and pinned wire dtype, and the COMPUTED_FIELDS allowlist derives from the table. Only those fields cross the wire on blank templates; every other Sample field never crosses — the driver overlays the wire fields onto deepcopies of its input sample and keeps its local value verbatim for the rest. Non-strict dtype conversions must be value-exact (a negative loss-mask entry wrapping into uint8 raises instead of shipping corrupt masks). The codec lives in samples/codec.py, depending only on Sample, NumPy, and safetensors (>=0.8.0, now an explicit dependency); deterministic assembly failures return 422 with the assertion text; empty replies carry a no_records/all_truncated discriminator preserving today's ABORTED semantics. The old records path (GET the full dump, GiB-scale under R3 with per-turn full-prefix arrays, parsed on the driver's single interpreter) is retired from the training path. collect_samples posts once via the new http_utils.post_bytes_no_retry (post() force-decodes json/text and blind-retries; a 5xx here means the owning instance died with the records, a 422 is deterministic), raises on non-2xx with the body text, raises on timeout instead of silently ABORTing the sample (a measured data-loss bug in the records path), and attempts the session DELETE on every path so failed sessions cannot accumulate. agentic_tool_call.generate shrinks to a shell: agent call -> collect_samples -> empty_reason mapping -> metadata application in today's order. Tests assert golden Sample field values derived from the records fixture (per-turn, merged, truncated), the wire-codec round-trip, malformed-payload contracts, the 422/404/empty_reason/route-order contracts, and the client's behavior deltas. Co-Authored-By: Claude Fable 5 --- docs/user-guide/rollout-endpoints.md | 5 +- .../rollout/generate_hub/agentic_tool_call.py | 70 ++-- .../generate_utils/openai_endpoint_utils.py | 53 ++- miles/rollout/session/core.py | 59 ++- miles/rollout/session/samples/codec.py | 202 ++++++++++ miles/rollout/session/samples/merge.py | 29 +- miles/rollout/session/sessions.py | 16 + miles/utils/http_utils.py | 22 ++ requirements.txt | 1 + tests/fast/fixtures/generation_fixtures.py | 16 + tests/fast/fixtures/rollout_fixtures.py | 1 + .../rollout/generate_hub/test_multi_turn.py | 8 +- .../test_openai_endpoint_utils.py | 156 +++++++- tests/fast/rollout/session/test_samples.py | 84 ++--- .../rollout/session/test_samples_codec.py | 230 ++++++++++++ tests/fast/router/test_session_samples_op.py | 348 ++++++++++++++++++ 16 files changed, 1140 insertions(+), 160 deletions(-) create mode 100644 miles/rollout/session/samples/codec.py create mode 100644 tests/fast/rollout/session/test_samples_codec.py create mode 100644 tests/fast/router/test_session_samples_op.py diff --git a/docs/user-guide/rollout-endpoints.md b/docs/user-guide/rollout-endpoints.md index 47a7b419b7..0bb7a42804 100644 --- a/docs/user-guide/rollout-endpoints.md +++ b/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. diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index 3cd1f7ef06..ee00f5a999 100644 --- a/miles/rollout/generate_hub/agentic_tool_call.py +++ b/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.", ) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 19338f8798..f8f15da3c5 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/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) diff --git a/miles/rollout/session/core.py b/miles/rollout/session/core.py index e6cfc341da..20c2c0a917 100644 --- a/miles/rollout/session/core.py +++ b/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,6 +15,7 @@ from starlette.responses import Response +from miles.rollout.generate_utils.sample_utils import merge_samples from miles.rollout.session.errors import ( MessageValidationError, SessionNotFoundError, @@ -21,6 +23,8 @@ 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: diff --git a/miles/rollout/session/samples/codec.py b/miles/rollout/session/samples/codec.py new file mode 100644 index 0000000000..c89293086d --- /dev/null +++ b/miles/rollout/session/samples/codec.py @@ -0,0 +1,202 @@ +"""Wire codec for the `POST /sessions/{id}/samples` reply. + +The wire contract is the single table `SAMPLES_VALUE_SPEC`: every computed field +maps to a `ValueSpec` naming its codec and wire dtype. Only these fields are +assembled from the records on the server and cross the samples wire; every other +`Sample` field never crosses — the driver overlay keeps its local deepcopy's +value verbatim. + +Payload layout — one safetensors container: + +- Tensor codecs ("tensor", "tensor_list"): one tensor per sample per non-null + field, named `{field}.{sample_index}`. A null field carries no tensor and is + listed in that sample's `nulls` marker instead. +- "json" fields ride `_samples_meta`, a rank-one uint8 tensor holding one UTF-8 + JSON document: per-sample json fields and null markers, the session metadata, + and the empty-reason discriminator. (`safetensors.numpy.load` exposes no + header metadata and `__metadata__` is reserved by the format, hence the + tensor carrier; malformed payloads still fail inside safetensors' validated + parser instead of hand-rolled framing checks.) + +- Server: `encode_samples_reply(samples, session_metadata, empty_reason)`. +- Driver: `decode_samples_reply(payload, input_sample)` — rebuilds each `Sample` + by overlaying the wire's computed fields onto a deepcopy of `input_sample`. +""" + +import dataclasses +import json +from copy import copy, deepcopy + +import numpy as np +import safetensors.numpy + +from miles.utils.types import Sample + + +@dataclasses.dataclass(frozen=True) +class ValueSpec: + """Wire contract of one computed field.""" + + codec: str # "tensor": ndarray on both sides | "tensor_list": tensor on the wire, list on decode | "json" + dtype: np.dtype | None = None # tensor codecs: pinned on both sides; a mismatch raises instead of converting + strict: bool = False # encode never converts, only validates (R3 replay tensors must arrive as int32) + null: object = None # decoded value for a null-marked field; copied per sample so no instance is shared + + +# tokens/loss_mask/logprobs are re-materialized as Python lists on decode, +# exactly like the legacy JSON path (int64/f64 round-trips are lossless). +SAMPLES_VALUE_SPEC: dict[str, ValueSpec] = { + "tokens": ValueSpec("tensor_list", np.dtype(np.int64), null=[]), + "response": ValueSpec("json"), + "response_length": ValueSpec("json"), + "loss_mask": ValueSpec("tensor_list", np.dtype(np.uint8)), + "rollout_log_probs": ValueSpec("tensor_list", np.dtype(np.float64)), + "rollout_routed_experts": ValueSpec("tensor", np.dtype(np.int32), strict=True), + "rollout_indexer_topk": ValueSpec("tensor", np.dtype(np.int32), strict=True), + "status": ValueSpec("json"), + "weight_versions": ValueSpec("json"), + "prefix_cache_info": ValueSpec("json"), +} + +# The wire allowlist, derived: only table fields cross the samples wire. +COMPUTED_FIELDS = tuple(SAMPLES_VALUE_SPEC) + +assert all( + spec.codec in ("tensor", "tensor_list", "json") for spec in SAMPLES_VALUE_SPEC.values() +), "unknown codec in SAMPLES_VALUE_SPEC" + +_TENSOR_FIELDS = frozenset(field for field, spec in SAMPLES_VALUE_SPEC.items() if spec.codec != "json") + +_SAMPLES_META_KEY = "_samples_meta" +_OPD_STUDENT_TOP_LOGPROBS_KEY = "opd_student_top_logprobs" + + +@dataclasses.dataclass +class SamplesReply: + """Decoded `POST /sessions/{id}/samples` reply.""" + + samples: list[Sample] + session_metadata: dict + empty_reason: str | None + + +def _asarray_wire(field: str, value, dtype: np.dtype) -> np.ndarray: + """Convert to the wire dtype, refusing any conversion that changes values + (e.g. a negative loss-mask entry wrapping into uint8).""" + arr = np.asarray(value) + if arr.dtype == dtype: + return arr + converted = arr.astype(dtype) + if not np.array_equal(converted, arr): + raise ValueError(f"{field} values do not fit wire dtype {dtype}") + return converted + + +def encode_samples_reply(samples: list[Sample], session_metadata: dict, empty_reason: str | None = None) -> bytes: + """Server side: pack assembled samples into one safetensors payload.""" + tensors: dict[str, np.ndarray] = {} + sample_metas = [] + for sample_index, sample in enumerate(samples): + sample_meta: dict = {} + nulls: list[str] = [] + for field, spec in SAMPLES_VALUE_SPEC.items(): + value = getattr(sample, field) + if spec.codec == "json": + if field == "status": + value = value.value + elif field == "prefix_cache_info": + value = value.to_dict() + sample_meta[field] = value + continue + if value is None: + nulls.append(field) + continue + if spec.strict: + arr = np.asarray(value) + if arr.dtype != spec.dtype: + raise ValueError(f"{field} must have dtype {spec.dtype}, got {arr.dtype}") + else: + arr = _asarray_wire(field, value, spec.dtype) + # ascontiguousarray is a correctness requirement: the numpy adapter + # serializes some non-contiguous views without raising, with wrong values. + tensors[f"{field}.{sample_index}"] = np.ascontiguousarray(arr) + sample_meta["nulls"] = nulls + sample_metas.append(sample_meta) + meta = {"samples": sample_metas, "session_metadata": session_metadata, "empty_reason": empty_reason} + # Compact separators: no reason to ship JSON padding on every reply. + meta_bytes = json.dumps(meta, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + tensors[_SAMPLES_META_KEY] = np.frombuffer(meta_bytes, dtype=np.uint8) + return safetensors.numpy.save(tensors) + + +def decode_samples_reply(payload: bytes, input_sample: Sample) -> SamplesReply: + """Driver side: overlay each wire sample's computed fields onto a deepcopy of `input_sample`.""" + tensors = safetensors.numpy.load(payload) # SafetensorError propagates: invalid container + meta_arr = tensors.pop(_SAMPLES_META_KEY) # KeyError propagates: missing meta is malformed + if meta_arr.ndim != 1 or meta_arr.dtype != np.uint8: + raise ValueError( + f"{_SAMPLES_META_KEY} must be a rank-one uint8 tensor, got {meta_arr.dtype} rank {meta_arr.ndim}" + ) + meta = json.loads(meta_arr.tobytes().decode("utf-8")) + if meta["samples"]: + _assert_overlay_template_defaults(input_sample) + samples = [] + for sample_index, sample_meta in enumerate(meta["samples"]): + sample = deepcopy(input_sample) + nulls = set(sample_meta["nulls"]) # KeyError propagates: missing null markers are malformed + if nulls - _TENSOR_FIELDS: + raise ValueError(f"null markers reference non-tensor fields: {sorted(nulls - _TENSOR_FIELDS)}") + for field, spec in SAMPLES_VALUE_SPEC.items(): + if spec.codec == "json": + value = sample_meta[field] + if field == "status": + value = Sample.Status(value) + elif field == "weight_versions": + value = list(value) + elif field == "prefix_cache_info": + value = Sample.PrefixCacheInfo.from_dict(value) + setattr(sample, field, value) + continue + if field in nulls: + setattr(sample, field, copy(spec.null)) + continue + arr = tensors.pop(f"{field}.{sample_index}") # KeyError propagates: a promised tensor must exist + if arr.dtype != spec.dtype: + raise ValueError(f"{field} must have dtype {spec.dtype}, got {arr.dtype}") + setattr(sample, field, arr.tolist() if spec.codec == "tensor_list" else arr) + samples.append(sample) + if tensors: + raise ValueError(f"payload carries unreferenced tensors: {sorted(tensors)}") + return SamplesReply(samples=samples, session_metadata=meta["session_metadata"], empty_reason=meta["empty_reason"]) + + +def _assert_overlay_template_defaults(input_sample: Sample) -> None: + """Overlay equivalence precondition (fail-loud). + + The legacy driver-side pipeline EVOLVED some fields of the input sample in + place (`weight_versions` append, `prefix_cache_info` accumulate, merge sums + `spec_info` across turns, `strip_last_output_tokens` trims + `teacher_log_probs`/`opd_reverse_kl`/`metadata["opd_student_top_logprobs"]`), + while the overlay REPLACES the computed fields and carries the template + verbatim. The two agree exactly when the input sample holds dataclass + defaults on those fields — true for every sample fresh from the data loader + (and `reset_for_retry` restores it on framework retries). + """ + assert input_sample.weight_versions == [], ( + f"input sample must not carry weight_versions (got {input_sample.weight_versions}); " + "the legacy pipeline appended to it, the samples-wire overlay replaces it" + ) + assert ( + input_sample.prefix_cache_info.to_dict() == Sample.PrefixCacheInfo().to_dict() + ), f"input sample must carry a default prefix_cache_info (got {input_sample.prefix_cache_info.to_dict()})" + assert ( + input_sample.spec_info.to_dict() == Sample.SpecInfo().to_dict() + ), f"input sample must carry a default spec_info (got {input_sample.spec_info.to_dict()})" + assert input_sample.teacher_log_probs is None and input_sample.opd_reverse_kl is None, ( + "input sample must not carry teacher_log_probs/opd_reverse_kl; " + "the legacy pipeline trimmed them per turn, the samples-wire overlay carries them verbatim" + ) + assert _OPD_STUDENT_TOP_LOGPROBS_KEY not in (input_sample.metadata or {}), ( + f"input sample metadata must not carry {_OPD_STUDENT_TOP_LOGPROBS_KEY!r}; " + "merge_samples gives it per-token semantics that only hold for per-turn values" + ) diff --git a/miles/rollout/session/samples/merge.py b/miles/rollout/session/samples/merge.py index 5ef9d9fe23..325c63b78e 100644 --- a/miles/rollout/session/samples/merge.py +++ b/miles/rollout/session/samples/merge.py @@ -1,22 +1,12 @@ -# doc-dev: docs/developer/session-server-sample-assembly.md -"""Training-sample assembly: session records -> `Sample` objects. - -Owned by the session package so the assembly can run inside the owning -session worker (records never have to leave the session server); the rollout -driver imports the same functions until the client-side switch lands, so -there is exactly one implementation either way. - -- Depends on `generate_utils.generate_endpoint_utils` for the R3 replay - decoders (accepted utils-level dependency: the decoders have other - consumers on the single-turn `/generate` path and must not fork). -- Order contract: `truncate_samples_by_total_tokens` runs BEFORE - `merge_samples` — truncation is a turn-level budget decision (which turns - survive; the overflowing turn is cut at a turn boundary, later turns are - dropped) and the turn structure only exists pre-merge. +"""Training-sample assembly: session records -> per-turn `Sample`s, truncated at turn boundaries. + +Owned by the session package so the assembly runs on the owning instance (records never have to leave the session server). The wire codec for the assembled reply lives in `codec`. + +- Depends on `generate_utils.generate_endpoint_utils` for the R3 replay decoders (accepted utils-level dependency: the decoders have other consumers on the single-turn `/generate` path and must not fork). +- Order contract: `truncate_samples_by_total_tokens` runs BEFORE `merge_samples` — truncation is a turn-level budget decision (which turns survive; the overflowing turn is cut at a turn boundary, later turns are dropped) and the turn structure only exists pre-merge. """ from argparse import Namespace -from copy import deepcopy from miles.rollout.generate_utils.generate_endpoint_utils import ( get_indexer_topk_from_response, @@ -29,7 +19,6 @@ def compute_samples_from_openai_records( args: Namespace, - input_sample: Sample, records: list[SessionRecord], tokenizer, accumulated_token_ids: list[int] | None = None, @@ -85,7 +74,7 @@ def compute_samples_from_openai_records( # Step 4: advance cursor past matched output to the next turn cursor += matched - sample = _compute_sample_from_openai_record(args, input_sample, record, tokenizer, trim_count) + sample = _compute_sample_from_openai_record(args, record, tokenizer, trim_count) attach_lifecycle_metadata(sample, record, records[i - 1] if i else None, turn=i + 1) if is_last and args.save_debug_trajectory_data is not None: sample.metadata["messages"] = record.request["messages"] + [record.response["choices"][0]["message"]] @@ -102,7 +91,7 @@ def compute_samples_from_openai_records( def _compute_sample_from_openai_record( - args: Namespace, input_sample: Sample, record: SessionRecord, tokenizer, trim_count: int = 0 + args: Namespace, record: SessionRecord, tokenizer, trim_count: int = 0 ) -> Sample: choice = record.response["choices"][0] @@ -113,7 +102,7 @@ def _compute_sample_from_openai_record( output_token_ids = [item[1] for item in choice["meta_info"]["output_token_logprobs"]] output_log_probs = [item[0] for item in choice["meta_info"]["output_token_logprobs"]] - sample = deepcopy(input_sample) + sample = Sample() sample.tokens = prompt_token_ids + output_token_ids sample.rollout_log_probs = output_log_probs sample.response = tokenizer.decode(output_token_ids) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index f3ad6b813f..e14e0ec452 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -4,6 +4,7 @@ ``SessionCore``. All session/TITO logic lives in ``core``. """ +import json import logging from fastapi import Request @@ -71,6 +72,21 @@ async def chat_completions(request: Request, session_id: str): body=body, ) + @app.post("/sessions/{session_id}/samples") + async def collect_samples(request: Request, session_id: str): + # Must stay registered BEFORE the catch-all session_proxy below: Starlette + # matches in registration order, and the catch-all would otherwise swallow + # this path and forward it to the inference backend. + # Request params are parsed here, OUTSIDE core.collect_samples's 422 lane: + # a malformed body is a protocol violation (500), not an assembly failure. + body = await request.body() + params = json.loads(body) if body else {} + return await core.collect_samples( + session_id, + multi_samples=bool(params.get("multi_samples", False)), + max_seq_len=params.get("max_seq_len"), + ) + @app.api_route("/sessions/{session_id}/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) async def session_proxy(request: Request, session_id: str, path: str): body = await request.body() diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index 9b27fdb613..6839808201 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -224,6 +224,28 @@ async def _post(client, url, payload, max_retries=60, action="post", headers=Non return output +async def post_bytes_no_retry(url: str, payload: dict, *, timeout: float) -> bytes: + """Single POST over the shared client: no retries, raw-bytes reply. + + For endpoints where a retry cannot help (the session samples endpoint: a + 5xx means the owning worker died and its state died with it, a 422 is a + deterministic assembly failure) and where the reply is a binary payload + that `post()`'s json()/text decoding would mangle. A non-2xx raises with + the response body text; `timeout` bounds the whole call via wait_for (the + shared client itself has timeout=None, and httpx timeouts are per-phase, + not total). + """ + assert _http_client is not None, "init_http_client() must run before post_bytes_no_retry()" + + async def _do() -> bytes: + response = await _http_client.post(url, json=payload) + if not (200 <= response.status_code < 300): + raise RuntimeError(f"POST {url} failed with {response.status_code}: {response.text}") + return response.content + + return await asyncio.wait_for(_do(), timeout=timeout) + + def init_http_client(args): """Initialize HTTP client and optionally enable distributed POST via Ray.""" global _http_client, _client_concurrency, _distributed_post_enabled diff --git a/requirements.txt b/requirements.txt index e696aa7d8b..37991feecc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -20,6 +20,7 @@ pyyaml qwen_vl_utils # for VLM ray[default] ring_flash_attn; platform_system == "Linux" +safetensors>=0.8.0 # samples-reply wire codec; malformed-payload exception contract validated on 0.8.0 sglang-router>=0.2.3 tensorboard torchft-nightly==2026.4.3; platform_system == "Linux" and platform_machine == "x86_64" diff --git a/tests/fast/fixtures/generation_fixtures.py b/tests/fast/fixtures/generation_fixtures.py index 95838f5c74..353dd7db00 100644 --- a/tests/fast/fixtures/generation_fixtures.py +++ b/tests/fast/fixtures/generation_fixtures.py @@ -157,6 +157,8 @@ def make_args( generate_execute_tool_function_path: str = "miles.utils.test_utils.mock_tools.execute_tool_call", rollout_max_context_len: int | None = None, chat_template_path: str | None = None, + num_layers: int | None = None, + moe_router_topk: int | None = None, ) -> Namespace: argv = [ "pytest", @@ -212,6 +214,14 @@ def make_args( with patch("sys.argv", argv): args = parse_args() + # R3 decode shape overrides — not CLI flags (derived from the model config + # in production). Applied here, before with_session_server copies args into + # the worker namespace, because sample assembly runs inside the worker. + if num_layers is not None: + args.num_layers = num_layers + if moe_router_topk is not None: + args.moe_router_topk = moe_router_topk + init_http_client(args) return args @@ -240,6 +250,12 @@ def with_session_server( tito_model=args.tito_model, tito_allowed_append_roles=args.tito_allowed_append_roles, use_rollout_routing_replay=args.use_rollout_routing_replay, + # Sample assembly runs inside the server, so the R3 decode shape args + # must reach the server namespace (set them via args_kwargs BEFORE the + # server starts; assigning to the driver args afterwards has no effect). + num_layers=getattr(args, "num_layers", None), + moe_router_topk=getattr(args, "moe_router_topk", None), + save_debug_trajectory_data=getattr(args, "save_debug_trajectory_data", None), session_server_instance_id=instance_id, ) session_server = SessionServer(server_args, backend_url=backend_url) diff --git a/tests/fast/fixtures/rollout_fixtures.py b/tests/fast/fixtures/rollout_fixtures.py index 6df6de97a3..6ef8d028d9 100644 --- a/tests/fast/fixtures/rollout_fixtures.py +++ b/tests/fast/fixtures/rollout_fixtures.py @@ -110,6 +110,7 @@ def _with_session_server(args: Namespace, backend_url: str) -> Iterator[UvicornT tito_model=getattr(args, "tito_model", "default"), tito_allowed_append_roles=getattr(args, "tito_allowed_append_roles", ["tool"]), use_rollout_routing_replay=getattr(args, "use_rollout_routing_replay", False), + save_debug_trajectory_data=getattr(args, "save_debug_trajectory_data", None), ) session_server = SessionServer(session_args, backend_url=backend_url) port = find_available_port(31000) diff --git a/tests/fast/rollout/generate_hub/test_multi_turn.py b/tests/fast/rollout/generate_hub/test_multi_turn.py index cef0f611ec..96c476bcc8 100644 --- a/tests/fast/rollout/generate_hub/test_multi_turn.py +++ b/tests/fast/rollout/generate_hub/test_multi_turn.py @@ -578,6 +578,10 @@ class TestRoutedExpertsMultiTurn: { "args_kwargs": { "use_rollout_routing_replay": True, + # Must be in args BEFORE the session server starts: the R3 + # decode now runs inside the worker during sample assembly. + "num_layers": 2, + "moe_router_topk": 4, } } ], @@ -585,9 +589,7 @@ class TestRoutedExpertsMultiTurn: ) def test_two_turns_routed_experts(self, variant, generation_env): S = TwoTurnStub - num_layers, moe_router_topk = 2, 4 - generation_env.args.num_layers = num_layers - generation_env.args.moe_router_topk = moe_router_topk + num_layers, moe_router_topk = generation_env.args.num_layers, generation_env.args.moe_router_topk if is_agentic_variant(variant): tito = get_tito_tokenizer( TOKENIZER, diff --git a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py index e39ebc6743..1d3aa1b3f2 100644 --- a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py +++ b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py @@ -1,14 +1,24 @@ """Tests for OpenAIEndpointTracer (session-server client side). The sample-assembly and TITO multi-turn merge tests live in -tests/fast/rollout/session/test_samples.py, next to the functions. +tests/fast/rollout/session/test_samples.py (assembly) and +test_samples_codec.py (wire codec), next to the functions. +The collect_samples tests here lock the client's HTTP behavior deltas vs the +old collect_records path: single POST with no retries, non-2xx raises with the +body text, timeout raises (instead of silently ABORTing), and the session +DELETE is attempted on every path. """ +import asyncio from types import SimpleNamespace import pytest +import miles.utils.http_utils as http_utils from miles.rollout.generate_utils.openai_endpoint_utils import OpenAIEndpointTracer +from miles.rollout.session.samples.codec import encode_samples_reply +from miles.utils.http_utils import post_bytes_no_retry +from miles.utils.types import Sample @pytest.mark.asyncio @@ -53,7 +63,7 @@ async def fake_post(url: str, payload: dict, action: str = "post"): @pytest.mark.asyncio async def test_create_distributes_sessions_across_port_range(monkeypatch): """With a multi-port range, sessions land on more than one instance, and every - request of a session (create, chat, GET, DELETE) hits the port chosen + request of a session (create, samples POST, DELETE) hits the port chosen at create time — the URL is the router.""" calls: list[tuple[str, str]] = [] @@ -61,9 +71,14 @@ async def fake_post(url: str, payload: dict, action: str = "post"): calls.append((action, url)) if action == "post" and url.endswith("/sessions"): return {"session_id": f"session-{len(calls)}"} - return {"session_id": url.rsplit("/", 1)[1], "records": [], "metadata": {}} + return {} + + async def fake_post_bytes(url, payload, *, timeout): + calls.append(("post_bytes", url)) + return encode_samples_reply([], {}, "no_records") monkeypatch.setattr("miles.rollout.generate_utils.openai_endpoint_utils.post", fake_post) + monkeypatch.setattr("miles.rollout.generate_utils.openai_endpoint_utils.post_bytes_no_retry", fake_post_bytes) ports = [12345, 12346, 12347, 12348] args = SimpleNamespace(session_server_ip="127.0.0.1", session_server_ports=ports) @@ -76,14 +91,145 @@ async def fake_post(url: str, payload: dict, action: str = "post"): assert port in ports chosen_ports.add(port) - await tracer.collect_records() + await tracer.collect_samples(Sample(), multi_samples=False, max_seq_len=None) prefix = f"http://127.0.0.1:{port}" assert [url for _, url in calls] == [ f"{prefix}/sessions", - tracer.base_url, + f"{tracer.base_url}/samples", tracer.base_url, ] assert tracer.base_url.startswith(f"{prefix}/sessions/") # 32 uniform picks over 4 ports miss a given port with p = (3/4)^32 ≈ 1e-4. assert len(chosen_ports) > 1 + + +# ── collect_samples client behavior ── + + +def _tracer() -> OpenAIEndpointTracer: + return OpenAIEndpointTracer(router_url="http://127.0.0.1:12345", session_id="sid-1") + + +def _computed_reply_payload() -> bytes: + sample = Sample() + sample.tokens = [1, 2, 10] + sample.response = "r" + sample.response_length = 1 + sample.loss_mask = [1] + sample.rollout_log_probs = [-0.5] + sample.status = Sample.Status.COMPLETED + return encode_samples_reply([sample], {"max_trim_tokens": 1}, None) + + +class _CollectCalls: + """Patches the two HTTP primitives collect_samples uses and records order.""" + + def __init__(self, monkeypatch, *, post_outcome, delete_outcome=None): + self.calls: list[str] = [] + + async def fake_post_bytes(url, payload, *, timeout): + self.calls.append(f"POST {url}") + assert payload == {"multi_samples": False, "max_seq_len": 7} + if isinstance(post_outcome, Exception): + raise post_outcome + return post_outcome + + async def fake_post(url, payload, action="post"): + assert action == "delete" + self.calls.append(f"DELETE {url}") + if isinstance(delete_outcome, Exception): + raise delete_outcome + return {} + + monkeypatch.setattr("miles.rollout.generate_utils.openai_endpoint_utils.post_bytes_no_retry", fake_post_bytes) + monkeypatch.setattr("miles.rollout.generate_utils.openai_endpoint_utils.post", fake_post) + + +@pytest.mark.asyncio +async def test_collect_samples_single_post_then_delete(monkeypatch): + calls = _CollectCalls(monkeypatch, post_outcome=_computed_reply_payload()) + result = await _tracer().collect_samples(Sample(), multi_samples=False, max_seq_len=7) + + assert calls.calls == [ + "POST http://127.0.0.1:12345/sessions/sid-1/samples", + "DELETE http://127.0.0.1:12345/sessions/sid-1", + ] + (sample,) = result.samples + assert sample.tokens == [1, 2, 10] and sample.status == Sample.Status.COMPLETED + assert result.session_metadata == {"max_trim_tokens": 1} + + +@pytest.mark.asyncio +async def test_collect_samples_non_2xx_raises_with_body_and_still_deletes(monkeypatch): + calls = _CollectCalls(monkeypatch, post_outcome=RuntimeError("422: trim_count 2 exceeds allowed=1")) + with pytest.raises(RuntimeError, match="trim_count 2 exceeds allowed=1"): + await _tracer().collect_samples(Sample(), multi_samples=False, max_seq_len=7) + assert calls.calls[-1] == "DELETE http://127.0.0.1:12345/sessions/sid-1" + + +@pytest.mark.asyncio +async def test_collect_samples_timeout_raises_and_still_deletes(monkeypatch): + # The old collect_records swallowed the timeout and returned empty records + # (silently ABORTing the sample); the samples path must raise it. + calls = _CollectCalls(monkeypatch, post_outcome=asyncio.TimeoutError()) + with pytest.raises(asyncio.TimeoutError): + await _tracer().collect_samples(Sample(), multi_samples=False, max_seq_len=7) + assert calls.calls[-1] == "DELETE http://127.0.0.1:12345/sessions/sid-1" + + +@pytest.mark.asyncio +async def test_collect_samples_delete_failure_is_tolerated(monkeypatch): + _CollectCalls(monkeypatch, post_outcome=_computed_reply_payload(), delete_outcome=RuntimeError("delete boom")) + result = await _tracer().collect_samples(Sample(), multi_samples=False, max_seq_len=7) + assert len(result.samples) == 1 + + +# ── post_bytes_no_retry primitive ── + + +class _FakeResponse: + def __init__(self, status_code: int, content: bytes = b"", text: str = ""): + self.status_code = status_code + self.content = content + self.text = text + + +class _FakeClient: + def __init__(self, responses): + self.responses = list(responses) + self.post_count = 0 + + async def post(self, url, json=None): + self.post_count += 1 + outcome = self.responses.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + +@pytest.mark.asyncio +async def test_post_bytes_no_retry_returns_raw_bytes(monkeypatch): + client = _FakeClient([_FakeResponse(200, content=b"\x00\x01binary")]) + monkeypatch.setattr(http_utils, "_http_client", client) + assert await post_bytes_no_retry("http://x/samples", {}, timeout=5) == b"\x00\x01binary" + assert client.post_count == 1 + + +@pytest.mark.asyncio +async def test_post_bytes_no_retry_does_not_retry_and_carries_body(monkeypatch): + # Two queued outcomes; a retrying client would consume both. It must not. + client = _FakeClient([_FakeResponse(422, text="cursor 3 != len(accumulated_token_ids) 4"), RuntimeError("late")]) + monkeypatch.setattr(http_utils, "_http_client", client) + with pytest.raises(RuntimeError, match="422.*cursor 3"): + await post_bytes_no_retry("http://x/samples", {}, timeout=5) + assert client.post_count == 1 + + +@pytest.mark.asyncio +async def test_post_bytes_no_retry_transport_error_propagates_once(monkeypatch): + client = _FakeClient([ConnectionError("boom"), RuntimeError("late")]) + monkeypatch.setattr(http_utils, "_http_client", client) + with pytest.raises(ConnectionError, match="boom"): + await post_bytes_no_retry("http://x/samples", {}, timeout=5) + assert client.post_count == 1 diff --git a/tests/fast/rollout/session/test_samples.py b/tests/fast/rollout/session/test_samples.py index 72d53d6235..58a6a2a0c4 100644 --- a/tests/fast/rollout/session/test_samples.py +++ b/tests/fast/rollout/session/test_samples.py @@ -26,22 +26,6 @@ def _mock_tokenizer(): return tok -def _make_input_sample(**overrides): - defaults = dict( - group_index=0, - index=0, - prompt="test prompt", - tokens=[], - response="", - response_length=0, - status=Sample.Status.PENDING, - label="test", - reward=1.0, - ) - defaults.update(overrides) - return Sample(**defaults) - - def _make_record( prompt_token_ids: list[int], output_token_ids: list[int], @@ -49,11 +33,14 @@ def _make_record( finish_reason: str = "stop", cached_tokens: int | None = None, prompt_tokens: int | None = None, + weight_version: str | None = None, + routed_experts: str | None = None, ) -> SessionRecord: """Build a minimal session record mimicking SGLang's response format. Token IDs and logprobs are stored in meta_info.output_token_logprobs as (logprob, token_id) tuples, matching the real SGLang response. + `routed_experts` is the base64 int32 buffer exactly as SGLang returns it. """ if output_log_probs is None: output_log_probs = [-0.1 * (i + 1) for i in range(len(output_token_ids))] @@ -70,6 +57,10 @@ def _make_record( meta_info["cached_tokens"] = cached_tokens if prompt_tokens is not None: meta_info["prompt_tokens"] = prompt_tokens + if weight_version is not None: + meta_info["weight_version"] = weight_version + if routed_experts is not None: + meta_info["routed_experts"] = routed_experts return SessionRecord( timestamp=0.0, method="POST", @@ -100,9 +91,8 @@ def test_single_record_builds_correct_sample(self): output_token_ids=[10, 11], output_log_probs=[-0.5, -0.6], ) - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, [record], tok) + samples = compute_samples_from_openai_records(_ARGS, [record], tok) assert len(samples) == 1 s = samples[0] @@ -118,9 +108,8 @@ def test_multiple_records_produce_multiple_samples(self): _make_record(prompt_token_ids=[1, 2], output_token_ids=[10]), _make_record(prompt_token_ids=[1, 2, 10, 20], output_token_ids=[30]), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) assert len(samples) == 2 assert samples[0].tokens == [1, 2, 10] @@ -133,7 +122,7 @@ def test_last_sample_carries_raw_conversation(self): _make_record(prompt_token_ids=[1, 2, 10, 20], output_token_ids=[30]), ] - samples = compute_samples_from_openai_records(_ARGS_RECORDING, _make_input_sample(), records, tok) + samples = compute_samples_from_openai_records(_ARGS_RECORDING, records, tok) assert "messages" not in (samples[0].metadata or {}) assert samples[1].metadata["messages"] == records[1].request["messages"] + [ @@ -147,7 +136,7 @@ def test_merge_keeps_last_conversation_snapshot(self): _make_record(prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], output_token_ids=[30, 31]), ] - samples = compute_samples_from_openai_records(_ARGS_RECORDING, _make_input_sample(), records, tok) + samples = compute_samples_from_openai_records(_ARGS_RECORDING, records, tok) merged = merge_samples(samples, tok) assert merged.metadata["messages"] == samples[-1].metadata["messages"] @@ -159,9 +148,8 @@ def test_finish_reason_length_gives_truncated(self): output_token_ids=[10], finish_reason="length", ) - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, [record], tok) + samples = compute_samples_from_openai_records(_ARGS, [record], tok) assert samples[0].status == Sample.Status.TRUNCATED @@ -197,9 +185,8 @@ def test_two_turn_merge_succeeds(self): output_log_probs=[-0.3, -0.4], ), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) merged = merge_samples(samples, tok) assert merged.tokens == [1, 2, 3, 10, 11, 20, 21, 30, 31] @@ -228,9 +215,8 @@ def test_three_turn_merge_succeeds(self): output_log_probs=[-0.3], ), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) merged = merge_samples(samples, tok) assert merged.tokens == [1, 2, 10, 20, 30, 40, 50] @@ -251,9 +237,8 @@ def test_prefix_mismatch_raises(self): output_token_ids=[30, 31], ), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) with pytest.raises(AssertionError, match="b.tokens must start with a.tokens"): merge_samples(samples, tok) @@ -271,8 +256,7 @@ def test_two_turn_merge_propagates_teacher_log_probs(self): output_log_probs=[-0.3, -0.4], ), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) # OPD attaches per-response-token teacher log-probs to each turn's sample. samples[0].teacher_log_probs = [-1.0, -1.1] @@ -297,8 +281,7 @@ def test_two_turn_merge_propagates_opd_student_top_logprobs_metadata(self): output_log_probs=[-0.3, -0.4], ), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) turn_0_top_logprobs = [[[-0.1, 101]], [[-0.2, 102]]] turn_1_top_logprobs = [[[-0.3, 103]], [[-0.4, 104]]] @@ -330,8 +313,7 @@ def test_two_turn_merge_teacher_log_probs_none_stays_none(self): _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11]), _make_record(prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], output_token_ids=[30, 31]), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) merged = merge_samples(samples, tok) @@ -345,8 +327,7 @@ def test_merge_raises_on_teacher_log_probs_length_mismatch(self): _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11]), _make_record(prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], output_token_ids=[30, 31]), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) samples[0].teacher_log_probs = [-1.0] # length 1 != response_length 2 @@ -427,11 +408,9 @@ def test_three_turn_trim_trailing_stop_tokens(self): _make_record(prompt_token_ids=[1, 2, 3, 10, 4, 5, 6, 20, 7, 8, 9], output_token_ids=[30, STOP]), ] accumulated = [1, 2, 3, 10, 4, 5, 6, 20, 7, 8, 9, 30, STOP] - input_sample = _make_input_sample() samples = compute_samples_from_openai_records( _ARGS, - input_sample, records, tok, accumulated_token_ids=accumulated, @@ -462,11 +441,9 @@ def test_no_trim_when_no_trailing_stop(self): _make_record(prompt_token_ids=[1, 2, 10, 11, 3, 4], output_token_ids=[20, 21]), ] accumulated = [1, 2, 10, 11, 3, 4, 20, 21] - input_sample = _make_input_sample() samples = compute_samples_from_openai_records( _ARGS, - input_sample, records, tok, accumulated_token_ids=accumulated, @@ -487,11 +464,9 @@ def test_single_turn_no_trim(self): _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11, STOP]), ] accumulated = [1, 2, 3, 10, 11, STOP] - input_sample = _make_input_sample() samples = compute_samples_from_openai_records( _ARGS, - input_sample, records, tok, accumulated_token_ids=accumulated, @@ -510,11 +485,9 @@ def test_no_accumulated_skips_trimming(self): _make_record(prompt_token_ids=[1, 2], output_token_ids=[10, STOP]), _make_record(prompt_token_ids=[1, 2, 10, STOP, 3, 4], output_token_ids=[20, STOP]), ] - input_sample = _make_input_sample() samples = compute_samples_from_openai_records( _ARGS, - input_sample, records, tok, accumulated_token_ids=None, @@ -537,12 +510,10 @@ def test_trim_exceeding_max_raises(self): _make_record(prompt_token_ids=[1, 2, 10, 3, 4], output_token_ids=[20]), ] accumulated = [1, 2, 10, 3, 4, 20] - input_sample = _make_input_sample() with pytest.raises(AssertionError, match="trim_count 2 exceeds allowed=1"): compute_samples_from_openai_records( _ARGS, - input_sample, records, tok, accumulated_token_ids=accumulated, @@ -560,12 +531,10 @@ def test_cursor_covers_entire_accumulated(self): ] # Missing the last token — accumulated should be [1,2,10,3,20] but we give [1,2,10,3,20,99] accumulated = [1, 2, 10, 3, 20, 99] - input_sample = _make_input_sample() with pytest.raises(AssertionError, match="cursor .* != len\\(accumulated_token_ids\\)"): compute_samples_from_openai_records( _ARGS, - input_sample, records, tok, accumulated_token_ids=accumulated, @@ -623,9 +592,8 @@ def test_thinking_tokens_break_prefix_chain(self): output_token_ids=[30, 31], ), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) # sample[0].tokens = [1,2,3] + thinking + [10,11] = [1,2,3, ,\n,42,43,\n,,\n, 10,11] # sample[1].tokens = [1,2,3, 10,11, 20,21, 30,31] @@ -648,9 +616,8 @@ def test_no_thinking_tokens_prefix_chain_holds(self): output_token_ids=[30, 31], ), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) merged = merge_samples(samples, tok) assert merged.tokens == [1, 2, 3, 10, 11, 20, 21, 30, 31] @@ -671,8 +638,7 @@ def test_single_record_with_cache_stats(self): cached_tokens=2, prompt_tokens=3, ) - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, [record], tok) + samples = compute_samples_from_openai_records(_ARGS, [record], tok) assert samples[0].prefix_cache_info.cached_tokens == 2 assert samples[0].prefix_cache_info.total_prompt_tokens == 3 @@ -696,8 +662,7 @@ def test_multi_turn_cache_stats_accumulate_after_merge(self): prompt_tokens=7, ), ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) + samples = compute_samples_from_openai_records(_ARGS, records, tok) merged = merge_samples(samples, tok) assert merged.prefix_cache_info.cached_tokens == 0 + 5 @@ -711,8 +676,7 @@ def test_missing_cache_fields_default_to_zero(self): prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11], ) - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, [record], tok) + samples = compute_samples_from_openai_records(_ARGS, [record], tok) assert samples[0].prefix_cache_info.cached_tokens == 0 assert samples[0].prefix_cache_info.total_prompt_tokens == 0 diff --git a/tests/fast/rollout/session/test_samples_codec.py b/tests/fast/rollout/session/test_samples_codec.py new file mode 100644 index 0000000000..f8fd49b920 --- /dev/null +++ b/tests/fast/rollout/session/test_samples_codec.py @@ -0,0 +1,230 @@ +"""Tests for the samples wire codec: encode on the worker, overlay on the driver. + +Covers safetensors-tensor round-trips, malformed-payload rejection, and the +overlay defaults guard (`_assert_overlay_template_defaults`). +""" + +import json + +import numpy as np +import pytest +import safetensors.numpy +from safetensors import SafetensorError + +from miles.rollout.session.samples.codec import decode_samples_reply, encode_samples_reply +from miles.utils.types import Sample + + +def _computed_sample(**overrides) -> Sample: + """A blank-template sample carrying every computed field, as the worker produces.""" + s = Sample() + s.tokens = [1, 2, 3, 10, 11] + s.response = "[10][11]" + s.response_length = 2 + s.loss_mask = [1, 1] + s.rollout_log_probs = [-0.5, -0.1234567891234567] + s.rollout_routed_experts = np.arange(24, dtype=np.int32).reshape(4, 3, 2) + s.rollout_indexer_topk = None + s.status = Sample.Status.COMPLETED + s.weight_versions = ["w1", "w2"] + s.prefix_cache_info = Sample.PrefixCacheInfo.from_dict({"cached_tokens": 2, "total_prompt_tokens": 3}) + for name, value in overrides.items(): + setattr(s, name, value) + return s + + +def _mutated_payload(payload: bytes, mutate) -> bytes: + """Re-pack a valid payload after `mutate(meta, tensors)` edits, for malformed-wire cases.""" + tensors = safetensors.numpy.load(payload) + meta = json.loads(tensors.pop("_samples_meta").tobytes().decode("utf-8")) + mutate(meta, tensors) + meta_bytes = json.dumps(meta, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + tensors["_samples_meta"] = np.frombuffer(meta_bytes, dtype=np.uint8) + return safetensors.numpy.save(tensors) + + +class TestSamplesWireCodec: + def test_round_trip_overlays_computed_and_keeps_template(self): + template = Sample( + group_index=7, + index=3, + prompt=[{"role": "user", "content": "hi"}], + label="lbl", + reward=1.5, + metadata={"task": "t"}, + routing_key="rk", + train_metadata={"loss": "ppo"}, + ) + payload = encode_samples_reply([_computed_sample()], {"max_trim_tokens": 1}, None) + reply = decode_samples_reply(payload, template) + + assert reply.empty_reason is None + assert reply.session_metadata == {"max_trim_tokens": 1} + (out,) = reply.samples + # computed fields overlaid, with exact types/values + assert out.tokens == [1, 2, 3, 10, 11] and type(out.tokens) is list + assert out.rollout_log_probs == [-0.5, -0.1234567891234567] + assert out.loss_mask == [1, 1] + assert out.response == "[10][11]" and out.response_length == 2 + assert out.status == Sample.Status.COMPLETED + assert out.weight_versions == ["w1", "w2"] + assert out.prefix_cache_info.to_dict() == {"cached_tokens": 2, "total_prompt_tokens": 3} + assert out.rollout_routed_experts.dtype == np.int32 + assert np.array_equal(out.rollout_routed_experts, np.arange(24, dtype=np.int32).reshape(4, 3, 2)) + assert out.rollout_indexer_topk is None + # template fields carried from the input sample, untouched + assert out.group_index == 7 and out.index == 3 + assert out.prompt == [{"role": "user", "content": "hi"}] + assert out.label == "lbl" and out.reward == 1.5 + assert out.metadata == {"task": "t"} and out.routing_key == "rk" + assert out.train_metadata == {"loss": "ppo"} + # the input template itself is never mutated + assert template.tokens == [] and template.metadata == {"task": "t"} + + def test_multi_sample_reply_keeps_per_sample_tensors(self): + a = _computed_sample() + b = _computed_sample( + tokens=[1, 2, 3, 10, 11, 20, 21, 30], + rollout_routed_experts=np.arange(100, 100 + 42, dtype=np.int32).reshape(7, 3, 2), + rollout_log_probs=[-1.0, -2.0], + ) + reply = decode_samples_reply(encode_samples_reply([a, b], {}, None), Sample()) + out_a, out_b = reply.samples + assert out_a.tokens == a.tokens and out_b.tokens == b.tokens + assert np.array_equal(out_a.rollout_routed_experts, a.rollout_routed_experts) + assert np.array_equal(out_b.rollout_routed_experts, b.rollout_routed_experts) + assert out_b.rollout_log_probs == [-1.0, -2.0] + + def test_empty_reply_round_trips_reason_and_skips_defaults_guard(self): + # The empty reply is decoded before the driver takes its ABORTED path, so + # the overlay-defaults guard must not fire on it — even for an input + # sample that would violate the guard. + evolved = Sample(weight_versions=["stale"]) + reply = decode_samples_reply(encode_samples_reply([], {"max_trim_tokens": 0}, "no_records"), evolved) + assert reply.samples == [] and reply.empty_reason == "no_records" + + def test_defaults_guard_rejects_evolved_template(self): + payload = encode_samples_reply([_computed_sample()], {}, None) + with pytest.raises(AssertionError, match="weight_versions"): + decode_samples_reply(payload, Sample(weight_versions=["stale"])) + with pytest.raises(AssertionError, match="teacher_log_probs"): + decode_samples_reply(payload, Sample(teacher_log_probs=[-1.0])) + with pytest.raises(AssertionError, match="opd_student_top_logprobs"): + decode_samples_reply(payload, Sample(metadata={"opd_student_top_logprobs": [[[-0.1, 1]]]})) + + def test_safetensors_container_round_trips_non_contiguous_replay_tensors(self): + routed = np.arange(24, dtype=np.int32).reshape(3, 4, 2).transpose(1, 0, 2) + indexer = np.arange(48, dtype=np.int32).reshape(8, 3, 2)[::2] + assert not routed.flags["C_CONTIGUOUS"] and not indexer.flags["C_CONTIGUOUS"] + sample = _computed_sample(rollout_routed_experts=routed, rollout_indexer_topk=indexer) + + payload = encode_samples_reply([sample], {}, None) + # the reply is a plain safetensors buffer: no Miles framing needed to open it + tensors = safetensors.numpy.load(payload) + assert set(tensors) == { + "_samples_meta", + "tokens.0", + "loss_mask.0", + "rollout_log_probs.0", + "rollout_routed_experts.0", + "rollout_indexer_topk.0", + } + assert tensors["_samples_meta"].dtype == np.uint8 and tensors["_samples_meta"].ndim == 1 + assert tensors["loss_mask.0"].dtype == np.uint8 + + (out,) = decode_samples_reply(payload, Sample()).samples + assert out.tokens == sample.tokens and type(out.tokens) is list + assert out.rollout_log_probs == sample.rollout_log_probs + assert out.rollout_routed_experts.dtype == np.int32 and out.rollout_routed_experts.shape == (4, 3, 2) + assert np.array_equal(out.rollout_routed_experts, routed) + assert out.rollout_indexer_topk.dtype == np.int32 and np.array_equal(out.rollout_indexer_topk, indexer) + + def test_zero_size_tensor_is_distinct_from_none(self): + sample = _computed_sample( + rollout_routed_experts=np.empty((0, 3, 2), dtype=np.int32), rollout_indexer_topk=None + ) + (out,) = decode_samples_reply(encode_samples_reply([sample], {}, None), Sample()).samples + assert isinstance(out.rollout_routed_experts, np.ndarray) and out.rollout_routed_experts.shape == (0, 3, 2) + assert out.rollout_indexer_topk is None + + def test_null_tokens_restore_fresh_empty_lists(self): + # A null marker restores per-sample fresh instances, never one shared list. + def null_out_tokens(meta, tensors): + for index, sample_meta in enumerate(meta["samples"]): + sample_meta["nulls"].append("tokens") + del tensors[f"tokens.{index}"] + + payload = _mutated_payload( + encode_samples_reply([_computed_sample(), _computed_sample()], {}, None), null_out_tokens + ) + out_a, out_b = decode_samples_reply(payload, Sample()).samples + assert out_a.tokens == [] and out_b.tokens == [] + assert out_a.tokens is not out_b.tokens + + def test_encode_rejects_non_int32_replay_dtype(self): + sample = _computed_sample(rollout_routed_experts=np.arange(24, dtype=np.int64).reshape(4, 3, 2)) + with pytest.raises(ValueError, match="rollout_routed_experts must have dtype int32"): + encode_samples_reply([sample], {}, None) + + def test_encode_rejects_loss_mask_values_that_wrap_in_uint8(self): + with pytest.raises(ValueError, match="loss_mask values do not fit wire dtype uint8"): + encode_samples_reply([_computed_sample(loss_mask=[1, -1])], {}, None) + + @pytest.mark.parametrize( + ("build_payload", "expected_error", "match"), + [ + pytest.param(lambda p: p[: len(p) - 100], SafetensorError, None, id="truncated-container"), + pytest.param(lambda p: b"", SafetensorError, None, id="empty-container"), + pytest.param( + lambda p: safetensors.numpy.save({"tokens.0": np.arange(3, dtype=np.int64)}), + KeyError, + "_samples_meta", + id="missing-samples-meta", + ), + pytest.param( + lambda p: safetensors.numpy.save({"_samples_meta": np.zeros(4, dtype=np.int32)}), + ValueError, + "rank-one uint8", + id="meta-wrong-dtype", + ), + pytest.param( + lambda p: safetensors.numpy.save({"_samples_meta": np.zeros((2, 2), dtype=np.uint8)}), + ValueError, + "rank-one uint8", + id="meta-wrong-rank", + ), + pytest.param( + lambda p: _mutated_payload(p, lambda meta, tensors: tensors.pop("tokens.0")), + KeyError, + "tokens.0", + id="promised-tensor-missing", + ), + pytest.param( + lambda p: _mutated_payload( + p, + lambda meta, tensors: tensors.update({"tokens.0": tensors["tokens.0"].astype(np.int32)}), + ), + ValueError, + "tokens must have dtype int64", + id="wire-dtype-contract-violated", + ), + pytest.param( + lambda p: _mutated_payload(p, lambda meta, tensors: meta["samples"][0]["nulls"].append("response")), + ValueError, + "non-tensor fields", + id="null-marker-for-non-tensor-field", + ), + pytest.param( + lambda p: _mutated_payload( + p, lambda meta, tensors: tensors.update(orphan=np.zeros(1, dtype=np.uint8)) + ), + ValueError, + "unreferenced tensors", + id="unreferenced-leftover-tensor", + ), + ], + ) + def test_malformed_safetensors_reply_fails_loudly(self, build_payload, expected_error, match): + valid = encode_samples_reply([_computed_sample()], {}, None) + with pytest.raises(expected_error, match=match): + decode_samples_reply(build_payload(valid), Sample()) diff --git a/tests/fast/router/test_session_samples_op.py b/tests/fast/router/test_session_samples_op.py new file mode 100644 index 0000000000..90ef04dbb4 --- /dev/null +++ b/tests/fast/router/test_session_samples_op.py @@ -0,0 +1,348 @@ +"""Samples-op tests: golden assembled Samples plus the op-level error/route contracts. + +Drives `SessionCore.collect_samples` in-process against a real tokenizer (the +`test_sessions.py` precedent), with records injected via the registry — the +broken-chain and R3 fixtures cannot be produced through the chat path. The +HTTP surface (route registration order, 404 mapping) is exercised through the +real `setup_session_routes` app with a `TestClient`. + +The golden tests assert the exact `Sample` field values derivable from the +two-turn records fixture — through `collect_samples` → `decode_samples_reply` +overlay → the driver-side metadata application `agentic_tool_call.generate` +performs — including the template-field overlay and the metadata application +order (agent metadata overrides the input's keys; session metadata, applied +last, overrides the agent's). +""" + +import json +import uuid +from types import SimpleNamespace + +import numpy as np +import pybase64 +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from tests.fast.rollout.session.test_samples import _make_record + +from miles.rollout.session.core import SessionCore +from miles.rollout.session.linear_trajectory import SessionRegistry +from miles.rollout.session.samples.codec import decode_samples_reply +from miles.rollout.session.sessions import setup_session_routes +from miles.utils.chat_template_utils import get_tito_tokenizer +from miles.utils.processing_utils import load_tokenizer +from miles.utils.types import Sample + +NUM_LAYERS = 3 +TOPK = 2 + +_ARGS = SimpleNamespace( + miles_router_timeout=30, + hf_checkpoint="Qwen/Qwen3-0.6B", + chat_template_path=None, + apply_chat_template_kwargs={"enable_thinking": False}, + tito_model="default", + tito_allowed_append_roles=["tool"], + session_server_instance_id=uuid.uuid4().hex, + num_layers=NUM_LAYERS, + moe_router_topk=TOPK, + save_debug_trajectory_data=None, +) + + +class _UnusedBackend: + """collect_samples never proxies; any backend call is a test bug.""" + + async def do_proxy(self, *args, **kwargs): + raise AssertionError("collect_samples must not touch the proxy backend") + + +def _build_core() -> SessionCore: + # Mirrors setup_session_routes (sessions.py): tokenizer + registry + core. + tokenizer = load_tokenizer( + _ARGS.hf_checkpoint, chat_template_path=_ARGS.chat_template_path, trust_remote_code=True + ) + tito_tokenizer = get_tito_tokenizer( + tokenizer, + tokenizer_type=_ARGS.tito_model, + chat_template_kwargs=_ARGS.apply_chat_template_kwargs, + allowed_append_roles=_ARGS.tito_allowed_append_roles, + ) + registry = SessionRegistry(_ARGS, tokenizer, tito_tokenizer=tito_tokenizer) + return SessionCore(_UnusedBackend(), registry, _ARGS, _ARGS.session_server_instance_id) + + +@pytest.fixture(scope="module") +def core(): + return _build_core() + + +# ── fixtures: a two-turn trajectory with R3 / cache stats / weight versions ── + + +def _r3_b64(num_tokens: int, seed: int) -> str: + arr = np.arange(seed, seed + num_tokens * NUM_LAYERS * TOPK, dtype=np.int32) + return pybase64.b64encode(arr.tobytes()).decode("ascii") + + +def _two_turn_records(): + # R3 buffer length per record = (len(prompt) + len(output) - 1) * layers * topk. + return [ + _make_record( + prompt_token_ids=[1, 2, 3], + output_token_ids=[10, 11], + output_log_probs=[-0.125, -0.25], + cached_tokens=0, + prompt_tokens=3, + weight_version="w1", + routed_experts=_r3_b64(4, seed=0), + ), + _make_record( + prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], + output_token_ids=[30, 31], + output_log_probs=[-0.5, -1.0], + cached_tokens=5, + prompt_tokens=7, + weight_version="w2", + routed_experts=_r3_b64(8, seed=100), + ), + ] + + +_ACCUMULATED = [1, 2, 3, 10, 11, 20, 21, 30, 31] + + +def _input_sample() -> Sample: + return Sample( + group_index=4, + index=9, + prompt=[{"role": "user", "content": "hi"}], + label="lbl", + reward=2.5, + metadata={"task": "t1", "shared_key": "from-input"}, + routing_key="routing-sid", + train_metadata={"loss": "ppo"}, + generate_function_path="gen.fn", + ) + + +# Overlapping keys lock the application order: agent overrides the input's +# shared_key; session_metadata (applied last) overrides the agent's +# max_trim_tokens plant. +_AGENT_METADATA = {"shared_key": "from-agent", "agent_only": 1, "max_trim_tokens": "agent-plant"} + + +async def _make_session(core, records, accumulated) -> str: + response = await core.create_session() + sid = json.loads(response.body)["session_id"] + session = core.registry.sessions[sid] + for record in records: + session.append_record(record) + if accumulated is not None: + session.trajectory_token_ids.append(list(accumulated)) + return sid + + +async def _collect_via_op(core, sid, *, multi_samples=False, max_seq_len=None): + response = await core.collect_samples(sid, multi_samples=multi_samples, max_seq_len=max_seq_len) + return response.status_code, response.body + + +def _new_pipeline(payload, input_sample, *, multi_samples): + """What collect_samples() does after the cutover: overlay + driver-side metadata.""" + reply = decode_samples_reply(payload, input_sample) + samples = reply.samples + for s in samples: + s.metadata.update(_AGENT_METADATA) + if samples: + if not multi_samples: + (merged,) = samples + merged.metadata.update(reply.session_metadata) + else: + samples[-1].metadata.update(reply.session_metadata) + return samples, reply + + +# ── golden assembly: exact expected Samples for the two-turn fixture ── + + +def _expected_r3(seed: int, num_tokens: int): + return np.arange(seed, seed + num_tokens * NUM_LAYERS * TOPK, dtype=np.int32).reshape(num_tokens, NUM_LAYERS, TOPK) + + +async def test_assembled_samples_golden_multi(core): + """multi_samples=True, no truncation: per-turn Samples carry exactly the + values derivable from the records fixture, template fields come from the + driver's input sample, and the metadata application order holds.""" + sid = await _make_session(core, _two_turn_records(), _ACCUMULATED) + status, payload = await _collect_via_op(core, sid, multi_samples=True) + assert status == 200 + samples, reply = _new_pipeline(payload, _input_sample(), multi_samples=True) + assert reply.empty_reason is None + tokenizer = core.registry.tokenizer + + s1, s2 = samples + assert s1.tokens == [1, 2, 3, 10, 11] + assert s1.response == tokenizer.decode([10, 11]) + assert s1.response_length == 2 + assert s1.loss_mask == [1, 1] + assert s1.rollout_log_probs == [-0.125, -0.25] + assert s1.status == Sample.Status.COMPLETED + assert s1.weight_versions == ["w1"] + assert np.array_equal(s1.rollout_routed_experts, _expected_r3(0, 4)) + assert s1.prefix_cache_info.to_dict() == {"cached_tokens": 0, "total_prompt_tokens": 3} + + assert s2.tokens == _ACCUMULATED + assert s2.response == tokenizer.decode([30, 31]) + assert s2.response_length == 2 + assert s2.loss_mask == [1, 1] + assert s2.rollout_log_probs == [-0.5, -1.0] + assert s2.status == Sample.Status.COMPLETED + assert s2.weight_versions == ["w2"] + assert np.array_equal(s2.rollout_routed_experts, _expected_r3(100, 8)) + assert s2.prefix_cache_info.to_dict() == {"cached_tokens": 5, "total_prompt_tokens": 7} + + # Overlay: template fields are the driver's, untouched by the wire. + for s in samples: + assert s.prompt == [{"role": "user", "content": "hi"}] + assert s.label == "lbl" + assert s.reward == 2.5 + assert s.routing_key == "routing-sid" + assert s.train_metadata == {"loss": "ppo"} + assert s.metadata["task"] == "t1" + # Metadata application order: agent overrides the input's shared_key on every + # sample; session_metadata (applied last, last sample only) overrides the + # agent's max_trim_tokens plant. + assert samples[0].metadata["shared_key"] == "from-agent" + assert samples[0].metadata["max_trim_tokens"] == "agent-plant" + assert samples[-1].metadata["max_trim_tokens"] == reply.session_metadata["max_trim_tokens"] + assert samples[-1].metadata["accumulated_token_ids"] == _ACCUMULATED + + +async def test_assembled_samples_golden_merged(core): + """multi_samples=False: turns merge into one trajectory Sample; the env + tokens between turns get zero loss/logprob; the last turn's R3 is kept.""" + sid = await _make_session(core, _two_turn_records(), _ACCUMULATED) + status, payload = await _collect_via_op(core, sid) + assert status == 200 + samples, reply = _new_pipeline(payload, _input_sample(), multi_samples=False) + (m,) = samples + tokenizer = core.registry.tokenizer + + assert m.tokens == _ACCUMULATED + assert m.response == tokenizer.decode([10, 11]) + tokenizer.decode([20, 21]) + tokenizer.decode([30, 31]) + assert m.response_length == 6 + assert m.loss_mask == [1, 1, 0, 0, 1, 1] + assert m.rollout_log_probs == [-0.125, -0.25, 0.0, 0.0, -0.5, -1.0] + assert m.status == Sample.Status.COMPLETED + assert m.weight_versions == ["w1", "w2"] + assert np.array_equal(m.rollout_routed_experts, _expected_r3(100, 8)) + assert m.prefix_cache_info.to_dict() == {"cached_tokens": 5, "total_prompt_tokens": 10} + assert m.metadata["shared_key"] == "from-agent" + assert m.metadata["accumulated_token_ids"] == _ACCUMULATED + + +@pytest.mark.parametrize("multi_samples", [False, True], ids=["merge", "multi"]) +async def test_truncation_golden(core, multi_samples): + """max_seq_len=8 strips one output token off the second turn (a turn-level + budget applied before merge): the final sample ends TRUNCATED at 8 tokens + with its per-token fields (including R3) trimmed in lockstep.""" + sid = await _make_session(core, _two_turn_records(), _ACCUMULATED) + status, payload = await _collect_via_op(core, sid, multi_samples=multi_samples, max_seq_len=8) + assert status == 200 + samples, _ = _new_pipeline(payload, _input_sample(), multi_samples=multi_samples) + + last = samples[-1] + assert last.status == Sample.Status.TRUNCATED + assert last.tokens == _ACCUMULATED[:8] + if multi_samples: + assert len(samples) == 2 + assert last.loss_mask == [1] + assert last.rollout_log_probs == [-0.5] + else: + assert last.loss_mask == [1, 1, 0, 0, 1] + assert last.rollout_log_probs == [-0.125, -0.25, 0.0, 0.0, -0.5] + assert np.array_equal(last.rollout_routed_experts, _expected_r3(100, 8)[:-1]) + + +async def test_session_metadata_matches_get_session(core): + """The samples reply and the records GET must expose the same metadata dict + (both are built by the extracted _session_metadata helper).""" + sid = await _make_session(core, _two_turn_records(), _ACCUMULATED) + _, payload = await _collect_via_op(core, sid) + reply = decode_samples_reply(payload, Sample()) + + response = await core.get_session(sid) + assert response.status_code == 200 + assert reply.session_metadata == json.loads(response.body)["metadata"] + assert reply.session_metadata["accumulated_token_ids"] == _ACCUMULATED + + +# ── empty_reason discriminator ── + + +async def test_no_records_reply(core): + sid = await _make_session(core, [], None) + status, payload = await _collect_via_op(core, sid) + assert status == 200 + reply = decode_samples_reply(payload, Sample()) + assert reply.samples == [] and reply.empty_reason == "no_records" + + +async def test_all_truncated_reply(core): + # max_seq_len=2 < the first turn's prompt+1: truncate_samples_by_total_tokens + # drops every turn -> empty samples with the all_truncated reason; the old + # pipeline returns [] on the same fixture (today's ABORTED path). + records = _two_turn_records() + sid = await _make_session(core, records, _ACCUMULATED) + status, payload = await _collect_via_op(core, sid, max_seq_len=2) + assert status == 200 + reply = decode_samples_reply(payload, Sample()) + assert reply.samples == [] and reply.empty_reason == "all_truncated" + + +# ── the 422 lane ── + + +async def test_broken_chain_returns_422_and_server_survives(core): + # The accumulated sequence carries one token the records never produced -> + # the cursor consistency assert fires -> 422 with the assertion text, and + # the server keeps serving (the failure never escapes as an unhandled 500). + sid = await _make_session(core, _two_turn_records(), _ACCUMULATED + [99]) + status, payload = await _collect_via_op(core, sid) + assert status == 422 + assert "cursor" in payload.decode() + + health = await core.health() + assert health.status_code == 200 + + +# ── the HTTP surface: route order and error mapping through the real app ── + + +@pytest.fixture(scope="module") +def app_client(): + app = FastAPI() + setup_session_routes(app, _UnusedBackend(), _ARGS) + with TestClient(app) as client: + yield client + + +def test_missing_session_returns_404(app_client): + response = app_client.post( + f"/sessions/{uuid.uuid4().hex}/samples", content=b'{"multi_samples":false,"max_seq_len":null}' + ) + assert response.status_code == 404 + assert "not found" in response.json()["error"] + + +def test_samples_route_registered_before_catch_all_proxy(app_client): + # The catch-all session_proxy would forward the request to the inference + # backend (_UnusedBackend raises); the samples route must win instead and + # answer with a decodable empty reply for a fresh session. + sid = app_client.post("/sessions").json()["session_id"] + response = app_client.post(f"/sessions/{sid}/samples", content=b'{"multi_samples":false,"max_seq_len":null}') + assert response.status_code == 200 + assert response.headers["content-type"] == "application/octet-stream" + reply = decode_samples_reply(response.content, Sample()) + assert reply.empty_reason == "no_records", "catch-all session_proxy swallowed the samples route"