(2/2) refactor(session): assemble training samples on the session server; records never leave it by guapisolo · Pull Request #1759 · radixark/miles

guapisolo · 2026-07-21T23:09:32Z

Depends on: #1916
Stacked on the parent PR's branch. Review / merge the parent first.

Summary

Assemble agentic training samples on the owning session server over HTTP.

Motivation

The rollout driver previously downloaded every session record before computing, truncating, and merging the training sample. Because each record contains the full trajectory prefix, that transfer grows quadratically with the number of turns and makes the driver serialize, transfer, and parse data whose only consumer is sample assembly.

This PR establishes HTTP as the complete baseline transport. It intentionally adds no Mooncake dependency or runtime assumption; a later PR can add Mooncake as an optional optimization without changing the assembly contract.

Before / After

Before / After: Replace driver-side records download with server-side compute → truncate → unconditional merge.
The endpoint returns one safetensors payload described by SAMPLES_VALUE_SPEC.
The driver overlays the computed fields onto a deepcopy of its input Sample, so input-only fields remain local.
Server-produced metadata crosses as a delta before downstream agent/session overlays.
OpenAIEndpointTracer.collect_samples performs one non-retrying binary HTTP POST.

Behavior Preservation

How we know: Golden tests compare exact post-merge Sample values, including truncation.
Metadata authority is input < server < agent < session.
Lifecycle/debug messages survive the wire.
no_records and all_truncated remain singleton ABORTED results.
POST timeouts, transport failures, or non-2xx replies raise instead of silently losing a sample.
Session DELETE is attempted on every path, while a DELETE failure cannot hide the POST result.

Verification

Existing test suite: 143 passed; 17 skipped across session assembly/codec, router samples-op, OpenAI endpoint client, generate-hub, plus rollout integration tests.
Pre-commit passed.

Review Focus

Scrutinize SessionCore.collect_samples: assembly stays synchronous on the event loop so its lock-free read sees one coherent records list.
Scrutinize SAMPLES_VALUE_SPEC: it is the sole field allowlist plus dtype/codec contract for the wire.
Confirm the transport boundary is HTTP-only in this PR; Mooncake remains a separate optional follow-up.

gemini-code-assist · 2026-07-21T23:09:35Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

guapisolo · 2026-07-22T00:40:26Z

