Skip to content

feat: track policy versions in fully async rollout - #2279

Open
0z5a wants to merge 1 commit into
THUDM:mainfrom
0z5a:codex/fully-async-policy-version-lifecycle
Open

feat: track policy versions in fully async rollout#2279
0z5a wants to merge 1 commit into
THUDM:mainfrom
0z5a:codex/fully-async-policy-version-lifecycle

Conversation

@0z5a

@0z5a 0z5a commented Aug 17, 2026

Copy link
Copy Markdown

Summary

This is the first, dependency-free slice of #1800. It establishes policy-version lifecycle semantics for fully async rollout without adding staleness control, eviction, or partial-resume policy yet.

  • wrap both initial and periodic actor-to-rollout weight updates with optional before_weight_update / after_weight_update hooks
  • advance the rollout policy version only after a successful weight publication; failed updates keep the previous version
  • let RolloutManager forward lifecycle hooks to the configured rollout module while preserving compatibility with rollout functions that do not define them
  • snapshot the current policy version when a fully async group is admitted, then preserve it in an immutable CompletedSampleRecord even if completions arrive out of order
  • expose the admission version on Sample.policy_version

Semantics

policy_version starts at zero in RolloutManager. A successful initial or periodic weight update increments it once. The async worker captures that published value at admission, so a weight update completing while generation is in flight cannot relabel an older trajectory. The existing behavior of recycling aborted groups remains unchanged.

Validation

  • python -m pytest tests/test_fully_async_rollout.py tests/test_sample.py tests/test_train_async_policy_version.py — 21 passed
  • pre-commit run --from-ref 41014d1f29e201137fdffce737bb8bac65bc5219 --to-ref HEAD — all hooks passed

The tests cover successful and failed weight updates, monotonic version publication, out-of-order completion across an update boundary, immutable completion records, queue backpressure, and Sample serialization.

Follow-ups

To keep review scope small, configurable staleness/requeue policy, partial rollout resume, and fixed-cardinality observability metrics will be submitted as follow-up PRs after this lifecycle foundation is reviewed.

Part of #1800.

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