Skip to content

refactor: share training batch production#3027

Draft
tim0120 wants to merge 1 commit into
feat/static-sftfrom
feat/shared-batch-producer
Draft

refactor: share training batch production#3027
tim0120 wants to merge 1 commit into
feat/static-sftfrom
feat/shared-batch-producer

Conversation

@tim0120

@tim0120 tim0120 commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

What

Extracts the trainer-bound batch lifecycle behind a source-independent interface:

  • TrainingBatchSource prepares examples and handles source-specific post-send work.
  • TrainingBatchProducer is the single loop that publishes TrainingBatch, advances the trainer-aligned step, and notifies the source.
  • The online rollout orchestrator and static dataset source both use that loop.

This is a stacked architectural draft on #2954 for examining whether this is the right invariant boundary.

Why

#2954's DatasetBatchProducer still owned its own send/step loop, making it structurally close to a renamed static orchestrator. The first interface shared by rollout training and static SFT is the prepared trainer batch, not environments, dataset rows, or rollouts.

This change shares the lifecycle at that boundary while leaving source-specific behavior on the appropriate side:

  • Rollout draining, eval handling, traces, and rollout metrics remain in Orchestrator.
  • Dataset rendering, cursor accounting, trainer acknowledgement, and static metrics remain in DatasetBatchSource.

No config, transport, trainer, normalization, or checkpoint format changes are intended.

Validation

  • Ruff check and format check passed for the three changed files.
  • Python compilation passed for the three changed files.
  • A focused in-memory producer check verified publication order, step advancement, and the post-send callback contract.
  • git diff --check passed.

No README changes or new test files.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant