[Feat] Support immediate rollout recovery by export hf asynchronously#1966
[Feat] Support immediate rollout recovery by export hf asynchronously#1966YanhuiDua wants to merge 1 commit into
Conversation
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| self._fit() | ||
| finally: | ||
| self._exp_tracker.close() |
There was a problem hiding this comment.
Claude: Bug (Critical): close_trace() has been removed from the finally block. The trace runtime (TraceRuntime.close()) flushes pending spans and stops local collector processes. Without it, trace data in-flight at end of training may be lost or delayed until the process exits (relying only on the atexit fallback).
The import of close_trace was also dropped (line 44). This appears to be an accidental deletion during refactoring — the executor shutdown and trace shutdown serve different purposes.
| self._fit() | |
| finally: | |
| self._exp_tracker.close() | |
| finally: | |
| self._exp_tracker.close() | |
| if self._hf_export_executor is not None: | |
| self._hf_export_executor.shutdown(wait=True) | |
| close_trace() |
(Same fix needed in the DisaggregatedRLTrainer.fit() at line ~1981.)
| previous_ready_hf_path is not None | ||
| and str(previous_ready_hf_path) not in self._meta.latest_exp.hf_checkpoint_list | ||
| ): | ||
| rmtree(previous_ready_hf_path, ignore_errors=True) |
There was a problem hiding this comment.
Claude: Race condition (Warning): There is a timing window between clear_ready_recovery_hf and the rmtree below (line 898) where the health manager's background restart thread may have already captured the old _ready_recovery_hf value in a local variable (ready_recovery_hf = self._ready_recovery_hf in _restart_worker_group) and started a reinit.remote(model_path=old_path). Deleting the path from disk while a Ray worker is loading from it would cause the restart to fail.
This isn't catastrophic (the health check will detect the failure and shut down the group), but it means the immediate-recovery path degrades to a restart failure for that group, which is counterproductive.
A safer approach: defer deletion of the old path until after the next successful weight sync (where all workers are confirmed healthy), or track in-flight restarts with a reference count before removing the directory.
| return None | ||
|
|
||
| self._ready_recovery_hf_path = finalized_hf_path | ||
| return finalized_hf_path |
There was a problem hiding this comment.
Claude: Thread safety (Warning): The closure wait_and_publish_recovery_hf writes self._ready_recovery_hf_path (line 938) from the background ThreadPoolExecutor thread. The main thread reads it in the next _maybe_save_recovery_hf call (line 881: previous_ready_hf_path = self._ready_recovery_hf_path).
In CPython this works because Future.result() (called in _wait_or_disable_immediate_recovery) synchronizes through internal condition variables and the GIL, establishing a happens-before relationship. But this relies on an implementation detail and is fragile for anyone maintaining this code. Consider either:
- Returning the path only through the Future's result value (which you already do) and reading it from
pending.result()instead of fromself._ready_recovery_hf_path, or - Protecting the attribute with an explicit
threading.Lock.
| pending = cast(Future[Path], self._pending_hf_export) | ||
| return pending.done() | ||
|
|
There was a problem hiding this comment.
Claude: Defensiveness (Nit): is_hf_export_done and wait_hf_export both use cast(Future[Path], self._pending_hf_export) without guarding against None. If either is ever called out of sequence (before start_hf_export), this will produce a confusing AttributeError: 'NoneType' object has no attribute 'done'.
Consider raising an explicit RuntimeError if self._pending_hf_export is None:
if self._pending_hf_export is None:
raise RuntimeError("No pending HF export — call start_hf_export first.")
return self._pending_hf_export.done()
| if pending is None: | ||
| return | ||
|
|
||
| disable_immediate_recovery = not pending.done() | ||
| if disable_immediate_recovery: | ||
| self._enable_immediate_recovery = False | ||
| self.logger.warning( | ||
| "Disable immediate recovery because the previous recovery HF " |
There was a problem hiding this comment.
Claude: Design concern (Warning): When the export doesn't finish before the next weight sync, _enable_immediate_recovery is permanently set to False — it cannot be re-enabled for the remainder of training.
If the slow export was a transient issue (e.g., I/O spike, one large model step), this permanently disables a useful optimization. Consider either:
- Making this a "cooldown" that re-enables after N successful weight syncs, or
- At minimum, documenting that this is a permanent disable (the log says "Disable" but doesn't say "for the rest of training").
| @@ -2007,12 +2119,16 @@ async def _fit(self): | |||
|
|
|||
| async def _sync_weights_and_save(self, model_step: int, step_timer_dict: dict): | |||
| # producer 已暂停；保持 save -> bind -> update 顺序。 | |||
There was a problem hiding this comment.
Claude: Behavioral change (Warning): This line adds restart_inactive_workers to the disaggregated trainer's weight-sync path. The old code had a TODO indicating this was intentionally omitted pending design evaluation.
Is this safe for all disaggregated configurations? The colocate trainer restarts inactive workers at the start of _sync_weights_and_save (before offload), but here it's inside the sync_weight timer, right before bind_train_rollout. If the restart takes significant time (with ROLLOUT_RAY_GET_TIMEOUT), it could delay weight updates. Also, since this runs inside timer("sync_weight"), the restart time inflates the sync_weight metric.
Please confirm this ordering is intentional and tested in the disaggregated setup.
|
Claude: ## Summary This PR adds an "immediate recovery" mode for rollout workers: instead of restarting crashed workers with empty weights and waiting for the next weight-sync, it exports an HF checkpoint asynchronously after each weight update and lets the health manager reload workers from that checkpoint immediately. The refactoring of IssuesCritical
Warning
Nit
VerdictREQUEST_CHANGES ProduceBatchResult impact: not directly affected — the recovery path rejoins the normal flow before batch accounting. RoutedExperts impact: not affected — no changes to routed-experts plumbing or object-ref ownership. |
1 participant
No description provided.