pass the engine weight version from the trainer instead of polling the router

[yueming-yuan]

FullyAsyncRolloutFn needs the current engine weight version to measure and bound sample staleness. Today it rediscovers it over HTTP: _CachedWeightVersion polls the router's /model_info once per drained group, behind a 1s TTL cache that also throttles failures so an unreachable router does not cost one 2s timeout per group.

That number is ours. The trainer stamps it onto the engines (update_weight_version(weight_version=str(self.weight_version))), so polling reads back what miles just wrote.

Change

weight_updater.weight_version → RolloutManager.set_weight_version → RolloutFnTrainInput.weight_version → the drain.

• Both training actors already hold self.rollout_manager and already call it at rank 0 inside update_weights (clear_updatable_has_new_engines); this adds one call right after the push completes, with the same precondition.
• RolloutFnTrainInput gains weight_version: int | None, mirroring RolloutFnEvalInput.
• _CachedWeightVersion and its HTTP/timeout/TTL machinery are deleted.

Same numbering on both sides: engines are stamped from this counter and sample.weight_versions come back from engine meta_info, so current - group_oldest_weight_version compares like with like — that is what the ci_test version assertion in the actors already checks.

None until the first weight update, which turns the staleness filter off rather than reporting staleness 0. train_async.py pushes weights once before the loop, so a real run has a version before its first drain.

Notes

• The notify lands after update_weights() returns, i.e. after engines resume, so a brief window under-reports staleness. Smaller than the 1s TTL it replaces.
• No longer degrades to None on router flakiness, and no longer needs to special-case the engines' pre-update "default" string.

Motivated by #2030, which moves staleness control into the fully-async data buffer and pays this HTTP round trip on every consumed group. Independent of it — this lands on main and #2030 picks it up on the next merge.

Tests

test_stale_group_recycled now passes the version on the input; test_weight_version_throttles_failed_queries is replaced by test_staleness_filter_off_before_the_first_weight_update, covering the None case.

⚠️ Not executed: the devbox was unreachable while preparing this. Verified locally with py_compile and pre-commit run --all-files only. Please run tests/fast/rollout/test_fully_async_rollout.py before merging.

🤖 Generated with Claude Code

[guapisolo]

[P1] Stamp LoRA updates before trusting the trainer counter

input.weight_version is only comparable with Sample.weight_versions when both are stamped from the same counter. However, the supported --fully-async single-LoRA path with the default broadcast transport calls load_lora_adapter_from_distributed without any weight_version (the tensor path is the same, and the actor version-equality check explicitly skips LoRA), while the trainer counter still increments. Samples can therefore retain "default" or a fixed base version: oldest_weight_version then either returns None and silently disables the configured --max-weight-staleness, or remains fixed until every group is recycled as stale. Please stamp this counter on the engine after a successful LoRA load, or reject this configuration, and add coverage for the LoRA path.

[guapisolo]

[cc] [P2] Cell failover regresses the published weight version; the staleness filter goes silently inert

RolloutManager.set_weight_version (rollout_manager.py:379) is fed from weight_updater.weight_version, a per-actor in-memory counter (= 0 in __init__, incremented at the top of every update_weights). It is not checkpointed and not restored from the engines, and RolloutManager.save/load does not persist the manager-side copy either.

Under FT (--indep-dp), RayTrainGroup.update_weights runs on the first alive cell (_execute_first_alive, wrapped in retry, group.py:264-267). When cell 0 dies after N updates, the next update runs on cell 1, whose own updater counter is still 0 — so the manager's version jumps backwards (e.g. 57 → 1). Queued fully-async groups still carry samples stamped with the old higher versions, so staleness = current - oldest in _drain (fully_async_rollout.py:206) goes negative. A negative value never exceeds max_weight_staleness, so exactly the most off-policy groups are admitted into training, and rollout/fully_async/{avg,max}_staleness are reported as negative numbers — a wrong metric rather than an error.

To be fair, the counter reset itself predates this PR (the same reset used to reach the filter via engine stamps + router polling), so this is inherited rather than introduced. But the new channel carries it forward, and the PR does add one new case: a RolloutManager actor restart resets weight_version to None, silently disabling the filter until the next update.

Minimal ask: make the regression observable — warn (or assert) in set_weight_version when the incoming value is lower than the stored one. This is safe within a manager lifetime: a fresh manager starts from None, so legitimate full-job restarts never trip it. Optionally also clamp/warn on negative staleness in _drain. Actually persisting/recovering the counter across failover can be follow-up work.