diff --git a/.github/workflows/gpu_tests.yaml b/.github/workflows/gpu_tests.yaml index cbab1b4f7d..16b90d3c16 100644 --- a/.github/workflows/gpu_tests.yaml +++ b/.github/workflows/gpu_tests.yaml @@ -121,3 +121,11 @@ jobs: PYTEST_OUTPUT_DIR: /tmp/outputs VLLM_WORKER_MULTIPROC_METHOD: spawn run: uv run pytest -vvs ${{ matrix.pytest_args }} + - name: Upload run logs on failure + if: failure() + uses: actions/upload-artifact@v4 + with: + name: ${{ matrix.test }}-logs + path: /tmp/outputs/**/*.log + retention-days: 7 + if-no-files-found: ignore diff --git a/deps/research-environments b/deps/research-environments index e40d32690f..f5d423c321 160000 --- a/deps/research-environments +++ b/deps/research-environments @@ -1 +1 @@ -Subproject commit e40d32690f74f4cdefff9dec4cdb654d63a59170 +Subproject commit f5d423c3210cd438b3923da4d4e2494745c06e00 diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index b99d2a3dc9..43ff4c6a1b 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -176,6 +176,8 @@ 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`` + self.version_advanced = asyncio.Event() self.wait_for_policy_time = 0.0 self.component_tasks = [] @@ -601,6 +603,35 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: f"({n_trainable / len(batch.rollouts):.1%}) — consider reviewing task difficulty / filter config" ) + # Ship batch ``step`` only once the trainer has published v{step-1-TARGET_LAG}. + # Without this, fast envs fill batches from buffered rollouts, the + # orchestrator finishes early, and its teardown strands the trainer inside + # an in-memory broadcast handshake that needs the live weight watcher. + # Always satisfiable: the trainer skips only the final TARGET_LAG+1 + # in-memory broadcasts. Bench runs have no trainer, so they ship freely. + required_version = step - 1 - TARGET_LAG + if not config.bench and self.policy.version < required_version: + get_logger().info( + f"Holding batch {step} until the trainer publishes policy v{required_version} " + f"(currently v{self.policy.version})" + ) + hold_start = time.perf_counter() + while True: + self.version_advanced.clear() + if self.policy.version >= required_version: + break + await self.version_advanced.wait() + self.wait_for_policy_time += time.perf_counter() - hold_start + + # Stamp each rollout's true staleness: batch ``step`` trains on policy + # v{step-1}, so a rollout generated from v{k} is (step-1)-k versions + # off-policy — queue time included, unlike the dispatcher's in-flight + # counter, which only sees weight updates during generation. Frozen- + # sourced rollouts stay 0 (their sampler doesn't follow the policy). + for r in batch.rollouts: + if self.train_envs.get(r.env_name).sampler.samples_from_live_policy: + r.off_policy_steps = (step - 1) - r.policy_version + # The effective (clean, trained-on) subset lands in the per-step ``effective`` trace file # at ship time; the full arrival window already streamed into ``all`` on arrival. # to_record drops the per-node training tensors — they're for training, not the rollout @@ -904,11 +935,14 @@ def update_dispatch_gate(self) -> 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.""" if self.model_express is not None: await asyncio.to_thread(self.model_express.set_status, p2p_pb2.SOURCE_STATUS_READY) + self.version_advanced.set() async def on_new_version(self, step: int) -> None: - """``VersionObserver`` hook: the watcher just advanced ``policy.version``; + """``VersionObserver`` hook: the weight update completed; re-evaluate the dispatch gate (may resume if the trainer caught up).""" if self.model_express is not None: await asyncio.to_thread(self.model_express.set_status, p2p_pb2.SOURCE_STATUS_INITIALIZING) diff --git a/src/prime_rl/orchestrator/watcher.py b/src/prime_rl/orchestrator/watcher.py index 5a818f5c51..d660cd80b4 100644 --- a/src/prime_rl/orchestrator/watcher.py +++ b/src/prime_rl/orchestrator/watcher.py @@ -64,6 +64,12 @@ async def start(self) -> None: async def stop(self) -> None: self.stopped.set() + # Let an in-flight apply finish before dying: the trainer blocks inside + # its in-memory broadcast until the apply completes, so cancelling + # mid-apply would strand it. The orchestrator's global teardown budget + # bounds this wait. + async with self.update_lock: + pass if self.task is not None: await safe_cancel(self.task) self.task = None @@ -103,6 +109,13 @@ async def apply_policy_update(self, next_step: int) -> None: await wait_for_path(stable_marker) 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. + self.ckpt_step = next_step + self.policy.version = next_step + # Drain off-policy rollouts BEFORE pausing the inference engines. # Aborting a rollout triggers vLLM's KV-connector cleanup (NIXL's # ``_reqs_not_processed``), which is only propagated to the workers @@ -128,8 +141,6 @@ async def apply_policy_update(self, next_step: int) -> None: self.update_count += 1 get_logger().debug(f"Updated weights to step {next_step} in {format_time(self.last_update_weights_time)}") - self.ckpt_step = next_step - self.policy.version = next_step if self.lora_name is not None: self.inference.update_model_name(self.lora_name) self.policy.model_name = self.lora_name diff --git a/src/prime_rl/trainer/rl/token_export.py b/src/prime_rl/trainer/rl/token_export.py index a2280dda81..6571d15ea9 100644 --- a/src/prime_rl/trainer/rl/token_export.py +++ b/src/prime_rl/trainer/rl/token_export.py @@ -10,6 +10,7 @@ from prime_rl.configs.trainer import DefaultLossConfig, TrainerConfig from prime_rl.trainer.rl.loss import compute_importance_ratio_and_mismatch_kl +from prime_rl.utils.logger import get_logger SCHEMA_VERSION = 1 @@ -93,7 +94,12 @@ def mark_stable(self, ready_run_ids: set[str] | None = None) -> None: ready_run_ids = ready_run_ids or set() for run_id in [rid for rid in self._pending_stable_dirs if rid is None or rid in ready_run_ids]: for stable_dir in self._pending_stable_dirs.pop(run_id): - (stable_dir / "STABLE").touch() + try: + (stable_dir / "STABLE").touch() + except FileNotFoundError: + # A multi-run run dir can be deleted while its tail steps are + # still flushing (mirrors the broadcast paths' guard) + get_logger().warning(f"Run dir for {run_id} is deleted, skipping token-export STABLE marker") def _export_dir(self, export_step: int, run_id: str | None) -> Path: if run_id is not None: diff --git a/tests/conftest.py b/tests/conftest.py index 413aeba3d3..6d0c2b1e79 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,4 @@ import os -import shutil import signal import socket import subprocess @@ -74,11 +73,11 @@ def branch_name() -> str: @pytest.fixture(scope="session") def output_dir(tmp_path_factory: pytest.TempPathFactory) -> Generator[Path, None, None]: - """Fixture for temporary output directory for tests with automatic cleanup""" + """Fixture for the tests' output directory. Never deleted — CI uploads run logs + from it on failure; runs start with ``--clean-output-dir`` for a fresh workspace.""" output_dir = Path(os.environ.get("PYTEST_OUTPUT_DIR", tmp_path_factory.mktemp("outputs"))) output_dir.mkdir(parents=True, exist_ok=True) yield output_dir - shutil.rmtree(output_dir, ignore_errors=True) @pytest.fixture(scope="session") diff --git a/tests/integration/test_alphabet_sort.py b/tests/integration/test_alphabet_sort.py index c8f3852f3c..5e81896fe7 100644 --- a/tests/integration/test_alphabet_sort.py +++ b/tests/integration/test_alphabet_sort.py @@ -38,6 +38,7 @@ def rl_process( "rl", "@", "configs/ci/integration/alphabet_sort.toml", + "--clean-output-dir", "--wandb.project", wandb_project, "--wandb.name", diff --git a/tests/integration/test_reverse_text_moe.py b/tests/integration/test_reverse_text_moe.py index db510837e7..d713b6ae7c 100644 --- a/tests/integration/test_reverse_text_moe.py +++ b/tests/integration/test_reverse_text_moe.py @@ -29,6 +29,7 @@ def rl_process( "rl", "@", "configs/ci/integration/reverse-text-moe/start.toml", + "--clean-output-dir", "--trainer.model.impl", "custom", "--wandb.project", diff --git a/tests/integration/test_reverse_text_multi_run.py b/tests/integration/test_reverse_text_multi_run.py index 4c9f323434..7ad782a4ea 100644 --- a/tests/integration/test_reverse_text_multi_run.py +++ b/tests/integration/test_reverse_text_multi_run.py @@ -20,6 +20,17 @@ ORCHESTRATOR_NAMES = ["alpha", "beta", "gamma"] +def remove_run_dir(run_dir: Path) -> None: + """Delete a run dir, tolerating the shared trainer still flushing its tail + (token exports recreate directories mid-rmtree -> 'Directory not empty').""" + for _ in range(10): + shutil.rmtree(run_dir, ignore_errors=True) + if not run_dir.exists(): + return + time.sleep(1) + shutil.rmtree(run_dir) + + def wait_for_file( file_path: Path, timeout: int = 600, @@ -274,7 +285,7 @@ def multi_run_result( print(f"Copied alpha checkpoint to {tmp_path / 'alpha_ckpt_step_10'}") # Remove alpha run directory - shutil.rmtree(output_dir / "run_alpha") + remove_run_dir(output_dir / "run_alpha") # =========================== # Queue alpha's resume proc @@ -305,7 +316,7 @@ def multi_run_result( shutil.copy(run_dir / "logs" / "orchestrator.log", log_dir / "beta_orchestrator.log") shutil.copytree(beta_ckpt_dir, tmp_path / "beta_ckpt_step_20") print(f"Copied {beta_ckpt_dir} to {tmp_path / 'beta_ckpt_step_20'}") - shutil.rmtree(run_dir) + remove_run_dir(run_dir) # ===================== # Queue beta's resume @@ -329,7 +340,7 @@ def multi_run_result( timeout=TIMEOUT, ) shutil.copy(output_dir / "run_gamma" / "logs" / "orchestrator.log", log_dir / "gamma_orchestrator.log") - shutil.rmtree(output_dir / "run_gamma") + remove_run_dir(output_dir / "run_gamma") # ================================================ # Wait for alpha_resume and beta_resume to finish diff --git a/uv.lock b/uv.lock index a9bc3cad73..36e4694c64 100644 --- a/uv.lock +++ b/uv.lock @@ -4632,7 +4632,7 @@ requires-dist = [ [[package]] name = "openswe-v1" -version = "0.1.2" +version = "0.1.3" source = { editable = "deps/research-environments/environments/swe/openswe_v1" } dependencies = [ { name = "datasets" }, @@ -4774,7 +4774,7 @@ wheels = [ [[package]] name = "openthoughts-tblite-v1" -version = "0.1.0" +version = "0.1.1" source = { editable = "deps/research-environments/environments/terminal/openthoughts_tblite_v1" } dependencies = [ { name = "verifiers" }, @@ -6446,7 +6446,7 @@ wheels = [ [[package]] name = "senior-swe-bench-v1" -version = "0.1.0" +version = "0.1.1" source = { editable = "deps/research-environments/environments/swe/senior_swe_bench_v1" } dependencies = [ { name = "verifiers", extra = ["harbor"] }, @@ -7091,7 +7091,7 @@ requires-dist = [{ name = "verifiers", specifier = ">=0.2.1" }] [[package]] name = "terminal-lego-v1" -version = "0.1.0" +version = "0.1.1" source = { editable = "deps/research-environments/environments/terminal/terminal_lego_v1" } dependencies = [ { name = "huggingface-hub" },