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

[guapisolo]

Summary

Assemble training samples on the session server and drop --generate-multi-samples.

Motivation

Each record's R3 buffer covers the full prefix, so a session's records dump grows quadratically with turns — measured ~3.15 GiB per full GET /sessions/{id} at production shape (an mp-stack figure; the magnitude carries to this stack) — then the rollout driver parsed every byte of it on one interpreter per session. The only consumer of records is Sample assembly, whose inputs all already live on the owning instance. Relocating the assembly also retires --generate-multi-samples: sample boundaries are a property of the trajectory (linear turns merge; only a genuine context discontinuity — future compaction/subagent support — starts a new sample), not a CLI switch, and per-turn siblings shared one group_index, polluting group baselines.

Before / After

• Before / After: the driver fetched the full records dump per session; now POST /sessions/{session_id}/samples runs compute → truncate → merge on the owning instance, synchronously on its event loop (the lock-free get_session invariant).
• The reply is one opaque safetensors payload; the scalar JSON rides as the rank-one uint8 tensor _samples_meta.
• What moved where: compute_samples_from_openai_records / truncate_samples_by_total_tokens moved as-is to miles/rollout/session/samples/merge.py (review with --find-copies); the wire codec is samples/codec.py; SessionCore.collect_samples, the route, http_utils.post_bytes_no_retry, plus the client cutover follow.
• agentic_tool_call.generate shrinks to a shell: agent call → collect_samples → empty_reason mapping → today's metadata application order. It now always returns list[Sample] (aborted path included), so a custom RM on this path is always called batched with a list[Sample]; multi_turn.generate always TITO-accumulates one scalar sample.

Behavior Preservation

• How we know: golden tests assert the exact assembled Sample field values (merged, truncated) derivable from a fixed two-turn records fixture.
• Only the ten COMPUTED_FIELDS cross the wire on blank templates; the driver overlays them onto deepcopies of its input sample; an import-time guard forces every new Sample field to be classified computed-or-template.
• Three deliberate client deltas, each locked by a test: a non-2xx raises with the server's body text (a 422 carries the assembly assertion verbatim); a collect timeout raises instead of silently ABORTing the sample (a measured data-loss bug in the records path); the session DELETE is attempted on every path.
• Deliberately removed: the per-turn sample mode. check_reward_nonzero_std now flattens nested groups; agentic + --group-rm / --partial-rollout / --recompute-logprobs-via-prefill is documented as unsupported.

Verification

• tests/fast/router/test_session_samples_op.py — golden field values; 422 / 404; empty_reason discriminators; route registered before the catch-all proxy.
• tests/fast/rollout/session/ — assembly goldens plus codec round-trip / drift-guard tests; client deltas in test_openai_endpoint_utils.py.
• Latest full loop: 73 passed (router op + session + client), 107 passed / 16 skipped (generate_hub), 36 passed (inference-rollout integration).

Review Focus

• Scrutinize collect_samples in core.py: assembly must stay synchronous on the loop — offloading to an executor without snapshotting records breaks the lock-free read.
• Scrutinize the COMPUTED_FIELDS / TEMPLATE_FIELDS split in samples/codec.py: a misclassified field silently keeps the driver's stale value at training time.
• Scrutinize the always-list return of agentic_tool_call.generate: downstream isinstance(sample, list) forks now always take the list branch on this path.

[gemini-code-assist]

Code Review

This pull request shifts training-sample assembly into the session server's owning worker, preventing raw session records from leaving the server and significantly reducing IPC payload sizes. It introduces a new POST /sessions/{session_id}/samples endpoint, moves the assembly logic to a new sample_assembly.py module, and implements a custom binary wire codec to overlay computed fields onto the driver's local input sample. Feedback on the changes highlights opportunities to improve robustness and performance: safely retrieving input_ids to avoid an uncaught KeyError (which would bypass the 422 error handler), resolving inconsistent dictionary access for meta_info, explicitly defining dtypes for all serialized segments to prevent type mismatches, and using memoryview for zero-copy buffer slicing to optimize memory usage during deserialization.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In compute_samples_from_openai_records, record.request["input_ids"] is accessed directly. If "input_ids" is missing from the request, this will raise a KeyError before reaching _compute_sample_from_openai_record. Since KeyError is not caught in collect_samples (which only catches AssertionError and ValueError), this will propagate as an uncaught exception and return a 502 Bad Gateway instead of a 422 response.

To ensure missing "input_ids" is correctly handled as a ValueError (which maps to a 422 response), retrieve it safely using .get() and raise a ValueError if it is missing.

Suggested change

	prompt_ids = record.request["input_ids"]
	prompt_ids = record.request.get("input_ids")
	if prompt_ids is None:
	raise ValueError("input_ids not found in request — the session server should populate it")

[gemini-code-assist]

Using choice.get("meta_info", {}) on line 145 and then directly accessing choice["meta_info"] on line 146 is inconsistent and redundant. If "meta_info" is missing, the direct access on line 146 will raise a KeyError.

To be consistent and safe, retrieve "meta_info" once using a safe accessor and reuse the local variable.

Suggested change

	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"])
	meta_info = choice.get("meta_info") or {}
	sample.prefix_cache_info.add(meta_info)
	if "weight_version" in meta_info:
	sample.weight_versions.append(meta_info["weight_version"])

References

1. Avoid using safe accessors (e.g., .get()) for a dictionary key if a prior direct access in the same code path already guarantees its existence. This practice is redundant and inconsistent.

[gemini-code-assist]

The _SEGMENT_DTYPES dictionary does not specify dtypes for rollout_routed_experts and rollout_indexer_topk. If these fields are empty lists [] or have other non-array types, np.asarray(value, dtype=None) will infer their dtypes (e.g., float64 for empty lists), which can cause type mismatches downstream. Specifying their dtypes explicitly as np.int32 ensures consistent serialization and deserialization.

Suggested change

	_SEGMENT_DTYPES = {"tokens": np.int64, "rollout_log_probs": np.float64}
	_SEGMENT_DTYPES = {
	"tokens": np.int64,
	"rollout_log_probs": np.float64,
	"rollout_routed_experts": np.int32,
	"rollout_indexer_topk": np.int32,
	}

[gemini-code-assist]

Slicing the body bytes object (body[start : start + nbytes]) creates a copy of the sliced bytes in memory. For large R3 arrays (which can be hundreds of megabytes in production), this slicing and copying can be a performance bottleneck and increase memory usage.

Using memoryview allows zero-copy slicing of the buffer, which is much more efficient.

Suggested change

	def _read_segment(body: bytes, segment_meta: dict | None) -> np.ndarray | None:
	if segment_meta is None:
	return None
	start = segment_meta["offset"]
	arr = np.frombuffer(body[start : start + segment_meta["nbytes"]], dtype=segment_meta["dtype"])
	return arr.reshape(segment_meta["shape"])
	def _read_segment(body: bytes, segment_meta: dict | None) -> np.ndarray | None:
	if segment_meta is None:
	return None
	start = segment_meta["offset"]
	arr = np.frombuffer(
	memoryview(body)[start : start + segment_meta["nbytes"]],
	dtype=segment_meta["dtype"],
	)
	return arr.reshape(segment_meta["shape"])

[guapisolo]

Bench repro on this exact base (rebased onto #1602 / da8870e) — same machine, same full production shape (32×50, R3 ≡ 25 layers × topk 10 int32, 16 workers). Core conclusions reproduce; the PR-body table stands.

metric	records-GET (repro / round 1)	samples-POST (repro / round 1)
collect per session avg	107.1 s / 122.7 s	11.8 s / 11.7 s
collect reply size avg	3281.0 MiB (byte-identical across runs)	98.3 MiB (byte-identical)
chat p99, collect-overlap, owning worker	198.1 s / 183.7 s	42.6 s* / 17.9 s
/health non-200 (all 503, all inside collect windows)	38/69 / 35/66 (~55%)	14/53 / 10/35 (~27%)
collect success / 422s	32/32, 0 / same	32/32, 0 × 422 / same
bench wall clock	336.6 s / 323.0 s	120.6 s† / 81.9 s

* single outlier in n=19 overlap samples (p95=p99=max — one tail sample); still 4.7× better than the GET path's owning p99; likely machine load noise (other work ran concurrently on the box), flagged for a re-run rather than papered over.
† +47% wall driven by chat-phase throughput (13.3 vs 19.5 turns/s) under the same suspected load noise; the GET leg re-ran at +4%.

Honest residual, same as the PR body: /health false-503s inside collect windows are reduced ~2× in ratio (55%→27%) but not zero — a single ~12 s synchronous assembly exceeds _HEALTH_TIMEOUT = 5.0 by design (no-await invariant). Follow-up candidates: treat a plain health-ping timeout as busy-but-alive (a production-verified one-liner exists, not yet landed), or offload assembly with a records snapshot.

Artifacts: bench-records-get-1602base.json / bench-samples-post-1602base.json (+ the round-1 pair and the uncommitted bench samples-scenario patch), archived in the ops log directory.

🤖 Generated with Claude Code

[guapisolo]

Superseded by the split stack (identical final tree, verified byte-for-byte): #1758 (1/3) → #1759 (2/3) → #1760 (3/3).