diff --git a/docs/user-guide/rollout-endpoints.md b/docs/user-guide/rollout-endpoints.md index 4cef87074d3..0e104e198b0 100644 --- a/docs/user-guide/rollout-endpoints.md +++ b/docs/user-guide/rollout-endpoints.md @@ -86,14 +86,14 @@ Helpers: - `compute_prompt_ids_from_sample` and `compute_request_payload` from `miles/rollout/generate_utils/generate_endpoint_utils.py` build `/generate` requests. -- For multi-sample outputs, set `--generate-multi-samples` and return a list. +- Returning a `list[Sample]` from a generate function is supported natively; no flag is needed. ### Reference generators - **`single_turn.py`**: single-turn generation via `/generate`. Text or multimodal prompts. - **`multi_turn.py`**: multi-turn tool calling via `/generate`. Adds CLI flags `--generate-max-turns`, `--generate-tool-specs-path`, `--generate-tool-call-parser`, - `--generate-execute-tool-function-path`, `--generate-multi-samples`. + `--generate-execute-tool-function-path`. - **`benchmarkers.py`**: forces random output sequence length for benchmarking. --- @@ -181,6 +181,18 @@ remain a `messages` list. SGLang handles templating server-side. + + +**Agentic output is a `list[Sample]`.** `agentic_tool_call.generate` always returns a list +(one merged TITO sample per linear run today). Consequences: + +- A custom reward model (`--custom-rm-path`) is called in batch form with a + `list[Sample]` argument; it must handle that shape. +- `--group-rm`, `--partial-rollout`, and `--recompute-logprobs-via-prefill` are not + supported in combination with the agentic generator. + + + ### Optional teardown: the `abort` hook The module named by `--custom-agent-function-path` may expose an optional `abort` @@ -217,8 +229,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. @@ -234,7 +247,7 @@ TITO needs two things from every SGLang response: By default, `build_chat_request_kwargs` sets both flags. The session middleware forwards raw `messages` to SGLang, which tokenizes the prompt and returns the response. `_compute_sample_from_openai_record` in -[`openai_endpoint_utils.py`](https://github.com/radixark/miles/blob/main/miles/rollout/generate_utils/openai_endpoint_utils.py) +[`samples.py`](https://github.com/radixark/miles/blob/main/miles/rollout/session/samples.py) extracts prompt and output ids from the response and concatenates them into `sample.tokens`. You don't need to provide `input_ids` yourself. diff --git a/examples/experimental/swe-agent-v2/README.md b/examples/experimental/swe-agent-v2/README.md index 1c8aebc54be..30149b68a53 100644 --- a/examples/experimental/swe-agent-v2/README.md +++ b/examples/experimental/swe-agent-v2/README.md @@ -289,8 +289,6 @@ merge_samples() -> logprobs: -------- [real] [0.0] [real] ``` -Without TITO, use `--generate-multi-samples` to skip merge and train on per-turn samples instead (current default in `run.sh`). - ## Troubleshooting ### Harbor containers can't reach Miles Router @@ -312,7 +310,9 @@ The task directory for the given `instance_id` doesn't exist under `HARBOR_TASKS ### `b.tokens must start with a.tokens` assertion error -Multi-turn merge fails due to BPE re-tokenization inconsistency. Use `--generate-multi-samples` (already default in `run.sh`) to skip merge and train on per-turn samples. +Multi-turn merge fails due to BPE re-tokenization inconsistency. The session-server TITO path +(pretokenized `input_ids`) avoids the re-tokenization entirely; check that the run goes through +`--use-session-server` and that the chat template round-trips (see the TITO docs). ### Trace-viewer shows no trajectories diff --git a/miles/rollout/filter_hub/dynamic_sampling_filters.py b/miles/rollout/filter_hub/dynamic_sampling_filters.py index 0f6685c674e..505a23b060f 100644 --- a/miles/rollout/filter_hub/dynamic_sampling_filters.py +++ b/miles/rollout/filter_hub/dynamic_sampling_filters.py @@ -6,17 +6,8 @@ __all__ = ["check_reward_nonzero_std", "check_no_aborted"] -def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): - rewards = [sample.get_reward_value(args) for sample in samples] - keep = torch.tensor(rewards, dtype=torch.float64).std() > 1e-8 - return DynamicFilterOutput( - keep=keep, - reason=None if keep else f"zero_std_{round(rewards[0], 1)}", - ) - - -def _flatten_samples(samples): - """Flatten samples that may contain nested lists (from --generate-multi-samples).""" +def _flatten_samples(samples: list[Sample | list[Sample]]): + """Flatten a group whose elements are `Sample` or `list[Sample]` (generate-function dependent).""" for s in samples: if isinstance(s, list): yield from s @@ -24,7 +15,16 @@ def _flatten_samples(samples): yield s -def check_no_aborted(args, samples: list[Sample], **kwargs): +def check_reward_nonzero_std(args, samples: list[Sample | list[Sample]], **kwargs): + rewards = [sample.get_reward_value(args) for sample in _flatten_samples(samples)] + keep = torch.tensor(rewards, dtype=torch.float64).std() > 1e-8 + return DynamicFilterOutput( + keep=keep, + reason=None if keep else f"zero_std_{round(rewards[0], 1)}", + ) + + +def check_no_aborted(args, samples: list[Sample | list[Sample]], **kwargs): """Reject entire group if any sample was aborted (e.g. env timeout, Docker crash).""" if any(s.status == Sample.Status.ABORTED for s in _flatten_samples(samples)): return DynamicFilterOutput(keep=False, reason="group_has_aborted") diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index fe0a1252c62..aa51122cea5 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( @@ -32,12 +33,7 @@ async def my_agent( from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput -from miles.rollout.generate_utils.openai_endpoint_utils import ( - OpenAIEndpointTracer, - compute_samples_from_openai_records, - truncate_samples_by_total_tokens, -) -from miles.rollout.generate_utils.sample_utils import merge_samples +from miles.rollout.generate_utils.openai_endpoint_utils import OpenAIEndpointTracer from miles.utils.misc import load_function from miles.utils.types import Sample @@ -84,69 +80,51 @@ 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, 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), - ) + 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 (applied to every returned sample). 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) + samples[-1].metadata.update(result.session_metadata) return GenerateFnOutput(samples=samples) def _add_arguments(parser: argparse.ArgumentParser): parser.add_argument("--custom-agent-function-path", type=str) - parser.add_argument("--generate-multi-samples", action="store_true", default=False) parser.add_argument( "--max-seq-len", type=int, 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_hub/multi_turn.py b/miles/rollout/generate_hub/multi_turn.py index 99bec2b1a9d..1abc4ed220b 100644 --- a/miles/rollout/generate_hub/multi_turn.py +++ b/miles/rollout/generate_hub/multi_turn.py @@ -36,8 +36,6 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: tool_specs = load_function(args.generate_tool_specs_path) tool_call_parser = create_tool_call_parser(tool_specs, args.generate_tool_call_parser) - multi_samples = [] - # ----------------------- Initial prompts ------------------------- prompt_tokens_ids = compute_prompt_ids_from_sample(input.state, sample, tools=tool_specs) @@ -50,19 +48,11 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: payload, halt_status = compute_request_payload(args, sample.tokens, input.sampling_params) if payload is None: sample.status = halt_status - if args.generate_multi_samples and multi_samples: - multi_samples[-1].status = halt_status break - if args.generate_multi_samples: - sample = deepcopy(input.sample) - output = await post(url, payload, headers=compute_routing_headers(args, sample)) await update_sample_from_response(args, sample, payload=payload, output=output, update_loss_mask=True) - if args.generate_multi_samples: - multi_samples.append(deepcopy(sample)) - if output["meta_info"]["finish_reason"]["type"] in ("abort", "length"): break @@ -75,7 +65,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: tool_messages = await execute_tool_calls(tool_calls, execute_tool_function) update_sample_with_tool_responses(sample, tool_messages, tokenizer=tokenizer) - return GenerateFnOutput(samples=multi_samples if args.generate_multi_samples else sample) + return GenerateFnOutput(samples=sample) def _add_arguments(parser: argparse.ArgumentParser): @@ -83,7 +73,6 @@ def _add_arguments(parser: argparse.ArgumentParser): parser.add_argument("--generate-tool-specs-path", type=str) parser.add_argument("--generate-tool-call-parser", type=str) parser.add_argument("--generate-execute-tool-function-path", type=str) - parser.add_argument("--generate-multi-samples", action="store_true") generate.add_arguments = _add_arguments diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 205f29b90a3..4e8c34ee6c9 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -6,14 +6,9 @@ import logging import random from argparse import Namespace -from copy import deepcopy -from miles.rollout.generate_utils.generate_endpoint_utils import ( - get_indexer_topk_from_response, - get_routed_experts_from_response, -) -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__) @@ -56,179 +51,31 @@ 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 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", + {"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}") - - return (records or []), metadata - - -def compute_samples_from_openai_records( - args: Namespace, - input_sample: Sample, - records: list[SessionRecord], - tokenizer, - accumulated_token_ids: list[int] | None = None, - max_trim_tokens: int = 0, -) -> list[Sample]: - """Convert per-turn session records into training Samples, aligning each - turn's output tokens against the TITO accumulated token sequence. - - Each record carries its own ``prompt_token_ids`` and ``output_token_ids`` - (with logprobs). We want to reuse those per-turn logprobs directly - instead of re-decoding, but we must first trim "trailing tokens" — stop - tokens the model emitted that the chat template also renders as the next - turn's delimiter — to avoid double-counting. - - See ``TestTITOTrailingTokenTrim`` in - ``tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py`` - for a concrete worked example with token-level walkthroughs. - """ - samples = [] - cursor = 0 - - for i, record in enumerate(records): - is_last = i == len(records) - 1 - prompt_ids = record.request["input_ids"] - output_ids = [t[1] for t in record.response["choices"][0]["meta_info"]["output_token_logprobs"]] - - trim_count = 0 - if accumulated_token_ids is not None: - # Step 1: position cursor right after this turn's prompt - cursor = len(prompt_ids) - - # Step 2: greedily match output_ids against accumulated[cursor:] - matched = 0 - for j in range(len(output_ids)): - idx = cursor + j - if idx < len(accumulated_token_ids) and output_ids[j] == accumulated_token_ids[idx]: - matched += 1 - else: - break - - # Step 3: unmatched trailing tokens were consumed by the next - # turn's template rendering (e.g. stop tokens that double as - # the next message delimiter) — strip them from the sample. - trim_count = len(output_ids) - matched - allowed = 0 if is_last else max_trim_tokens - assert trim_count <= allowed, ( - f"trim_count {trim_count} exceeds allowed={allowed} " - f"(is_last={is_last}, max_trim_tokens={max_trim_tokens}); " - f"output_ids[-3:]={output_ids[-3:]}, " - f"accumulated[cursor:cursor+3]={accumulated_token_ids[cursor:cursor+3]}" - ) - - # 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) - samples.append(sample) - - if accumulated_token_ids is not None: - # Step 5: verify the entire accumulated sequence was consumed - assert cursor == len(accumulated_token_ids), ( - f"cursor {cursor} != len(accumulated_token_ids) {len(accumulated_token_ids)} " - f"after processing all {len(records)} records" - ) - - return samples - - -def _compute_sample_from_openai_record( - args: Namespace, input_sample: Sample, record: SessionRecord, tokenizer, trim_count: int = 0 -) -> Sample: - choice = record.response["choices"][0] - - prompt_token_ids = record.request.get("input_ids") - if prompt_token_ids is None: - raise ValueError("input_ids not found in request — the session server should populate it") - - 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.tokens = prompt_token_ids + output_token_ids - sample.rollout_log_probs = output_log_probs - sample.response = tokenizer.decode(output_token_ids) - sample.response_length = len(output_token_ids) - sample.loss_mask = [1] * len(output_token_ids) - sample.rollout_routed_experts = get_routed_experts_from_response(args, choice, sample) - sample.rollout_indexer_topk = get_indexer_topk_from_response(args, choice, sample) - - if trim_count > 0: - sample.strip_last_output_tokens(trim_count, tokenizer) - - # TODO unify with Sample.update_from_meta_info - match choice["finish_reason"]: - case "stop" | "tool_calls": - sample.status = Sample.Status.COMPLETED - case "length": - sample.status = Sample.Status.TRUNCATED - case "abort": - sample.status = Sample.Status.ABORTED - - sample.prefix_cache_info.add(choice.get("meta_info", {})) - if "weight_version" in choice["meta_info"]: - sample.weight_versions.append(choice["meta_info"]["weight_version"]) - - return sample - - -def truncate_samples_by_total_tokens( - samples: list[Sample], - max_seq_len: int, - tokenizer, -) -> list[Sample]: - """Truncate samples so the total token count (prompt + output, including - env responses) does not exceed ``max_seq_len``. - """ - result: list[Sample] = [] - - for sample in samples: - total = len(sample.tokens) - if total <= max_seq_len: - result.append(sample) - continue - - overshoot = total - max_seq_len - allowed_output = sample.response_length - overshoot - if allowed_output <= 0: - break - - sample.strip_last_output_tokens(overshoot, tokenizer) - sample.status = Sample.Status.TRUNCATED - result.append(sample) - break + except Exception as e: + logger.warning(f"Failed to delete session {self.session_id} after collecting samples: {e}") - return result + return decode_samples_reply(payload, input_sample) diff --git a/miles/rollout/session/core.py b/miles/rollout/session/core.py index e1471ef6d20..8b81828769f 100644 --- a/miles/rollout/session/core.py +++ b/miles/rollout/session/core.py @@ -4,6 +4,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. +- ``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 @@ -13,6 +14,7 @@ from starlette.responses import Response +from miles.rollout.generate_utils.sample_utils import merge_samples from miles.rollout.session.errors import ( MessageValidationError, SessionNotFoundError, @@ -20,6 +22,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__) @@ -44,6 +48,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") @@ -107,8 +116,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) @@ -119,11 +130,59 @@ 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, 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")) + # Sample boundaries are a property of the trajectory, not a caller + # choice: consecutive linear turns fold into one TITO sample. The + # registry rejects non-linear appends today, so this is always one + # merged run; compaction/subagent support will partition the turns + # into multiple runs here (the reply is already a list on the wire). + 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/__init__.py b/miles/rollout/session/samples/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/rollout/session/samples/codec.py b/miles/rollout/session/samples/codec.py new file mode 100644 index 00000000000..8dbeccb0c85 --- /dev/null +++ b/miles/rollout/session/samples/codec.py @@ -0,0 +1,222 @@ +"""Black-box sample packing for the `POST /sessions/{id}/samples` reply; only the +input/output contract matters, the wire format is an implementation detail. + +- Server: `encode_samples_reply(samples, session_metadata, empty_reason)` — assembled + `Sample`s in, one opaque payload (`bytes`) out. +- Driver: `decode_samples_reply(payload, input_sample)` — that payload in, `SamplesReply` + out, each `Sample` rebuilt by overlaying the wire's computed fields onto a deepcopy of + the driver's `input_sample`. +""" + +import dataclasses +import json +from collections.abc import Callable +from copy import deepcopy + +import numpy as np +import safetensors.numpy + +from miles.utils.types import Sample + +# Every Sample field is either COMPUTED by assembly from the records (crosses +# the wire) or belongs to the driver's input-sample TEMPLATE (never crosses; +# the driver overlay keeps its local deepcopy's value). Adding a Sample field +# without classifying it here fails at import time, not silently at training. +COMPUTED_FIELDS = ( + "tokens", + "response", + "response_length", + "loss_mask", + "rollout_log_probs", + "rollout_routed_experts", + "rollout_indexer_topk", + "status", + "weight_versions", + "prefix_cache_info", +) +TEMPLATE_FIELDS = ( + "group_index", + "index", + "prompt", + "multimodal_inputs", + "multimodal_train_inputs", + "label", + "reward", + "remove_sample", + "teacher_log_probs", + "opd_reverse_kl", + "metadata", + "generate_function_path", + "train_metadata", + "routing_key", + "non_generation_time", + "spec_info", +) + +_SAMPLE_FIELDS = {f.name for f in dataclasses.fields(Sample)} +assert set(COMPUTED_FIELDS) | set(TEMPLATE_FIELDS) == _SAMPLE_FIELDS and not set(COMPUTED_FIELDS) & set( + TEMPLATE_FIELDS +), ( + "Sample fields drifted: every field must be classified as COMPUTED (crosses the samples wire) " + "or TEMPLATE (stays on the driver's input sample). " + f"Unclassified: {sorted(_SAMPLE_FIELDS - set(COMPUTED_FIELDS) - set(TEMPLATE_FIELDS))}, " + f"unknown: {sorted((set(COMPUTED_FIELDS) | set(TEMPLATE_FIELDS)) - _SAMPLE_FIELDS)}, " + f"overlap: {sorted(set(COMPUTED_FIELDS) & set(TEMPLATE_FIELDS))}" +) + + +@dataclasses.dataclass(frozen=True) +class _TensorSpec: + """Wire contract of one COMPUTED tensor field; the codec is a loop over `_TENSOR_SPECS`.""" + + normalize_dtype: np.dtype | None # np.asarray target on encode; None keeps the caller's dtype + wire_dtype: np.dtype # pinned on both sides; a mismatch raises instead of silently converting + restore_list: bool # decode returns .tolist() (legacy JSON-path types) instead of the ndarray + null_factory: Callable[[], object] # decode value for JSON null; factory so no instance is shared + + +# Token ids and logprobs are re-materialized as Python lists on decode, exactly +# like the legacy JSON path (int64/f64 round-trips are lossless for both). The +# R3 replay fields must arrive as int32 and are never converted. +_TENSOR_SPECS = { + "tokens": _TensorSpec(np.dtype(np.int64), np.dtype(np.int64), True, list), + "rollout_log_probs": _TensorSpec(np.dtype(np.float64), np.dtype(np.float64), True, lambda: None), + "rollout_routed_experts": _TensorSpec(None, np.dtype(np.int32), False, lambda: None), + "rollout_indexer_topk": _TensorSpec(None, np.dtype(np.int32), False, lambda: None), +} +_SCALAR_FIELDS = ("response", "response_length", "loss_mask", "status", "weight_versions", "prefix_cache_info") + +assert set(_TENSOR_SPECS) | set(_SCALAR_FIELDS) == set(COMPUTED_FIELDS) and not set(_TENSOR_SPECS) & set( + _SCALAR_FIELDS +), ( + "every COMPUTED field needs exactly one wire representation (tensor spec or scalar JSON); " + f"uncovered: {sorted(set(COMPUTED_FIELDS) - set(_TENSOR_SPECS) - set(_SCALAR_FIELDS))}, " + f"unknown: {sorted((set(_TENSOR_SPECS) | set(_SCALAR_FIELDS)) - set(COMPUTED_FIELDS))}, " + f"overlap: {sorted(set(_TENSOR_SPECS) & set(_SCALAR_FIELDS))}" +) + +_SAMPLES_META_KEY = "_samples_meta" +_OPD_STUDENT_TOP_LOGPROBS_KEY = "opd_student_top_logprobs" + + +def _tensor_name(sample_index: int, field: str) -> str: + return f"sample.{sample_index}.{field}" + + +@dataclasses.dataclass +class SamplesReply: + """Decoded `POST /sessions/{id}/samples` reply.""" + + samples: list[Sample] + session_metadata: dict + empty_reason: str | None + + +def encode_samples_reply(samples: list[Sample], session_metadata: dict, empty_reason: str | None = None) -> bytes: + """Worker side: pack assembled samples into one safetensors payload.""" + tensors: dict[str, np.ndarray] = {} + sample_metas = [] + for sample_index, sample in enumerate(samples): + tensor_meta = {} + for field, spec in _TENSOR_SPECS.items(): + value = getattr(sample, field) + if value is None: + tensor_meta[field] = None + continue + arr = np.asarray(value, dtype=spec.normalize_dtype) + if arr.dtype != spec.wire_dtype: + raise ValueError(f"{field} must have dtype {spec.wire_dtype}, got {arr.dtype}") + name = _tensor_name(sample_index, field) + # ascontiguousarray is a correctness requirement: the numpy adapter + # serializes some non-contiguous views without raising, with wrong values. + tensors[name] = np.ascontiguousarray(arr) + tensor_meta[field] = name + scalar_meta = {} + for field in _SCALAR_FIELDS: + value = getattr(sample, field) + if field == "status": + value = value.value + elif field == "prefix_cache_info": + value = value.to_dict() + scalar_meta[field] = value + sample_metas.append({**scalar_meta, "tensors": tensor_meta}) + meta = {"samples": sample_metas, "session_metadata": session_metadata, "empty_reason": empty_reason} + # Compact separators are load-bearing: the default ", "/": " padding costs + # ~1 byte per loss_mask/token entry, ~100KB on production-sized replies. + 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) + tensor_meta = sample_meta["tensors"] + for field, spec in _TENSOR_SPECS.items(): + name = tensor_meta[field] + if name is None: + setattr(sample, field, spec.null_factory()) + continue + expected = _tensor_name(sample_index, field) + if name != expected: + raise ValueError(f"{field} references tensor {name!r}, expected {expected!r}") + arr = tensors.pop(name) # KeyError propagates: a referenced tensor must exist + if arr.dtype != spec.wire_dtype: + raise ValueError(f"{field} must have dtype {spec.wire_dtype}, got {arr.dtype}") + setattr(sample, field, arr.tolist() if spec.restore_list else arr) + for field in _SCALAR_FIELDS: + 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) + 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 new file mode 100644 index 00000000000..48992462964 --- /dev/null +++ b/miles/rollout/session/samples/merge.py @@ -0,0 +1,155 @@ +"""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 miles.rollout.generate_utils.generate_endpoint_utils import ( + get_indexer_topk_from_response, + get_routed_experts_from_response, +) +from miles.rollout.session.types import SessionRecord +from miles.utils.types import Sample + + +def compute_samples_from_openai_records( + args: Namespace, + records: list[SessionRecord], + tokenizer, + accumulated_token_ids: list[int] | None = None, + max_trim_tokens: int = 0, +) -> list[Sample]: + """Convert per-turn session records into training Samples, aligning each + turn's output tokens against the TITO accumulated token sequence. + + Each record carries its own ``prompt_token_ids`` and ``output_token_ids`` + (with logprobs). We want to reuse those per-turn logprobs directly + instead of re-decoding, but we must first trim "trailing tokens" — stop + tokens the model emitted that the chat template also renders as the next + turn's delimiter — to avoid double-counting. + + See ``TestTITOTrailingTokenTrim`` in + ``tests/fast/rollout/session/test_samples.py`` + for a concrete worked example with token-level walkthroughs. + """ + samples = [] + cursor = 0 + + for i, record in enumerate(records): + is_last = i == len(records) - 1 + prompt_ids = record.request["input_ids"] + output_ids = [t[1] for t in record.response["choices"][0]["meta_info"]["output_token_logprobs"]] + + trim_count = 0 + if accumulated_token_ids is not None: + # Step 1: position cursor right after this turn's prompt + cursor = len(prompt_ids) + + # Step 2: greedily match output_ids against accumulated[cursor:] + matched = 0 + for j in range(len(output_ids)): + idx = cursor + j + if idx < len(accumulated_token_ids) and output_ids[j] == accumulated_token_ids[idx]: + matched += 1 + else: + break + + # Step 3: unmatched trailing tokens were consumed by the next + # turn's template rendering (e.g. stop tokens that double as + # the next message delimiter) — strip them from the sample. + trim_count = len(output_ids) - matched + allowed = 0 if is_last else max_trim_tokens + assert trim_count <= allowed, ( + f"trim_count {trim_count} exceeds allowed={allowed} " + f"(is_last={is_last}, max_trim_tokens={max_trim_tokens}); " + f"output_ids[-3:]={output_ids[-3:]}, " + f"accumulated[cursor:cursor+3]={accumulated_token_ids[cursor:cursor+3]}" + ) + + # Step 4: advance cursor past matched output to the next turn + cursor += matched + + sample = _compute_sample_from_openai_record(args, record, tokenizer, trim_count) + samples.append(sample) + + if accumulated_token_ids is not None: + # Step 5: verify the entire accumulated sequence was consumed + assert cursor == len(accumulated_token_ids), ( + f"cursor {cursor} != len(accumulated_token_ids) {len(accumulated_token_ids)} " + f"after processing all {len(records)} records" + ) + + return samples + + +def _compute_sample_from_openai_record( + args: Namespace, record: SessionRecord, tokenizer, trim_count: int = 0 +) -> Sample: + choice = record.response["choices"][0] + + prompt_token_ids = record.request.get("input_ids") + if prompt_token_ids is None: + raise ValueError("input_ids not found in request — the session server should populate it") + + 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 = Sample() + sample.tokens = prompt_token_ids + output_token_ids + sample.rollout_log_probs = output_log_probs + sample.response = tokenizer.decode(output_token_ids) + sample.response_length = len(output_token_ids) + sample.loss_mask = [1] * len(output_token_ids) + sample.rollout_routed_experts = get_routed_experts_from_response(args, choice, sample) + sample.rollout_indexer_topk = get_indexer_topk_from_response(args, choice, sample) + + if trim_count > 0: + sample.strip_last_output_tokens(trim_count, tokenizer) + + # TODO unify with Sample.update_from_meta_info + match choice["finish_reason"]: + case "stop" | "tool_calls": + sample.status = Sample.Status.COMPLETED + case "length": + sample.status = Sample.Status.TRUNCATED + case "abort": + sample.status = Sample.Status.ABORTED + + sample.prefix_cache_info.add(choice.get("meta_info", {})) + if "weight_version" in choice["meta_info"]: + sample.weight_versions.append(choice["meta_info"]["weight_version"]) + + return sample + + +def truncate_samples_by_total_tokens( + samples: list[Sample], + max_seq_len: int, + tokenizer, +) -> list[Sample]: + """Truncate samples so the total token count (prompt + output, including + env responses) does not exceed ``max_seq_len``. + """ + result: list[Sample] = [] + + for sample in samples: + total = len(sample.tokens) + if total <= max_seq_len: + result.append(sample) + continue + + overshoot = total - max_seq_len + allowed_output = sample.response_length - overshoot + if allowed_output <= 0: + break + + sample.strip_last_output_tokens(overshoot, tokenizer) + sample.status = Sample.Status.TRUNCATED + result.append(sample) + break + + return result diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index f3ad6b813f2..7d41ba355d0 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,17 @@ 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, 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 9b27fdb6131..68398082018 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 74b84a2977f..e20019db703 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,6 +19,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/e2e/sglang/utils/logprob_verify_generate.py b/tests/e2e/sglang/utils/logprob_verify_generate.py deleted file mode 100644 index b92aade7bbc..00000000000 --- a/tests/e2e/sglang/utils/logprob_verify_generate.py +++ /dev/null @@ -1,411 +0,0 @@ -"""Custom generate function: agentic flow + re-prefill logprob verification. - -Design rationale -~~~~~~~~~~~~~~~~ -This file needs to run the *exact same* pipeline as the production -``agentic_tool_call.generate``, but insert a verification step between -"collect session records" and "convert records to training samples". -There are two ways to achieve this: - -1. **Modify production code** — split ``generate()`` into reusable - ``_generate_core()`` / ``_finalize()`` helpers and import them here. -2. **Inline the production flow** — copy the pipeline steps verbatim - from ``agentic_tool_call.generate`` into this file, and insert - verification in the middle. - -We chose option 2 to avoid modifying production code for test purposes. -The tradeoff is a maintenance burden: if ``agentic_tool_call.generate`` -changes, this file must be updated in lockstep. - -Equivalence guarantee -~~~~~~~~~~~~~~~~~~~~~ -Steps 1 and 4 in ``generate()`` below are a *verbatim* copy of the -production ``agentic_tool_call.generate`` (as of the commit that -introduced this test). Specifically: - -- **Step 1** (lines up to ``collect_records``) mirrors production lines - that create the tracer, load the agent, build metadata, run the agent, - and collect records. Every function call, argument, and branch is - identical. -- **Step 4** (from ``compute_samples_from_openai_records`` onward) - mirrors the production finalization: sample computation, metadata - merge, truncation, and multi-sample merge. Same calls, same order. - -Steps 2-3 are *test-only* assertions and the re-prefill verification. -They are read-only checks that do not mutate ``records`` or -``session_metadata``, so they cannot affect the finalization output. - -To detect drift, ``grep`` for the functions called in steps 1/4 -(``OpenAIEndpointTracer.create``, ``build_chat_request_kwargs``, -``compute_samples_from_openai_records``, ``truncate_samples_by_total_tokens``, -``merge_samples``) — if their signatures change in production, this file -will fail to compile or produce assertion errors. - -Verification approach -~~~~~~~~~~~~~~~~~~~~~ -After the multi-turn agent session completes, the TITO session server -exposes the full ``accumulated_token_ids`` — the token sequence built -incrementally across all turns. We send this to ``/generate`` with -``max_new_tokens=0`` for a single prefill pass and compare the resulting -``input_token_logprobs`` against the per-turn ``output_token_logprobs`` -from the session's decode phase. Any token ID mismatch is fatal (TITO -tokenization bug); logprob values must match within a tight tolerance -(prefill vs decode numerical differences). - -When ``use_rollout_routing_replay`` is enabled (MoE models), per-turn -``routed_experts`` arrays are also compared against the re-prefill. -""" - -import logging -import statistics -from collections.abc import Callable -from copy import deepcopy - -import numpy as np -import pybase64 - -from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput -from miles.rollout.generate_hub.agentic_tool_call import build_chat_request_kwargs -from miles.rollout.generate_utils.openai_endpoint_utils import ( - OpenAIEndpointTracer, - compute_samples_from_openai_records, - truncate_samples_by_total_tokens, -) -from miles.rollout.generate_utils.sample_utils import merge_samples -from miles.utils.http_utils import post -from miles.utils.misc import load_function -from miles.utils.types import Sample - -logger = logging.getLogger(__name__) - -LOGPROB_ABS_TOL = 1e-8 # deterministic inference → prefill and decode must be bit-identical -LOGPROB_WARN_TOL = 0.0 # any nonzero diff is worth logging - - -async def generate(input: GenerateFnInput) -> GenerateFnOutput: - """Run production agentic flow with mid-pipeline logprob verification. - - Steps 1 and 4 are a verbatim copy of ``agentic_tool_call.generate``. - Steps 2-3 are test-only read-only checks inserted between record - collection and sample finalization. See module docstring for the - equivalence rationale. - """ - - # ── Step 1: core agentic flow ────────────────────────────────────── - # Verbatim from agentic_tool_call.generate — do NOT diverge. - tracer = await OpenAIEndpointTracer.create(input.args) - - custom_agent_function: Callable = load_function(input.args.custom_agent_function_path) - assert ( - custom_agent_function is not None - ), f"Custom agent function {input.args.custom_agent_function_path} not found" - - max_seq_len = getattr(input.args, "max_seq_len", None) - - metadata = input.sample.metadata - if max_seq_len is not None: - metadata = {**metadata, "max_seq_len": max_seq_len} - - agent_metadata = await custom_agent_function( - base_url=tracer.base_url, - prompt=input.sample.prompt, - request_kwargs=build_chat_request_kwargs(input.sampling_params), - metadata=metadata, - ) - - records, session_metadata = await tracer.collect_records() - - if not records: - logger.warning("No model calls recorded for sample") - sample = deepcopy(input.sample) - sample.status = Sample.Status.ABORTED - return GenerateFnOutput(samples=sample) - # ── End step 1 ───────────────────────────────────────────────────── - - # ── Step 2: test-only session-level precondition checks ──────────── - # These are stricter than production (which silently tolerates missing - # fields). We enforce them because the re-prefill comparison below - # is meaningless without a valid accumulated token sequence. - mismatch = session_metadata.get("tito_session_mismatch") - assert mismatch == [], f"tito_session_mismatch is not empty: {mismatch}" - - accumulated = session_metadata.get("accumulated_token_ids") - assert accumulated and len(accumulated) > 0, "accumulated_token_ids is empty" - - max_trim_tokens = session_metadata.get("max_trim_tokens", 0) - - assert len(records) >= 2, f"Expected at least 2 turns for TITO verification, got {len(records)}" - - # ── Step 3: re-prefill logprob verification (test-only) ──────────── - sglang_url = f"http://{input.args.sglang_router_ip}:{input.args.sglang_router_port}" - use_r3 = getattr(input.args, "use_rollout_routing_replay", False) - - await _verify_logprobs_via_reprefill( - sglang_url, - records, - accumulated, - max_trim_tokens=max_trim_tokens, - use_r3=use_r3, - ) - - logger.info( - "Logprob equivalence verified: %d turns, %d accumulated tokens", - len(records), - len(accumulated), - ) - - # ── Step 4: finalize ─────────────────────────────────────────────── - # Verbatim from agentic_tool_call.generate — do NOT diverge. - samples = compute_samples_from_openai_records( - input.args, - input.sample, - records, - input.state.tokenizer, - accumulated_token_ids=accumulated, - max_trim_tokens=max_trim_tokens, - ) - - for s in samples: - s.metadata.update(agent_metadata or {}) - - 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) - return GenerateFnOutput(samples=samples) - # ── End step 4 ───────────────────────────────────────────────────── - - -# Reuse the same CLI arguments as agentic_tool_call so that the -# framework registers --custom-agent-function-path, --max-seq-len, etc. -generate.add_arguments = None # set below after import - - -def _init_add_arguments(): - from miles.rollout.generate_hub.agentic_tool_call import generate as _base_generate - - generate.add_arguments = _base_generate.add_arguments - - -_init_add_arguments() - - -# --------------------------------------------------------------------------- -# Verification helpers -# --------------------------------------------------------------------------- - - -def _match_output_tokens( - output_ids: list[int], - accumulated_token_ids: list[int], - cursor: int, -) -> int: - """Count how many leading output tokens match accumulated[cursor:]. - - Shared by logprob and routed-expert verification. Mirrors the - greedy-match loop in ``compute_samples_from_openai_records`` - (openai_endpoint_utils.py) — intentionally duplicated here to avoid - modifying production code for test purposes. - """ - matched = 0 - for j in range(len(output_ids)): - idx = cursor + j - if idx < len(accumulated_token_ids) and output_ids[j] == accumulated_token_ids[idx]: - matched += 1 - else: - break - return matched - - -async def _verify_logprobs_via_reprefill( - sglang_url: str, - records: list, - accumulated_token_ids: list[int], - max_trim_tokens: int, - use_r3: bool, -) -> None: - """Re-prefill the full accumulated sequence and compare logprobs. - - Sends ``accumulated_token_ids`` to ``/generate`` with - ``max_new_tokens=0, return_logprob=True``. Compares the resulting - ``input_token_logprobs`` (single prefill pass) against per-turn - ``output_token_logprobs`` (incremental decode) from session records. - """ - first_prompt_len = len(records[0].request["input_ids"]) - - # ── A: send re-prefill request ── - payload = { - "input_ids": accumulated_token_ids, - "sampling_params": {"max_new_tokens": 0, "temperature": 0}, - "return_logprob": True, - "logprob_start_len": first_prompt_len, - } - if use_r3: - payload["return_routed_experts"] = True - - reprefill_resp = await post(f"{sglang_url}/generate", payload) - reprefill_logprobs = reprefill_resp["meta_info"]["input_token_logprobs"] - - expected_len = len(accumulated_token_ids) - first_prompt_len - assert len(reprefill_logprobs) == expected_len, ( - f"Re-prefill returned {len(reprefill_logprobs)} input_token_logprobs, " - f"expected {expected_len} (accumulated={len(accumulated_token_ids)}, " - f"first_prompt={first_prompt_len})" - ) - - # ── B: walk records and compare per-turn ── - all_diffs: list[float] = [] - # Per-turn (cursor, matched) pairs — reused by R3 verification to - # avoid recomputing the same greedy match. - turn_matches: list[tuple[int, int]] = [] - - for i, record in enumerate(records): - is_last = i == len(records) - 1 - choice = record.response["choices"][0] - prompt_ids = record.request["input_ids"] - session_output_logprobs = choice["meta_info"]["output_token_logprobs"] - output_ids = [t[1] for t in session_output_logprobs] - - if not output_ids: - logger.warning("Turn %d: no output tokens, skipping", i) - turn_matches.append((len(prompt_ids), 0)) - continue - - cursor = len(prompt_ids) - matched = _match_output_tokens(output_ids, accumulated_token_ids, cursor) - turn_matches.append((cursor, matched)) - - trim_count = len(output_ids) - matched - allowed = 0 if is_last else max_trim_tokens - assert trim_count <= allowed, f"Turn {i}: trim_count {trim_count} exceeds allowed={allowed}" - - # Compare matched tokens against re-prefill - turn_reprefill_start = cursor - first_prompt_len - mismatches = [] - warnings = [] - - for j in range(matched): - rp_entry = reprefill_logprobs[turn_reprefill_start + j] # (logprob, token_id, text) - sp_entry = session_output_logprobs[j] # (logprob, token_id) - - rp_tid = rp_entry[1] - sp_tid = sp_entry[1] - assert rp_tid == sp_tid, f"Turn {i}, token {j}: token_id mismatch — reprefill={rp_tid} vs session={sp_tid}" - - rp_lp = rp_entry[0] - sp_lp = sp_entry[0] - if rp_lp is None or sp_lp is None: - continue - - diff = abs(rp_lp - sp_lp) - all_diffs.append(diff) - - if diff > LOGPROB_ABS_TOL: - mismatches.append(f" token {j}: prefill={rp_lp:.8f} decode={sp_lp:.8f} diff={diff:.4f}") - elif diff > LOGPROB_WARN_TOL: - warnings.append(f" token {j}: prefill={rp_lp:.8f} decode={sp_lp:.8f} diff={diff:.4f}") - - if warnings: - logger.warning( - "Turn %d: %d tokens with diff > %.4f (but within tolerance):\n%s", - i, - len(warnings), - LOGPROB_WARN_TOL, - "\n".join(warnings[:10]), - ) - - assert ( - not mismatches - ), f"Turn {i}: {len(mismatches)} logprob differences exceed " f"tolerance {LOGPROB_ABS_TOL}:\n" + "\n".join( - mismatches - ) - - logger.info("Turn %d: verified %d output tokens (trimmed %d)", i, matched, trim_count) - - # ── C: R3 routed-expert comparison ── - if use_r3: - _verify_routed_experts(records, reprefill_resp, accumulated_token_ids, first_prompt_len, turn_matches) - - # ── D: summary statistics ── - if all_diffs: - sorted_diffs = sorted(all_diffs) - p99_idx = min(int(len(sorted_diffs) * 0.99), len(sorted_diffs) - 1) - logger.info( - "Logprob diff stats: mean=%.6f, max=%.6f, p99=%.6f, count=%d", - statistics.mean(all_diffs), - max(all_diffs), - sorted_diffs[p99_idx], - len(all_diffs), - ) - - -def _verify_routed_experts( - records: list, - reprefill_resp: dict, - accumulated_token_ids: list[int], - first_prompt_len: int, - turn_matches: list[tuple[int, int]], -) -> None: - """Compare per-turn routed_experts from session decode vs re-prefill. - - Reuses ``turn_matches`` (cursor, matched) computed by the logprob - verification pass to avoid recomputing the greedy token match. - """ - reprefill_re_b64 = reprefill_resp["meta_info"].get("routed_experts") - if reprefill_re_b64 is None: - logger.warning("Re-prefill response missing routed_experts, skipping R3 check") - return - - reprefill_re_flat = np.frombuffer(pybase64.b64decode(reprefill_re_b64.encode("ascii")), dtype=np.int32) - - for i, record in enumerate(records): - choice = record.response["choices"][0] - prompt_ids = record.request["input_ids"] - output_logprobs = choice["meta_info"]["output_token_logprobs"] - output_ids = [t[1] for t in output_logprobs] - - session_re_b64 = choice["meta_info"].get("routed_experts") - if session_re_b64 is None: - logger.warning("Turn %d: session missing routed_experts, skipping", i) - continue - - cursor, matched = turn_matches[i] - if matched == 0: - continue - - session_re = np.frombuffer(pybase64.b64decode(session_re_b64.encode("ascii")), dtype=np.int32) - - # routed_experts shape: [total_tokens - 1, num_layers, top_k]. - # Entry k corresponds to the routing decision for token k+1. - total_tokens = len(prompt_ids) + len(output_ids) - if total_tokens <= 1 or len(session_re) == 0: - continue - - per_token_size = len(session_re) // (total_tokens - 1) - - # Session output starts at position P (=len(prompt_ids)), - # so its expert entries start at index (P-1) in the flat array. - session_start = (len(prompt_ids) - 1) * per_token_size - session_slice = session_re[session_start : session_start + matched * per_token_size] - - # Re-prefill covers all of accumulated_token_ids. - # Token at accumulated[cursor] → expert entry at index (cursor-1). - rp_offset = (cursor - 1) * per_token_size - rp_slice = reprefill_re_flat[rp_offset : rp_offset + matched * per_token_size] - - np.testing.assert_array_equal( - session_slice, - rp_slice, - err_msg=f"Turn {i}: routed_experts mismatch", - ) - logger.info("Turn %d: routed_experts match (%d entries)", i, len(session_slice)) diff --git a/tests/fast/fixtures/generation_fixtures.py b/tests/fast/fixtures/generation_fixtures.py index 95838f5c748..3dff6775588 100644 --- a/tests/fast/fixtures/generation_fixtures.py +++ b/tests/fast/fixtures/generation_fixtures.py @@ -31,10 +31,8 @@ VARIANT_TO_GENERATE_FN_PATH = { "old_sglang_rollout": "miles.rollout.sglang_rollout.generate", "single_turn": "miles.rollout.generate_hub.single_turn.generate", - "multi_turn_single_sample": "miles.rollout.generate_hub.multi_turn.generate", - "multi_turn_multi_samples": "miles.rollout.generate_hub.multi_turn.generate", - "agentic_tool_call_single_sample": "miles.rollout.generate_hub.agentic_tool_call.generate", - "agentic_tool_call_multi_samples": "miles.rollout.generate_hub.agentic_tool_call.generate", + "multi_turn": "miles.rollout.generate_hub.multi_turn.generate", + "agentic_tool_call": "miles.rollout.generate_hub.agentic_tool_call.generate", } @@ -53,7 +51,7 @@ def extra_argv_for_variant( custom_generate_function_path or VARIANT_TO_GENERATE_FN_PATH[variant], ] - if variant in ("multi_turn_single_sample", "multi_turn_multi_samples"): + if variant == "multi_turn": argv += [ "--generate-max-turns", str(generate_max_turns), @@ -63,13 +61,9 @@ def extra_argv_for_variant( generate_execute_tool_function_path, ] argv += ["--generate-tool-call-parser", generate_tool_call_parser] - if variant == "multi_turn_multi_samples": - argv.append("--generate-multi-samples") - elif variant in ("agentic_tool_call_single_sample", "agentic_tool_call_multi_samples"): + elif variant == "agentic_tool_call": argv += ["--custom-agent-function-path", custom_agent_function_path] argv += ["--use-session-server", "--tito-model", "qwen3", "--tito-allowed-append-roles", "tool"] - if variant == "agentic_tool_call_multi_samples": - argv.append("--generate-multi-samples") return argv @@ -157,6 +151,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 +208,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 +244,11 @@ 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), session_server_instance_id=instance_id, ) session_server = SessionServer(server_args, backend_url=backend_url) diff --git a/tests/fast/rollout/generate_hub/test_multi_turn.py b/tests/fast/rollout/generate_hub/test_multi_turn.py index cbfad5079af..e79c1ea3e60 100644 --- a/tests/fast/rollout/generate_hub/test_multi_turn.py +++ b/tests/fast/rollout/generate_hub/test_multi_turn.py @@ -22,7 +22,7 @@ def is_agentic_variant(variant: str) -> bool: - return variant in ("agentic_tool_call_single_sample", "agentic_tool_call_multi_samples") + return variant == "agentic_tool_call" # ------------------------------------ fixtures and consts ---------------------------------------- @@ -33,14 +33,7 @@ def is_agentic_variant(variant: str) -> bool: TOKENIZER = load_tokenizer(MODEL_NAME, trust_remote_code=True) -@pytest.fixture( - params=[ - "multi_turn_single_sample", - "multi_turn_multi_samples", - "agentic_tool_call_single_sample", - "agentic_tool_call_multi_samples", - ] -) +@pytest.fixture(params=["multi_turn", "agentic_tool_call"]) def variant(request): return request.param @@ -233,41 +226,21 @@ def test_two_turns_with_tool_call(self, variant, generation_env): expected_request(S.FIRST_PROMPT_TOKEN_IDS), expected_request(S.SECOND_PROMPT_TOKEN_IDS), ] - if variant in ("multi_turn_single_sample", "agentic_tool_call_single_sample"): - full_response = S.FIRST_RESPONSE + S.FIRST_TOOL_RESPONSE + S.SECOND_RESPONSE - expected = [ - ExpectedSampleInfo( - chunks=[ - expected_chunk(S.FIRST_RESPONSE, 1), - expected_chunk(S.FIRST_TOOL_RESPONSE, 0), - expected_chunk(S.SECOND_RESPONSE, 1), - ], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=full_response, - response_length=token_len(full_response), - ), - ), - ] - else: - expected = [ - ExpectedSampleInfo( - chunks=[expected_chunk(S.FIRST_RESPONSE, 1)], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=S.FIRST_RESPONSE, - response_length=token_len(S.FIRST_RESPONSE), - ), + full_response = S.FIRST_RESPONSE + S.FIRST_TOOL_RESPONSE + S.SECOND_RESPONSE + expected = [ + ExpectedSampleInfo( + chunks=[ + expected_chunk(S.FIRST_RESPONSE, 1), + expected_chunk(S.FIRST_TOOL_RESPONSE, 0), + expected_chunk(S.SECOND_RESPONSE, 1), + ], + partial_sample=expected_partial_sample( + prompt=S.PROMPT, + response=full_response, + response_length=token_len(full_response), ), - ExpectedSampleInfo( - chunks=[expected_chunk(S.SECOND_RESPONSE, 1)], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=S.SECOND_RESPONSE, - response_length=token_len(S.SECOND_RESPONSE), - ), - ), - ] + ), + ] verify_samples(result.sample, expected) @@ -347,7 +320,7 @@ def test_max_turns_reached(self, variant, generation_env): assert _strip_pretokenized(result.requests) == [expected_openai_request(S.OPENAI_MESSAGES_FIRST_TURN)] else: assert result.requests == [expected_request(S.FIRST_PROMPT_TOKEN_IDS)] - if variant == "multi_turn_single_sample": + if variant == "multi_turn": expected = [ ExpectedSampleInfo( chunks=[ @@ -384,17 +357,14 @@ def test_prompt_exceeds_max_context_len_returns_truncated(self, variant, generat pytest.skip("TODO: implement") result = _run_generate(variant, generation_env, make_sample(prompt=SINGLE_TURN_PROMPT)) assert result.requests == [] - if variant == "multi_turn_single_sample": - expected = [ - ExpectedSampleInfo( - chunks=[], - partial_sample=expected_partial_sample( - prompt=SINGLE_TURN_PROMPT, response="", response_length=0, status=Sample.Status.TRUNCATED - ), - ) - ] - else: - expected = [] + expected = [ + ExpectedSampleInfo( + chunks=[], + partial_sample=expected_partial_sample( + prompt=SINGLE_TURN_PROMPT, response="", response_length=0, status=Sample.Status.TRUNCATED + ), + ) + ] verify_samples(result.sample, expected) @pytest.mark.parametrize( @@ -419,34 +389,21 @@ def test_second_turn_exceeds_max_context_len_returns_truncated(self, variant, ge result = _run_generate(variant, generation_env, make_sample(prompt=S.PROMPT)) assert result.requests == [expected_request(S.FIRST_PROMPT_TOKEN_IDS)] - if variant == "multi_turn_single_sample": - partial_response = S.FIRST_RESPONSE + S.FIRST_TOOL_RESPONSE - expected = [ - ExpectedSampleInfo( - chunks=[ - expected_chunk(S.FIRST_RESPONSE, 1), - expected_chunk(S.FIRST_TOOL_RESPONSE, 0), - ], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=partial_response, - response_length=token_len(partial_response), - status=Sample.Status.TRUNCATED, - ), - ), - ] - else: - expected = [ - ExpectedSampleInfo( - chunks=[expected_chunk(S.FIRST_RESPONSE, 1)], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=S.FIRST_RESPONSE, - response_length=token_len(S.FIRST_RESPONSE), - status=Sample.Status.TRUNCATED, - ), + partial_response = S.FIRST_RESPONSE + S.FIRST_TOOL_RESPONSE + expected = [ + ExpectedSampleInfo( + chunks=[ + expected_chunk(S.FIRST_RESPONSE, 1), + expected_chunk(S.FIRST_TOOL_RESPONSE, 0), + ], + partial_sample=expected_partial_sample( + prompt=S.PROMPT, + response=partial_response, + response_length=token_len(partial_response), + status=Sample.Status.TRUNCATED, ), - ] + ), + ] verify_samples(result.sample, expected) @pytest.mark.parametrize( @@ -497,57 +454,25 @@ def test_three_turns_with_sequential_tool_calls(self, variant, generation_env): expected_request(S.SECOND_PROMPT_TOKEN_IDS), expected_request(S.THIRD_PROMPT_TOKEN_IDS), ] - if variant in ("multi_turn_single_sample", "agentic_tool_call_single_sample"): - full_response = ( - S.FIRST_RESPONSE - + S.FIRST_TOOL_RESPONSE - + S.SECOND_RESPONSE - + S.SECOND_TOOL_RESPONSE - + S.THIRD_RESPONSE - ) - expected = [ - ExpectedSampleInfo( - chunks=[ - expected_chunk(S.FIRST_RESPONSE, 1), - expected_chunk(S.FIRST_TOOL_RESPONSE, 0), - expected_chunk(S.SECOND_RESPONSE, 1), - expected_chunk(S.SECOND_TOOL_RESPONSE, 0), - expected_chunk(S.THIRD_RESPONSE, 1), - ], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=full_response, - response_length=token_len(full_response), - ), - ), - ] - else: - expected = [ - ExpectedSampleInfo( - chunks=[expected_chunk(S.FIRST_RESPONSE, 1)], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=S.FIRST_RESPONSE, - response_length=token_len(S.FIRST_RESPONSE), - ), - ), - ExpectedSampleInfo( - chunks=[expected_chunk(S.SECOND_RESPONSE, 1)], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=S.SECOND_RESPONSE, - response_length=token_len(S.SECOND_RESPONSE), - ), - ), - ExpectedSampleInfo( - chunks=[expected_chunk(S.THIRD_RESPONSE, 1)], - partial_sample=expected_partial_sample( - prompt=S.PROMPT, - response=S.THIRD_RESPONSE, - response_length=token_len(S.THIRD_RESPONSE), - ), + full_response = ( + S.FIRST_RESPONSE + S.FIRST_TOOL_RESPONSE + S.SECOND_RESPONSE + S.SECOND_TOOL_RESPONSE + S.THIRD_RESPONSE + ) + expected = [ + ExpectedSampleInfo( + chunks=[ + expected_chunk(S.FIRST_RESPONSE, 1), + expected_chunk(S.FIRST_TOOL_RESPONSE, 0), + expected_chunk(S.SECOND_RESPONSE, 1), + expected_chunk(S.SECOND_TOOL_RESPONSE, 0), + expected_chunk(S.THIRD_RESPONSE, 1), + ], + partial_sample=expected_partial_sample( + prompt=S.PROMPT, + response=full_response, + response_length=token_len(full_response), ), - ] + ), + ] verify_samples(result.sample, expected) @@ -558,6 +483,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, } } ], @@ -565,9 +494,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, @@ -645,7 +572,7 @@ def process_fn(prompt: str) -> ProcessResult: assert len(sample.tokens) - 1 == second_routed_experts.shape[0] -_AGENTIC_VARIANTS = ["agentic_tool_call_single_sample", "agentic_tool_call_multi_samples"] +_AGENTIC_VARIANTS = ["agentic_tool_call"] _AGENT_METADATA = {"reward": 1.0, "exit_status": "Submitted", "eval_report": {"passed": True}} diff --git a/tests/fast/rollout/generate_hub/test_single_turn.py b/tests/fast/rollout/generate_hub/test_single_turn.py index 12a51385774..b16a7613d12 100644 --- a/tests/fast/rollout/generate_hub/test_single_turn.py +++ b/tests/fast/rollout/generate_hub/test_single_turn.py @@ -30,7 +30,7 @@ DEFAULT_MAX_NEW_TOKENS = SAMPLING_PARAMS["max_new_tokens"] -@pytest.fixture(params=["old_sglang_rollout", "single_turn", "multi_turn_single_sample", "multi_turn_multi_samples"]) +@pytest.fixture(params=["old_sglang_rollout", "single_turn", "multi_turn"]) def variant(request): return request.param @@ -49,9 +49,9 @@ def expected_request( "sampling_params": sampling_params or SAMPLING_PARAMS, "return_logprob": True, } - if variant in ("single_turn", "multi_turn_single_sample", "multi_turn_multi_samples") or return_routed_experts: + if variant in ("single_turn", "multi_turn") or return_routed_experts: result["return_routed_experts"] = return_routed_experts - if variant in ("single_turn", "multi_turn_single_sample", "multi_turn_multi_samples") or return_indexer_topk: + if variant in ("single_turn", "multi_turn") or return_indexer_topk: result["return_indexer_topk"] = return_indexer_topk if image_data is not None: result["image_data"] = image_data @@ -85,11 +85,7 @@ def expected_sample( ) -> Sample: actual_response_length = response_length if response_length is not None else len(RESPONSE_TOKENS) if isinstance(loss_mask, _Unset): - loss_mask = ( - [1] * actual_response_length - if variant in ("multi_turn_single_sample", "multi_turn_multi_samples") - else None - ) + loss_mask = [1] * actual_response_length if variant == "multi_turn" else None return Sample( group_index=None, @@ -143,7 +139,7 @@ def test_basic_generation(self, variant, generation_env): class TestResumedSingleTurn: def test_two_consecutive_calls_on_same_sample(self, variant, generation_env): - if variant in ("multi_turn_single_sample", "multi_turn_multi_samples"): + if variant == "multi_turn": pytest.skip("not tested yet") partial_text = "\\boxed" partial_tokens = [59, 79075] @@ -279,7 +275,7 @@ def test_allowed_statuses(self, variant, generation_env, status): @pytest.mark.parametrize("status", [Sample.Status.COMPLETED, Sample.Status.TRUNCATED]) def test_rejected_statuses(self, variant, generation_env, status): - if variant in ("multi_turn_single_sample", "multi_turn_multi_samples"): + if variant == "multi_turn": pytest.skip("not tested yet") with pytest.raises(AssertionError): _run_generate(variant, generation_env, _make_sample(status=status)) @@ -298,7 +294,7 @@ def test_sampling_params_passed_through(self, variant, generation_env): class TestBoundaryConditions: def test_max_new_tokens_zero_returns_truncated(self, variant, generation_env): - if variant in ("multi_turn_single_sample", "multi_turn_multi_samples"): + if variant == "multi_turn": pytest.skip("not tested yet") existing_tokens = [1, 2, 3, 4, 5, 6, 7] + list(range(100, 110)) sample = _make_sample(tokens=existing_tokens, response="x" * 10, response_length=10) @@ -319,11 +315,9 @@ def test_max_new_tokens_zero_returns_truncated(self, variant, generation_env): def test_prompt_exceeds_max_context_len_returns_truncated(self, variant, generation_env): if variant == "old_sglang_rollout": pytest.skip("old_sglang_rollout does not support rollout_max_context_len") - if variant == "multi_turn_multi_samples": - pytest.skip("multi_turn_multi_samples returns empty list when first turn fails") result = _run_generate(variant, generation_env) assert result.requests == [] - tokens = PROMPT_TOKENS if variant in ("multi_turn_single_sample", "multi_turn_multi_samples") else [] + tokens = PROMPT_TOKENS if variant == "multi_turn" else [] assert listify(result.sample) == [ expected_sample( variant, @@ -333,7 +327,7 @@ def test_prompt_exceeds_max_context_len_returns_truncated(self, variant, generat rollout_log_probs=None, status=Sample.Status.TRUNCATED, prompt_tokens=0, - loss_mask=None if variant == "multi_turn_single_sample" else _UNSET, + loss_mask=None if variant == "multi_turn" else _UNSET, ) ] @@ -363,11 +357,9 @@ def test_moderate_length_input_adjusts_max_new_tokens(self, variant, generation_ def test_adjusted_max_new_tokens_zero_returns_truncated(self, variant, generation_env): if variant == "old_sglang_rollout": pytest.skip("old_sglang_rollout does not support rollout_max_context_len") - if variant == "multi_turn_multi_samples": - pytest.skip("multi_turn_multi_samples returns empty list when first turn fails") result = _run_generate(variant, generation_env) assert result.requests == [] - tokens = PROMPT_TOKENS if variant == "multi_turn_single_sample" else [] + tokens = PROMPT_TOKENS if variant == "multi_turn" else [] assert listify(result.sample) == [ expected_sample( variant, @@ -377,7 +369,7 @@ def test_adjusted_max_new_tokens_zero_returns_truncated(self, variant, generatio rollout_log_probs=None, status=Sample.Status.TRUNCATED, prompt_tokens=0, - loss_mask=None if variant == "multi_turn_single_sample" else _UNSET, + loss_mask=None if variant == "multi_turn" else _UNSET, ) ] @@ -398,7 +390,7 @@ def test_empty_response(self, variant, generation_env): class TestMultimodal: @pytest.mark.parametrize("generation_env", [{"args_kwargs": {"model_name": VLM_MODEL_NAME}}], indirect=True) def test_multimodal_inputs_processed(self, variant, generation_env): - if variant in ("multi_turn_single_sample", "multi_turn_multi_samples"): + if variant == "multi_turn": pytest.skip("not tested yet") test_image = Image.new("RGB", (64, 64), color="red") multimodal_inputs = {"images": [test_image]} 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 9a40f126dec..e5253528ef3 100644 --- a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py +++ b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py @@ -1,95 +1,25 @@ -"""Tests for compute_samples_from_openai_records and TITO multi-turn merge workflow. - -Validates the contract between session records, sample construction, -and merge_samples — the core of the TITO (Token In Token Out) pipeline. +"""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 (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 -from unittest.mock import MagicMock import pytest -from miles.rollout.generate_utils.openai_endpoint_utils import ( - OpenAIEndpointTracer, - compute_samples_from_openai_records, -) -from miles.rollout.generate_utils.sample_utils import merge_samples -from miles.rollout.session.types import SessionRecord +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 -# ── helpers ────────────────────────────────────────────────────────── - -_ARGS = SimpleNamespace() - - -def _mock_tokenizer(): - tok = MagicMock() - tok.decode = lambda ids: "".join(f"[{i}]" for i in ids) - 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], - output_log_probs: list[float] | None = None, - finish_reason: str = "stop", - cached_tokens: int | None = None, - prompt_tokens: int | 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. - """ - if output_log_probs is None: - output_log_probs = [-0.1 * (i + 1) for i in range(len(output_token_ids))] - - output_token_logprobs = [(lp, tid) for tid, lp in zip(output_token_ids, output_log_probs, strict=True)] - logprobs_content = [ - {"logprob": lp, "token": f"t{tid}"} for tid, lp in zip(output_token_ids, output_log_probs, strict=True) - ] - meta_info = { - "output_token_logprobs": output_token_logprobs, - "completion_tokens": len(output_token_ids), - } - if cached_tokens is not None: - meta_info["cached_tokens"] = cached_tokens - if prompt_tokens is not None: - meta_info["prompt_tokens"] = prompt_tokens - return SessionRecord( - timestamp=0.0, - method="POST", - path="/v1/chat/completions", - status_code=200, - request={"messages": [{"role": "user", "content": "hello"}], "input_ids": prompt_token_ids}, - response={ - "choices": [ - { - "message": {"role": "assistant", "content": "response"}, - "finish_reason": finish_reason, - "logprobs": {"content": logprobs_content}, - "meta_info": meta_info, - } - ] - }, - ) - @pytest.mark.asyncio async def test_create_reads_session_server_instance_id_from_args(monkeypatch): @@ -133,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]] = [] @@ -141,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) @@ -156,11 +91,11 @@ 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(), 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/") @@ -169,604 +104,132 @@ async def fake_post(url: str, payload: dict, action: str = "post"): assert len(chosen_ports) > 1 -# ── test: compute_samples_from_openai_records ──────────────────────── - - -class TestComputeSamplesFromRecords: - def test_single_record_builds_correct_sample(self): - tok = _mock_tokenizer() - record = _make_record( - prompt_token_ids=[1, 2, 3], - 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) - - assert len(samples) == 1 - s = samples[0] - assert s.tokens == [1, 2, 3, 10, 11] - assert s.rollout_log_probs == [-0.5, -0.6] - assert s.response_length == 2 - assert s.loss_mask == [1, 1] - assert s.status == Sample.Status.COMPLETED - - def test_multiple_records_produce_multiple_samples(self): - tok = _mock_tokenizer() - records = [ - _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) - - assert len(samples) == 2 - assert samples[0].tokens == [1, 2, 10] - assert samples[1].tokens == [1, 2, 10, 20, 30] - - def test_finish_reason_length_gives_truncated(self): - tok = _mock_tokenizer() - record = _make_record( - prompt_token_ids=[1, 2], - output_token_ids=[10], - finish_reason="length", - ) - input_sample = _make_input_sample() - - samples = compute_samples_from_openai_records(_ARGS, input_sample, [record], tok) - - assert samples[0].status == Sample.Status.TRUNCATED - - -# ── test: multi-turn prefix chain (merge_samples integration) ──────── - - -class TestMultiTurnPrefixChain: - """Validate that session records from a well-behaved multi-turn - conversation satisfy the prefix chain required by merge_samples. - - The contract: for consecutive records r[i] and r[i+1], - r[i+1].prompt_token_ids must start with r[i].prompt_token_ids + r[i].output_token_ids. - This is because the agent includes the previous response in the next prompt. - """ - - def test_two_turn_merge_succeeds(self): - """Normal two-turn conversation: samples merge without error.""" - tok = _mock_tokenizer() - - # Turn 1: prompt=[1,2,3], model outputs [10,11] - # Turn 2: prompt=[1,2,3, 10,11, 20,21], model outputs [30,31] - # (tokens 20,21 are the tool/observation tokens added by the environment) - records = [ - _make_record( - prompt_token_ids=[1, 2, 3], - output_token_ids=[10, 11], - output_log_probs=[-0.1, -0.2], - ), - _make_record( - prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], - output_token_ids=[30, 31], - output_log_probs=[-0.3, -0.4], - ), - ] - input_sample = _make_input_sample() - - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) - merged = merge_samples(samples, tok) - - assert merged.tokens == [1, 2, 3, 10, 11, 20, 21, 30, 31] - assert merged.response_length == 2 + 2 + 2 # resp1 + obs + resp2 - assert merged.loss_mask == [1, 1, 0, 0, 1, 1] - assert merged.status == Sample.Status.COMPLETED - - def test_three_turn_merge_succeeds(self): - """Three-turn conversation: prefix chain holds across all turns.""" - tok = _mock_tokenizer() - - records = [ - _make_record( - prompt_token_ids=[1, 2], - output_token_ids=[10], - output_log_probs=[-0.1], - ), - _make_record( - prompt_token_ids=[1, 2, 10, 20], - output_token_ids=[30], - output_log_probs=[-0.2], - ), - _make_record( - prompt_token_ids=[1, 2, 10, 20, 30, 40], - output_token_ids=[50], - output_log_probs=[-0.3], - ), - ] - input_sample = _make_input_sample() - - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) - merged = merge_samples(samples, tok) - - assert merged.tokens == [1, 2, 10, 20, 30, 40, 50] - assert merged.response_length == 1 + 1 + 1 + 1 + 1 # 3 responses + 2 obs - - def test_prefix_mismatch_raises(self): - """When the prefix chain is broken, merge_samples must fail.""" - tok = _mock_tokenizer() - - # Turn 2's prompt does NOT start with turn 1's full tokens - records = [ - _make_record( - prompt_token_ids=[1, 2, 3], - output_token_ids=[10, 11], - ), - _make_record( - prompt_token_ids=[1, 2, 3, 99, 99, 20, 21], # 99,99 != 10,11 - output_token_ids=[30, 31], - ), - ] - input_sample = _make_input_sample() - - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) - - with pytest.raises(AssertionError, match="b.tokens must start with a.tokens"): - merge_samples(samples, tok) +# ── collect_samples client behavior ── - def test_two_turn_merge_propagates_teacher_log_probs(self): - """OPD teacher_log_probs merge like rollout_log_probs: per-turn values - concatenated with zeros over the injected observation span.""" - tok = _mock_tokenizer() - records = [ - _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11], output_log_probs=[-0.1, -0.2]), - _make_record( - prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], - output_token_ids=[30, 31], - output_log_probs=[-0.3, -0.4], - ), - ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) - - # OPD attaches per-response-token teacher log-probs to each turn's sample. - samples[0].teacher_log_probs = [-1.0, -1.1] - samples[1].teacher_log_probs = [-1.2, -1.3] - - merged = merge_samples(samples, tok) - - # resp1 (2) + obs (2 zeros) + resp2 (2) - assert merged.teacher_log_probs == [-1.0, -1.1, 0.0, 0.0, -1.2, -1.3] - assert len(merged.teacher_log_probs) == merged.response_length - merged.validate() # the new teacher_log_probs length assertion must hold - - def test_two_turn_merge_propagates_opd_student_top_logprobs_metadata(self): - """Top-k OPD student top-logprobs are per-token metadata, not equal metadata.""" - tok = _mock_tokenizer() - - records = [ - _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11], output_log_probs=[-0.1, -0.2]), - _make_record( - prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], - output_token_ids=[30, 31], - output_log_probs=[-0.3, -0.4], - ), - ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) - - turn_0_top_logprobs = [[[-0.1, 101]], [[-0.2, 102]]] - turn_1_top_logprobs = [[[-0.3, 103]], [[-0.4, 104]]] - samples[0].metadata = { - "opd_student_top_logprobs": turn_0_top_logprobs, - "shared_metadata": "same", - } - samples[1].metadata = { - "opd_student_top_logprobs": turn_1_top_logprobs, - "shared_metadata": "same", - } - - merged = merge_samples(samples, tok) - - assert merged.metadata["shared_metadata"] == "same" - assert merged.metadata["opd_student_top_logprobs"] == [ - *turn_0_top_logprobs, - [], - [], - *turn_1_top_logprobs, - ] - assert len(merged.metadata["opd_student_top_logprobs"]) == merged.response_length +def _tracer() -> OpenAIEndpointTracer: + return OpenAIEndpointTracer(router_url="http://127.0.0.1:12345", session_id="sid-1") - def test_two_turn_merge_teacher_log_probs_none_stays_none(self): - """Non-OPD runs leave teacher_log_probs unset; merge must keep it None.""" - tok = _mock_tokenizer() - records = [ - _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) +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) - merged = merge_samples(samples, tok) - assert merged.teacher_log_probs is None +class _CollectCalls: + """Patches the two HTTP primitives collect_samples uses and records order.""" - def test_merge_raises_on_teacher_log_probs_length_mismatch(self): - """validate() guards teacher_log_probs length (surfaced via merge_samples).""" - tok = _mock_tokenizer() + def __init__(self, monkeypatch, *, post_outcome, delete_outcome=None): + self.calls: list[str] = [] - records = [ - _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) + async def fake_post_bytes(url, payload, *, timeout): + self.calls.append(f"POST {url}") + assert payload == {"max_seq_len": 7} + if isinstance(post_outcome, Exception): + raise post_outcome + return post_outcome - samples[0].teacher_log_probs = [-1.0] # length 1 != response_length 2 + 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 {} - with pytest.raises(AssertionError, match="teacher_log_probs length"): - merge_samples(samples, tok) + 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) -# ── test: TITO trailing token trimming ──────────────────────────────── +@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(), max_seq_len=7) -STOP = 99 # stands for <|observation|> stop token + 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} -class TestTITOTrailingTokenTrim: - """Validate trailing-token trimming via ``accumulated_token_ids``. +@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(), max_seq_len=7) + assert calls.calls[-1] == "DELETE http://127.0.0.1:12345/sessions/sid-1" - Worked example — agentic tool-call retries - ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - An agent makes three turns. The model's tool call fails to parse on - turns 1 and 2, so the agent feeds back an error and retries. +@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(), max_seq_len=7) + assert calls.calls[-1] == "DELETE http://127.0.0.1:12345/sessions/sid-1" - The session server sees three request/response pairs (records). Each - record's response is an independent inference re-stitched via - pretokenized prefix reuse:: - record 0 prompt_token_ids: [<|sys|>, aaa, <|user|>, bbb, <|asst|>] - output_token_ids: [ccc, <|obs|>] ← model stopped with <|obs|> +@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(), max_seq_len=7) + assert len(result.samples) == 1 - record 1 prompt_token_ids: [<|sys|>, aaa, <|user|>, bbb, <|asst|>, ccc, <|sys|>, ddd, <|asst|>] - output_token_ids: [eee, <|obs|>] - record 2 prompt_token_ids: [..., eee, <|sys|>, fff, <|asst|>] - output_token_ids: [ggg, <|obs|>] +# ── post_bytes_no_retry primitive ── - ``accumulated_token_ids`` = record 2's prompt + output:: - [<|sys|>, aaa, <|user|>, bbb, <|asst|>, ccc, <|sys|>, ddd, - <|asst|>, eee, <|sys|>, fff, <|asst|>, ggg, <|obs|>] +class _FakeResponse: + def __init__(self, status_code: int, content: bytes = b"", text: str = ""): + self.status_code = status_code + self.content = content + self.text = text - Note: there is NO ``<|obs|>`` between ``ccc`` and ``<|sys|>`` in the - accumulated sequence — the stop token the model emitted at turn 1 was - consumed by the chat template when rendering turn 2's prompt. - The algorithm walks ``accumulated_token_ids`` with a cursor:: +class _FakeClient: + def __init__(self, responses): + self.responses = list(responses) + self.post_count = 0 - Record 0: cursor = len(prompt_0) → points to "ccc" - Match output [ccc, <|obs|>] against accumulated[cursor:]: - ccc OK, <|obs|> MISMATCH (accumulated has <|sys|> here) - → trim_count=1, strip <|obs|>; cursor advances past "ccc" + async def post(self, url, json=None): + self.post_count += 1 + outcome = self.responses.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome - Record 1: cursor = len(prompt_1) → points to "eee" - Match [eee, <|obs|>]: eee OK, <|obs|> MISMATCH - → trim_count=1; cursor advances past "eee" - Record 2: cursor = len(prompt_2) → points to "ggg" - Match [ggg, <|obs|>]: ggg OK, <|obs|> OK (last turn) - → trim_count=0; cursor reaches end +@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 - Result: three Samples with output tokens [ccc], [eee], [ggg, <|obs|>], - each carrying original per-turn logprobs. - The tests below encode this example (and variants) with concrete - token IDs. We use ``STOP = 99`` to represent ``<|observation|>``. - """ +@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 - def test_three_turn_trim_trailing_stop_tokens(self): - """Three-turn retry: non-final turns have 1 trailing stop token trimmed.""" - tok = _mock_tokenizer() - # prompt: [1, 2, 3] output: [10, STOP] - # prompt: [1, 2, 3, 10, 4, 5, 6] output: [20, STOP] - # prompt: [1, 2, 3, 10, 4, 5, 6, 20, 7, 8, 9] output: [30, STOP] - # accumulated (no intermediate STOPs): - # [1, 2, 3, 10, 4, 5, 6, 20, 7, 8, 9, 30, STOP] - records = [ - _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, STOP]), - _make_record(prompt_token_ids=[1, 2, 3, 10, 4, 5, 6], output_token_ids=[20, STOP]), - _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, - max_trim_tokens=1, - ) - - assert len(samples) == 3 - # Turn 0: [10, STOP] → trim 1 → response_length=1 - assert samples[0].tokens == [1, 2, 3, 10] - assert samples[0].response_length == 1 - # Turn 1: [20, STOP] → trim 1 → response_length=1 - assert samples[1].tokens == [1, 2, 3, 10, 4, 5, 6, 20] - assert samples[1].response_length == 1 - # Turn 2 (last): [30, STOP] → trim 0 → response_length=2 - assert samples[2].tokens == [1, 2, 3, 10, 4, 5, 6, 20, 7, 8, 9, 30, STOP] - assert samples[2].response_length == 2 - - def test_no_trim_when_no_trailing_stop(self): - """When output tokens fully match accumulated, trim_count=0 for all turns.""" - tok = _mock_tokenizer() - - # Two turns, no trailing stop tokens — output aligns perfectly - # prompt: [1, 2] output: [10, 11] - # prompt: [1, 2, 10, 11, 3, 4] output: [20, 21] - # accumulated: [1, 2, 10, 11, 3, 4, 20, 21] - records = [ - _make_record(prompt_token_ids=[1, 2], output_token_ids=[10, 11]), - _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, - max_trim_tokens=1, - ) - - assert len(samples) == 2 - assert samples[0].tokens == [1, 2, 10, 11] - assert samples[0].response_length == 2 - assert samples[1].tokens == [1, 2, 10, 11, 3, 4, 20, 21] - assert samples[1].response_length == 2 - - def test_single_turn_no_trim(self): - """Single turn: last turn never trims, even with accumulated_token_ids.""" - tok = _mock_tokenizer() - - records = [ - _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, - max_trim_tokens=1, - ) - - assert len(samples) == 1 - assert samples[0].tokens == [1, 2, 3, 10, 11, STOP] - assert samples[0].response_length == 3 - - def test_no_accumulated_skips_trimming(self): - """Without accumulated_token_ids, no trimming is performed at all.""" - tok = _mock_tokenizer() - - records = [ - _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, - ) - - assert len(samples) == 2 - # No trimming — STOP is kept for both turns - assert samples[0].tokens == [1, 2, 10, STOP] - assert samples[0].response_length == 2 - assert samples[1].tokens == [1, 2, 10, STOP, 3, 4, 20, STOP] - assert samples[1].response_length == 2 - - def test_trim_exceeding_max_raises(self): - """If trailing tokens exceed max_trim_tokens, assert fires.""" - tok = _mock_tokenizer() - - # Output has 2 trailing tokens that don't match, but max_trim_tokens=1 - records = [ - _make_record(prompt_token_ids=[1, 2], output_token_ids=[10, STOP, STOP]), - _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, - max_trim_tokens=1, - ) - - def test_cursor_covers_entire_accumulated(self): - """After processing all records, cursor must equal len(accumulated).""" - tok = _mock_tokenizer() - - # accumulated is shorter than what records imply — cursor won't reach end - records = [ - _make_record(prompt_token_ids=[1, 2], output_token_ids=[10, STOP]), - _make_record(prompt_token_ids=[1, 2, 10, 3], output_token_ids=[20]), - ] - # 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, - max_trim_tokens=1, - ) - - -# ── test: thinking token issue (documents known failure mode) ──────── - - -class TestThinkingTokenPrefixBreak: - """Documents the known issue where model-generated ... - tokens break the prefix chain. - - When a model (e.g. Qwen3) generates reasoning before - the actual response, agents strip the thinking content from conversation - history. This causes the next turn's prompt to not include the thinking - tokens, breaking the prefix assumption in merge_samples. - - This is a MODEL-LEVEL issue — the fix should be at the model/serving - config level (disable thinking mode), not in the merge logic. - """ - - THINK_TOKEN = 151667 # in Qwen3 - END_THINK_TOKEN = 151668 # in Qwen3 - NEWLINE_TOKEN = 198 # \n - - def test_thinking_tokens_break_prefix_chain(self): - """Demonstrates the failure: model outputs ..., but the agent - strips it from history, so the next prompt doesn't include those tokens.""" - tok = _mock_tokenizer() - - # Turn 1: model generates \nreasoning\n\n then actual response - thinking_tokens = [ - self.THINK_TOKEN, - self.NEWLINE_TOKEN, - 42, - 43, - self.NEWLINE_TOKEN, - self.END_THINK_TOKEN, - self.NEWLINE_TOKEN, - ] - response_tokens = [10, 11] - all_output = thinking_tokens + response_tokens - - records = [ - _make_record( - prompt_token_ids=[1, 2, 3], - output_token_ids=all_output, - ), - # Turn 2: agent only included the actual response [10, 11] in history - # (stripped thinking tokens), plus observation [20, 21] - _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) - - # 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] - # sample[1] does NOT start with sample[0] — prefix chain broken - with pytest.raises(AssertionError, match="b.tokens must start with a.tokens"): - merge_samples(samples, tok) - - def test_no_thinking_tokens_prefix_chain_holds(self): - """When thinking is disabled, the same conversation merges fine.""" - tok = _mock_tokenizer() - - # Same conversation but model output has no thinking prefix - records = [ - _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) - merged = merge_samples(samples, tok) - - assert merged.tokens == [1, 2, 3, 10, 11, 20, 21, 30, 31] - - -# ── test: prefix cache info population ──────────────────────────────── - - -class TestPrefixCacheInfo: - """Validate that prefix cache statistics from meta_info are collected.""" - - def test_single_record_with_cache_stats(self): - """cached_tokens and prompt_tokens from meta_info populate prefix_cache_info.""" - tok = _mock_tokenizer() - record = _make_record( - prompt_token_ids=[1, 2, 3], - output_token_ids=[10, 11], - cached_tokens=2, - prompt_tokens=3, - ) - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, [record], tok) - - assert samples[0].prefix_cache_info.cached_tokens == 2 - assert samples[0].prefix_cache_info.total_prompt_tokens == 3 - - def test_multi_turn_cache_stats_accumulate_after_merge(self): - """After merge_samples, prefix_cache_info sums across turns.""" - tok = _mock_tokenizer() - records = [ - _make_record( - prompt_token_ids=[1, 2, 3], - output_token_ids=[10, 11], - output_log_probs=[-0.1, -0.2], - cached_tokens=0, - prompt_tokens=3, - ), - _make_record( - prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], - output_token_ids=[30, 31], - output_log_probs=[-0.3, -0.4], - cached_tokens=5, - prompt_tokens=7, - ), - ] - input_sample = _make_input_sample() - samples = compute_samples_from_openai_records(_ARGS, input_sample, records, tok) - merged = merge_samples(samples, tok) - - assert merged.prefix_cache_info.cached_tokens == 0 + 5 - assert merged.prefix_cache_info.total_prompt_tokens == 3 + 7 - assert merged.prefix_cache_info.prefix_cache_hit_rate == 5 / 10 - - def test_missing_cache_fields_default_to_zero(self): - """Records without cached_tokens/prompt_tokens give zero prefix_cache_info (regression).""" - tok = _mock_tokenizer() - record = _make_record( - 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) - - assert samples[0].prefix_cache_info.cached_tokens == 0 - assert samples[0].prefix_cache_info.total_prompt_tokens == 0 +@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/inference_rollout/integration/test_agent_metadata.py b/tests/fast/rollout/inference_rollout/integration/test_agent_metadata.py index 7de04cbbe3d..78f82bec5ca 100644 --- a/tests/fast/rollout/inference_rollout/integration/test_agent_metadata.py +++ b/tests/fast/rollout/inference_rollout/integration/test_agent_metadata.py @@ -9,7 +9,7 @@ TWO_TURN_DATA_ROWS = [{"input": [{"role": "user", "content": TwoTurnStub.USER_QUESTION}], "label": "2008"}] -_AGENTIC_VARIANTS = ["agentic_tool_call_single_sample", "agentic_tool_call_multi_samples"] +_AGENTIC_VARIANTS = ["agentic_tool_call"] _METADATA_RM_EXTRA_ARGV = [ "--rollout-batch-size", diff --git a/tests/fast/rollout/inference_rollout/integration/test_multi_turn.py b/tests/fast/rollout/inference_rollout/integration/test_multi_turn.py index c41d7139919..16236d69065 100644 --- a/tests/fast/rollout/inference_rollout/integration/test_multi_turn.py +++ b/tests/fast/rollout/inference_rollout/integration/test_multi_turn.py @@ -11,12 +11,7 @@ TWO_TURN_DATA_ROWS = [{"input": [{"role": "user", "content": TwoTurnStub.USER_QUESTION}], "label": "2008"}] -_VARIANT_NAMES = [ - "multi_turn_single_sample", - "multi_turn_multi_samples", - "agentic_tool_call_single_sample", - "agentic_tool_call_multi_samples", -] +_VARIANT_NAMES = ["multi_turn", "agentic_tool_call"] _BASE_EXTRA_ARGV = [ "--rollout-batch-size", @@ -60,38 +55,26 @@ def test_rollout(rollout_env, variant, test_type): def _verify_samples(variant: str, samples: list[Any]): - is_multi_samples = variant in ("multi_turn_multi_samples", "agentic_tool_call_multi_samples") - - if is_multi_samples: + if variant == "agentic_tool_call": if len(samples) > 0 and isinstance(samples[0], list): - # Train mode: list[list[Sample]], grouped by prompt + # Train mode: list[list[Sample]] — one singleton list (merged TITO sample) per generate assert len(samples) == 2, f"n_samples_per_prompt=2, so group should have 2 samples, got {len(samples)}" for group_sample in samples: - assert isinstance(group_sample, list), "multi_samples variant should return list[Sample] per generate" - _verify_group_samples(group_sample) + assert isinstance(group_sample, list), "agentic_tool_call returns list[Sample] per generate" + assert len(group_sample) == 1, "linear trajectory merges into exactly one sample" + _verify_sample(group_sample[0]) else: - # Eval mode: list[Sample], flattened - # n_samples_per_eval_prompt=2, and each generate returns 2 turns, so 2*2=4 samples - assert ( - len(samples) == 4 - ), f"n_samples_per_eval_prompt=2, each generate returns 2 turns, so should have 4 samples, got {len(samples)}" - # Group samples by prompt (every 2 samples form a group) - group_samples_list = [samples[i : i + 2] for i in range(0, len(samples), 2)] - for group_samples in group_samples_list: - _verify_group_samples(group_samples) + # Eval mode: list[Sample], flattened (one merged sample per generate) + assert len(samples) == 2, f"n_samples_per_eval_prompt=2, so should have 2 samples, got {len(samples)}" + for sample in samples: + _verify_sample(sample) else: assert len(samples) == 2, f"n_samples_per_prompt=2, so group should have 2 samples, got {len(samples)}" for sample in samples: - assert isinstance(sample, Sample), "single_sample variant should return Sample, not list" + assert isinstance(sample, Sample), "multi_turn returns a scalar Sample per generate" _verify_sample(sample) -def _verify_group_samples(group_samples: list[Sample], expected_count: int = 2): - assert len(group_samples) == expected_count, f"Group should have {expected_count} samples (one per turn)" - for i, sample in enumerate(group_samples): - _verify_sample(sample, expect_answer=(i == len(group_samples) - 1)) - - def _verify_sample(sample: Sample, expected_reward: float = 1.0, expect_answer: bool = True): assert sample.status == Sample.Status.COMPLETED assert sample.reward == expected_reward, f"Sample should have reward={expected_reward}" @@ -101,13 +84,8 @@ def _verify_sample(sample: Sample, expected_reward: float = 1.0, expect_answer: async def _simple_reward_function(args, samples: Sample | list[Sample]) -> float | list[float]: if isinstance(samples, list): - # For multi_samples variants, use the last sample's reward - if getattr(args, "generate_multi_samples", False): - return [_check_reward(samples[-1])] * len(samples) - else: - return [_check_reward(sample) for sample in samples] - else: - return _check_reward(samples) + return [_check_reward(sample) for sample in samples] + return _check_reward(samples) def _check_reward(sample: Sample) -> float: diff --git a/tests/fast/rollout/session/__init__.py b/tests/fast/rollout/session/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/rollout/session/test_samples.py b/tests/fast/rollout/session/test_samples.py new file mode 100644 index 00000000000..20a4951973d --- /dev/null +++ b/tests/fast/rollout/session/test_samples.py @@ -0,0 +1,655 @@ +"""Tests for compute_samples_from_openai_records and TITO multi-turn merge workflow. + +Validates the contract between session records, sample construction, +and merge_samples — the core of the TITO (Token In Token Out) pipeline. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from miles.rollout.generate_utils.sample_utils import merge_samples +from miles.rollout.session.samples.merge import compute_samples_from_openai_records +from miles.rollout.session.types import SessionRecord +from miles.utils.types import Sample + +# ── helpers ────────────────────────────────────────────────────────── + +_ARGS = SimpleNamespace() + + +def _mock_tokenizer(): + tok = MagicMock() + tok.decode = lambda ids: "".join(f"[{i}]" for i in ids) + return tok + + +def _make_record( + prompt_token_ids: list[int], + output_token_ids: list[int], + output_log_probs: list[float] | None = None, + 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))] + + output_token_logprobs = [(lp, tid) for tid, lp in zip(output_token_ids, output_log_probs, strict=True)] + logprobs_content = [ + {"logprob": lp, "token": f"t{tid}"} for tid, lp in zip(output_token_ids, output_log_probs, strict=True) + ] + meta_info = { + "output_token_logprobs": output_token_logprobs, + "completion_tokens": len(output_token_ids), + } + if cached_tokens is not None: + 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", + path="/v1/chat/completions", + status_code=200, + request={"messages": [{"role": "user", "content": "hello"}], "input_ids": prompt_token_ids}, + response={ + "choices": [ + { + "message": {"role": "assistant", "content": "response"}, + "finish_reason": finish_reason, + "logprobs": {"content": logprobs_content}, + "meta_info": meta_info, + } + ] + }, + ) + + +# ── test: compute_samples_from_openai_records ──────────────────────── + + +class TestComputeSamplesFromRecords: + def test_single_record_builds_correct_sample(self): + tok = _mock_tokenizer() + record = _make_record( + prompt_token_ids=[1, 2, 3], + output_token_ids=[10, 11], + output_log_probs=[-0.5, -0.6], + ) + + samples = compute_samples_from_openai_records(_ARGS, [record], tok) + + assert len(samples) == 1 + s = samples[0] + assert s.tokens == [1, 2, 3, 10, 11] + assert s.rollout_log_probs == [-0.5, -0.6] + assert s.response_length == 2 + assert s.loss_mask == [1, 1] + assert s.status == Sample.Status.COMPLETED + + def test_multiple_records_produce_multiple_samples(self): + tok = _mock_tokenizer() + records = [ + _make_record(prompt_token_ids=[1, 2], output_token_ids=[10]), + _make_record(prompt_token_ids=[1, 2, 10, 20], output_token_ids=[30]), + ] + + samples = compute_samples_from_openai_records(_ARGS, records, tok) + + assert len(samples) == 2 + assert samples[0].tokens == [1, 2, 10] + assert samples[1].tokens == [1, 2, 10, 20, 30] + + def test_finish_reason_length_gives_truncated(self): + tok = _mock_tokenizer() + record = _make_record( + prompt_token_ids=[1, 2], + output_token_ids=[10], + finish_reason="length", + ) + + samples = compute_samples_from_openai_records(_ARGS, [record], tok) + + assert samples[0].status == Sample.Status.TRUNCATED + + +# ── test: multi-turn prefix chain (merge_samples integration) ──────── + + +class TestMultiTurnPrefixChain: + """Validate that session records from a well-behaved multi-turn + conversation satisfy the prefix chain required by merge_samples. + + The contract: for consecutive records r[i] and r[i+1], + r[i+1].prompt_token_ids must start with r[i].prompt_token_ids + r[i].output_token_ids. + This is because the agent includes the previous response in the next prompt. + """ + + def test_two_turn_merge_succeeds(self): + """Normal two-turn conversation: samples merge without error.""" + tok = _mock_tokenizer() + + # Turn 1: prompt=[1,2,3], model outputs [10,11] + # Turn 2: prompt=[1,2,3, 10,11, 20,21], model outputs [30,31] + # (tokens 20,21 are the tool/observation tokens added by the environment) + records = [ + _make_record( + prompt_token_ids=[1, 2, 3], + output_token_ids=[10, 11], + output_log_probs=[-0.1, -0.2], + ), + _make_record( + prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], + output_token_ids=[30, 31], + output_log_probs=[-0.3, -0.4], + ), + ] + + 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] + assert merged.response_length == 2 + 2 + 2 # resp1 + obs + resp2 + assert merged.loss_mask == [1, 1, 0, 0, 1, 1] + assert merged.status == Sample.Status.COMPLETED + + def test_three_turn_merge_succeeds(self): + """Three-turn conversation: prefix chain holds across all turns.""" + tok = _mock_tokenizer() + + records = [ + _make_record( + prompt_token_ids=[1, 2], + output_token_ids=[10], + output_log_probs=[-0.1], + ), + _make_record( + prompt_token_ids=[1, 2, 10, 20], + output_token_ids=[30], + output_log_probs=[-0.2], + ), + _make_record( + prompt_token_ids=[1, 2, 10, 20, 30, 40], + output_token_ids=[50], + output_log_probs=[-0.3], + ), + ] + + samples = compute_samples_from_openai_records(_ARGS, records, tok) + merged = merge_samples(samples, tok) + + assert merged.tokens == [1, 2, 10, 20, 30, 40, 50] + assert merged.response_length == 1 + 1 + 1 + 1 + 1 # 3 responses + 2 obs + + def test_prefix_mismatch_raises(self): + """When the prefix chain is broken, merge_samples must fail.""" + tok = _mock_tokenizer() + + # Turn 2's prompt does NOT start with turn 1's full tokens + records = [ + _make_record( + prompt_token_ids=[1, 2, 3], + output_token_ids=[10, 11], + ), + _make_record( + prompt_token_ids=[1, 2, 3, 99, 99, 20, 21], # 99,99 != 10,11 + output_token_ids=[30, 31], + ), + ] + + 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) + + def test_two_turn_merge_propagates_teacher_log_probs(self): + """OPD teacher_log_probs merge like rollout_log_probs: per-turn values + concatenated with zeros over the injected observation span.""" + tok = _mock_tokenizer() + + records = [ + _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11], output_log_probs=[-0.1, -0.2]), + _make_record( + prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], + output_token_ids=[30, 31], + output_log_probs=[-0.3, -0.4], + ), + ] + 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] + samples[1].teacher_log_probs = [-1.2, -1.3] + + merged = merge_samples(samples, tok) + + # resp1 (2) + obs (2 zeros) + resp2 (2) + assert merged.teacher_log_probs == [-1.0, -1.1, 0.0, 0.0, -1.2, -1.3] + assert len(merged.teacher_log_probs) == merged.response_length + merged.validate() # the new teacher_log_probs length assertion must hold + + def test_two_turn_merge_propagates_opd_student_top_logprobs_metadata(self): + """Top-k OPD student top-logprobs are per-token metadata, not equal metadata.""" + tok = _mock_tokenizer() + + records = [ + _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11], output_log_probs=[-0.1, -0.2]), + _make_record( + prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], + output_token_ids=[30, 31], + output_log_probs=[-0.3, -0.4], + ), + ] + 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]]] + samples[0].metadata = { + "opd_student_top_logprobs": turn_0_top_logprobs, + "shared_metadata": "same", + } + samples[1].metadata = { + "opd_student_top_logprobs": turn_1_top_logprobs, + "shared_metadata": "same", + } + + merged = merge_samples(samples, tok) + + assert merged.metadata["shared_metadata"] == "same" + assert merged.metadata["opd_student_top_logprobs"] == [ + *turn_0_top_logprobs, + [], + [], + *turn_1_top_logprobs, + ] + assert len(merged.metadata["opd_student_top_logprobs"]) == merged.response_length + + def test_two_turn_merge_teacher_log_probs_none_stays_none(self): + """Non-OPD runs leave teacher_log_probs unset; merge must keep it None.""" + tok = _mock_tokenizer() + + records = [ + _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]), + ] + samples = compute_samples_from_openai_records(_ARGS, records, tok) + + merged = merge_samples(samples, tok) + + assert merged.teacher_log_probs is None + + def test_merge_raises_on_teacher_log_probs_length_mismatch(self): + """validate() guards teacher_log_probs length (surfaced via merge_samples).""" + tok = _mock_tokenizer() + + records = [ + _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]), + ] + samples = compute_samples_from_openai_records(_ARGS, records, tok) + + samples[0].teacher_log_probs = [-1.0] # length 1 != response_length 2 + + with pytest.raises(AssertionError, match="teacher_log_probs length"): + merge_samples(samples, tok) + + +# ── test: TITO trailing token trimming ──────────────────────────────── + +STOP = 99 # stands for <|observation|> stop token + + +class TestTITOTrailingTokenTrim: + """Validate trailing-token trimming via ``accumulated_token_ids``. + + Worked example — agentic tool-call retries + ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + An agent makes three turns. The model's tool call fails to parse on + turns 1 and 2, so the agent feeds back an error and retries. + + The session server sees three request/response pairs (records). Each + record's response is an independent inference re-stitched via + pretokenized prefix reuse:: + + record 0 prompt_token_ids: [<|sys|>, aaa, <|user|>, bbb, <|asst|>] + output_token_ids: [ccc, <|obs|>] ← model stopped with <|obs|> + + record 1 prompt_token_ids: [<|sys|>, aaa, <|user|>, bbb, <|asst|>, ccc, <|sys|>, ddd, <|asst|>] + output_token_ids: [eee, <|obs|>] + + record 2 prompt_token_ids: [..., eee, <|sys|>, fff, <|asst|>] + output_token_ids: [ggg, <|obs|>] + + ``accumulated_token_ids`` = record 2's prompt + output:: + + [<|sys|>, aaa, <|user|>, bbb, <|asst|>, ccc, <|sys|>, ddd, + <|asst|>, eee, <|sys|>, fff, <|asst|>, ggg, <|obs|>] + + Note: there is NO ``<|obs|>`` between ``ccc`` and ``<|sys|>`` in the + accumulated sequence — the stop token the model emitted at turn 1 was + consumed by the chat template when rendering turn 2's prompt. + + The algorithm walks ``accumulated_token_ids`` with a cursor:: + + Record 0: cursor = len(prompt_0) → points to "ccc" + Match output [ccc, <|obs|>] against accumulated[cursor:]: + ccc OK, <|obs|> MISMATCH (accumulated has <|sys|> here) + → trim_count=1, strip <|obs|>; cursor advances past "ccc" + + Record 1: cursor = len(prompt_1) → points to "eee" + Match [eee, <|obs|>]: eee OK, <|obs|> MISMATCH + → trim_count=1; cursor advances past "eee" + + Record 2: cursor = len(prompt_2) → points to "ggg" + Match [ggg, <|obs|>]: ggg OK, <|obs|> OK (last turn) + → trim_count=0; cursor reaches end + + Result: three Samples with output tokens [ccc], [eee], [ggg, <|obs|>], + each carrying original per-turn logprobs. + + The tests below encode this example (and variants) with concrete + token IDs. We use ``STOP = 99`` to represent ``<|observation|>``. + """ + + def test_three_turn_trim_trailing_stop_tokens(self): + """Three-turn retry: non-final turns have 1 trailing stop token trimmed.""" + tok = _mock_tokenizer() + + # prompt: [1, 2, 3] output: [10, STOP] + # prompt: [1, 2, 3, 10, 4, 5, 6] output: [20, STOP] + # prompt: [1, 2, 3, 10, 4, 5, 6, 20, 7, 8, 9] output: [30, STOP] + # accumulated (no intermediate STOPs): + # [1, 2, 3, 10, 4, 5, 6, 20, 7, 8, 9, 30, STOP] + records = [ + _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, STOP]), + _make_record(prompt_token_ids=[1, 2, 3, 10, 4, 5, 6], output_token_ids=[20, STOP]), + _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] + + samples = compute_samples_from_openai_records( + _ARGS, + records, + tok, + accumulated_token_ids=accumulated, + max_trim_tokens=1, + ) + + assert len(samples) == 3 + # Turn 0: [10, STOP] → trim 1 → response_length=1 + assert samples[0].tokens == [1, 2, 3, 10] + assert samples[0].response_length == 1 + # Turn 1: [20, STOP] → trim 1 → response_length=1 + assert samples[1].tokens == [1, 2, 3, 10, 4, 5, 6, 20] + assert samples[1].response_length == 1 + # Turn 2 (last): [30, STOP] → trim 0 → response_length=2 + assert samples[2].tokens == [1, 2, 3, 10, 4, 5, 6, 20, 7, 8, 9, 30, STOP] + assert samples[2].response_length == 2 + + def test_no_trim_when_no_trailing_stop(self): + """When output tokens fully match accumulated, trim_count=0 for all turns.""" + tok = _mock_tokenizer() + + # Two turns, no trailing stop tokens — output aligns perfectly + # prompt: [1, 2] output: [10, 11] + # prompt: [1, 2, 10, 11, 3, 4] output: [20, 21] + # accumulated: [1, 2, 10, 11, 3, 4, 20, 21] + records = [ + _make_record(prompt_token_ids=[1, 2], output_token_ids=[10, 11]), + _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] + + samples = compute_samples_from_openai_records( + _ARGS, + records, + tok, + accumulated_token_ids=accumulated, + max_trim_tokens=1, + ) + + assert len(samples) == 2 + assert samples[0].tokens == [1, 2, 10, 11] + assert samples[0].response_length == 2 + assert samples[1].tokens == [1, 2, 10, 11, 3, 4, 20, 21] + assert samples[1].response_length == 2 + + def test_single_turn_no_trim(self): + """Single turn: last turn never trims, even with accumulated_token_ids.""" + tok = _mock_tokenizer() + + records = [ + _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11, STOP]), + ] + accumulated = [1, 2, 3, 10, 11, STOP] + + samples = compute_samples_from_openai_records( + _ARGS, + records, + tok, + accumulated_token_ids=accumulated, + max_trim_tokens=1, + ) + + assert len(samples) == 1 + assert samples[0].tokens == [1, 2, 3, 10, 11, STOP] + assert samples[0].response_length == 3 + + def test_no_accumulated_skips_trimming(self): + """Without accumulated_token_ids, no trimming is performed at all.""" + tok = _mock_tokenizer() + + records = [ + _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]), + ] + + samples = compute_samples_from_openai_records( + _ARGS, + records, + tok, + accumulated_token_ids=None, + ) + + assert len(samples) == 2 + # No trimming — STOP is kept for both turns + assert samples[0].tokens == [1, 2, 10, STOP] + assert samples[0].response_length == 2 + assert samples[1].tokens == [1, 2, 10, STOP, 3, 4, 20, STOP] + assert samples[1].response_length == 2 + + def test_trim_exceeding_max_raises(self): + """If trailing tokens exceed max_trim_tokens, assert fires.""" + tok = _mock_tokenizer() + + # Output has 2 trailing tokens that don't match, but max_trim_tokens=1 + records = [ + _make_record(prompt_token_ids=[1, 2], output_token_ids=[10, STOP, STOP]), + _make_record(prompt_token_ids=[1, 2, 10, 3, 4], output_token_ids=[20]), + ] + accumulated = [1, 2, 10, 3, 4, 20] + + with pytest.raises(AssertionError, match="trim_count 2 exceeds allowed=1"): + compute_samples_from_openai_records( + _ARGS, + records, + tok, + accumulated_token_ids=accumulated, + max_trim_tokens=1, + ) + + def test_cursor_covers_entire_accumulated(self): + """After processing all records, cursor must equal len(accumulated).""" + tok = _mock_tokenizer() + + # accumulated is shorter than what records imply — cursor won't reach end + records = [ + _make_record(prompt_token_ids=[1, 2], output_token_ids=[10, STOP]), + _make_record(prompt_token_ids=[1, 2, 10, 3], output_token_ids=[20]), + ] + # 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] + + with pytest.raises(AssertionError, match="cursor .* != len\\(accumulated_token_ids\\)"): + compute_samples_from_openai_records( + _ARGS, + records, + tok, + accumulated_token_ids=accumulated, + max_trim_tokens=1, + ) + + +# ── test: thinking token issue (documents known failure mode) ──────── + + +class TestThinkingTokenPrefixBreak: + """Documents the known issue where model-generated ... + tokens break the prefix chain. + + When a model (e.g. Qwen3) generates reasoning before + the actual response, agents strip the thinking content from conversation + history. This causes the next turn's prompt to not include the thinking + tokens, breaking the prefix assumption in merge_samples. + + This is a MODEL-LEVEL issue — the fix should be at the model/serving + config level (disable thinking mode), not in the merge logic. + """ + + THINK_TOKEN = 151667 # in Qwen3 + END_THINK_TOKEN = 151668 # in Qwen3 + NEWLINE_TOKEN = 198 # \n + + def test_thinking_tokens_break_prefix_chain(self): + """Demonstrates the failure: model outputs ..., but the agent + strips it from history, so the next prompt doesn't include those tokens.""" + tok = _mock_tokenizer() + + # Turn 1: model generates \nreasoning\n\n then actual response + thinking_tokens = [ + self.THINK_TOKEN, + self.NEWLINE_TOKEN, + 42, + 43, + self.NEWLINE_TOKEN, + self.END_THINK_TOKEN, + self.NEWLINE_TOKEN, + ] + response_tokens = [10, 11] + all_output = thinking_tokens + response_tokens + + records = [ + _make_record( + prompt_token_ids=[1, 2, 3], + output_token_ids=all_output, + ), + # Turn 2: agent only included the actual response [10, 11] in history + # (stripped thinking tokens), plus observation [20, 21] + _make_record( + prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], + output_token_ids=[30, 31], + ), + ] + + 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] + # sample[1] does NOT start with sample[0] — prefix chain broken + with pytest.raises(AssertionError, match="b.tokens must start with a.tokens"): + merge_samples(samples, tok) + + def test_no_thinking_tokens_prefix_chain_holds(self): + """When thinking is disabled, the same conversation merges fine.""" + tok = _mock_tokenizer() + + # Same conversation but model output has no thinking prefix + records = [ + _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], + ), + ] + + 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] + + +# ── test: prefix cache info population ──────────────────────────────── + + +class TestPrefixCacheInfo: + """Validate that prefix cache statistics from meta_info are collected.""" + + def test_single_record_with_cache_stats(self): + """cached_tokens and prompt_tokens from meta_info populate prefix_cache_info.""" + tok = _mock_tokenizer() + record = _make_record( + prompt_token_ids=[1, 2, 3], + output_token_ids=[10, 11], + cached_tokens=2, + prompt_tokens=3, + ) + 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 + + def test_multi_turn_cache_stats_accumulate_after_merge(self): + """After merge_samples, prefix_cache_info sums across turns.""" + tok = _mock_tokenizer() + records = [ + _make_record( + prompt_token_ids=[1, 2, 3], + output_token_ids=[10, 11], + output_log_probs=[-0.1, -0.2], + cached_tokens=0, + prompt_tokens=3, + ), + _make_record( + prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], + output_token_ids=[30, 31], + output_log_probs=[-0.3, -0.4], + cached_tokens=5, + prompt_tokens=7, + ), + ] + samples = compute_samples_from_openai_records(_ARGS, records, tok) + merged = merge_samples(samples, tok) + + assert merged.prefix_cache_info.cached_tokens == 0 + 5 + assert merged.prefix_cache_info.total_prompt_tokens == 3 + 7 + assert merged.prefix_cache_info.prefix_cache_hit_rate == 5 / 10 + + def test_missing_cache_fields_default_to_zero(self): + """Records without cached_tokens/prompt_tokens give zero prefix_cache_info (regression).""" + tok = _mock_tokenizer() + record = _make_record( + prompt_token_ids=[1, 2, 3], + output_token_ids=[10, 11], + ) + 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 00000000000..2d948f43ff9 --- /dev/null +++ b/tests/fast/rollout/session/test_samples_codec.py @@ -0,0 +1,241 @@ +"""Tests for the samples wire codec: encode on the worker, overlay on the driver. + +Covers the COMPUTED/TEMPLATE field partition, safetensors-tensor round-trips, +malformed-payload rejection, and the overlay defaults guard +(`_assert_overlay_template_defaults`). +""" + +import dataclasses +import json + +import numpy as np +import pytest +import safetensors.numpy +from safetensors import SafetensorError + +from miles.rollout.session.samples.codec import ( + COMPUTED_FIELDS, + TEMPLATE_FIELDS, + 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_field_partition_is_total_and_disjoint(self): + all_fields = {f.name for f in dataclasses.fields(Sample)} + assert set(COMPUTED_FIELDS) | set(TEMPLATE_FIELDS) == all_fields + assert not set(COMPUTED_FIELDS) & set(TEMPLATE_FIELDS) + + 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", + "sample.0.tokens", + "sample.0.rollout_log_probs", + "sample.0.rollout_routed_experts", + "sample.0.rollout_indexer_topk", + } + assert tensors["_samples_meta"].dtype == np.uint8 and tensors["_samples_meta"].ndim == 1 + + (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): + # JSON null 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["tensors"]["tokens"] = None + del tensors[f"sample.{index}.tokens"] + + 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) + + @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({"sample.0.tokens": 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("sample.0.tokens")), + KeyError, + "sample.0.tokens", + id="referenced-tensor-missing", + ), + pytest.param( + lambda p: _mutated_payload( + p, + lambda meta, tensors: tensors.update( + {"sample.0.tokens": tensors["sample.0.tokens"].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]["tensors"].update(tokens="sample.0.rollout_log_probs"), + ), + ValueError, + "references tensor", + id="reference-name-mismatch", + ), + 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 00000000000..c5ee3ec84c3 --- /dev/null +++ b/tests/fast/router/test_session_samples_op.py @@ -0,0 +1,296 @@ +"""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, +) + + +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, *, max_seq_len=None): + response = await core.collect_samples(sid, max_seq_len=max_seq_len) + return response.status_code, response.body + + +def _new_pipeline(payload, input_sample): + """What the driver does with the reply: 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: + 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_merged(core): + """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()) + (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} + # Overlay: template fields are the driver's, untouched by the wire. + assert m.prompt == [{"role": "user", "content": "hi"}] + assert m.label == "lbl" + assert m.reward == 2.5 + assert m.routing_key == "routing-sid" + assert m.train_metadata == {"loss": "ppo"} + assert m.metadata["task"] == "t1" + # Metadata application order: the agent overrides the input's shared_key; + # session_metadata (applied last) overrides the agent's max_trim_tokens plant. + assert m.metadata["shared_key"] == "from-agent" + assert m.metadata["max_trim_tokens"] == reply.session_metadata["max_trim_tokens"] + assert m.metadata["accumulated_token_ids"] == _ACCUMULATED + + +async def test_truncation_golden(core): + """max_seq_len=8 strips one output token off the second turn (a turn-level + budget applied before merge): the merged 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, max_seq_len=8) + assert status == 200 + samples, _ = _new_pipeline(payload, _input_sample()) + + (last,) = samples + assert last.status == Sample.Status.TRUNCATED + assert last.tokens == _ACCUMULATED[:8] + 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'{"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'{"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" diff --git a/tests/manual/session/bench_session_server_overhead.py b/tests/manual/session/bench_session_server_overhead.py index 3e69b5cbe00..92b8ca5dab2 100644 --- a/tests/manual/session/bench_session_server_overhead.py +++ b/tests/manual/session/bench_session_server_overhead.py @@ -338,7 +338,6 @@ def _build_server_args( apply_chat_template_kwargs=chat_template_kwargs, tito_model=bench_args.tito_model, tito_allowed_append_roles=bench_args.allowed_append_roles, - generate_multi_samples=False, use_rollout_routing_replay=True, use_rollout_indexer_replay=False, miles_router_timeout=600.0,