Skip to content

perf(reassembler): batch group staging reads and expose fetch timings - #4072

Draft
pthombre wants to merge 2 commits into
mainfrom
pranav/reassembler-group-fetch
Draft

perf(reassembler): batch group staging reads and expose fetch timings#4072
pthombre wants to merge 2 commits into
mainfrom
pranav/reassembler-group-fetch

Conversation

@pthombre

@pthombre pthombre commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Token-capture reassembly currently fetches staging rows once per rollout and combines data-plane latency with verification time. Fetch the deduplicated staging-key union once per group, preserve per-rollout failure isolation, and expose separate fetch and verification timings.

  • Healthy groups use one base read and at most one targeted route-payload read in direct mode. Failed batches subdivide to isolate missing, malformed, or unready rows.
  • Retain existing aggregate metrics and public Gym/source interfaces.
  • Prefer the creating node for reassembler actors, with soft affinity and explicit spill when that node is busy.

Issues

Closes #4010. Follow-up to #3837.

PR stack

This is the first PR in the stack: mainpranav/reassembler-group-fetchpranav/capture-orchestration.

Next in stack: #4073 (capture orchestration).

The capture-orchestration PR builds on this branch. Merge this PR first, then rebase/retarget the dependent PR to main.

Validation

  • Slurm job 3656922, final commit 364119fa60553475f3b8022242c7f40a873d032b: completed with exit 0:0.
  • 44 live TQ/Gym tests + 20 actor tests passed, zero failures or skips; includes 18 new regression cases.
  • Scoped pyrefly: zero errors; Ruff, import sorting, formatting, and diff checks passed.
  • Used the existing rl-gym container and its Gym environment. Live storage coverage uses TQ's simple backend; full training throughput and multi-node placement were not benchmarked.

Before ready for review

  • Followed contributor guidelines; commits include sign-off trailers.
  • Added and ran relevant unit/regression tests through Slurm.
  • Updated affected API docstrings and metric semantics.
  • Run a representative full training/functional workload.

This PR is intentionally a draft.

Pranav Thombre added 2 commits September 9, 2026 23:09
Signed-off-by: Pranav Thombre <pthombre@aws-cmh-slurm-1-vscode-01.cm.cluster>
Signed-off-by: Pranav Thombre <pthombre@aws-cmh-slurm-1-vscode-01.cm.cluster>
@pthombre pthombre added the CI:L1 Run doctests, unit tests, and functional tests label Sep 10, 2026
@copy-pr-bot

copy-pr-bot Bot commented Sep 10, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L1 Run doctests, unit tests, and functional tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Token capture finalizer: batch the group staging fetch, and give the pool a placement policy + a splittable timer

1 participant