(2/3) refactor(session): assemble training samples on the session server; records never leave it#1759
Draft
guapisolo wants to merge 1 commit into
Draft
(2/3) refactor(session): assemble training samples on the session server; records never leave it#1759guapisolo wants to merge 1 commit into
guapisolo wants to merge 1 commit into
Conversation
guapisolo
requested review from
Shi-Dong,
Zhichenzzz,
fzyzcjy,
jybsuper,
maocheng23 and
yueming-yuan
as code owners
July 21, 2026 23:09
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
guapisolo
commented
Jul 22, 2026
| @@ -0,0 +1,222 @@ | |||
| """Black-box sample packing for the `POST /sessions/{id}/samples` reply; only the | |||
Collaborator
Author
There was a problem hiding this comment.
This is a vibed file, but verified by e2e CIs. No tito mismatch rate regression.
guapisolo
force-pushed
the
refactor/session-samples-move
branch
from
July 23, 2026 19:40
0d64438 to
8c55e0e
Compare
guapisolo
force-pushed
the
refactor/session-samples-op
branch
from
July 23, 2026 19:40
95a27be to
8fdc5d5
Compare
guapisolo
marked this pull request as draft
July 24, 2026 20:39
guapisolo
force-pushed
the
refactor/session-samples-op
branch
from
July 24, 2026 20:59
8fdc5d5 to
a70ff78
Compare
guapisolo
force-pushed
the
refactor/session-samples-move
branch
from
July 24, 2026 21:13
8c55e0e to
b82b089
Compare
…ds never leave it
POST /sessions/{session_id}/samples (registered before the catch-all proxy, which would otherwise forward it to the inference backend): the owning instance runs compute -> truncate -> merge synchronously on its event loop (the lock-free get_session invariant) and replies with one safetensors payload — per-sample tensors named {field}.{index} for the array fields (loss_mask crosses as uint8, one byte per token), with the json-codec fields riding as the rank-one uint8 tensor _samples_meta (safetensors.numpy.load exposes no header metadata and __metadata__ is reserved by the format), so malformed payloads fail inside safetensors' validated parser instead of hand-rolled framing checks. The wire contract is one declarative table, SAMPLES_VALUE_SPEC: every computed field maps to a ValueSpec naming its codec ("tensor" | "tensor_list" | "json") and pinned wire dtype, and the COMPUTED_FIELDS allowlist derives from the table. Only those fields cross the wire on blank templates; every other Sample field never crosses — the driver overlays the wire fields onto deepcopies of its input sample and keeps its local value verbatim for the rest. Non-strict dtype conversions must be value-exact (a negative loss-mask entry wrapping into uint8 raises instead of shipping corrupt masks). The codec lives in samples/codec.py, depending only on Sample, NumPy, and safetensors (>=0.8.0, now an explicit dependency); deterministic assembly failures return 422 with the assertion text; empty replies carry a no_records/all_truncated discriminator preserving today's ABORTED semantics.
The old records path (GET the full dump, GiB-scale under R3 with per-turn full-prefix arrays, parsed on the driver's single interpreter) is retired from the training path. collect_samples posts once via the new http_utils.post_bytes_no_retry (post() force-decodes json/text and blind-retries; a 5xx here means the owning instance died with the records, a 422 is deterministic), raises on non-2xx with the body text, raises on timeout instead of silently ABORTing the sample (a measured data-loss bug in the records path), and attempts the session DELETE on every path so failed sessions cannot accumulate. agentic_tool_call.generate shrinks to a shell: agent call -> collect_samples -> empty_reason mapping -> metadata application in today's order.
Tests assert golden Sample field values derived from the records fixture (per-turn, merged, truncated), the wire-codec round-trip, malformed-payload contracts, the 422/404/empty_reason/route-order contracts, and the client's behavior deltas.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
guapisolo
force-pushed
the
refactor/session-samples-op
branch
from
July 24, 2026 21:24
a70ff78 to
c95ee80
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Assemble training samples on the session server; records never leave the owning instance.
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. What stays the same: the assembly pipeline (compute, truncate, merge) runs unchanged, with the assembled values locked by golden tests.Before / After
POST /sessions/{session_id}/samplesruns compute → truncate → merge on the owning instance, synchronously on its event loop (the lock-freeget_sessioninvariant).SAMPLES_VALUE_SPEC(field → codec/dtype): per-sample tensors named{field}.{index}(loss_maskcrosses as uint8), json-codec fields riding as the rank-one uint8 tensor_samples_meta.miles/rollout/session/samples/codec.py, depending only onSample, NumPy, safetensors;SessionCore.collect_samples, the route,http_utils.post_bytes_no_retry, plus the client cutover land together.agentic_tool_call.generateshrinks to a shell: agent call →collect_samples→empty_reasonmapping → today's metadata application order.Behavior Preservation
Samplefield values (per-turn, merged, truncated) derivable from a fixed two-turn records fixture.COMPUTED_FIELDScross the wire on blank templates; the driver overlays them onto deepcopies of its input sample; every otherSamplefield never crosses — the overlay keeps the driver deepcopy's value verbatim.Verification
tests/fast/router/test_session_samples_op.py— golden field values; 422 / 404;empty_reasondiscriminators; route registered before the catch-all proxy.tests/fast/rollout/session/— assembly goldens plus codec round-trip / malformed-payload tests; client deltas intest_openai_endpoint_utils.py.Review Focus
collect_samplesincore.py: assembly must stay synchronous on the loop — offloading to an executor without snapshotting records breaks the lock-free read.COMPUTED_FIELDSallowlist insamples/codec.py: a computed field missing from it silently keeps the driver's stale value at training time.