diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index 2f1cdec11f..c325986e01 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -471,9 +471,6 @@ def auto_setup_weight_broadcast(self): "PEFT-shaped directory on disk (LoRAModel.from_local_checkpoint) - in-memory transports " "have no disk artifact to load from." ) - # The final version v{max_steps} is broadcast iff something consumes it: - # training never samples from it, but a configured final eval measures it. - broadcast_final = self.orchestrator.eval is not None if self.weight_broadcast.type in ("nccl", "nixl"): inference_world_size = ( self.inference.vllm.data_parallel_size * self.inference.vllm.tensor_parallel_size @@ -485,7 +482,6 @@ def auto_setup_weight_broadcast(self): port=self.weight_broadcast.port, timeout=self.weight_broadcast.timeout, inference_world_size=inference_world_size, - broadcast_final=broadcast_final, ) if self.weight_broadcast.type == "nccl": transport_config = dict( @@ -501,10 +497,10 @@ def auto_setup_weight_broadcast(self): self.orchestrator.weight_broadcast = orchestrator_config_type(**common_config, **transport_config) elif self.weight_broadcast.type == "filesystem": self.trainer.weight_broadcast = TrainerFileSystemWeightBroadcastConfig( - timeout=self.weight_broadcast.timeout, broadcast_final=broadcast_final + timeout=self.weight_broadcast.timeout ) self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig( - timeout=self.weight_broadcast.timeout, broadcast_final=broadcast_final + timeout=self.weight_broadcast.timeout ) if self.inference is not None: self.inference.weight_broadcast = InferenceWeightBroadcastConfig(type=self.weight_broadcast.type) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 24efcaea17..28b497a2fb 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -42,12 +42,6 @@ class BaseWeightBroadcastConfig(BaseConfig): """Timeout in seconds for the broadcast handshake and transfer. The trainer fails the run when no consumer acknowledges an offered version in time.""" - broadcast_final: bool = True - """Internal - stamped by the RL/SFT launcher, never set by users: whether the - trainer broadcasts the final version v{max_steps}. True iff something consumes - it (a configured final eval) - training itself never samples from the final - version.""" - class RunConfig(BaseConfig): name: str | None = None diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 0cc31fb2c4..d779f0c6f6 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -210,6 +210,8 @@ def __init__( # *why* — the orchestrator toggles this based on step / policy lead. self.dispatch_allowed = asyncio.Event() self.dispatch_allowed.set() + self.policy_update_pending = False + self.scheduling_lock = asyncio.Lock() self.stopped = asyncio.Event() self.task: asyncio.Task | None = None @@ -371,6 +373,12 @@ async def on_version_pending(self, step: int) -> None: the resulting aborts are processed while the engine is still stepping — otherwise the orphaned KV transfers crash the decode engine on resume (see ``WeightWatcher.apply_policy_update``).""" + self.policy_update_pending = True + # Wait for a scheduling call that started before the pending update. + # No rollout can cross the inference weight swap after this barrier. + async with self.scheduling_lock: + pass + if self.train_envs is None or self.progress is None: return min_version = min_fresh_version(self.progress.step, self.max_off_policy_steps) @@ -392,7 +400,8 @@ async def on_version_pending(self, step: int) -> None: ) async def on_new_version(self, step: int) -> None: - """No-op: the dispatcher drains in ``on_version_pending`` (pre-pause).""" + """Resume rollout scheduling after inference applies the new policy.""" + self.policy_update_pending = False async def fill_inflight(self) -> None: """Schedule new rollouts up to ``max_inflight``, honoring @@ -401,28 +410,33 @@ async def fill_inflight(self) -> None: respects it. When ``PREFER_EVAL``'s source exhausts we flip back to ``PREFER_TRAIN`` so the eval tail drains alongside fresh train.""" while True: + if self.policy_update_pending: + return if self.available_permits <= 0 or self.admission_budget() <= 0: return - if self.mode == DispatcherMode.PREFER_EVAL: - # PREFER_EVAL is only entered when the orchestrator triggers - # eval, which requires ``eval_source`` to be configured - assert self.eval_source is not None - if not self.eval_has_work: - # Eval source + all eval groups fully dispatched. Flip - # to PREFER_TRAIN so any remaining permits go to train - # while the in-flight eval tail completes naturally - self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained") - continue - scheduled = await self.try_schedule("eval") - if not scheduled: - return - else: # PREFER_TRAIN — respects the orchestrator's dispatch gate - if not self.dispatch_allowed.is_set(): - return - scheduled = await self.try_schedule("train") - if not scheduled: + async with self.scheduling_lock: + if self.policy_update_pending: return + if self.mode == DispatcherMode.PREFER_EVAL: + # PREFER_EVAL is only entered when the orchestrator triggers + # eval, which requires ``eval_source`` to be configured + assert self.eval_source is not None + if not self.eval_has_work: + # Eval source + all eval groups fully dispatched. Flip + # to PREFER_TRAIN so any remaining permits go to train + # while the in-flight eval tail completes naturally + self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained") + continue + scheduled = await self.try_schedule("eval") + if not scheduled: + return + else: # PREFER_TRAIN — respects the orchestrator's dispatch gate + if not self.dispatch_allowed.is_set(): + return + scheduled = await self.try_schedule("train") + if not scheduled: + return def switch_mode(self, new_mode: DispatcherMode, *, reason: str) -> None: if new_mode == self.mode: diff --git a/src/prime_rl/orchestrator/eval_source.py b/src/prime_rl/orchestrator/eval_source.py index ca3b33aa6a..18e191c601 100644 --- a/src/prime_rl/orchestrator/eval_source.py +++ b/src/prime_rl/orchestrator/eval_source.py @@ -1,7 +1,7 @@ """EvalSource: trigger-driven, finite-per-epoch pull of eval examples. -The orchestrator pokes ``trigger(step)`` after each ship + once at -startup; the dispatcher pulls via ``next_task()`` until +The policy watcher calls ``trigger(step)`` after each applied policy, +including startup. The dispatcher pulls via ``next_task()`` until ``bool(source) == False``. Constructed only when eval is configured.""" from __future__ import annotations @@ -37,8 +37,8 @@ def __init__( self.queue: deque[TaskRequest] = deque() - # On resume we skip the startup eval; on fresh start the first - # trigger fires every env (subject to ``skip_first_step``) + # A fresh run evaluates the base policy. Resumed runs apply interval + # rules to the loaded checkpoint and later policies. self.first_trigger = not is_resumed def trigger(self, step: int, *, force: bool = False) -> list[str]: diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 1fc1645d68..6509741551 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -78,11 +78,7 @@ from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.logger import format_time, get_logger, setup_logger from prime_rl.utils.pathing import get_broadcast_dir -from prime_rl.utils.utils import ( - clean_exit, - final_broadcast_version, - resolve_latest_ckpt_step, -) +from prime_rl.utils.utils import clean_exit, resolve_latest_ckpt_step monkey_patch_oai_iterable_types() monkey_patch_chat_completion_logprobs() @@ -166,9 +162,10 @@ def __init__(self, config: OrchestratorConfig) -> None: self.eval_triggered_at = {} self.consecutive_empty_batches = 0 self.gate_closed_at = None - # Pulsed by the version hooks so a held ship can re-check ``policy.version`` + # Pulsed after inference applies a policy so held work can re-check it. self.version_advanced = asyncio.Event() self.wait_for_policy_time = 0.0 + self.eval_triggered_steps: set[int] = set() self.component_tasks = [] # Always assigned by ``setup()``; None-initialized so teardown can run @@ -320,14 +317,9 @@ async def setup(self) -> None: # scratch). The startup broadcast is always coming, so wait for it rather # than failing immediately when it is not there yet. sync_version = self.resume_step if self.resume_step is not None else 0 - get_logger().info(f"Syncing inference to the trainer's startup broadcast (v{sync_version})") - t0 = time.perf_counter() wait_timeout = (config.ckpt.wait_for_weights_timeout if config.ckpt else None) or ( STARTUP_WEIGHT_WAIT_TIMEOUT_S ) - await self.receiver.sync_startup(sync_version, timeout=wait_timeout) - self.policy.version = sync_version - get_logger().debug(f"Synced inference to policy v{sync_version} in {format_time(time.perf_counter() - t0)}") self.eval_source: EvalSource | None = ( EvalSource( @@ -390,9 +382,12 @@ async def setup(self) -> None: self.watcher = WeightWatcher( self.receiver, policy=self.policy, - observers=[self.dispatcher, self], - ckpt_step=self.policy.version, + observers=[self.dispatcher], + ckpt_step=sync_version, ) + if self.eval_source is not None: + self.watcher.on_update(self.trigger_eval) + self.watcher.on_update(self.on_policy_update) # Single periodic logger for the whole pipeline. It's the only # consumer of ``dispatcher.metrics.drained()`` (which clears on read) self.lag_monitor = EventLoopLagMonitor() @@ -419,6 +414,11 @@ async def setup(self) -> None: wandb_enabled=wandb_enabled, ) + get_logger().info(f"Syncing inference to the trainer's startup broadcast (v{sync_version})") + t0 = time.perf_counter() + await self.watcher.sync_startup(sync_version, timeout=wait_timeout) + get_logger().debug(f"Synced inference to policy v{sync_version} in {format_time(time.perf_counter() - t0)}") + async def start(self) -> None: """Run the orchestrator until shutdown. Drives setup, spawns the background tasks, runs the main loop in this task, then cleans up.""" @@ -438,14 +438,6 @@ async def start(self) -> None: asyncio.create_task(self.watcher.start(), name="watcher"), ] - # Base-model eval (policy v0) — fires before any train rollouts, logged at the first - # step, unless ``eval.skip_first_step=True``. On resume, defaults to assuming a clean - # exit (evals already completed); set ``eval.retrigger_on_resume=True`` to also re-fire - # interval-aligned evals at the checkpoint step (e.g. after a crash). - if config.eval is not None and config.eval.retrigger_on_resume and self.resume_step is not None: - self.maybe_trigger_eval(self.resume_step) - self.maybe_trigger_eval(self.progress.step) - # Anchor step-time clock so the first step measures startup → first batch self.last_batch_at = time.perf_counter() @@ -484,18 +476,11 @@ async def start(self) -> None: get_logger().warning("Orchestrator cleanup complete (forced)") trim_process_memory() - @property - def final_version(self) -> int | None: - """Newest policy version the trainer will ever broadcast.""" - if self.config.max_steps is None: - return None - return final_broadcast_version(self.config.max_steps, self.config.weight_broadcast.broadcast_final) - async def wait_for_version(self, version: int, reason: str) -> None: """Bounded wait until the watcher has applied v{version}.""" if self.policy.version >= version: return - get_logger().info(f"Waiting for the trainer to broadcast v{version} {reason}") + get_logger().info(f"Waiting for inference to apply policy v{version} {reason}") async def wait() -> None: while self.policy.version < version: @@ -514,15 +499,15 @@ async def wait() -> None: try: await asyncio.wait_for(wait(), timeout=timeout) except asyncio.TimeoutError: - get_logger().warning(f"Trainer did not broadcast v{version} within {timeout}s — proceeding anyway") + get_logger().warning(f"Inference did not apply policy v{version} within {timeout}s — proceeding anyway") async def wait_for_final_broadcast(self) -> None: """Stay alive for the trainer's last broadcast. Every broadcast is a blocking rendezvous — tearing down the watcher before it would strand the trainer inside the handshake.""" - if self.final_version is None: + if self.config.max_steps is None: return - await self.wait_for_version(self.final_version, reason="before shutdown") + await self.wait_for_version(self.config.max_steps, reason="before shutdown") async def main_loop(self) -> None: """Consume dispatcher results and route them to the train / eval sink. @@ -644,16 +629,16 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: f"({n_trainable / effective.num_traces:.1%}) — consider reviewing task difficulty" ) - # Ship batch ``step`` only once the trainer has published v{step-1-TARGET_LAG}. + # Ship batch ``step`` only once inference has applied v{step-1-TARGET_LAG}. # Without this, fast envs fill batches from buffered rollouts and the # orchestrator races arbitrarily far ahead of the trainer. Always - # satisfiable: the trainer broadcasts every version except v{max_steps}, - # and ``wait_for_final_broadcast`` keeps the watcher alive through the - # last rendezvous after the pipeline drains. + # satisfiable: the trainer broadcasts every version, and + # ``wait_for_final_broadcast`` keeps the watcher alive through the last + # rendezvous after the pipeline drains. required_version = step - 1 - TARGET_LAG if self.policy.version < required_version: get_logger().info( - f"Holding batch {step} until the trainer publishes policy v{required_version} " + f"Holding batch {step} until inference applies policy v{required_version} " f"(currently v{self.policy.version})" ) hold_start = time.perf_counter() @@ -757,14 +742,8 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: self.log_train_batch(batch, step=step, step_time=step_time) - # The final eval must measure the final weights: hold until the - # trainer's last broadcast (v{max_steps} when evals consume it) has - # been applied before triggering it. Satisfiable — the trainer - # broadcasts right after consuming the batch this call just shipped. - if config.eval is not None and config.max_steps is not None and step >= config.max_steps: - assert self.final_version is not None - await self.wait_for_version(self.final_version, reason="for the final eval") - self.maybe_trigger_eval(self.progress.step) + if config.max_steps is not None and step >= config.max_steps: + await self.wait_for_version(step, reason="before shutdown") # Drain right after shipping the final batch. Waiting for a further # batch to fill would burn inference on data that can never train — # and with a tight ``max_off_policy_steps`` it never fills at all (the @@ -784,14 +763,18 @@ async def start_draining(self, reason: str) -> None: f"train episode(s); any in-flight evals will complete)" ) - def maybe_trigger_eval(self, step: int) -> None: + async def trigger_eval(self, step: int) -> None: """Fire eligible eval epochs and flip to ``PREFER_EVAL`` if anything fires. No-op when eval is not configured.""" - if self.eval_source is None: + if self.eval_source is None or step in self.eval_triggered_steps: return - fired = self.eval_source.trigger(step) + if self.resume_step == step and self.config.eval is not None and not self.config.eval.retrigger_on_resume: + return + is_final = self.config.max_steps is not None and step >= self.config.max_steps + fired = self.eval_source.trigger(step, force=is_final) if not fired: return + self.eval_triggered_steps.add(step) reason = f"eval was triggered at step {step}" self.dispatcher.switch_mode(DispatcherMode.PREFER_EVAL, reason=reason) now = time.perf_counter() @@ -942,10 +925,6 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: policy_versions = {span.start for span in policy_spans if span is not None} policy_versions.update(failure.policy_version for failure in batch.failures) policy_version = min(policy_versions) - if len(policy_versions) > 1: - get_logger().warning( - f"Eval {batch.env_name} step {batch.step} had mixed policy versions: {sorted(policy_versions)}" - ) # Episode metrics over {all,effective} (eval batches are per-env, so no `agg` axis). # ``effective`` = non-errored; pass@k / pass^k only over the effective set. episodes = batch.episodes @@ -969,7 +948,7 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: triggered_at = self.eval_triggered_at.pop((batch.env_name, batch.step), None) elapsed = (time.perf_counter() - triggered_at) if triggered_at is not None else 0.0 get_logger().success( - f"Evaluated {batch.env_name} (Step {batch.step}) | " + f"Evaluated {batch.env_name} | " f"Policy v{policy_version} | {format_time(elapsed):>7} | Reward {eff.reward.mean():.4f} | " f"Turns {eff.num_turns.mean():.1f} | Branches {eff.num_branches.mean():.1f} | " f"Error {full.has_error.mean():.1%} | Truncation {eff.is_truncated.mean():.1%}" @@ -1006,7 +985,7 @@ def update_dispatch_gate(self) -> None: if lead > TARGET_LAG: if was_set: get_logger().info( - f"Pausing dispatcher until the trainer publishes policy v{self.progress.step - 1 - TARGET_LAG} " + f"Pausing dispatcher until inference applies policy v{self.progress.step - 1 - TARGET_LAG} " f"(currently v{self.policy.version})" ) self.gate_closed_at = time.perf_counter() @@ -1019,15 +998,10 @@ def update_dispatch_gate(self) -> None: self.gate_closed_at = None gate.set() - async def on_version_pending(self, step: int) -> None: - """``VersionObserver`` hook, fired at publish confirmation (pre-apply): - ``policy.version`` already carries the new version, so wake a held ship.""" - self.version_advanced.set() - - async def on_new_version(self, step: int) -> None: - """``VersionObserver`` hook: the weight update completed; - re-evaluate the dispatch gate (may resume if the trainer caught up).""" + async def on_policy_update(self, _step: int) -> None: + """Refresh policy-dependent state after inference applies new weights.""" self.update_dispatch_gate() + self.version_advanced.set() async def stop(self) -> None: """Bounded best-effort teardown of all components. Has a global diff --git a/src/prime_rl/orchestrator/watcher.py b/src/prime_rl/orchestrator/watcher.py index 968bd7d67d..adf1f9d961 100644 --- a/src/prime_rl/orchestrator/watcher.py +++ b/src/prime_rl/orchestrator/watcher.py @@ -6,6 +6,7 @@ import asyncio import time +from collections.abc import Awaitable, Callable from prime_rl.orchestrator.types import Policy, VersionObserver from prime_rl.transports.weights import WeightReceiver @@ -38,6 +39,19 @@ def __init__( self.task: asyncio.Task | None = None self.update_lock = asyncio.Lock() self.stopped = asyncio.Event() + self.update_hooks: list[Callable[[int], Awaitable[None]]] = [] + + def on_update(self, hook: Callable[[int], Awaitable[None]]) -> None: + """Register an async hook that runs after inference applies new weights.""" + self.update_hooks.append(hook) + + async def sync_startup(self, step: int, timeout: float) -> None: + """Apply the startup policy and notify the registered update hooks.""" + async with self.update_lock: + await self.receiver.sync_startup(step, timeout) + self.ckpt_step = step + self.policy.version = step + await self._notify_update(step) async def start(self) -> None: self.task = asyncio.current_task() @@ -72,12 +86,9 @@ async def apply_policy_update(self, next_step: int) -> None: await self.receiver.wait_published(next_step, cancelled=self.stopped.is_set) self.last_wait_for_ckpt_time = time.perf_counter() - t0 - # Publish confirmed: the policy version advances here, before the - # apply — inference pauses during the update, so nothing can generate - # under the new number from the old weights, and a ship held on this - # version releases without waiting out the inference weight reload. + # Record the published version before notifying pending observers. + # ``policy.version`` advances only after inference applies it. self.ckpt_step = next_step - self.policy.version = next_step # Drain stale rollouts BEFORE pausing the inference engines. # Aborting a rollout triggers vLLM's KV-connector cleanup (NIXL's @@ -102,17 +113,22 @@ async def apply_policy_update(self, next_step: int) -> None: await self.receiver.receive(next_step) self.last_update_weights_time = time.perf_counter() - t1 self.update_count += 1 + self.policy.version = next_step get_logger().debug( f"Updated inference weights to policy v{next_step} in {format_time(self.last_update_weights_time)}" ) - for observer in self.observers: - try: - await observer.on_new_version(next_step) - except Exception as exc: - get_logger().warning( - f"Observer {type(observer).__name__}.on_new_version({next_step}) raised: {exc!r}" - ) + await self._notify_update(next_step) + + async def _notify_update(self, step: int) -> None: + for observer in self.observers: + try: + await observer.on_new_version(step) + except Exception as exc: + get_logger().warning(f"Observer {type(observer).__name__}.on_new_version({step}) raised: {exc!r}") + + for hook in self.update_hooks: + await hook(step) def gauges(self) -> dict[str, float]: return { diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index fe63de8d47..ef387113ec 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -67,7 +67,7 @@ from prime_rl import monitors from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title -from prime_rl.utils.utils import clean_exit, final_broadcast_version, resolve_latest_ckpt_step +from prime_rl.utils.utils import clean_exit, resolve_latest_ckpt_step from ring_flash_attn import substitute_hf_flash_attn @@ -590,31 +590,23 @@ def train(config: TrainerConfig): current_lr = optimizer.param_groups[0]["lr"] forward_backward_time = time.perf_counter() - forward_backward_start_time - # Broadcast the model just produced (policy v{progress.step}) so the orchestrator can - # sample its next step from it. Every broadcast is a handshake with the consumer, so - # versions past the last consumed one are skipped (``final_broadcast_version``: - # training never samples v{max_steps}, but a configured final eval measures it). + # Broadcast the model just produced (policy v{progress.step}). The final + # broadcast keeps inference synchronized with the completed trainer. if weight_sender is None: broadcast_weights_time = 0 else: - broadcast_unused = config.max_steps is not None and progress.step > final_broadcast_version( - config.max_steps, config.weight_broadcast.broadcast_final - ) - if not broadcast_unused: - broadcast_weights_start_time = time.perf_counter() - # The per-layer gather + fp8 conversion peaks ~50 GiB above the - # resident weights; release cached blocks (incl. offload-stream - # pools) so the broadcast gets the full headroom. Drain all - # pending work first: empty_cache returns blocks to the driver, - # so a still-running kernel holding a cached block (e.g. the - # optimizer step's tail) faults with an illegal memory access - # once its block is freed under it. - torch.cuda.synchronize() - torch.cuda.empty_cache() - weight_sender.broadcast(model, step=progress.step) - broadcast_weights_time = time.perf_counter() - broadcast_weights_start_time - else: - broadcast_weights_time = 0 + broadcast_weights_start_time = time.perf_counter() + # The per-layer gather + fp8 conversion peaks ~50 GiB above the + # resident weights; release cached blocks (incl. offload-stream + # pools) so the broadcast gets the full headroom. Drain all + # pending work first: empty_cache returns blocks to the driver, + # so a still-running kernel holding a cached block (e.g. the + # optimizer step's tail) faults with an illegal memory access + # once its block is freed under it. + torch.cuda.synchronize() + torch.cuda.empty_cache() + weight_sender.broadcast(model, step=progress.step) + broadcast_weights_time = time.perf_counter() - broadcast_weights_start_time # Checkpoint the step we just finished (model = policy v{progress.step}). if ( diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index 9192e56399..9992e9ae80 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -236,17 +236,17 @@ def validate_run_dir( def clean_future_steps(output_dir: Path, resume_step: int) -> None: - """Remove stale rollouts, broadcasts, and traces past ``resume_step``. + """Remove stale rollouts past ``resume_step`` and broadcasts from it onward. Pass ``resume_step=-1`` to wipe every step directory (fresh runs). """ - dirs = [ - get_rollout_dir(output_dir), - get_broadcast_dir(output_dir), + cleanup_rules = [ + (get_rollout_dir(output_dir), lambda step: step > resume_step), + (get_broadcast_dir(output_dir), lambda step: step >= resume_step), ] - for directory in dirs: - steps_to_delete = [step for step in get_all_ckpt_steps(directory) if step > resume_step] + for directory, should_delete in cleanup_rules: + steps_to_delete = [step for step in get_all_ckpt_steps(directory) if should_delete(step)] if not steps_to_delete: continue get_logger().info( diff --git a/src/prime_rl/utils/utils.py b/src/prime_rl/utils/utils.py index f02e9d742b..047fea4cab 100644 --- a/src/prime_rl/utils/utils.py +++ b/src/prime_rl/utils/utils.py @@ -163,13 +163,6 @@ def get_free_port() -> int: return port -def final_broadcast_version(max_steps: int, broadcast_final: bool) -> int: - """Newest policy version the trainer will ever broadcast: v{max_steps} when - something consumes it (a configured final eval), else v{max_steps - 1} - - training never samples from the final version.""" - return max_steps if broadcast_final else max_steps - 1 - - @contextmanager def default_dtype(dtype): prev = torch.get_default_dtype() diff --git a/tests/unit/utils/test_pathing.py b/tests/unit/utils/test_pathing.py index 09c73f9a74..6e9a5c7eb5 100644 --- a/tests/unit/utils/test_pathing.py +++ b/tests/unit/utils/test_pathing.py @@ -1,6 +1,12 @@ import pytest -from prime_rl.utils.pathing import validate_run_dir +from prime_rl.utils.pathing import ( + clean_future_steps, + get_broadcast_dir, + get_rollout_dir, + get_step_path, + validate_run_dir, +) def test_nonexistent_dir_passes(tmp_path): @@ -88,3 +94,20 @@ def test_clean_outside_output_dir_raises(tmp_path): with pytest.raises(ValueError, match="remain under output_dir"): validate_run_dir(escaped, output_dir=tmp_path / "outputs", resuming=False, clean=True) assert escaped.exists() + + +def test_clean_future_steps_rebuilds_resume_broadcast(tmp_path): + rollout_dir = get_rollout_dir(tmp_path) + broadcast_dir = get_broadcast_dir(tmp_path) + for parent in (rollout_dir, broadcast_dir): + for step in (1, 2, 3): + get_step_path(parent, step).mkdir(parents=True) + + clean_future_steps(tmp_path, resume_step=2) + + assert get_step_path(rollout_dir, 1).exists() + assert get_step_path(rollout_dir, 2).exists() + assert not get_step_path(rollout_dir, 3).exists() + assert get_step_path(broadcast_dir, 1).exists() + assert not get_step_path(broadcast_dir, 2).exists() + assert not get_step_path(broadcast_dir, 3).exists() diff --git a/tests/utils.py b/tests/utils.py index c1587adb27..048df9d0b9 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -108,7 +108,7 @@ def check_loss_goes_down(lines: list[str]): def check_final_eval_reward_above(lines: list[str], env_name: str, min_threshold: float): - """Assert the LAST `Evaluated {env_name} (Step N) | ... | Reward X.XXXX` + """Assert the last `Evaluated {env_name} | ... | Reward X.XXXX` line reports a reward above ``min_threshold``. Robust for short distill smokes: until the policy converges, eval reward is