decouple trainer step completion from weight broadcast

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -747,7 +747,7 @@ def auto_setup_bench(self):
 
     @model_validator(mode="after")
     def validate_no_inference_broadcast(self):
-        """Static SFT uses the filesystem marker as its trainer handshake."""
+        """NCCL broadcast requires a running inference deployment."""
         if not self.needs_inference and self.weight_broadcast.type != "filesystem":
             raise ValueError(
                 "dataset-backed SFT requires filesystem weight broadcast — there is no "

src/prime_rl/orchestrator/static_sft.py

@@ -21,7 +21,7 @@
 from prime_rl.utils.config import to_toml_dict
 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, get_step_path, wait_for_path
+from prime_rl.utils.pathing import get_step_path, get_trainer_step_dir, wait_for_path
 from prime_rl.utils.utils import resolve_latest_ckpt_step
 
 if TYPE_CHECKING:
@@ -260,7 +260,7 @@ async def start(self) -> None:
                 samples, attempts = await source.build_batch(config.batch_size, config.token_batch_size)
                 await self.sender.send(TrainingBatch(examples=samples, step=step))
 
-                stable = get_step_path(get_broadcast_dir(config.output_dir), step) / "STABLE"
+                stable = get_step_path(get_trainer_step_dir(config.output_dir), step) / "STABLE"
                 await wait_for_path(stable)
 
                 num_tokens = sum(len(sample.token_ids) for sample in samples)

src/prime_rl/trainer/rl/broadcast/filesystem.py

@@ -37,6 +37,9 @@ def __init__(
 
     def broadcast_weights(self, model: nn.Module, step: int) -> None:
         """Broadcast weights by saving a HF-compatible checkpoint to shared filesystem and notifies the orchestrator."""
+        ready_idxs = self.multi_run_manager.ready_to_broadcast_idxs
+        if not ready_idxs:
+            return
         self.logger.debug("Starting broadcasting weights to inference engine via shared filesystem")
         start_time = time.perf_counter()
         adapter_only = self.lora_config is not None
@@ -50,7 +53,7 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None:
 
                 state_dict = revert_weight_conversion(model, state_dict)
 
-        for idx in self.multi_run_manager.ready_to_update_idxs:
+        for idx in ready_idxs:
             self.logger.debug(
                 f"Broadcasting weights for run {idx} (ready_to_update={self.multi_run_manager.ready_to_update[idx]})"
             )
@@ -99,8 +102,6 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None:
                     self.logger.warning(f"Run {idx} is deleted, skipping")
                 except Exception as e:
                     self.logger.error(f"Error broadcasting weights for run {idx}: {e}")
-                finally:
-                    self.multi_run_manager.ready_to_update[idx] = False
 
         if self.world.is_master:
             self.logger.debug(f"Weights broadcasted in {time.perf_counter() - start_time:.2f}s")

src/prime_rl/trainer/rl/broadcast/nccl.py

@@ -207,6 +207,8 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None:
         # the orchestrator has paused inference, the collectives sit unmatched
         # until NCCL's watchdog kills the process after 10 min.
         notified_runs = self._compute_notified_runs()
+        if not notified_runs:
+            return
         if self.world.is_master:
             self._notify_orchestrator(notified_runs)
             self._wait_for_nccl_ready(notified_runs)
@@ -223,9 +225,7 @@ def _compute_notified_runs(self) -> list[tuple[int, Path]]:
         non-master ranks agree on which NCCL_READY markers to wait for.
         """
         notified_runs: list[tuple[int, Path]] = []
-        for idx in self.multi_run_manager.used_idxs:
-            if not self.multi_run_manager.ready_to_update[idx]:
-                continue
+        for idx in self.multi_run_manager.ready_to_broadcast_idxs:
             try:
                 # pack() already advanced progress to the next step, so the model we just
                 # trained — policy v(step-1) — broadcasts to broadcasts/step_{step-1}.
@@ -241,10 +241,10 @@ def _compute_notified_runs(self) -> list[tuple[int, Path]]:
         return notified_runs
 
     def _notify_orchestrator(self, notified_runs: list[tuple[int, Path]]) -> None:
-        """Create STABLE markers for each notified run and clear their ready flags.
+        """Create STABLE markers for each notified run.
 
-        Master-only side effects (filesystem writes + state mutation). Called
-        after `_compute_notified_runs`; non-master ranks skip this entirely.
+        Called after `_compute_notified_runs`; non-master ranks skip this
+        filesystem side effect entirely.
         """
         for idx, save_dir in notified_runs:
             try:
@@ -255,8 +255,6 @@ def _notify_orchestrator(self, notified_runs: list[tuple[int, Path]]) -> None:
                 self.logger.warning(f"Run {idx} is deleted, skipping")
             except Exception as e:
                 self.logger.error(f"Error broadcasting weights for run {idx}: {e}")
-            finally:
-                self.multi_run_manager.ready_to_update[idx] = False
 
     def _wait_for_nccl_ready(self, notified_runs: list[tuple[int, Path]]):
         """Wait for inference workers to signal they are ready to receive NCCL broadcast."""

src/prime_rl/trainer/rl/train.py

@@ -590,9 +590,8 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
                     weight_broadcast.maybe_clean(interval_to_keep)
             else:
                 broadcast_weights_time = 0
-                # Usually the broadcast will set this. If broadcast is skipped, we need to reset this here.
-                for idx in multi_run_manager.used_idxs:
-                    multi_run_manager.ready_to_update[idx] = False
+
+        multi_run_manager.mark_ready_steps_complete()
 
         # Checkpoint the step we just finished (model = policy v{progress.step}).
         if config.max_concurrent_runs > 1:

src/prime_rl/trainer/runs.py

@@ -12,7 +12,7 @@
 from prime_rl.configs.trainer import LoRAConfig
 from prime_rl.trainer.world import get_world
 from prime_rl.utils.logger import get_logger
-from prime_rl.utils.pathing import get_all_ckpt_steps, get_stable_ckpt_steps
+from prime_rl.utils.pathing import get_all_ckpt_steps, get_stable_ckpt_steps, get_step_path, get_trainer_step_dir
 
 if TYPE_CHECKING:
     from prime_rl.configs.orchestrator import OrchestratorConfig
@@ -486,6 +486,24 @@ def used_idxs(self):
     def ready_to_update_idxs(self):
         return [idx for idx, ready in enumerate(self.ready_to_update) if ready]
 
+    @property
+    def ready_to_broadcast_idxs(self) -> list[int]:
+        return [idx for idx in self.used_idxs if self.ready_to_update[idx] and self.config[idx].needs_inference]
+
+    def mark_ready_steps_complete(self) -> None:
+        """Publish trainer-step completion and release every run updated this step."""
+        completed_idxs = [idx for idx in self.used_idxs if self.ready_to_update[idx]]
+        if self.world.is_master:
+            for idx in completed_idxs:
+                run_dir = self.get_run_dir(idx)
+                if not run_dir.exists():
+                    continue
+                step_dir = get_step_path(get_trainer_step_dir(run_dir), self.progress[idx].step - 1)
+                step_dir.mkdir(parents=True, exist_ok=True)
+                (step_dir / "STABLE").touch()
+        for idx in completed_idxs:
+            self.ready_to_update[idx] = False
+
     def run_dirs(self) -> list[Path]:
         return [self.output_dir / run_id for run_id in self.id_2_idx.keys()]
 

src/prime_rl/utils/pathing.py

@@ -88,6 +88,10 @@ def get_broadcast_dir(output_dir: Path) -> Path:
     return output_dir / "broadcasts"
 
 
+def get_trainer_step_dir(output_dir: Path) -> Path:
+    return output_dir / "trainer_steps"
+
+
 def get_step_path(path: Path, step: int) -> Path:
     return path / f"step_{step}"
 
@@ -157,7 +161,7 @@ def validate_output_dir(output_dir: Path, *, resuming: bool, clean: bool, ckpt_o
 
 
 def clean_future_steps(output_dir: Path, resume_step: int) -> None:
-    """Remove stale rollouts, broadcasts, and traces past ``resume_step``.
+    """Remove stale rollouts, broadcasts, trainer markers, and traces past ``resume_step``.
 
     Pass ``resume_step=-1`` to wipe every step directory (fresh runs).
     """
@@ -166,6 +170,7 @@ def clean_future_steps(output_dir: Path, resume_step: int) -> None:
         get_rollout_dir(output_dir),
         get_rollout_dir(run_default),
         get_broadcast_dir(run_default),
+        get_trainer_step_dir(run_default),
     ]
 
     for directory in dirs:

tests/unit/train/test_runs.py

@@ -1,11 +1,12 @@
 from pathlib import Path
+from types import SimpleNamespace
 from typing import Generator
 
 import pytest
 import tomli_w
 import torch.distributed as dist
 
-from prime_rl.trainer.runs import MultiRunManager
+from prime_rl.trainer.runs import MultiRunManager, Progress
 
 
 @pytest.fixture(autouse=True, scope="module")
@@ -64,6 +65,29 @@ def test_initial_state(tmp_path: Path) -> None:
     assert multi_run_manager.run_dirs() == []
 
 
+def test_step_completion_is_independent_from_weight_broadcast(tmp_path: Path) -> None:
+    multi_run_manager = MultiRunManager.__new__(MultiRunManager)
+    multi_run_manager.output_dir = tmp_path
+    multi_run_manager.world = SimpleNamespace(is_master=True)
+    multi_run_manager.idx_2_id = {0: "run_static", 1: "run_policy"}
+    multi_run_manager.config = {
+        0: SimpleNamespace(needs_inference=False),
+        1: SimpleNamespace(needs_inference=True),
+    }
+    multi_run_manager.progress = {0: Progress(step=4), 1: Progress(step=7)}
+    multi_run_manager.ready_to_update = [True, True]
+    for run_id in multi_run_manager.idx_2_id.values():
+        (tmp_path / run_id).mkdir()
+
+    assert multi_run_manager.ready_to_broadcast_idxs == [1]
+
+    multi_run_manager.mark_ready_steps_complete()
+
+    assert (tmp_path / "run_static" / "trainer_steps" / "step_3" / "STABLE").exists()
+    assert (tmp_path / "run_policy" / "trainer_steps" / "step_6" / "STABLE").exists()
+    assert multi_run_manager.ready_to_update == [False, False]
+
+
 def test_detect_new_runs(tmp_path: Path) -> None:
     """Test that new runs are detected correctly."""
     multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5)