@@ -0,0 +1,222 @@
+"""Black-box sample packing for the `POST /sessions/{id}/samples` reply; only the


This is a vibed file, but verified by e2e CIs. No tito mismatch rate regression.

…server v2 The v5 multi-lineage design lands as an OPT-IN second implementation instead of replacing v1: --use-session-server grows an optional value (bare flag or "v1" = the untouched linear server, "v2" = tree serving), and the whole tree stack lives in miles/rollout/session/v2/. The v1 modules stay byte-identical to their pre-branch state. v2 package (carried from the v5 line, imports adjusted only): - session_state: always-branch serving over the forest — deepest attach point, suffix becomes the branch delta, non-extensions grow siblings or new roots; retry semantics moves entirely to the samples-op picker. Branch suffixes may carry client assistants (compaction); snapshot splice with zero-inheritance-root fallback on canonical-prefix divergence. Truncated-path extension is 409 (TruncatedGenerationError, defined here — v1 never raises it). --session-strict-append-only is the single-chain guard: branch shapes fail loud with attach diagnostics. - assembly: per-leaf fold (compute -> truncate -> fold), default pick (temporal-supersession retry trim; roots never trimmed; non-retry-shaped trees 422) and default merge (exactly-once completion masking over the surviving set, rewards keyed by response id, fold < agent < server metadata layering); both replaceable via --session-sample-picker-path / --session-merge-function-path (sync-only, loaded in-process). - core: v2 twin of SessionCore reusing v1's HTTP plumbing unchanged. Shared seams, each inert for v1 by construction: - sessions.py dispatches registry/core on the flag and forwards the collect body's "metadata" (agent semantic layer) only under v2. - codec: encode/decode parameterized by a fields tuple; v1 default keeps the wire byte-identical, v2 adds reward + per-sample metadata (COMPUTED_FIELDS_V2) with conditional overlay semantics on decode. - arguments: flag upgrade (nargs="?"), unknown values rejected, and the three v2 flags require --use-session-server v2. Tests: v2 HTTP matrix (tree pins incl. deep/root divergence branching, strict guard, truncation + compaction), session_state unit matrix, and the samples-op suite (golden, multi-leaf trim/masking/rewards, hook lanes) — all against a v2-flagged server; the restored v1 suites and the byte-exact A-list (git diff against the pre-branch base is empty) pin that the default path did not move. Rebased onto the tito/session PR stack (#1628/#1777/#1778 + #1759/#1760), with the stack's landed shape taken as authoritative: - supports_midpath_assistant_rerender (True on the base class) moves in here with its consumer, the v2 suffix gate. No family sets it False anymore: the vendored DeepSeek-V3.2 encoder renders position-independently under drop_thinking=False, so the gate text no longer names it; a stub family pins the fail-loud branch. - v2 serving fills SessionRecord.request_timestamp like v1, so folded samples keep their lifecycle req_ts across the wire. - codec: COMPUTED_FIELDS_V2 is re-expressed over the ValueSpec table (SAMPLES_VALUE_SPEC_V2 superset), same wire semantics. - v2 unit tests pin allowed_append_roles=["tool"] explicitly; the ctor default is tool+user since the append-roles PR. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Shi-Dong · 2026-07-29T19:03:28Z

-        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.


Something I don't understand: what if TimeoutError is raised here? Wouldn't the error be propagated out and crash the rollout group?

let me carefully check this.

…ds never leave it POST /sessions/{session_id}/samples lets the owning session-server instance run record assembly, turn-boundary truncation, and the unconditional linear merge before returning one safetensors payload. SAMPLES_VALUE_SPEC defines the tensor/json wire fields and pinned dtypes; the driver overlays them onto a deepcopy of its input Sample. Server-produced metadata crosses as a delta, so lifecycle and debug messages survive while input-only task metadata remains local. The training path no longer downloads full session records to the rollout driver. OpenAIEndpointTracer performs one non-retrying binary POST, raises on timeout and non-2xx replies, and still attempts DELETE on every path. agentic_tool_call keeps its stable list[Sample] contract, maps empty reasons to a singleton ABORTED result, then applies agent and session metadata in order. Tests cover merged and truncated golden samples, metadata authority and isolation, malformed safetensors payloads, 422/404/empty-reason/route-order behavior, client failure semantics, and the HTTP-only generate happy path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Rename the codec entry points to distinguish wire encoding from driver-side Sample assembly. Simplify their local documentation while preserving the payload and merge behavior.

Keep the samples collection and assembly documentation focused on stable behavior without retaining implementation-history detail.

Map collect_samples timeouts to a copied ABORTED sample at the agentic generate boundary so one transient collection timeout does not fail the rollout group. Keep OpenAIEndpointTracer fail-loud and continue propagating non-timeout collection errors.

Document that collect_samples leaves request timeouts fail-loud after the session cleanup attempt so the agentic generate boundary can map them to an ABORTED sample.

Shi-Dong

LGTM!

guapisolo requested review from Shi-Dong, Zhichenzzz, fzyzcjy, jybsuper, maocheng23 and yueming-yuan as code owners July 21, 2026 23:09

This was referenced Jul 21, 2026

(3/3) refactor(rollout): drop --generate-multi-samples and its per-turn sample semantics #1760

Merged

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

Closed

guapisolo added the run-ci-sglang label Jul 22, 2026

guapisolo commented Jul 22, 2026

View reviewed changes

guapisolo mentioned this pull request Jul 23, 2026

feat(session): session server v2 — opt-in trajectory-tree serving; v1 preserved byte-exact #1771

Closed

guapisolo force-pushed the refactor/session-samples-move branch from 0d64438 to 8c55e0e Compare July 23, 2026 19:40

guapisolo force-pushed the refactor/session-samples-op branch from 95a27be to 8fdc5d5 Compare July 23, 2026 19:40

guapisolo removed the run-ci-sglang label Jul 24, 2026

guapisolo marked this pull request as draft July 24, 2026 20:39

guapisolo force-pushed the refactor/session-samples-op branch from 8fdc5d5 to a70ff78 Compare July 24, 2026 20:59

guapisolo force-pushed the refactor/session-samples-move branch from 8c55e0e to b82b089 Compare July 24, 2026 21:13

Base automatically changed from refactor/session-samples-move to main July 24, 2026 21:16

guapisolo force-pushed the refactor/session-samples-op branch from a70ff78 to c95ee80 Compare July 24, 2026 21:24

guapisolo mentioned this pull request Jul 28, 2026

[rollout] Agentic server roadmap #712

Open

4 tasks

guapisolo force-pushed the refactor/session-samples-op branch from c95ee80 to 98b4e42 Compare July 28, 2026 07:47

guapisolo marked this pull request as ready for review July 28, 2026 18:12

guapisolo force-pushed the refactor/session-samples-op branch from 98b4e42 to 8875bc6 Compare July 28, 2026 18:56

guapisolo changed the base branch from main to refactor/remove-generate-multi-samples July 28, 2026 18:58

guapisolo changed the title ~~(2/3) refactor(session): assemble training samples on the session server; records never leave it~~ (2/2) refactor(session): assemble training samples on the session server; records never leave it Jul 28, 2026

guapisolo added the run-ci-sglang label Jul 28, 2026

guapisolo force-pushed the refactor/session-samples-op branch from d2ac3c5 to 51b62ac Compare July 29, 2026 00:27

Base automatically changed from refactor/remove-generate-multi-samples to main July 29, 2026 18:44

Shi-Dong reviewed Jul 29, 2026

View reviewed changes

guapisolo and others added 3 commits July 29, 2026 20:46

refactor(session): name sample codec assembly explicitly

2cf3ae8

Rename the codec entry points to distinguish wire encoding from driver-side Sample assembly. Simplify their local documentation while preserving the payload and merge behavior.

docs(session): simplify samples path comments

21b59dd

Keep the samples collection and assembly documentation focused on stable behavior without retaining implementation-history detail.

guapisolo force-pushed the refactor/session-samples-op branch from 51b62ac to 21b59dd Compare July 29, 2026 21:17

guapisolo removed the run-ci-sglang label Jul 29, 2026

guapisolo added 2 commits July 29, 2026 22:26

docs(rollout): clarify timeout propagation

2730619

Document that collect_samples leaves request timeouts fail-loud after the session cleanup attempt so the agentic generate boundary can map them to an ABORTED sample.

Shi-Dong approved these changes Jul 30, 2026

View reviewed changes

guapisolo merged commit 0425ef0 into main Jul 30, 2026
9 checks passed

guapisolo deleted the refactor/session-samples-op branch July 30, 2026 00:12

Navigation Menu

Uh oh!

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

(2/2) refactor(session): assemble training samples on the session server; records never leave it#1759
guapisolo merged 5 commits into
mainfrom
refactor/session-samples-op

guapisolo commented Jul 21, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

guapisolo Jul 22, 2026 •

edited

Loading

Uh oh!

Shi-Dong Jul 29, 2026

Uh oh!

guapisolo Jul 29, 2026

Uh oh!

Shi-Dong left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -0,0 +1,222 @@
		"""Black-box sample packing for the `POST /sessions/{id}/samples` reply; only the

Uh oh!

Conversation

guapisolo commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Motivation

Before / After

Behavior Preservation

Verification

Review Focus

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

guapisolo Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Shi-Dong Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

guapisolo Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

Shi-Dong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

guapisolo commented Jul 21, 2026 •

edited

Loading

guapisolo Jul 22, 2026 •

edited

Loading