From cae3aa745c75ed00c79442c72a5bef2b11d63938 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Sun, 2 Aug 2026 13:55:32 +0800 Subject: [PATCH 1/4] fix(rollout): route async weight pushes through the engine version ledger Closes the weight_version accounting gap documented as a follow-up in #287: ARTrainer.evaluate() pushed weights via _prepare_rollout without advancing the driver-side counter, leaving the engine on unnumbered weights until the next interval boundary. The async trainer families also disagreed on eval policy: AR synced-without-bump, diffusion explicitly passed sync_weights=False, agentic has no eval. Structure: both driver-side engines replace bump_weight_version() with sync_weights(weight_sync) - one call that pushes and advances the ledger, so the pairing stops being call-site discipline. The batch engine also enforces the quiesce invariant (raises before pushing if any generation is in flight). All async sync sites (boundary, resume, agentic drive) route through it; bump_weight_version drops to zero callers and is removed. Policy: async eval becomes read-only. ARTrainer.evaluate() gains sync_weights: bool = True (mirroring DiffusionTrainer.evaluate); the async trainers override evaluate() with async-correct defaults (no push; diffusion also defaults sleep_after=False), so bare evaluate() calls are safe outside train() too. A pre-train explicit push raises a lifecycle RuntimeError instead of an incidental AttributeError. Deliberate behavior changes: (1) async-AR eval no longer pushes weights, so BOTH the eval series and the training rollouts launched between an eval and the next interval boundary change - they previously ran under eval-synced fresher weights; deployment cadence is now governed solely by weight_sync_interval. At interval=1, or when eval_interval is a multiple of the sync interval, training is point-identical to main. (2) The scored eval policy is 1..interval optimizer steps old (exactly 1 at interval=1). (3) The resume-time push advances the ledger 0->1 (metric offset only - eviction math is relative; the launch ceiling is computed from rollout_id). --- unirl/rollout/README.md | 8 +++++-- unirl/rollout/engine/asynchronous.py | 25 ++++++++++++++++---- unirl/trainer/agentic_async.py | 9 ++++---- unirl/trainer/agentic_partial.py | 3 +-- unirl/trainer/ar.py | 10 +++++--- unirl/trainer/async_ar.py | 34 ++++++++++++++++++++++------ unirl/trainer/async_diffusion.py | 11 ++++++--- 7 files changed, 74 insertions(+), 26 deletions(-) diff --git a/unirl/rollout/README.md b/unirl/rollout/README.md index 7ea2ea69..18542b2c 100644 --- a/unirl/rollout/README.md +++ b/unirl/rollout/README.md @@ -94,8 +94,12 @@ implements its weight-receive method and a matching `sync:` handler in engine also can't live on a `layout: separate` slab — `_build_rollout` raises. - **Quiesce before weight sync / eval / checkpoint on the batch async path** — `AsyncBatchRolloutEngine.quiesce()` drains every in-flight generation; a - weight + KV update corrupts one mid-flight. The agentic quiesce is a - turn-boundary `abort` + final poll, folded into + weight + KV update corrupts one mid-flight. Weight pushes go through + `engine.sync_weights(weight_sync)` — one call that pushes and advances + `weight_version`, and (batch engine) raises if anything is still in flight. + Async eval never pushes weights (`evaluate(..., sync_weights=False)`): it + scores the engine-resident policy so the ledger stays exact. The agentic + quiesce is a turn-boundary `abort` + final poll, folded into `AsyncAgenticRolloutEngine.quiesce()`. Reap-vs-launch ordering is trainer statement order (diffusion polls before topping up; see its `_next_step`). - **Reward/advantage methods are not engine code** — `Part.compute_advantages` and diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py index a930f9fd..9ebd1313 100644 --- a/unirl/rollout/engine/asynchronous.py +++ b/unirl/rollout/engine/asynchronous.py @@ -12,7 +12,8 @@ - :class:`InflightPool` — non-blocking pool of distributed ``generate`` calls. Engines, sharing one consumer surface (``poll`` / ``drain_freshest`` / -``pop_evicted`` / ``quiesce`` + engine-owned ``weight_version``): +``pop_evicted`` / ``quiesce`` + engine-owned ``weight_version``, advanced +only by ``sync_weights`` — every weight push goes through the ledger): - :class:`AsyncBatchRolloutEngine` — batch granularity over a single-turn engine slab; one ``submit`` is one non-blocking distributed ``generate``. @@ -239,7 +240,16 @@ def __init__( def weight_version(self) -> int: return self._weight_version - def bump_weight_version(self) -> int: + def sync_weights(self, weight_sync: Any) -> int: + """Push train weights via *weight_sync* and advance the version ledger. + + The only sanctioned weight-push path — pairing the push with the bump + is what keeps the ledger truthful. Raises if any generation is in + flight (a weight + KV update corrupts it); drain via ``quiesce`` first. + """ + if len(self._pool): + raise RuntimeError(f"sync_weights with {len(self._pool)} generations in flight; quiesce() first") + weight_sync.sync() self._weight_version += 1 return self._weight_version @@ -351,7 +361,14 @@ def __init__(self, rollout: Any, *, group_size: int, start_gen_id: int = 0) -> N def weight_version(self) -> int: return self._weight_version - def bump_weight_version(self) -> int: + def sync_weights(self, weight_sync: Any) -> int: + """Push train weights via *weight_sync* and advance the version ledger. + + The only sanctioned weight-push path. The facade cannot see the + coordinator's drive state, so the quiesce contract stays with the + caller: sync only while the drive is finalized/quiesced (decode-idle). + """ + weight_sync.sync() self._weight_version += 1 return self._weight_version @@ -380,7 +397,7 @@ def pop_evicted(self) -> List[List["Sample"]]: def quiesce(self) -> List["Sample"]: """Turn-boundary stop: abort, then one final poll for trajectories that completed DURING the quiesce (before the next ``submit`` resets worker - buffers). Call before ``bump_weight_version`` so those groups carry the + buffers). Call before ``sync_weights`` so those groups carry the version they completed under.""" carried = self._rollout.abort()[0] self.poll() diff --git a/unirl/trainer/agentic_async.py b/unirl/trainer/agentic_async.py index 4e78ae1a..bf8232fd 100644 --- a/unirl/trainer/agentic_async.py +++ b/unirl/trainer/agentic_async.py @@ -21,8 +21,8 @@ ``buffer_max_staleness``), reward + GRPO advantage + one optimizer step (reusing :class:`AgenticTrainer`'s helpers), then **quiesce + sync**: ``abort`` the in-flight tail at a turn boundary, apply the configured ``tail_policy`` (carry only when the - environment can resume from the ``Sample``; otherwise drop), ``weight_sync.sync()``, - bump the version. + environment can resume from the ``Sample``; otherwise drop), then + ``engine.sync_weights`` (one call: push + version bump). ONE single-threaded loop (the ``AsyncARTrainer`` shape): with disjoint slabs the rollout slab keeps generating in the background (the engine's per-worker drain) while @@ -404,7 +404,7 @@ def train( self._engine = AsyncAgenticRolloutEngine(self.rollout, group_size=self._n, start_gen_id=start_rollout) if start_rollout < num_rollouts and start_rollout and self.weight_sync is not None: - self.weight_sync.sync() # push restored weights into the fresh engine + self._engine.sync_weights(self.weight_sync) # push restored weights into the fresh engine if start_rollout < num_rollouts: self._submit_drive(carried=[], rollout_id=start_rollout) # prime the first drive @@ -437,8 +437,7 @@ def train( save_mode=save_mode, ) if need_sync: - self.weight_sync.sync() - self._engine.bump_weight_version() + self._engine.sync_weights(self.weight_sync) if step < num_rollouts: self._submit_drive(carried=carried, rollout_id=step) # resume safe tails + fresh finally: diff --git a/unirl/trainer/agentic_partial.py b/unirl/trainer/agentic_partial.py index b0e1d7a5..ff581e84 100644 --- a/unirl/trainer/agentic_partial.py +++ b/unirl/trainer/agentic_partial.py @@ -181,8 +181,7 @@ def _drive_partial(self, rollout_id: int, sync_weights: bool, stale: int) -> Lis # Sync at the top, AWAKE + decode-idle (barrier parity): TensorWeightSync writes the live # SRT weight pool, which a full sleep() would release. if sync_weights and self.weight_sync is not None: - self.weight_sync.sync() - self._engine.bump_weight_version() + self._engine.sync_weights(self.weight_sync) tasks = self._build_tasks(self._carried, rollout_id) # fresh + carried self._carried = [] # consumed into this drive self._engine.submit(tasks) diff --git a/unirl/trainer/ar.py b/unirl/trainer/ar.py index 9af9eeed..bc3651c2 100644 --- a/unirl/trainer/ar.py +++ b/unirl/trainer/ar.py @@ -493,9 +493,13 @@ def train_step( ) return result, mean_reward - def evaluate(self, rollout_id: int) -> float: + def evaluate(self, rollout_id: int, *, sync_weights: bool = True) -> float: """Periodic eval — ``avg@k`` accuracy on the eval prompt set. + ``sync_weights=False`` skips the pre-eval weight push and scores the + engine-resident policy (the async trainers' mode; keeps the driver-side + version ledger exact). + Mirrors :meth:`train_step`'s rollout+reward path but skips advantage/backward: iterate up to ``eval_num_prompts`` prompts from ``run.eval_data_path`` in ``eval_batch_size``-sized batches, expand @@ -538,13 +542,13 @@ def evaluate(self, rollout_id: int) -> float: ) train_state_offloaded = False if not anchored: - train_state_offloaded = self._prepare_rollout(sync_weights=self.weight_sync is not None) + train_state_offloaded = self._prepare_rollout(sync_weights=sync_weights and self.weight_sync is not None) # Anchored eval keeps FSDP offloaded and vLLM awake for the entire eval # set. Training still uses one _anchored_rollout_session per rollout in # train_step(), so its sleep/wake and onload/offload lifecycle is unchanged. eval_session = ( self._anchored_rollout_session( - sync_weights=self.weight_sync is not None, + sync_weights=sync_weights and self.weight_sync is not None, restore_backend=False, ) if anchored diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 5bda5f24..88d8e163 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -205,10 +205,31 @@ def __init__( if self.weight_sync is not None: self._connect_separate(sync_cfg) + def evaluate(self, rollout_id: int, *, sync_weights: bool = False) -> float: + """Resident-policy eval: the async default flips to ``sync_weights=False``. + + Eval is read-only — weight deployment is governed solely by + ``weight_sync_interval`` (pushes go through ``engine.sync_weights``). + The scored policy is the one already on the rollout slab, 1..interval + optimizer steps old (exactly 1 at ``weight_sync_interval=1``). + """ + return super().evaluate(rollout_id, sync_weights=sync_weights) + def _prepare_rollout(self, *, sync_weights: bool) -> bool: - """Sync a resident separate-slab engine without colocate handoffs.""" + """Prepare the resident separate-slab engine without colocate handoffs. + + An explicit weight push stays on the ledger and is only possible inside + ``train()``, where the async engine exists. + """ if sync_weights and self.weight_sync is not None: - self.weight_sync.sync() + engine = getattr(self, "_async_engine", None) + if engine is None: + raise RuntimeError( + "AsyncARTrainer: weight push outside train() — the async engine owns the " + "version ledger and exists only inside train(); use evaluate(sync_weights=False) " + "for a standalone resident-policy eval." + ) + engine.sync_weights(self.weight_sync) return False def _finish_rollout(self, *, train_state_offloaded: bool) -> None: @@ -360,9 +381,9 @@ def train( ) if resumed and self.weight_sync is not None: - self.weight_sync.sync() # push restored weights into the fresh engine + self._async_engine.sync_weights(self.weight_sync) # push restored weights into the fresh engine if self.eval_interval > 0: - self.evaluate(rollout_id=-1) # baseline; engine quiescent + self.evaluate(rollout_id=-1) # baseline; resident-policy eval, engine quiescent try: for rollout_id in range(start_rollout, num_rollouts): @@ -380,7 +401,7 @@ def train( step = rollout_id + 1 if self.eval_interval > 0 and step % self.eval_interval == 0: self._drain_all() # eval shares the engine - self.evaluate(rollout_id=rollout_id) + self.evaluate(rollout_id=rollout_id) # resident-policy eval; ledger stays exact if save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts): self._drain_all() # consistent engine + deterministic resume self.maybe_save_checkpoint( @@ -388,8 +409,7 @@ def train( ) if step % interval == 0 and self.weight_sync is not None: self._drain_all() # MANDATORY: weight/KV update corrupts in-flight generations - self.weight_sync.sync() - self._async_engine.bump_weight_version() + self._async_engine.sync_weights(self.weight_sync) finally: # Match BaseTrainer._finish_wandb: cleanup failures must not mask # the exception that caused teardown. diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 83ac178f..3a63d18d 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -111,6 +111,12 @@ def _score_completed(self, gen_id: int, completed: Sample) -> List[Sample]: self._drop_decoded(scored, rollout_id=gen_id) return scored.split() + def evaluate(self, step: int, *, sync_weights: bool = False, sleep_after: bool = False) -> float: + """Resident-policy eval — async defaults: no weight push (the ledger stays + exact; deployment cadence belongs to ``weight_sync_interval`` alone) and + no post-eval sleep (the disaggregated engine stays resident).""" + return super().evaluate(step, sync_weights=sync_weights, sleep_after=sleep_after) + def _drain_all(self) -> None: """Finish + buffer EVERY in-flight generation (the single-threaded quiesce). @@ -197,7 +203,7 @@ def train( ) if resumed and self.weight_sync is not None: - self.weight_sync.sync() # push restored weights into the fresh engine + self._async_engine.sync_weights(self.weight_sync) # push restored weights into the fresh engine if self.eval_interval > 0: # Evaluate the policy already resident on the rollout slab. Eval must # neither advance the async weight version nor offload this engine. @@ -227,8 +233,7 @@ def train( ) if step % interval == 0 and self.weight_sync is not None: self._drain_all() # MANDATORY: weight/KV update corrupts in-flight generations - self.weight_sync.sync() - self._async_engine.bump_weight_version() + self._async_engine.sync_weights(self.weight_sync) finally: # Cleanup failures must not mask the exception that stopped training. active_exception = sys.exc_info()[0] is not None From 97329c339570e9bd102361740bd1342b53d1be40 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Sun, 2 Aug 2026 21:12:54 +0800 Subject: [PATCH 2/4] fix(rollout): enforce decode-idle on agentic sync_weights; log eval/weight_version --- unirl/rollout/README.md | 14 ++++++++------ unirl/rollout/engine/asynchronous.py | 13 ++++++++++--- unirl/trainer/async_ar.py | 9 +++++++-- unirl/trainer/async_diffusion.py | 9 +++++++-- 4 files changed, 32 insertions(+), 13 deletions(-) diff --git a/unirl/rollout/README.md b/unirl/rollout/README.md index 18542b2c..a237be2e 100644 --- a/unirl/rollout/README.md +++ b/unirl/rollout/README.md @@ -96,12 +96,14 @@ implements its weight-receive method and a matching `sync:` handler in `AsyncBatchRolloutEngine.quiesce()` drains every in-flight generation; a weight + KV update corrupts one mid-flight. Weight pushes go through `engine.sync_weights(weight_sync)` — one call that pushes and advances - `weight_version`, and (batch engine) raises if anything is still in flight. - Async eval never pushes weights (`evaluate(..., sync_weights=False)`): it - scores the engine-resident policy so the ledger stays exact. The agentic - quiesce is a turn-boundary `abort` + final poll, folded into - `AsyncAgenticRolloutEngine.quiesce()`. Reap-vs-launch ordering is trainer - statement order (diffusion polls before topping up; see its `_next_step`). + `weight_version`, raising unless decode-idle (batch: generations still in + flight; agentic: a submitted drive not yet finalized/quiesced). Async eval + never pushes weights (`evaluate(..., sync_weights=False)`): it scores the + engine-resident policy, logged as `eval/weight_version`, so the ledger + stays exact. The agentic quiesce is a turn-boundary `abort` + final poll, + folded into `AsyncAgenticRolloutEngine.quiesce()`. Reap-vs-launch ordering + is trainer statement order (diffusion polls before topping up; see its + `_next_step`). - **Reward/advantage methods are not engine code** — `Part.compute_advantages` and `Sample.propagate_rewards` are called by the trainer after scoring. An engine fills generation fields such as `segment`, `conditions`, `primitive`, and diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py index 9ebd1313..f4b162c4 100644 --- a/unirl/rollout/engine/asynchronous.py +++ b/unirl/rollout/engine/asynchronous.py @@ -356,6 +356,7 @@ def __init__(self, rollout: Any, *, group_size: int, start_gen_id: int = 0) -> N self._buffer: VersionedBuffer[List["Sample"]] = VersionedBuffer() self._gen_id = int(start_gen_id) self._weight_version = 0 + self._drive_active = False @property def weight_version(self) -> int: @@ -364,16 +365,20 @@ def weight_version(self) -> int: def sync_weights(self, weight_sync: Any) -> int: """Push train weights via *weight_sync* and advance the version ledger. - The only sanctioned weight-push path. The facade cannot see the - coordinator's drive state, so the quiesce contract stays with the - caller: sync only while the drive is finalized/quiesced (decode-idle). + The only sanctioned weight-push path — pairing the push with the bump + is what keeps the ledger truthful. Raises while a drive is active (a + weight push must be decode-idle); a joined ``finalize_if_drained`` or + ``quiesce`` ends the drive. """ + if self._drive_active: + raise RuntimeError("sync_weights with a drive active; finalize or quiesce() first") weight_sync.sync() self._weight_version += 1 return self._weight_version def submit(self, tasks: List["Sample"]) -> None: """Fire a background drive over a flat task list (fresh siblings + carried partials).""" + self._drive_active = True # before the call: a failed submit must still block sync self._rollout.submit(tasks) def poll(self) -> int: @@ -386,6 +391,7 @@ def finalize_if_drained(self) -> Optional[int]: completed = self._rollout.finalize_if_drained()[0] if completed is None: return None + self._drive_active = False return self._ingest(completed) def drain_freshest(self, n: int, *, max_staleness: int) -> Optional[List[List["Sample"]]]: @@ -401,6 +407,7 @@ def quiesce(self) -> List["Sample"]: version they completed under.""" carried = self._rollout.abort()[0] self.poll() + self._drive_active = False return carried def discard_roots(self, roots: Iterable[str]) -> int: diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 88d8e163..df63d5e4 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -211,9 +211,14 @@ def evaluate(self, rollout_id: int, *, sync_weights: bool = False) -> float: Eval is read-only — weight deployment is governed solely by ``weight_sync_interval`` (pushes go through ``engine.sync_weights``). The scored policy is the one already on the rollout slab, 1..interval - optimizer steps old (exactly 1 at ``weight_sync_interval=1``). + optimizer steps old (exactly 1 at ``weight_sync_interval=1``); its + ledger version is logged as ``eval/weight_version``. """ - return super().evaluate(rollout_id, sync_weights=sync_weights) + acc = super().evaluate(rollout_id, sync_weights=sync_weights) + engine = getattr(self, "_async_engine", None) + if engine is not None: # absent only for a standalone eval outside train() + self.wandb_logger.log_eval(rollout_id + 1, {"weight_version": engine.weight_version}) + return acc def _prepare_rollout(self, *, sync_weights: bool) -> bool: """Prepare the resident separate-slab engine without colocate handoffs. diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 3a63d18d..ce7d99b1 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -114,8 +114,13 @@ def _score_completed(self, gen_id: int, completed: Sample) -> List[Sample]: def evaluate(self, step: int, *, sync_weights: bool = False, sleep_after: bool = False) -> float: """Resident-policy eval — async defaults: no weight push (the ledger stays exact; deployment cadence belongs to ``weight_sync_interval`` alone) and - no post-eval sleep (the disaggregated engine stays resident).""" - return super().evaluate(step, sync_weights=sync_weights, sleep_after=sleep_after) + no post-eval sleep (the disaggregated engine stays resident). The scored + ledger version is logged as ``eval/weight_version``.""" + result = super().evaluate(step, sync_weights=sync_weights, sleep_after=sleep_after) + engine = getattr(self, "_async_engine", None) + if engine is not None: # absent only for a standalone eval outside train() + self.wandb_logger.log_eval(step, {"weight_version": engine.weight_version}) + return result def _drain_all(self) -> None: """Finish + buffer EVERY in-flight generation (the single-threaded quiesce). From c4cbf9ba0d19dbdf578868fa36af4be5b8f1ac4b Mon Sep 17 00:00:00 2001 From: haonan3 Date: Mon, 3 Aug 2026 14:30:20 +0800 Subject: [PATCH 3/4] fix(trainer): deterministic diffusion eval sends a pure-ODE request MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit eval_eta=0 previously rode along with the training-resolved sde_indices — a contradictory request the central kernel silently degrades to ODE but BAGEL's worker-resident scheduler refuses (RuntimeError at the first gated step). Clearing the gate at eval_eta<=0 makes the request say what eval means; SD3-family trajectories are unchanged (they already ran ODE). --- unirl/trainer/diffusion.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/unirl/trainer/diffusion.py b/unirl/trainer/diffusion.py index aaf57c5d..3553b68f 100644 --- a/unirl/trainer/diffusion.py +++ b/unirl/trainer/diffusion.py @@ -531,8 +531,9 @@ def evaluate( Mirrors :meth:`train_step`'s rollout+reward path but skips advantage/backward. Generates at the deterministic best-quality setting (``cfg_text_scale= - eval_cfg_text_scale``, ``eta=eval_eta``; ``eval_samples_per_prompt`` x_T per - prompt) and scores. The training reward plus every shared-set + eval_cfg_text_scale``, ``eta=eval_eta`` — at ``eval_eta=0`` the SDE gate + is also cleared, so the request is pure ODE; ``eval_samples_per_prompt`` + x_T per prompt) and scores. The training reward plus every shared-set ``eval_rewards`` suite scores the SAME generated images over the default eval set (``run.eval_data_path``, ``eval_num_prompts`` prompts); each own-set suite then gets its own generation pass over its own prompts. @@ -555,6 +556,11 @@ def evaluate( samples_per_prompt=self.eval_samples_per_prompt, eta=self.eval_eta, ) + if self.eval_eta <= 0.0: + # Deterministic eval must also clear the SDE gate: eta=0 with gated + # steps is a contradictory request — the central kernel degrades such + # steps to ODE, but worker-resident schedulers (BAGEL) refuse the pair. + replace_kwargs.update(sde_indices=[], scheduler=None) if "cfg_text_scale" in {f.name for f in dataclasses.fields(base_diffusion)}: replace_kwargs["cfg_text_scale"] = self.eval_cfg_text_scale else: From 2ece8755c4e4e49757547fe943af2e538a2ceed2 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Mon, 3 Aug 2026 14:30:24 +0800 Subject: [PATCH 4/4] fix(rollout): log each ledger push at the engine The actor-side [LoRA-SYNC] lines never reach the driver log, leaving weight deployments invisible in stdout. sync_weights is now the single push path, so one driver-side INFO line covers every async push. --- unirl/rollout/engine/asynchronous.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py index f4b162c4..c2f95eb1 100644 --- a/unirl/rollout/engine/asynchronous.py +++ b/unirl/rollout/engine/asynchronous.py @@ -251,6 +251,7 @@ def sync_weights(self, weight_sync: Any) -> int: raise RuntimeError(f"sync_weights with {len(self._pool)} generations in flight; quiesce() first") weight_sync.sync() self._weight_version += 1 + logger.info("sync_weights: pushed train weights; weight_version -> %d", self._weight_version) return self._weight_version @property @@ -374,6 +375,7 @@ def sync_weights(self, weight_sync: Any) -> int: raise RuntimeError("sync_weights with a drive active; finalize or quiesce() first") weight_sync.sync() self._weight_version += 1 + logger.info("sync_weights: pushed train weights; weight_version -> %d", self._weight_version) return self._weight_version def submit(self, tasks: List["Sample"]) -> None: