From 3374057831410677b1584f3824003b743034f7f0 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Fri, 23 Oct 2026 11:40:32 +0000 Subject: [PATCH 01/12] unify the prime monitor and share the sft run One PrimeMonitor and one PrimeMonitorConfig replace the train and eval pair. A training run's processes log to one Prime training run: the rl orchestrator or sft trainer registers it, and an sft online-eval process attaches to the trainer's run instead of opening one Prime evaluation per epoch. PRL_RUN_FINISHER names the process that finishes a shared run; the sft launcher sets it to the online-eval process, which drains its final evals after the trainer exits. The other processes flush and drop the SDK's crash hook, which also stops a paused run from being reported as crashed. A standalone eval still opens one evaluation per source and epoch. Monitors share more of the base class: one MonitorContext replaces per-monitor init arguments, every hook but log_metrics defaults to a no-op, and the fan-out helpers share one error path. Episodes are logged once, on arrival; the unused effective-cohort log calls are gone, since batch membership already travels as annotations. Co-Authored-By: Claude Opus 5.5 --- docs/eval.md | 2 +- docs/training.md | 2 +- .../src/prime_rl/configs/eval.py | 15 +- .../src/prime_rl/configs/monitors.py | 25 +- .../src/prime_rl/configs/orchestrator.py | 4 +- .../src/prime_rl/configs/rl.py | 4 +- .../src/prime_rl/configs/sft.py | 4 +- .../src/prime_rl/configs/trainer.py | 4 +- src/prime_rl/entrypoints/dashboard.py | 4 +- src/prime_rl/entrypoints/eval.py | 4 +- src/prime_rl/entrypoints/sft.py | 18 +- src/prime_rl/eval/runner.py | 5 +- src/prime_rl/monitors/__init__.py | 113 ++--- src/prime_rl/monitors/base.py | 55 +- src/prime_rl/monitors/file/monitor.py | 16 +- src/prime_rl/monitors/prime.py | 477 ++++++++---------- src/prime_rl/monitors/wandb/monitor.py | 26 +- src/prime_rl/orchestrator/orchestrator.py | 10 +- .../templates/multi_node_sft.sbatch.j2 | 2 + src/prime_rl/utils/links.py | 4 +- src/prime_rl/utils/process.py | 4 + tests/unit/train/sft/test_prime_monitor.py | 376 ++++---------- tests/unit/utils/test_links.py | 17 +- 23 files changed, 443 insertions(+), 748 deletions(-) diff --git a/docs/eval.md b/docs/eval.md index 27ce9f7368b..5585c421c80 100644 --- a/docs/eval.md +++ b/docs/eval.md @@ -116,7 +116,7 @@ Each source's epoch streams as an evaluation on the Prime platform. The launcher uv run eval gsm8k -n 32 -r 4 -c 8 ``` -The monitor opens one platform evaluation per env when its epoch starts, streams every episode to it as it lands, and closes it with the epoch's aggregates. The run leaves its platform identity in `monitors/prime/run.json` (the training run's id and URL, or one entry per eval epoch), and the dashboard's top bar shows a `View on Prime` link from the moment the evaluation opens: a direct link for a training run or a single evaluation, a menu when a run has several. +The monitor opens one Prime evaluation per env when its epoch starts, streams every episode to it as it lands, and closes it with the epoch's aggregates. The run leaves its Prime identity in `monitors/prime/run.json` (one entry per evaluation), and the dashboard's top bar shows a `View on Prime` link from the moment the evaluation opens: a direct link for a training run or a single evaluation, a menu when a run has several. ## Metrics diff --git a/docs/training.md b/docs/training.md index 41590e3d7c9..7e6a7d2cb99 100644 --- a/docs/training.md +++ b/docs/training.md @@ -453,7 +453,7 @@ Name the Prime run in TOML: name = "my-experiment" ``` -On `sft` the monitor registers the run with the dataset-batched training fields (`max_steps`, `batch_size`, `seq_len`) and streams the trainer's per-step metrics and validation losses; SFT has no rollouts, so no episodes are uploaded. +On `sft` the monitor registers the run with the dataset-batched training fields (`max_steps`, `batch_size`, `seq_len`) and streams the trainer's per-step metrics and validation losses. Online evals log to the same run, like an `rl` run's: the online-eval process attaches to the trainer's run, uploads its eval metrics and every eval episode, and finishes the run after its final evals. The platform run needs `PRIME_API_KEY` (`prime login` or the env var) and a team (`PRIME_TEAM_ID`, or the team selected with `prime login`) enabled for external runs. Like the W&B monitor, a run without a key, or one whose registration fails (no team, a team outside the allowlist), fails at launch: authenticate, select a team with `prime switch`, or pass `--no-monitors.prime`. `PRIME_RUNS_MODE=disabled` keeps the monitor configured but opens no platform run; `RUN_ID=` attaches to an external run a launcher already created instead of registering a new one. Currently internal-only. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/eval.py b/packages/prime-rl-configs/src/prime_rl/configs/eval.py index ef2d62429ac..54b29f1c5ed 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/eval.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/eval.py @@ -4,7 +4,7 @@ import verifiers.v1 as vf from pydantic import AliasChoices, Field, model_validator -from prime_rl.configs.monitors import EvalMonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig from prime_rl.configs.orchestrator import ConcurrencyConfig, EvalSourcesConfig, ScheduledEvalConfig from prime_rl.configs.shared import ( ClientConfig, @@ -103,7 +103,7 @@ class EvalConfig(ServedEvalConfig): log: LogConfig = LogConfig() - monitors: EvalMonitorsConfig = EvalMonitorsConfig() + monitors: MonitorsConfig = MonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" @property @@ -167,14 +167,9 @@ class SFTOnlineEvalConfig(ScheduledEvalConfig, ServedEvalConfig): log: LogConfig = LogConfig() - monitors: EvalMonitorsConfig = EvalMonitorsConfig() - """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``). - - ``EvalMonitorsConfig``, not the base ``MonitorsConfig``: the launcher - converts the trainer's ``[monitors.prime]`` into a - ``PrimeEvalMonitorConfig`` for the online-eval process, and the base - type would drop (or forbid) the prime block at the eval.json - dump/re-parse boundary.""" + monitors: MonitorsConfig = MonitorsConfig() + """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``), the + trainer's: the online-eval process logs to the trainer's W&B and Prime runs.""" @model_validator(mode="after") def auto_setup_broadcasts_dir(self): diff --git a/packages/prime-rl-configs/src/prime_rl/configs/monitors.py b/packages/prime-rl-configs/src/prime_rl/configs/monitors.py index de3cd327ec5..1358cfa3e09 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/monitors.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/monitors.py @@ -43,17 +43,14 @@ class FileMonitorConfig(BaseConfig): not the records, so this only trades record size against overlay precision.""" -class PrimeTrainMonitorConfig(BaseConfig): +class PrimeMonitorConfig(BaseConfig): name: str | None = None - """Run name shown on the platform. Inherits ``run.name`` when unset.""" + """Run or evaluation name shown on Prime. Inherits ``run.name`` when unset.""" -class PrimeEvalMonitorConfig(BaseConfig): - name: str | None = None - """Evaluation name shown on the platform. Inherits ``run.name`` when unset.""" - +class TrainerMonitorsConfig(BaseConfig): + """The RL trainer's monitors: its metrics reach Prime through the orchestrator's run.""" -class MonitorsConfig(BaseConfig): wandb: WandbMonitorConfig | None = None """Log metrics to Weights & Biases. Off by default; enable with ``--monitors.wandb``.""" @@ -61,13 +58,7 @@ class MonitorsConfig(BaseConfig): """Log metrics and episode traces to the run's output directory. On by default; disable with ``--no-monitors.file``.""" -class TrainMonitorsConfig(MonitorsConfig): - prime: PrimeTrainMonitorConfig | None = PrimeTrainMonitorConfig() - """Log metrics and episodes to the Prime Intellect platform. On by default and - needs a Prime API key; disable with ``--no-monitors.prime``.""" - - -class EvalMonitorsConfig(MonitorsConfig): - prime: PrimeEvalMonitorConfig | None = PrimeEvalMonitorConfig() - """Upload every eval source's epoch as an evaluation on the Prime Intellect platform. On by - default and needs a Prime API key; disable with ``--no-monitors.prime``.""" +class MonitorsConfig(TrainerMonitorsConfig): + prime: PrimeMonitorConfig | None = PrimeMonitorConfig() + """Log metrics and episodes to Prime. On by default and needs a Prime API key; + disable with ``--no-monitors.prime``.""" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 0419a48e24e..f991974be05 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -12,7 +12,7 @@ AlgoConfig, GRPOAlgoConfig, ) -from prime_rl.configs.monitors import TrainMonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig from prime_rl.configs.shared import ( BaseModelConfig, ControlConfig, @@ -494,7 +494,7 @@ class OrchestratorConfig(BaseConfig): env_vars: EnvVars = {} """Extra environment variables for the orchestrator process(es). Merged on top of the launcher defaults.""" - monitors: TrainMonitorsConfig = TrainMonitorsConfig() + monitors: MonitorsConfig = MonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" collect_inference_metrics: bool = True 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 83b4c44bebc..053493071d9 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -8,7 +8,7 @@ from prime_rl.configs.inference import InferenceConfig from prime_rl.configs.inference import WeightBroadcastConfig as InferenceWeightBroadcastConfig -from prime_rl.configs.monitors import FileMonitorConfig, PrimeTrainMonitorConfig +from prime_rl.configs.monitors import FileMonitorConfig, PrimeMonitorConfig from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.configs.shared import ( EnvVars, @@ -78,7 +78,7 @@ class SharedMonitorsConfig(BaseConfig): file: FileMonitorConfig | None = None """Shared local JSONL metric sink. If set, enables ``/monitors/file/metrics.jsonl`` on both trainer and orchestrator.""" - prime: PrimeTrainMonitorConfig | None = None + prime: PrimeMonitorConfig | None = None """Prime platform monitor. Propagated to the orchestrator only — the trainer has no platform integration.""" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 7e578e5f0e0..5f62ce989b7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -11,7 +11,7 @@ from prime_rl.configs.eval import SFTOnlineEvalConfig from prime_rl.configs.inference import InferenceConfig from prime_rl.configs.inference import WeightBroadcastConfig as InferenceWeightBroadcastConfig -from prime_rl.configs.monitors import TrainMonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig from prime_rl.configs.shared import ( FileSystemWeightBroadcastConfig, NCCLWeightBroadcastConfig, @@ -236,7 +236,7 @@ class SFTConfig(BaseTrainerConfig): resume: ResumeConfig | None = None """Resume the run from a checkpoint (point at it with the previous run's ``run.name``). Without ``[ckpt]`` the run loads the checkpoint but saves no new ones. If None, does not resume.""" - monitors: TrainMonitorsConfig = TrainMonitorsConfig() + monitors: MonitorsConfig = MonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" run: RunConfig = Field(default_factory=RunConfig) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 222aca108f7..bccefa10097 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -5,7 +5,7 @@ from pydantic import Field, field_validator, model_validator -from prime_rl.configs.monitors import MonitorsConfig +from prime_rl.configs.monitors import TrainerMonitorsConfig from prime_rl.configs.shared import ( BaseModelConfig, ControlConfig, @@ -740,7 +740,7 @@ class TrainerConfig(BaseTrainerConfig): rollout_transport: TransportConfig = ZMQTransportConfig() """Transport used to ship rollouts from orchestrator to trainer.""" - monitors: MonitorsConfig = MonitorsConfig() + monitors: TrainerMonitorsConfig = TrainerMonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``).""" output_dir: Path = Field(default_factory=default_output_dir) diff --git a/src/prime_rl/entrypoints/dashboard.py b/src/prime_rl/entrypoints/dashboard.py index 3b7d0e6204a..7292dc503a6 100644 --- a/src/prime_rl/entrypoints/dashboard.py +++ b/src/prime_rl/entrypoints/dashboard.py @@ -28,7 +28,7 @@ from pathlib import Path from typing import Callable -from prime_rl.utils.links import PRIME_EVAL_LINKS, PRIME_LINKS, WANDB_LINKS +from prime_rl.utils.links import PRIME_LINKS, WANDB_LINKS from prime_rl.utils.logger import get_logger from prime_rl.utils.pathing import CACHE_DIR @@ -118,7 +118,7 @@ def ensure_dashboard(output_dir: Path, logger) -> str | None: return None -LINK_NAMES = {PRIME_LINKS: "Prime", PRIME_EVAL_LINKS: "Prime", WANDB_LINKS: "W&B"} +LINK_NAMES = {PRIME_LINKS: "Prime", WANDB_LINKS: "W&B"} def log_run_links(dashboard_url: str | None, links: Callable[[], dict[str, dict[str, str]]] | None = None) -> None: diff --git a/src/prime_rl/entrypoints/eval.py b/src/prime_rl/entrypoints/eval.py index 5121ba07864..9dc3aa68bb7 100644 --- a/src/prime_rl/entrypoints/eval.py +++ b/src/prime_rl/entrypoints/eval.py @@ -121,7 +121,7 @@ def main(): sys.argv = [sys.argv[0], *argv] from prime_rl.entrypoints.dashboard import ensure_dashboard, log_run_links, run_links - from prime_rl.utils.links import PRIME_EVAL_LINKS, WANDB_LINKS + from prime_rl.utils.links import PRIME_LINKS, WANDB_LINKS from prime_rl.utils.logger import setup_logger from prime_rl.utils.pathing import ( format_config_message, @@ -195,7 +195,7 @@ def sigterm_handler(signum, frame): # Each source's Prime evaluation opens once the eval counts its tasks, so the # banner waits for the links next to the running eval. eval_done = threading.Event() - monitors = {PRIME_EVAL_LINKS: config.monitors.prime, WANDB_LINKS: config.monitors.wandb} + monitors = {PRIME_LINKS: config.monitors.prime, WANDB_LINKS: config.monitors.wandb} links = run_links(config.run_dir, log_dir, monitors, eval_done.is_set) banner = threading.Thread(target=log_run_links, args=(dashboard_url, links), daemon=True) banner.start() diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 9ece644aa8f..dea87611034 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -5,7 +5,6 @@ from pathlib import Path from prime_rl.configs.eval import SFTOnlineEvalConfig -from prime_rl.configs.monitors import EvalMonitorsConfig, PrimeEvalMonitorConfig, TrainMonitorsConfig from prime_rl.configs.orchestrator import OnlineEvalSourceConfig from prime_rl.configs.sft import SFTConfig from prime_rl.configs.shared import LogConfig @@ -32,6 +31,7 @@ DEFAULT_COMMON_ENV_VARS, DEFAULT_INFERENCE_ENV_VARS, DEFAULT_TRAINER_ENV_VARS, + RUN_FINISHER_VAR, ProcessGroup, partition_gpus, set_proc_title, @@ -70,16 +70,6 @@ def resolve_resume_step(config: SFTConfig) -> int | None: return resolve_latest_ckpt_step(get_ckpt_dir(get_ckpt_base(config))) -def build_online_eval_monitors(monitors: TrainMonitorsConfig) -> EvalMonitorsConfig: - """The online-eval process's monitors: same W&B and file config as the trainer's, - with the platform monitor switched to the evaluation flavor under the same name.""" - return EvalMonitorsConfig( - wandb=monitors.wandb, - file=monitors.file, - prime=PrimeEvalMonitorConfig(name=monitors.prime.name) if monitors.prime is not None else None, - ) - - def build_online_eval_config(config: SFTConfig) -> SFTOnlineEvalConfig: """The online-eval process's config: the resolved ``[eval]`` block with the run-level fields filled from the SFT config. The launcher spawns the env servers itself; the @@ -100,7 +90,7 @@ def build_online_eval_config(config: SFTConfig) -> SFTOnlineEvalConfig: # by directory separation. output_dir=config.run_dir, log=LogConfig(level=config.log.level, json_logging=config.log.json_logging), - monitors=build_online_eval_monitors(config.monitors), + monitors=config.monitors, ) return SFTOnlineEvalConfig(**{**eval_config.model_dump(exclude=set(run_fields)), **run_fields}) @@ -353,6 +343,8 @@ def sft_local(config: SFTConfig): "WANDB_PROGRAM": "uv run sft", "WANDB_ARGS": json.dumps(sys.argv), } + # The online-eval process outlives the trainer, so it finishes the shared Prime run. + shared_run_env = {RUN_FINISHER_VAR: "online-eval"} if config.eval is not None else {} with ProcessGroup() as processes: # Optionally, start the inference server for online evals @@ -398,6 +390,7 @@ def sft_local(config: SFTConfig): "PRL_ATTEMPT_LOG_DIR": str(log_dir), "PRL_LOG_DIR": str(log_dir), **wandb_shared_env, + **shared_run_env, "WANDB_SHARED_LABEL": "online-eval", }, log_path=log_dir / "eval.log", @@ -411,6 +404,7 @@ def sft_local(config: SFTConfig): **DEFAULT_TRAINER_ENV_VARS, **config.env_vars, **wandb_shared_env, + **shared_run_env, **launcher_attempt_env(config_dir, log_dir), } if config.eval is not None: diff --git a/src/prime_rl/eval/runner.py b/src/prime_rl/eval/runner.py index 783f81ab939..9bf5c38f263 100644 --- a/src/prime_rl/eval/runner.py +++ b/src/prime_rl/eval/runner.py @@ -246,7 +246,7 @@ async def land(self, episode: vf.Episode, pending: set[str]) -> None: """One episode of the epoch, arrived or restored: through the monitors and into its env's batch, which is finalized once the epoch is complete.""" step = eval_work(episode).step - await monitors.log([episode], step, "eval", "all") + await monitors.log([episode], step, "eval") eval_batch = self.eval_sink.add(episode) if eval_batch is not None: await self.finalize_eval_batch(eval_batch) @@ -254,13 +254,12 @@ async def land(self, episode: vf.Episode, pending: set[str]) -> None: async def finalize_eval_batch(self, batch: EvalBatch) -> None: """Persist + log one completed eval epoch through the monitors, mirroring the - orchestrator: effective episodes plus the ``eval/{env}/...`` metric dict.""" + orchestrator: the effective subset's annotations plus the ``eval/{env}/...`` metric dict.""" if not batch.episodes and not batch.failures and not batch.cancelled: get_logger().warning(f"Eval @ step={batch.step} env={batch.env_name}: no attempts returned, skipping log") return if batch.episodes.effective: - await monitors.log(batch.episodes.effective.vf_episodes, batch.step, "eval", "effective") await monitors.log_annotations(stamp_batch(batch.episodes.effective.vf_episodes, batch.step)) await monitors.log_eval_epoch(batch.env_name, batch.step, batch.episodes.vf_episodes) diff --git a/src/prime_rl/monitors/__init__.py b/src/prime_rl/monitors/__init__.py index 53ce9413ae0..83c90092d9d 100644 --- a/src/prime_rl/monitors/__init__.py +++ b/src/prime_rl/monitors/__init__.py @@ -13,15 +13,10 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, overload -from prime_rl.configs.monitors import ( - FileMonitorConfig, - PrimeEvalMonitorConfig, - PrimeTrainMonitorConfig, - WandbMonitorConfig, -) -from prime_rl.monitors.base import Kind, Monitor, Subset +from prime_rl.configs.monitors import FileMonitorConfig, PrimeMonitorConfig, WandbMonitorConfig +from prime_rl.monitors.base import Kind, Monitor, MonitorContext from prime_rl.monitors.file import FileMonitor -from prime_rl.monitors.prime import PrimeEvalMonitor, PrimeTrainMonitor +from prime_rl.monitors.prime import PrimeMonitor from prime_rl.monitors.wandb import WandbMonitor from prime_rl.utils.config import BaseConfig from prime_rl.utils.logger import format_time, get_logger @@ -31,9 +26,9 @@ __all__ = [ "Monitor", + "MonitorContext", "WandbMonitor", - "PrimeTrainMonitor", - "PrimeEvalMonitor", + "PrimeMonitor", "FileMonitor", "setup", "get", @@ -47,7 +42,7 @@ async def setup( wandb: WandbMonitorConfig | None = None, - prime: PrimeTrainMonitorConfig | PrimeEvalMonitorConfig | None = None, + prime: PrimeMonitorConfig | None = None, file: FileMonitorConfig | None = None, *, output_dir: Path, @@ -68,35 +63,24 @@ async def setup( if rank != 0: return - monitors: list[tuple[str, Monitor, dict[str, Any]]] = [] - if isinstance(prime, PrimeTrainMonitorConfig): - monitors.append(("prime", PrimeTrainMonitor(prime), dict(config=run_config, output_dir=output_dir))) - elif isinstance(prime, PrimeEvalMonitorConfig): - monitors.append(("prime", PrimeEvalMonitor(prime), dict(config=run_config, output_dir=output_dir))) - if wandb is not None: - monitors.append( - ( - "wandb", - WandbMonitor(wandb), - dict( - output_dir=output_dir, - config=run_config, - train_env_names=train_env_names, - eval_env_names=eval_env_names, - overview_flavor=overview_flavor, - ), - ) - ) - if file is not None: - monitors.append(("file", FileMonitor(file), dict(output_dir=output_dir, producer=producer))) + context = MonitorContext( + output_dir=output_dir, + producer=producer, + run_config=run_config, + train_env_names=train_env_names or [], + eval_env_names=eval_env_names or [], + overview_flavor=overview_flavor, + ) + configured = [(PrimeMonitor, prime), (WandbMonitor, wandb), (FileMonitor, file)] + monitors = [cls(config, context) for cls, config in configured if config is not None] if not monitors: return - get_logger().info(f"Initializing monitors ({', '.join(name for name, _, _ in monitors)})") + get_logger().info(f"Initializing monitors ({', '.join(type(monitor).__name__ for monitor in monitors)})") t0 = time.perf_counter() - for name, monitor, init_kwargs in monitors: - get_logger().info(f"Initializing {name} monitor ({monitor.config})") - await monitor.init(**init_kwargs) + for monitor in monitors: + get_logger().info(f"Initializing {type(monitor).__name__} ({monitor.config})") + await monitor.init() MONITORS.append(monitor) get_logger().debug(f"Initialized monitors in {format_time(time.perf_counter() - t0)}") @@ -107,69 +91,48 @@ def get(monitor_cls: type[Monitor]) -> Monitor | None: return next((monitor for monitor in MONITORS if isinstance(monitor, monitor_cls)), None) +async def _fan_out(hook: str, *args: Any, **kwargs: Any) -> None: + for monitor in MONITORS: + try: + await getattr(monitor, hook)(*args, **kwargs) + except Exception as e: + get_logger().warning(f"Failed to {hook} on {type(monitor).__name__}: {e}") + + @overload async def log(data: dict[str, Any], step: int | None) -> None: ... @overload -async def log(data: vf.Episode | list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: ... +async def log(data: vf.Episode | list[vf.Episode], step: int, kind: Kind) -> None: ... -async def log( - data: dict[str, Any] | vf.Episode | list[vf.Episode], - step: int | None, - kind: Kind = "train", - subset: Subset = "effective", -) -> None: - """Log to all registered monitors: a dict of scalar metrics, or episodes - with their cohort coordinates (train/eval x all/effective).""" - for monitor in MONITORS: - try: - await monitor.log(data, step=step, kind=kind, subset=subset) - except Exception as e: - get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") +async def log(data: dict[str, Any] | vf.Episode | list[vf.Episode], step: int | None, kind: Kind = "train") -> None: + """Log to all registered monitors: a dict of scalar metrics, or episodes as they + arrive.""" + await _fan_out("log", data, step=step, kind=kind) async def log_annotations(updates: list[dict[str, Any]]) -> None: """Log trace updates to all registered monitors.""" - for monitor in MONITORS: - try: - await monitor.log_annotations(updates) - except Exception as e: - get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") + await _fan_out("log_annotations", updates) async def log_eval_plan(env_name: str, step: int, expected: int) -> None: """Log an eval epoch's expected episode count to all registered monitors.""" - for monitor in MONITORS: - try: - await monitor.log_eval_plan(env_name, step, expected) - except Exception as e: - get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") + await _fan_out("log_eval_plan", env_name, step, expected) async def log_live(events: list[dict[str, Any]]) -> None: """Log live-trace events to all registered monitors.""" - for monitor in MONITORS: - try: - await monitor.log_live(events) - except Exception as e: - get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") + await _fan_out("log_live", events) async def log_eval_epoch(env_name: str, step: int, episodes: list[vf.Episode]) -> None: """Log one finished eval epoch to all registered monitors.""" - for monitor in MONITORS: - try: - await monitor.log_eval_epoch(env_name, step, episodes) - except Exception as e: - get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") + await _fan_out("log_eval_epoch", env_name, step, episodes) async def finalize(paused: bool = False) -> None: """Finalize the run on all registered monitors. A paused run is not reported as completed.""" - for monitor in MONITORS: - try: - await monitor.finalize(paused=paused) - except Exception as e: - get_logger().warning(f"Failed to finalize {monitor.__class__.__name__}: {e}") + await _fan_out("finalize", paused=paused) diff --git a/src/prime_rl/monitors/base.py b/src/prime_rl/monitors/base.py index 260cc0e56dd..f9d88e9a30a 100644 --- a/src/prime_rl/monitors/base.py +++ b/src/prime_rl/monitors/base.py @@ -1,6 +1,8 @@ from __future__ import annotations from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, overload from prime_rl.utils.config import BaseConfig @@ -10,66 +12,81 @@ import verifiers.v1 as vf Kind = Literal["train", "eval"] -Subset = Literal["all", "effective"] + + +@dataclass(frozen=True) +class MonitorContext: + """What every monitor of a process knows about the run it logs.""" + + output_dir: Path + """The run directory.""" + producer: str | None = None + """The process logging: ``orchestrator``, ``trainer``, ``online-eval`` or ``eval``.""" + run_config: BaseConfig | None = None + """The process's resolved config.""" + train_env_names: list[str] = field(default_factory=list) + eval_env_names: list[str] = field(default_factory=list) + overview_flavor: Literal["rl", "sft", "eval"] = "rl" class Monitor(ABC): - """Base class for monitors.""" + """Base class for monitors. Every hook but ``log_metrics`` defaults to a no-op, so a + monitor implements only what it carries.""" - def __init__(self, config: BaseConfig): + def __init__(self, config: BaseConfig, context: MonitorContext): self.config = config + self.context = context self.logger = get_logger() - async def init(self, **kwargs: Any) -> None: - """Initialize run. Overrides name their own kwargs.""" + async def init(self) -> None: + """Open the monitor's run. Network or file I/O, before the first log.""" @overload async def log(self, data: dict[str, Any], step: int | None) -> None: ... @overload - async def log(self, data: vf.Episode | list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: ... + async def log(self, data: vf.Episode | list[vf.Episode], step: int, kind: Kind) -> None: ... async def log( self, data: dict[str, Any] | vf.Episode | list[vf.Episode], step: int | None, kind: Kind = "train", - subset: Subset = "effective", ) -> None: """Log scalar metrics, or episodes.""" if isinstance(data, dict): await self.log_metrics(data, step=step) - else: - episodes = data if isinstance(data, list) else [data] + return + episodes = data if isinstance(data, list) else [data] + if episodes: assert step is not None - await self.log_episodes(episodes, step=step, kind=kind, subset=subset) + await self.log_episodes(episodes, step=step, kind=kind) @abstractmethod async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: """Log scalar metrics. ``step=None`` logs a time-keyed row (e.g. inference metrics, which are sampled on wall time rather than the training step).""" - @abstractmethod - async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: - """Log episodes.""" + async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind) -> None: + """Log episodes as they arrive: every episode exactly once, errored and filtered + ones included. What the run learns about them later (batch membership, + advantages) arrives through ``log_annotations``.""" async def log_annotations(self, updates: list[dict[str, Any]]) -> None: - """Log trace updates — post-hoc facts about traces this run already logged. - Monitors that only carry scalars ignore them.""" + """Log trace updates — post-hoc facts about traces this run already logged.""" async def log_live(self, events: list[dict[str, Any]]) -> None: """Log the env servers' stream of in-flight traces: ``{"delta", "dispatch"}`` for - each delta as it arrived, ``{"done": trace_id}`` when a trace's episode finished. - Monitors that only carry finished work ignore it.""" + each delta as it arrived, ``{"done": trace_id}`` when a trace's episode finished.""" async def log_eval_plan(self, env_name: str, step: int, expected: int) -> None: """Log how many episodes the eval epoch of ``env_name`` at ``step`` will produce, - known once its tasks are counted. Monitors that only carry results ignore it.""" + known once its tasks are counted.""" async def log_eval_epoch(self, env_name: str, step: int, episodes: list[vf.Episode]) -> None: """Log one finished eval epoch: every episode ``env_name`` produced for ``step``, errored ones included. Fires once per epoch, after the episodes streamed through - ``log``. Monitors that carry episodes as they arrive ignore it.""" + ``log``.""" async def finalize(self, paused: bool = False) -> None: """Finalize run. A paused run stopped to resume later and did not complete.""" diff --git a/src/prime_rl/monitors/file/monitor.py b/src/prime_rl/monitors/file/monitor.py index 21d54d2bdf7..24b37220fe7 100644 --- a/src/prime_rl/monitors/file/monitor.py +++ b/src/prime_rl/monitors/file/monitor.py @@ -11,7 +11,7 @@ from verifiers.v1.trace import EXCLUDE_FIELDS from prime_rl.configs.monitors import FileMonitorConfig -from prime_rl.monitors.base import Kind, Monitor, Subset +from prime_rl.monitors.base import Kind, Monitor from prime_rl.monitors.file.traces import get_annotations_dir, get_index_path, get_trace_stream from prime_rl.monitors.file.traces.chunks import ChunkedJsonl from prime_rl.monitors.file.traces.index import index_row @@ -36,9 +36,9 @@ class FileMonitor(Monitor): config: FileMonitorConfig file: TextIO - async def init(self, output_dir: Path, producer: str | None = None) -> None: - self.output_dir = output_dir - self.producer = producer + async def init(self) -> None: + output_dir = self.output_dir = self.context.output_dir + self.producer = self.context.producer index = get_index_path(get_trace_stream(output_dir)) # a relaunch appends to the stream it finds, so the numbering carries on self._logged = sum(1 for _ in index.open("rb")) if index.is_file() else 0 @@ -138,14 +138,12 @@ def _stream(self, directory: Path) -> tuple[ChunkedJsonl, BinaryIO]: self._streams[directory] = (stream, open(get_index_path(directory), "ab")) return self._streams[directory] - async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: + async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind) -> None: """Append each episode to the trace stream as it completes — every episode is serialized exactly once, in arrival order, whatever its kind, so an in-progress run can be tailed. Episode-level failures are preserved even when no trace was - produced. The shipped cohort writes no second copy: what it learns arrives as - annotations that readers fold back onto the stream.""" - if subset == "effective": - return + produced. What the shipped cohort learns arrives as annotations that readers fold + back onto the stream.""" def write() -> None: stream, index = self._stream(get_trace_stream(self.output_dir)) diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index cb76a2d6b17..62f3538b946 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -3,18 +3,21 @@ import asyncio import atexit import os +import time from pathlib import Path from typing import TYPE_CHECKING, Any, Callable import prime_runs as pr from prime_sandboxes import Config as PrimeConfig -from prime_rl.configs.monitors import PrimeEvalMonitorConfig, PrimeTrainMonitorConfig +from prime_rl.configs.eval import EvalConfig, SFTOnlineEvalConfig +from prime_rl.configs.monitors import PrimeMonitorConfig from prime_rl.configs.sft import SFTConfig -from prime_rl.monitors.base import Kind, Monitor, Subset +from prime_rl.monitors.base import Kind, Monitor from prime_rl.utils.config import BaseConfig -from prime_rl.utils.links import PRIME_EVAL_LINKS, PRIME_LINKS, publish_link, register_links +from prime_rl.utils.links import ATTEMPT_VAR, PRIME_LINKS, publish_link, register_links from prime_rl.utils.pathing import get_platform_run_path +from prime_rl.utils.process import RUN_FINISHER_VAR from prime_rl.utils.record import read_record, update_record, write_record from prime_rl.utils.utils import sanitize @@ -27,10 +30,15 @@ # default (300 s) is sized for eval sample batches; a crashed training process should # not linger that long, and a clean finish rarely has more than the last step queued. FINISH_TIMEOUT = 60.0 +# How long an SFT online-eval process waits for the trainer to register the run: the +# trainer starts with it, but registers only after its imports and config setup. +TRAIN_RUN_TIMEOUT_S = 1200.0 +TRAIN_RUN_POLL_S = 1.0 +OFF_HINT = "turn the monitor off with --no-monitors.prime" def write_platform_record(output_dir: Path, record: dict[str, Any]) -> None: - """Leave the run's platform identity in the run directory for the dashboard's + """Leave the run's Prime identity in the run directory for the dashboard's "View on Prime" link.""" write_record(get_platform_run_path(output_dir), record) @@ -39,27 +47,16 @@ def read_platform_record(output_dir: Path) -> dict[str, Any] | None: return read_record(get_platform_run_path(output_dir)) -def _merge_platform_record( - output_dir: Path, - merge: Callable[[dict[str, Any]], dict[str, Any]], -) -> None: - """Locked read-modify-write of the platform record: the trainer and the - online-eval process merge into the SAME run.json. Best-effort, see - ``update_record``.""" +def _merge_platform_record(output_dir: Path, merge: Callable[[dict[str, Any]], dict[str, Any]]) -> None: + """Locked read-modify-write of the platform record. Best-effort, see ``update_record``.""" update_record(get_platform_run_path(output_dir), merge) def update_platform_record(output_dir: Path, update: dict[str, Any]) -> None: - """Best-effort locked read-modify-write of the platform record: `update` - overlays the current record (preserving keys it doesn't set). Use for - every writer that shares the file — a naive overwrite drops concurrent - writers' keys.""" + """Overlay ``update`` on the platform record, keeping the keys it doesn't set.""" _merge_platform_record(output_dir, lambda record: {**record, **update}) -OFF_HINT = "turn the monitor off with --no-monitors.prime" - - def prime_mode() -> str: """$PRIME_RUNS_MODE when set, else online. Like the W&B monitor, a configured monitor without credentials fails the run instead of leaving it untracked.""" @@ -79,324 +76,264 @@ def _base_url() -> str | None: return base.rstrip("/").removesuffix("/rft") if base else None -class PrimeTrainMonitor(Monitor): - """Logs metrics and episodes to the Prime platform through ``prime_runs``. +def _registration(name: str | None, config: BaseConfig | None) -> dict[str, Any]: + """``pr.init`` arguments that register a new training run for ``config``.""" + if isinstance(config, SFTConfig): + # SFT trains on a dataset, not on environments, and has no rollouts. + return dict( + name=name, + model=config.model.name, + environments=[], + training=pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.data.batch_size, + seq_len=config.data.seq_len, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + ), + config=config.model_dump(exclude_none=True, mode="json"), + ) + if config is not None: + return dict( + name=name, + model=config.model.name, + environments=[env.env_id for env in config.train.source], + training=pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.batch_size, + rollouts_per_example=max( + (source.group_size for source in config.train.source), default=config.train.group_size + ), + seq_len=config.seq_len, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + ), + config=config.model_dump(exclude_none=True, mode="json"), + ) + # The RFT API requires a base model. + return {"name": name, "model": "unknown"} - The run handle owns the RFT lifecycle (register or attach, finalize), the - per-step metrics POSTs, the episode uploads to Prime Traces, and the - terminal status. Delivery is BEST-EFFORT: the SDK's atexit hook - reports a process that exits without finalizing, but an abrupt kill - (SIGKILL, OOM, node loss) skips the hook and a lost finalize is not - retried here — the platform owns the terminal state of attached - ($RUN_ID) and managed runs. `FINISH_TIMEOUT` bounds the queued-upload - drain, not the SDK's finalize HTTP retries. - ``init``/``finish`` do network I/O and run in worker threads; the log - calls are queue puts onto the SDK's uploader thread, which owns retries - and backpressure, so they never stall the loop. +def wait_for_train_run(output_dir: Path, timeout: float = TRAIN_RUN_TIMEOUT_S) -> str: + """The id of the training run this attempt's trainer registered, from the platform + record it leaves in the shared run dir.""" + attempt = os.getenv(ATTEMPT_VAR) + deadline = time.monotonic() + timeout + while True: + try: + record = read_platform_record(output_dir) or {} + except (OSError, ValueError): + record = {} + if record.get("kind") == "train" and record.get("attempt") == attempt and record.get("id"): + return record["id"] + if time.monotonic() >= deadline: + raise RuntimeError(f"The trainer registered no Prime run within {timeout:.0f}s; {OFF_HINT}") + time.sleep(TRAIN_RUN_POLL_S) + + +class PrimeMonitor(Monitor): + """Logs a run to Prime through ``prime_runs``. + + A training run's processes (the rl orchestrator, the sft trainer and its online-eval + process) log to one Prime training run: metrics, and every episode as it arrives. + The process with its own config registers the run, or attaches to the one a + hosted launch created ($RUN_ID); an SFT online-eval process attaches to its + trainer's. ``$PRL_RUN_FINISHER`` names the process that finishes a shared run. + + A standalone eval opens one Prime evaluation per env and epoch when the epoch + starts (``log_eval_plan``), streams its episodes into it, and finishes it with the + epoch's aggregates (``log_eval_epoch``). + + Delivery is best-effort: the SDK's uploader thread owns retries and backpressure, + and its atexit hook reports a process that exits without finishing. An abrupt kill + skips the hook, and the platform owns the terminal state of attached runs. """ - config: PrimeTrainMonitorConfig - run: pr.Run - - async def init(self, config: BaseConfig | None = None, output_dir: Path | None = None) -> None: - init_kwargs: dict[str, Any] - if run_id := os.getenv("RUN_ID"): - # A managed launch pre-created the platform run and injected its id - - # attach instead of registering a duplicate. The platform owns the - # attached run's failure/completion marking; a clean finish() is - # delivered best-effort and does not by itself mark it completed. - init_kwargs = {"id": run_id} - elif isinstance(config, SFTConfig): - # SFT trains on a dataset, not on environments, and has no rollouts - the - # platform run is registered without both. - init_kwargs = dict( - name=self.config.name, - model=config.model.name, - environments=[], - training=pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.data.batch_size, - seq_len=config.data.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, - ), - config=config.model_dump(exclude_none=True, mode="json"), - ) - elif config is not None: - init_kwargs = dict( - name=self.config.name, - model=config.model.name, - environments=[env.env_id for env in config.train.source], - training=pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.batch_size, - rollouts_per_example=max( - (source.group_size for source in config.train.source), default=config.train.group_size - ), - seq_len=config.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, - ), - config=config.model_dump(exclude_none=True, mode="json"), - ) + config: PrimeMonitorConfig + + async def init(self) -> None: + config = self.context.run_config + self.mode = prime_mode() + self.run: pr.Run | None = None + # (env, step) -> its open evaluation; None once opening it failed, so the epoch's + # episodes do not retry Prime on every arrival + self.evaluations: dict[tuple[str, int], pr.Run | None] = {} + if isinstance(config, EvalConfig): + await self._init_evaluations(config) + return + online_eval = isinstance(config, SFTOnlineEvalConfig) + self.run = await asyncio.to_thread(self._open_run, config, online_eval) + if not self.run.url: + self.logger.info("Prime run disabled") else: - # The RFT API requires a base model; "unknown" is what the - # pre-SDK register sent for a config-less init. - init_kwargs = {"name": self.config.name, "model": "unknown"} + attached = " (attached)" if self.run.attached else "" + self.logger.info(f"Logging metrics and episodes to Prime run {self.run.id} ({self.run.url}){attached}") + if not online_eval: + # Off the event loop: the flock retry must not stall it. + await asyncio.to_thread(self._record_run) - mode = prime_mode() + def _open_run(self, config: BaseConfig | None, online_eval: bool) -> pr.Run: + """Register the training run, or attach to it. Blocking: runs in a worker thread.""" + run_id = os.getenv("RUN_ID") + if online_eval and not run_id and self.mode == "online": + run_id = wait_for_train_run(self.context.output_dir) + kwargs = {"id": run_id} if run_id else _registration(self.config.name, config) try: - self.run = await asyncio.to_thread( - pr.init, kind="train", mode=mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT, **init_kwargs - ) + return pr.init(kind="train", mode=self.mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT, **kwargs) except (pr.ConfigurationError, pr.APIError) as e: # Registering fails without a team, or for a team outside the external-runs # allowlist. An attached run's error is the launcher's to explain. - if "id" in init_kwargs: + if run_id: raise raise RuntimeError( f"Failed to register the Prime run ({type(e).__name__}: {e}). Select a team with " f"`prime switch` (or set PRIME_TEAM_ID), or {OFF_HINT}" ) from e - if self.run.url: - attached = " (attached via $RUN_ID)" if self.run.attached else "" - self.logger.info(f"Logging metrics and episodes to Prime run {self.run.id} ({self.run.url}){attached}") - else: - self.logger.info("Prime run disabled") - if output_dir is not None: - # Off the event loop: the flock retry must not stall it. - await asyncio.to_thread(self._record, output_dir) - def _record(self, output_dir: Path) -> None: + def _record_run(self) -> None: + output_dir = self.context.output_dir register_links(output_dir, PRIME_LINKS, 1 if self.run.url else 0) if self.run.url: - # Merge, not overwrite: a concurrent eval process (SFT online evals share - # the trainer's run dir) may already have written its evaluations record. - update_platform_record(output_dir, {"kind": "train", "id": self.run.id, "url": self.run.url}) + update_platform_record( + output_dir, + {"kind": "train", "id": self.run.id, "url": self.run.url, "attempt": os.getenv(ATTEMPT_VAR)}, + ) publish_link(output_dir, PRIME_LINKS, "run", self.run.url) - async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: - # The SDK also drops non-finite values, but silently; sanitize first so - # the dropped paths are named in the log. - metrics, dropped = sanitize(metrics) - if dropped: - self.logger.warning(f"Dropping {len(dropped)} non-finite metric value(s): {', '.join(dropped[:5])}") - # A queue put that can block briefly under backpressure - off the loop. The SDK - # stamps `_timestamp` on every row, so step=None rows keep a time anchor. - await asyncio.to_thread(self.run.log_metrics, metrics, step=step) - - async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: - """Every episode ships whole as it arrives (the ``all`` subset), errored and - filtered ones included, online evals too. The SDK keys each to its dispatch - step and work type off ``run.work`` - the ``TrainRunInfo`` the dispatcher - stamps at emit time.""" - if subset != "all" or not episodes: - return - # A queue put that can block briefly under backpressure - off the loop. - await asyncio.to_thread(self.run.log_episodes, episodes) - - async def finalize(self, paused: bool = False) -> None: - # Drains queued uploads so the final step's metrics and episodes land, - # then finalizes (idempotent on the platform side); an attached run's - # failure marking stays with the launcher. A paused run only drains and - # stays open: it continues on resume, and its owner sets its status. - await asyncio.to_thread(self.run.flush if paused else self.run.finish) - - -class PrimeEvalMonitor(Monitor): - """Streams each eval epoch to the Prime platform through ``prime_runs``: one - evaluation per env and step, opened when the epoch starts (``log_eval_plan``), fed - every episode as it lands (``log_episodes``) and finished with the epoch's aggregates - (``log_eval_epoch``), so the platform page follows the run from its first episode. - - Under a managed launch ($RUN_ID) the epoch's ``eval/{env}/...`` metrics ALSO go to - the parent training run, which is what the run dashboard groups into its `eval` - section. The orchestrator's in-process evals reach it through the train monitor, but - an eval process of its own (SFT online evals) has no other way there, so without - this the run shows every metric group except the one its evals produced.""" - - config: PrimeEvalMonitorConfig - - async def init(self, config: BaseConfig | None = None, output_dir: Path | None = None) -> None: - self.mode = prime_mode() - self.model: str = config.model if config is not None else "unknown" - self.sources = {source.resolved_name: source for source in config.source} if config is not None else {} - self.run_id = os.getenv("PRL_RUN_ID") - self.output_dir = output_dir - # (env, step) -> its open evaluation; None once opening it failed, so the epoch's - # episodes do not retry the platform on every arrival - self.runs: dict[tuple[str, int], pr.Run | None] = {} - self.parent: pr.Run | None = None - self._lock = asyncio.Lock() + async def _init_evaluations(self, config: EvalConfig) -> None: + self.model = config.model + self.sources = {source.resolved_name: source for source in config.source} self.evaluation_id = os.getenv(EVAL_ID_VAR) if self.evaluation_id and len(self.sources) != 1: raise ValueError( - f"${EVAL_ID_VAR} names one platform evaluation, so the run needs " + f"${EVAL_ID_VAR} names one Prime evaluation, so the run needs " f"exactly one eval source (got {len(self.sources)})" ) - if self.mode == "online": - self.logger.info("Streaming eval epochs to the Prime platform") - self.parent = await asyncio.to_thread(self._attach_parent) - if output_dir is not None: - # Merge, not overwrite: when the eval process shares the - # trainer's run dir (SFT online evals), the train record's - # platform link (kind/id/url) must survive; a standalone - # eval keeps the plain eval-shaped record. Locked RMW. - def _eval_init_update(record: dict[str, Any]) -> dict[str, Any]: - if record.get("kind") == "train": - record.setdefault("evaluations", {}) - record["run_id"] = self.run_id - return record - return {"kind": "eval", "run_id": self.run_id, "evaluations": {}} - - # Off the event loop: the flock retry must not stall it. - await asyncio.to_thread(_merge_platform_record, output_dir, _eval_init_update) - else: - self.logger.info("Prime evaluations disabled") - if output_dir is not None: - expected = len(self.sources) if self.mode == "online" else 0 - await asyncio.to_thread(register_links, output_dir, PRIME_EVAL_LINKS, expected) - - def _attach_parent(self) -> pr.Run | None: - """The parent training run this eval process belongs to, for metrics only. - Blocking: runs in a worker thread. - - Attaching is a local handle, not a request, so a run that turns out to be - unreachable costs the eval nothing. The caller owns the run's lifecycle: we - drop the SDK's atexit crash hook, because an eval process that exits while the - trainer is still going would otherwise report the WHOLE run failed.""" - run_id = os.getenv("RUN_ID") - if not run_id: - return None - run = pr.init( - kind="train", - mode=self.mode, - base_url=_base_url(), - id=run_id, - finish_timeout=FINISH_TIMEOUT, - ) - atexit.unregister(run._atexit_hook) # the SDK's own unregister handle - self.logger.info(f"Logging eval metrics to parent training run {run.id}") - return run + online = self.mode == "online" + self.logger.info("Streaming eval epochs to Prime" if online else "Prime evaluations disabled") + + def record() -> None: + if online: + write_platform_record(self.context.output_dir, {"kind": "eval", "evaluations": {}}) + register_links(self.context.output_dir, PRIME_LINKS, len(self.sources) if online else 0) + + await asyncio.to_thread(record) async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: - """An epoch's aggregates (``eval/{env}/...``) onto the parent training run. - Only there: a standalone eval has no parent, and the per-epoch evaluation - carries its own aggregates through ``log_eval_epoch``.""" - if self.parent is None: + """Onto the training run. A standalone eval's evaluations carry their own + aggregates through ``log_eval_epoch``.""" + if self.run is None: return + # The SDK also drops non-finite values, but silently; sanitize first so the + # dropped paths are named in the log. metrics, dropped = sanitize(metrics) if dropped: self.logger.warning(f"Dropping {len(dropped)} non-finite metric value(s): {', '.join(dropped[:5])}") - # A queue put that can block briefly under backpressure - off the loop. - await asyncio.to_thread(self.parent.log_metrics, metrics, step=step) + # A queue put that can block briefly under backpressure - off the loop. The SDK + # stamps `_timestamp` on every row, so step=None rows keep a time anchor. + await asyncio.to_thread(self.run.log_metrics, metrics, step=step) + + async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind) -> None: + """Every episode ships whole as it arrives, keyed by the SDK to its work type and + dispatch step off ``run.work``.""" + if self.run is not None: + # A queue put that can block briefly under backpressure - off the loop. + await asyncio.to_thread(self.run.log_episodes, episodes) + return + by_env: dict[str, list[vf.Episode]] = {} + for episode in episodes: + if episode.env.name is not None: + by_env.setdefault(episode.env.name, []).append(episode) + for env_name, batch in by_env.items(): + evaluation = await self._evaluation(env_name, step) + if evaluation is not None: + await asyncio.to_thread(evaluation.log_episodes, batch) - def open(self, env_name: str, step: int, expected: int | None) -> pr.Run: - """Open the platform evaluation of one epoch. Blocking: runs in a worker thread.""" + def _open_evaluation(self, env_name: str, step: int, expected: int | None) -> pr.Run: + """Open the Prime evaluation of one epoch. Blocking: runs in a worker thread.""" + common = dict(kind="eval", mode=self.mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT) if self.evaluation_id: - # A hosted launch pre-created the platform evaluation and injected its id - - # attach instead of registering a duplicate. The backend owns its failure - # marking then; a clean finish still completes it. - if self.runs: + # A hosted launch pre-created the evaluation and injected its id. + if self.evaluations: raise RuntimeError(f"${EVAL_ID_VAR} holds one epoch, and it already took one") - return pr.init( - kind="eval", - mode=self.mode, - base_url=_base_url(), - id=self.evaluation_id, - finish_timeout=FINISH_TIMEOUT, - ) + return pr.init(id=self.evaluation_id, **common) source = self.sources[env_name] - name = self.config.name if len(self.sources) == 1 else f"{self.config.name}--{env_name}" return pr.init( - kind="eval", - mode=self.mode, - base_url=_base_url(), - name=name, + name=self.config.name if len(self.sources) == 1 else f"{self.config.name}--{env_name}", environments=[source.env.taskset.id], model=self.model, framework="prime-rl", - finish_timeout=FINISH_TIMEOUT, config={ "model": self.model, "step": step, - "run_id": self.run_id, + "run_id": os.getenv("PRL_RUN_ID"), "num_examples": expected // source.group_size if expected else None, "rollouts_per_example": source.group_size, }, + **common, ) - async def run_for(self, env_name: str, step: int, expected: int | None = None) -> pr.Run | None: + async def _evaluation(self, env_name: str, step: int, expected: int | None = None) -> pr.Run | None: """The epoch's evaluation, opened on first use.""" key = (env_name, step) - async with self._lock: - if key in self.runs: - return self.runs[key] - try: - run = await asyncio.to_thread(self.open, env_name, step, expected) - except Exception as e: - self.logger.warning(f"Failed to open the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") - self.runs[key] = None - if self.output_dir is not None: - await asyncio.to_thread(publish_link, self.output_dir, PRIME_EVAL_LINKS, env_name, None) - return None - self.runs[key] = run - if run.url: - attached = f" (attached via ${EVAL_ID_VAR})" if run.attached else "" - self.logger.info(f"Streaming {env_name} (Step {step}) evaluation - {run.url}{attached}") - if self.output_dir is not None: - # Off the event loop: the flock retry must not stall it. - await asyncio.to_thread(self._record, env_name, step, run) - return run - - def _record(self, env_name: str, step: int, run: pr.Run) -> None: - if run.url: + if key in self.evaluations: + return self.evaluations[key] + try: + evaluation = await asyncio.to_thread(self._open_evaluation, env_name, step, expected) + except Exception as e: + self.logger.warning(f"Failed to open the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") + evaluation = None + if key in self.evaluations: # another arrival opened it meanwhile + return self.evaluations[key] + self.evaluations[key] = evaluation + if evaluation is not None and evaluation.url: + attached = f" (attached via ${EVAL_ID_VAR})" if evaluation.attached else "" + self.logger.info(f"Streaming {env_name} (Step {step}) evaluation - {evaluation.url}{attached}") + await asyncio.to_thread(self._record_evaluation, env_name, step, evaluation) + return evaluation + + def _record_evaluation(self, env_name: str, step: int, evaluation: pr.Run | None) -> None: + url = evaluation.url if evaluation is not None else None + if url: def update(record: dict[str, Any]) -> dict[str, Any]: - record.setdefault("kind", "eval") - record.setdefault("run_id", self.run_id) - record.setdefault("evaluations", {})[env_name] = {"step": step, "id": run.id, "url": run.url} + record.setdefault("evaluations", {})[env_name] = {"step": step, "id": evaluation.id, "url": url} return record - _merge_platform_record(self.output_dir, update) - publish_link(self.output_dir, PRIME_EVAL_LINKS, env_name, run.url) + _merge_platform_record(self.context.output_dir, update) + publish_link(self.context.output_dir, PRIME_LINKS, env_name, url) async def log_eval_plan(self, env_name: str, step: int, expected: int) -> None: - await self.run_for(env_name, step, expected) - - async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: - """Every eval episode as it lands (the ``all`` subset is the arrival stream).""" - if kind != "eval" or subset != "all": - return - by_env: dict[str, list[vf.Episode]] = {} - for episode in episodes: - if episode.env.name is not None: - by_env.setdefault(episode.env.name, []).append(episode) - for env_name, batch in by_env.items(): - run = await self.run_for(env_name, step) - if run is not None: - await asyncio.to_thread(run.log_episodes, batch) # a queue put, off the loop + if self.run is None: + await self._evaluation(env_name, step, expected) async def log_eval_epoch(self, env_name: str, step: int, episodes: list[vf.Episode]) -> None: - run = await self.run_for(env_name, step) - if run is None: + if self.run is not None: + return + evaluation = await self._evaluation(env_name, step) + if evaluation is None: return try: - await asyncio.to_thread(run.finish, pr.metrics.from_episodes(episodes)) + await asyncio.to_thread(evaluation.finish, pr.metrics.from_episodes(episodes)) except Exception as e: self.logger.warning(f"Failed to finish the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") return - self.logger.info(f"Uploaded {env_name} (Step {step}) evaluation - {run.url}") + self.logger.info(f"Uploaded {env_name} (Step {step}) evaluation - {evaluation.url}") async def finalize(self, paused: bool = False) -> None: # An epoch the run did not finish leaves its evaluation open: close it as cancelled. - for (env_name, step), run in self.runs.items(): - if run is None or run.finished: + for (env_name, step), evaluation in self.evaluations.items(): + if evaluation is None or evaluation.finished: continue try: - await asyncio.to_thread(run.finish, status=pr.RunStatus.CANCELLED, error="interrupted") + await asyncio.to_thread(evaluation.finish, status=pr.RunStatus.CANCELLED, error="interrupted") except Exception as e: self.logger.warning(f"Failed to close the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") - if self.parent is not None: - # Drain only. The parent run belongs to the trainer and the platform: an - # eval process must never report its status, not even a clean completion. - await asyncio.to_thread(self.parent.flush, FINISH_TIMEOUT) + if self.run is None: + return + producer = self.context.producer + if paused or os.getenv(RUN_FINISHER_VAR, producer) != producer: + # Drain only: a paused run continues on resume, and a shared run is finished + # by the process that outlives this one. A clean exit is not a crash. + await asyncio.to_thread(self.run.flush, FINISH_TIMEOUT) + atexit.unregister(self.run._atexit_hook) + return + await asyncio.to_thread(self.run.finish) diff --git a/src/prime_rl/monitors/wandb/monitor.py b/src/prime_rl/monitors/wandb/monitor.py index c85bf8c7baf..11eb32be611 100644 --- a/src/prime_rl/monitors/wandb/monitor.py +++ b/src/prime_rl/monitors/wandb/monitor.py @@ -6,21 +6,20 @@ import sys import time from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any import wandb from wandb.errors import CommError, UsageError from wandb.sdk.mailbox.mailbox_handle import ServerResponseError from prime_rl.configs.monitors import WandbMonitorConfig -from prime_rl.monitors.base import Kind, Monitor, Subset +from prime_rl.monitors.base import Monitor from prime_rl.monitors.wandb.overview import ensure_overview_view -from prime_rl.utils.config import BaseConfig from prime_rl.utils.links import WANDB_LINKS, publish_link, register_links from prime_rl.utils.logger import format_time if TYPE_CHECKING: - import verifiers.v1 as vf + pass class WandbMonitor(Monitor): @@ -28,14 +27,8 @@ class WandbMonitor(Monitor): config: WandbMonitorConfig - async def init( - self, - output_dir: Path, - config: BaseConfig | None = None, - train_env_names: list[str] | None = None, - eval_env_names: list[str] | None = None, - overview_flavor: Literal["rl", "sft", "eval"] = "rl", - ) -> None: + async def init(self) -> None: + output_dir, config = self.context.output_dir, self.context.run_config # W&B reads the start command off sys.argv; the launcher passes the original # command to subprocesses via $WANDB_ARGS. wandb_args = os.environ.get("WANDB_ARGS") @@ -135,9 +128,9 @@ def init_wandb(max_retries: int): url = ensure_overview_view( self.wandb.entity, self.wandb.project, - train_envs=train_env_names or [], - eval_envs=eval_env_names or [], - flavor=overview_flavor, + train_envs=self.context.train_env_names, + eval_envs=self.context.eval_env_names, + flavor=self.context.overview_flavor, ) if url: self.logger.info(f"Created W&B overview view - {url}") @@ -166,9 +159,6 @@ async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: else: wandb.log({**metrics, "step": step, "_timestamp": time.time()}) - async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: - pass - async def finalize(self, paused: bool = False) -> None: self.logger.info(f"Finalizing W&B run {self.wandb.id}") t0 = time.perf_counter() diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 01fcfe4d36a..33698566c44 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -558,7 +558,7 @@ async def main_loop(self) -> None: kind = episode.run.work.type step = episode.run.work.step if kind == "eval" else self.progress.step stamp_arrival([episode], kind, step) - await monitors.log([episode], step, kind, "all") + await monitors.log([episode], step, kind) if kind == "eval": assert self.eval_sink is not None # eval rollouts only emitted when eval is configured @@ -679,8 +679,6 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: # records against each trace's arrival record - membership, advantages, the step # it shipped at - never a second episode copy. effective_episodes = await asyncio.to_thread(lambda: effective.vf_episodes) - # Monitors may retain the list for asynchronous upload after publication. - await monitors.log(effective_episodes.copy(), step, "train", "effective") annotations = await asyncio.to_thread(stamp_batch, effective_episodes, step) await monitors.log_annotations(annotations) @@ -911,11 +909,9 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: get_logger().warning(f"Eval @ step={batch.step} env={batch.env_name}: no attempts returned, skipping log") return - # The non-errored subset is logged on epoch completion (multiple eval envs share the - # step's trace file — each epoch appends its cohort once, and every record carries - # ``env_name``); the full returned cohort already streamed into ``all`` on arrival. + # Every episode was logged on arrival; the non-errored subset is annotated on epoch + # completion. if batch.episodes.effective: - await monitors.log(batch.episodes.effective.vf_episodes, batch.step, "eval", "effective") await monitors.log_annotations(stamp_batch(batch.episodes.effective.vf_episodes, batch.step)) policy_spans = [eval_work(episode).policy for episode in batch.episodes] if any(span is None for span in policy_spans): diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index 32027951d71..3579a093b61 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -58,6 +58,8 @@ ln -sfn inference/node_0.log $LOG_DIR/inference.log export TRAINER_DONE=$LAUNCHER_DIR/.sft_trainer_done_$SLURM_JOB_ID {% if online_eval -%} export EVAL_DONE=$LAUNCHER_DIR/.sft_eval_done_$SLURM_JOB_ID +# The online-eval process outlives the trainer, so it finishes the shared Prime run. +export PRL_RUN_FINISHER=online-eval {%- endif %} mkdir -p "$LAUNCHER_DIR" rm -f "$TRAINER_DONE"{% if online_eval %} "$EVAL_DONE"{% endif %} diff --git a/src/prime_rl/utils/links.py b/src/prime_rl/utils/links.py index 1c7b10f245c..f2304d8f910 100644 --- a/src/prime_rl/utils/links.py +++ b/src/prime_rl/utils/links.py @@ -16,10 +16,8 @@ # The launcher exports it to every component: entries stamped with another attempt are # a previous launch's. ATTEMPT_VAR = "PRL_ATTEMPT_LOG_DIR" -# Monitor keys in the links record. An SFT run's trainer (train monitor) and -# online-eval process (eval monitor) share one run dir, so the Prime monitors differ. +# Monitor keys in the links record. PRIME_LINKS = "prime" -PRIME_EVAL_LINKS = "prime/eval" WANDB_LINKS = "wandb" LINK_TIMEOUT_S = 300.0 LINK_POLL_S = 0.5 diff --git a/src/prime_rl/utils/process.py b/src/prime_rl/utils/process.py index 29eab2cb0d7..bebbe9b4274 100644 --- a/src/prime_rl/utils/process.py +++ b/src/prime_rl/utils/process.py @@ -26,6 +26,10 @@ "TRITON_CACHE_DIR": f"/tmp/triton-{os.getuid()}", } +# The producer that finishes a run several processes log to, e.g. the SFT online-eval +# process, which drains its final evals after the trainer exits. +RUN_FINISHER_VAR = "PRL_RUN_FINISHER" + DEFAULT_TRAINER_ENV_VARS: dict[str, str] = { "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", # NCCL flight recorder: on a collective timeout every rank dumps its last 2000 collectives diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index 4049eb3f159..113a4b913a9 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -1,57 +1,64 @@ -"""The SFT trainer's platform wiring: registration shape and monitor fan-out. - -``monitors.prime`` on ``SFTConfig`` must register the run with the dataset-batched -training fields (no environments, no rollouts) and flow the trainer's metrics and -finalize through the registered ``PrimeTrainMonitor`` like the RL orchestrator does. -""" +"""The prime monitor's run wiring: SFT registration, online evals sharing the trainer's +run, and the platform record.""" import asyncio +import errno +import fcntl +import os from types import SimpleNamespace import pytest -from prime_rl import monitors -from prime_rl.configs.monitors import ( - EvalMonitorsConfig, - FileMonitorConfig, - PrimeEvalMonitorConfig, - PrimeTrainMonitorConfig, - TrainMonitorsConfig, - WandbMonitorConfig, -) +import prime_rl.monitors.prime as prime_module +import prime_rl.utils.record as record_module +from prime_rl.configs.monitors import PrimeMonitorConfig from prime_rl.configs.sft import SFTConfig -from prime_rl.monitors.prime import PrimeTrainMonitor +from prime_rl.monitors.base import MonitorContext +from prime_rl.monitors.prime import PrimeMonitor, read_platform_record, update_platform_record +from prime_rl.utils.links import ATTEMPT_VAR +from prime_rl.utils.pathing import get_platform_run_path +from prime_rl.utils.process import RUN_FINISHER_VAR +from prime_rl.utils.record import record_lock -@pytest.fixture -def prime_init(monkeypatch): - """Replace the ``prime_runs`` SDK entry point with one that records its kwargs.""" - calls = [] +class FakeRun(SimpleNamespace): + def __init__(self, **kwargs): + super().__init__(id=kwargs.get("id", "run-123"), url="https://prime/training/run-123", attached="id" in kwargs) + self.finished = False + self.flushed = False + self.kwargs = kwargs - def fake_init(**kwargs): - calls.append(kwargs) - return SimpleNamespace( - id="run-123", - url="https://platform/training/run-123", - attached=False, - finished=False, - finish=lambda *args, **kwargs_: None, - log_metrics=lambda *args, **kwargs_: None, - log_episodes=lambda *args, **kwargs_: None, - ) - - monkeypatch.setattr("prime_rl.monitors.prime.pr.init", fake_init) - monkeypatch.delenv("RUN_ID", raising=False) - monkeypatch.delenv("PRIME_RUNS_MODE", raising=False) - monkeypatch.setenv("PRIME_API_KEY", "test-key") - return calls + def _atexit_hook(self): + pass + + def finish(self, *args, **kwargs): + self.finished = True + + def flush(self, *args, **kwargs): + self.flushed = True + + def log_metrics(self, *args, **kwargs): + pass + + def log_episodes(self, *args, **kwargs): + pass @pytest.fixture -def clear_monitors(): - """``monitors.setup`` registers into a module-global list; empty it after each test.""" - yield - monitors.MONITORS.clear() +def prime_init(monkeypatch, tmp_path): + """Replace the ``prime_runs`` SDK entry point with one that records its runs.""" + runs: list[FakeRun] = [] + + def fake_init(**kwargs): + runs.append(FakeRun(**kwargs)) + return runs[-1] + + monkeypatch.setattr(prime_module.pr, "init", fake_init) + for var in ("RUN_ID", "PRIME_RUNS_MODE", RUN_FINISHER_VAR): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("PRIME_API_KEY", "test-key") + monkeypatch.setenv(ATTEMPT_VAR, str(tmp_path / "logs" / "attempt_1")) + return runs def _sft_config(overrides: dict | None = None) -> SFTConfig: @@ -60,262 +67,72 @@ def _sft_config(overrides: dict | None = None) -> SFTConfig: "data": {"type": "fake", "batch_size": 8, "seq_len": 64}, "max_steps": 10, } - if overrides: - config.update(overrides) + config.update(overrides or {}) return SFTConfig.model_validate(config) -def test_sft_prime_monitor_name_inherits_run_name(): - config = _sft_config({"monitors": {"prime": {}}}) - - assert config.monitors.prime.name == config.run.name - - -def test_sft_prime_monitor_keeps_explicit_name(): - config = _sft_config({"monitors": {"prime": {"name": "my-experiment"}}}) - - assert config.monitors.prime.name == "my-experiment" - - -def test_sft_registers_dataset_training_fields(prime_init): - config = _sft_config({"monitors": {"prime": {}}}) - - asyncio.run(PrimeTrainMonitor(config.monitors.prime).init(config=config, output_dir=None)) - - (call,) = prime_init - assert call["kind"] == "train" - assert call["name"] == config.run.name - assert call["model"] == "PrimeIntellect/Qwen3-0.6B" - # SFT trains on a dataset: no environments, no rollouts per example - assert call["environments"] == [] - assert call["training"].max_steps == 10 - assert call["training"].batch_size == 8 - assert call["training"].seq_len == 64 - assert call["training"].rollouts_per_example is None - assert call["config"]["model"]["name"] == "PrimeIntellect/Qwen3-0.6B" - - -def test_sft_registers_wandb_project(prime_init): - config = _sft_config( - { - "monitors": { - "wandb": {"project": "sft-tests"}, - "prime": {}, - } - } - ) - - asyncio.run(PrimeTrainMonitor(config.monitors.prime).init(config=config, output_dir=None)) - - assert prime_init[0]["training"].wandb_project == "sft-tests" - - -def test_online_eval_monitors_switch_prime_flavor(): - """The online-eval process parses EvalMonitorsConfig, whose ``prime`` field is the - evaluation flavor - the launcher converts the trainer's TrainMonitorsConfig.""" - from prime_rl.entrypoints.sft import build_online_eval_monitors +def _monitor(tmp_path, config, producer: str) -> PrimeMonitor: + return PrimeMonitor(PrimeMonitorConfig(name="my-run"), MonitorContext(tmp_path, producer, config)) - monitors = build_online_eval_monitors( - TrainMonitorsConfig( - wandb=WandbMonitorConfig(project="sft-tests"), - file=FileMonitorConfig(), - prime=PrimeTrainMonitorConfig(name="my-run"), - ) - ) - assert isinstance(monitors, EvalMonitorsConfig) - assert isinstance(monitors.prime, PrimeEvalMonitorConfig) - assert monitors.prime.name == "my-run" - assert monitors.wandb.project == "sft-tests" +def test_sft_registers_dataset_training_fields(prime_init, tmp_path): + config = _sft_config({"monitors": {"wandb": {"project": "sft-tests"}}}) + asyncio.run(_monitor(tmp_path, config, "trainer").init()) - # Without a platform run the evals process gets no platform monitor either - bare = build_online_eval_monitors(TrainMonitorsConfig(prime=None)) - assert bare.prime is None + (run,) = prime_init + assert run.kwargs["kind"] == "train" + assert run.kwargs["environments"] == [] + training = run.kwargs["training"] + assert (training.max_steps, training.batch_size, training.seq_len) == (10, 8, 64) + assert training.rollouts_per_example is None + assert training.wandb_project == "sft-tests" + assert read_platform_record(tmp_path)["id"] == "run-123" -def test_sft_setup_registers_prime_monitor(prime_init, tmp_path, clear_monitors): - config = _sft_config({"monitors": {"prime": {}}}) - - asyncio.run( - monitors.setup( - prime=config.monitors.prime, - output_dir=tmp_path, - run_config=config, - ) - ) - - assert monitors.get(PrimeTrainMonitor) is not None - - # The trainer's per-step metrics fan out to the platform run, and finalize - # closes it out (drains the SDK's upload queue). - asyncio.run(monitors.log({"loss/mean": 1.0, "step": 1}, step=1)) - asyncio.run(monitors.log({"val/loss": 0.5, "step": 1}, step=1)) - asyncio.run(monitors.finalize()) +def test_online_evals_share_the_trainers_run(prime_init, tmp_path, monkeypatch): + """The online-eval process attaches to the run the trainer registered, and only it + finishes the run: it drains its final evals after the trainer exits.""" + from prime_rl.entrypoints.sft import build_online_eval_config + config = _sft_config({"eval": {"interval": 5, "source": [{"env": {"taskset": {"id": "harbor"}}}]}, "inference": {}}) + eval_config = build_online_eval_config(config) + assert eval_config.monitors.prime == config.monitors.prime + monkeypatch.setenv(RUN_FINISHER_VAR, "online-eval") -def test_online_eval_config_keeps_prime_monitor(): - """SFTOnlineEvalConfig.monitors must be EvalMonitorsConfig so the - launcher-converted PrimeEvalMonitorConfig survives the eval.json - dump/re-parse boundary (base MonitorsConfig drops/forbids prime).""" - from prime_rl.configs.eval import SFTOnlineEvalConfig - from prime_rl.configs.monitors import EvalMonitorsConfig, PrimeEvalMonitorConfig - from prime_rl.entrypoints.sft import build_online_eval_monitors + trainer = _monitor(tmp_path, config, "trainer") + asyncio.run(trainer.init()) + online_eval = _monitor(tmp_path, eval_config, "online-eval") + asyncio.run(online_eval.init()) + assert online_eval.run.kwargs["id"] == trainer.run.id - field = SFTOnlineEvalConfig.model_fields["monitors"] - assert field.annotation is EvalMonitorsConfig + asyncio.run(trainer.finalize()) + assert trainer.run.flushed and not trainer.run.finished + asyncio.run(online_eval.finalize()) + assert online_eval.run.finished - train_monitors = TrainMonitorsConfig.model_validate({"prime": {"name": "r1"}}) - eval_monitors = build_online_eval_monitors(train_monitors) - assert isinstance(eval_monitors.prime, PrimeEvalMonitorConfig) - assert eval_monitors.prime.name == "r1" - # The online-eval process passes the prime monitor to monitors.setup - # (src/prime_rl/eval/online.py) — same wiring as the standalone eval - # entrypoint. Read the source text (importing the module pulls torch, - # which this torch-free test env deliberately lacks). - from pathlib import Path - - import prime_rl - - online_src = Path(prime_rl.__path__[0], "eval", "online.py").read_text() - assert "prime=config.monitors.prime" in online_src - - -def test_online_eval_shares_trainer_run_dir_and_merges_records(tmp_path): - """The online-eval process shares the trainer's run dir (the dashboard - reads its file-monitor artifacts from there), and the platform records - MERGE instead of clobbering each other: the trainer's kind="train" - link survives with the eval process's evaluations riding along.""" - from prime_rl.entrypoints.sft import build_online_eval_config - from prime_rl.monitors.prime import ( - read_platform_record, - write_platform_record, - ) - - config = _sft_config( - { - "eval": { - "interval": 50, - "source": [ - { - "name": "rev", - "env": {"taskset": {"id": "harbor"}}, - } - ], - }, - "inference": {}, - } - ) - eval_config = build_online_eval_config(config) - # Same run dir: dashboard visibility for the eval artifacts. - assert eval_config.output_dir == config.run_dir - assert eval_config.broadcasts_dir != eval_config.output_dir - - # Simulate the trainer's record, then the eval process merging on top. - write_platform_record(tmp_path, {"kind": "train", "id": "run-1", "url": "https://x/run-1"}) - record = read_platform_record(tmp_path) or {} - if record.get("kind") == "train": - record.setdefault("evaluations", {}) - record["run_id"] = "run-1" - else: - record = {"kind": "eval", "run_id": "run-1", "evaluations": {}} - record.setdefault("evaluations", {})["rev"] = {"step": 5, "id": "ev-1", "url": "https://x/ev-1"} - write_platform_record(tmp_path, record) - merged = read_platform_record(tmp_path) - assert merged["kind"] == "train" # the train link survives - assert merged["id"] == "run-1" - assert merged["evaluations"]["rev"]["id"] == "ev-1" +def test_online_eval_ignores_a_previous_attempts_run(prime_init, tmp_path, monkeypatch): + update_platform_record(tmp_path, {"kind": "train", "id": "old", "attempt": "attempt_0"}) + with pytest.raises(RuntimeError, match="registered no Prime run"): + prime_module.wait_for_train_run(tmp_path, timeout=0.0) def test_platform_record_lock_keeps_a_stable_inode(tmp_path): - """The record lock must NEVER unlink its lock file: a writer waiting on - the flock can still hold the unlinked inode's lock while a later writer - creates and locks a fresh file - two simultaneous "exclusive" locks and - a lost merge. The lock file must persist with the SAME inode across - acquisitions, and the locked update helper must merge two sequential - writers into one record. (Interprocess contention is exercised by the - two-process update_platform_record probe; here the writers are - sequential.)""" - import os - - from prime_rl.monitors.prime import update_platform_record - from prime_rl.utils.pathing import get_platform_run_path - from prime_rl.utils.record import record_lock - + """The record lock never unlinks its lock file: a waiting writer could hold the + unlinked inode's lock while a later writer locks a fresh file.""" lock_path = get_platform_run_path(tmp_path).with_suffix(".json.lock") inodes = [] for _ in range(2): with record_lock(get_platform_run_path(tmp_path)): - # The lock file exists DURING the critical section, always. - assert lock_path.is_file() inodes.append(os.stat(lock_path).st_ino) + assert lock_path.is_file() and len(set(inodes)) == 1 - # The lock file SURVIVES release, and every acquisition locked the same inode. - assert lock_path.is_file() - assert len(set(inodes)) == 1 - - # Two sequential locked updates merge instead of clobbering. - update_platform_record(tmp_path, {"kind": "train", "id": "run-1", "url": "https://x/run-1"}) - update_platform_record(tmp_path, {"evaluations": {"rev": {"step": 5, "id": "ev-1", "url": "https://x/ev-1"}}}) - from prime_rl.monitors.prime import read_platform_record - - record = read_platform_record(tmp_path) - assert record["kind"] == "train" - assert record["id"] == "run-1" - assert record["evaluations"]["rev"]["id"] == "ev-1" - - -def test_record_persistence_runs_off_the_event_loop(prime_init, tmp_path, monkeypatch): - """The record merge (bounded flock retry included) must run in a worker - thread via asyncio.to_thread, not on the event loop thread: contention - for the full lock bound would otherwise freeze the loop.""" - import threading - from pathlib import Path - - import prime_rl.monitors.prime as prime_module - - seen_threads = [] - real_merge = prime_module._merge_platform_record - - def spy(output_dir, merge): - seen_threads.append(threading.get_ident()) - return real_merge(output_dir, merge) - - monkeypatch.setattr(prime_module, "_merge_platform_record", spy) - - config = _sft_config({"monitors": {"prime": {}}}) - asyncio.run(PrimeTrainMonitor(config.monitors.prime).init(config=config, output_dir=tmp_path)) - - assert seen_threads, "the record merge never ran" - assert threading.get_ident() not in seen_threads, "record merge ran on the caller thread" - - # All three async write sites (train init, eval init, eval epoch) go - # through to_thread; only the sync helpers stay callable directly. - import re - - import prime_rl - - src = Path(prime_rl.__path__[0], "monitors", "prime.py").read_text() - assert len(re.findall(r"asyncio\.to_thread\(\s*_merge_platform_record", src)) == 1 - assert len(re.findall(r"asyncio\.to_thread\(\s*self\._record", src)) == 2 - - from prime_rl.monitors.prime import read_platform_record - - record = read_platform_record(tmp_path) - assert record["kind"] == "train" - assert record["id"] == "run-123" + update_platform_record(tmp_path, {"kind": "train", "id": "run-1"}) + update_platform_record(tmp_path, {"url": "https://x/run-1"}) + assert read_platform_record(tmp_path) == {"kind": "train", "id": "run-1", "url": "https://x/run-1"} def test_platform_record_lock_acquisition_is_bounded(tmp_path, monkeypatch): - """flock(LOCK_EX) must never block the event-loop thread indefinitely: - while another process holds the lock, acquisition retries - non-blockingly and raises TimeoutError past the bound instead of - hanging.""" - import fcntl - - import prime_rl.utils.record as record_module - from prime_rl.utils.pathing import get_platform_run_path - monkeypatch.setattr(record_module, "RECORD_LOCK_TIMEOUT", 0.2) path = get_platform_run_path(tmp_path) lock_path = path.with_suffix(".json.lock") @@ -323,23 +140,16 @@ def test_platform_record_lock_acquisition_is_bounded(tmp_path, monkeypatch): with open(lock_path, "a+b") as held: fcntl.flock(held.fileno(), fcntl.LOCK_EX) with pytest.raises(TimeoutError, match="still held"): - with record_module.record_lock(path): + with record_lock(path): pass def test_platform_record_persistence_is_best_effort(tmp_path, monkeypatch): - """A storage failure on the OPTIONAL dashboard record (e.g. ENOLCK on a - filesystem without flock support) must not abort training: the write is - skipped with a warning, and no record is left behind.""" - import errno - - import prime_rl.monitors.prime as prime_module - import prime_rl.utils.record as record_module + """A storage failure on the record (e.g. ENOLCK) is skipped, never fatal.""" def enolck(fd, op): raise OSError(errno.ENOLCK, "No locks available") monkeypatch.setattr(record_module.fcntl, "flock", enolck) - # Must not raise: the platform record is best-effort. - prime_module.update_platform_record(tmp_path, {"kind": "train", "id": "run-1", "url": "https://x/run-1"}) - assert prime_module.read_platform_record(tmp_path) is None + update_platform_record(tmp_path, {"kind": "train", "id": "run-1"}) + assert read_platform_record(tmp_path) is None diff --git a/tests/unit/utils/test_links.py b/tests/unit/utils/test_links.py index 48ae6045929..d58043b27e8 100644 --- a/tests/unit/utils/test_links.py +++ b/tests/unit/utils/test_links.py @@ -5,8 +5,9 @@ import pytest import prime_rl.monitors.prime as prime_module -from prime_rl.configs.monitors import PrimeTrainMonitorConfig -from prime_rl.monitors.prime import PrimeTrainMonitor +from prime_rl.configs.monitors import PrimeMonitorConfig +from prime_rl.monitors.base import MonitorContext +from prime_rl.monitors.prime import PrimeMonitor from prime_rl.utils.links import ATTEMPT_VAR, publish_link, register_links, wait_for_links @@ -19,16 +20,16 @@ def attempt(tmp_path, monkeypatch): def test_waits_for_every_registered_link(tmp_path, attempt): register_links(tmp_path, "prime", 0) - register_links(tmp_path, "prime/eval", 2) + register_links(tmp_path, "other", 2) register_links(tmp_path, "wandb", 1) publish_link(tmp_path, "wandb", "run", "https://wandb.ai/run") - publish_link(tmp_path, "prime/eval", "gsm8k", "https://prime/ev-1") + publish_link(tmp_path, "other", "gsm8k", "https://prime/ev-1") # A page that failed to open publishes None, which resolves the wait without a link. - publish_link(tmp_path, "prime/eval", "aime", None) + publish_link(tmp_path, "other", "aime", None) start = time.monotonic() - assert wait_for_links(tmp_path, attempt, ["prime", "prime/eval", "wandb"], lambda: False, timeout=5.0) == { + assert wait_for_links(tmp_path, attempt, ["prime", "other", "wandb"], lambda: False, timeout=5.0) == { "prime": {}, - "prime/eval": {"gsm8k": "https://prime/ev-1"}, + "other": {"gsm8k": "https://prime/ev-1"}, "wandb": {"run": "https://wandb.ai/run"}, } assert time.monotonic() - start < 1.0 @@ -48,4 +49,4 @@ def test_prime_monitor_without_credentials_fails_informatively(tmp_path, monkeyp monkeypatch.delenv("RUN_ID", raising=False) monkeypatch.setattr(prime_module, "PrimeConfig", lambda: SimpleNamespace(api_key=None)) with pytest.raises(RuntimeError, match=r"prime login.*--no-monitors\.prime"): - asyncio.run(PrimeTrainMonitor(PrimeTrainMonitorConfig()).init(config=None, output_dir=tmp_path)) + asyncio.run(PrimeMonitor(PrimeMonitorConfig(), MonitorContext(tmp_path)).init()) From a9e9b462133b39844ec62050436cf486bb5c112b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sat, 24 Oct 2026 14:36:23 +0000 Subject: [PATCH 02/12] keep the prime monitor change localized Drop the monitor base refactor (MonitorContext, the subset removal, the fan-out helper, the config class renames) from this change: the files it touched go back to main. PrimeMonitor keeps main's monitor API and the layout of the old train monitor, with the standalone-eval methods folded in; the configs keep their class names, with one PrimeMonitorConfig in place of the train and eval pair. Co-Authored-By: Claude Opus 5.5 --- .../src/prime_rl/configs/eval.py | 13 +- .../src/prime_rl/configs/monitors.py | 17 +- .../src/prime_rl/configs/orchestrator.py | 4 +- .../src/prime_rl/configs/sft.py | 4 +- .../src/prime_rl/configs/trainer.py | 4 +- src/prime_rl/entrypoints/sft.py | 9 +- src/prime_rl/eval/runner.py | 5 +- src/prime_rl/monitors/__init__.py | 105 +++-- src/prime_rl/monitors/base.py | 55 +-- src/prime_rl/monitors/file/monitor.py | 16 +- src/prime_rl/monitors/prime.py | 385 ++++++++++-------- src/prime_rl/monitors/wandb/monitor.py | 26 +- src/prime_rl/orchestrator/orchestrator.py | 10 +- tests/unit/train/sft/test_prime_monitor.py | 361 ++++++++++++---- tests/unit/utils/test_links.py | 13 +- 15 files changed, 649 insertions(+), 378 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/eval.py b/packages/prime-rl-configs/src/prime_rl/configs/eval.py index 43e069a2612..3269993002f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/eval.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/eval.py @@ -4,7 +4,7 @@ import verifiers.v1 as vf from pydantic import AliasChoices, Field, model_validator -from prime_rl.configs.monitors import MonitorsConfig +from prime_rl.configs.monitors import EvalMonitorsConfig from prime_rl.configs.orchestrator import ConcurrencyConfig, EvalSourcesConfig, ScheduledEvalConfig from prime_rl.configs.shared import ( ClientConfig, @@ -108,7 +108,7 @@ class EvalConfig(ServedEvalConfig): log: LogConfig = LogConfig() - monitors: MonitorsConfig = MonitorsConfig() + monitors: EvalMonitorsConfig = EvalMonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" @property @@ -172,9 +172,12 @@ class SFTOnlineEvalConfig(ScheduledEvalConfig, ServedEvalConfig): log: LogConfig = LogConfig() - monitors: MonitorsConfig = MonitorsConfig() - """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``), the - trainer's: the online-eval process logs to the trainer's W&B and Prime runs.""" + monitors: EvalMonitorsConfig = EvalMonitorsConfig() + """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``). + + ``EvalMonitorsConfig``, not the base ``MonitorsConfig``: the online-eval + process logs to the trainer's Prime run, and the base type would drop (or + forbid) the prime block at the eval.json dump/re-parse boundary.""" @model_validator(mode="after") def auto_setup_broadcasts_dir(self): diff --git a/packages/prime-rl-configs/src/prime_rl/configs/monitors.py b/packages/prime-rl-configs/src/prime_rl/configs/monitors.py index 1358cfa3e09..cc3cde5f9f0 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/monitors.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/monitors.py @@ -48,9 +48,7 @@ class PrimeMonitorConfig(BaseConfig): """Run or evaluation name shown on Prime. Inherits ``run.name`` when unset.""" -class TrainerMonitorsConfig(BaseConfig): - """The RL trainer's monitors: its metrics reach Prime through the orchestrator's run.""" - +class MonitorsConfig(BaseConfig): wandb: WandbMonitorConfig | None = None """Log metrics to Weights & Biases. Off by default; enable with ``--monitors.wandb``.""" @@ -58,7 +56,14 @@ class TrainerMonitorsConfig(BaseConfig): """Log metrics and episode traces to the run's output directory. On by default; disable with ``--no-monitors.file``.""" -class MonitorsConfig(TrainerMonitorsConfig): +class TrainMonitorsConfig(MonitorsConfig): + prime: PrimeMonitorConfig | None = PrimeMonitorConfig() + """Log metrics and episodes to the Prime Intellect platform. On by default and + needs a Prime API key; disable with ``--no-monitors.prime``.""" + + +class EvalMonitorsConfig(MonitorsConfig): prime: PrimeMonitorConfig | None = PrimeMonitorConfig() - """Log metrics and episodes to Prime. On by default and needs a Prime API key; - disable with ``--no-monitors.prime``.""" + """Upload every eval source's epoch as an evaluation on the Prime Intellect platform (an + SFT online-eval process logs to the trainer's run instead). On by default and needs a + Prime API key; disable with ``--no-monitors.prime``.""" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index c56c86748e6..42513f247ee 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -12,7 +12,7 @@ AlgoConfig, GRPOAlgoConfig, ) -from prime_rl.configs.monitors import MonitorsConfig +from prime_rl.configs.monitors import TrainMonitorsConfig from prime_rl.configs.shared import ( BaseModelConfig, ControlConfig, @@ -495,7 +495,7 @@ class OrchestratorConfig(BaseConfig): env_vars: EnvVars = {} """Extra environment variables for the orchestrator process(es). Merged on top of the launcher defaults.""" - monitors: MonitorsConfig = MonitorsConfig() + monitors: TrainMonitorsConfig = TrainMonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" collect_inference_metrics: bool = True diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 5f62ce989b7..7e578e5f0e0 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -11,7 +11,7 @@ from prime_rl.configs.eval import SFTOnlineEvalConfig from prime_rl.configs.inference import InferenceConfig from prime_rl.configs.inference import WeightBroadcastConfig as InferenceWeightBroadcastConfig -from prime_rl.configs.monitors import MonitorsConfig +from prime_rl.configs.monitors import TrainMonitorsConfig from prime_rl.configs.shared import ( FileSystemWeightBroadcastConfig, NCCLWeightBroadcastConfig, @@ -236,7 +236,7 @@ class SFTConfig(BaseTrainerConfig): resume: ResumeConfig | None = None """Resume the run from a checkpoint (point at it with the previous run's ``run.name``). Without ``[ckpt]`` the run loads the checkpoint but saves no new ones. If None, does not resume.""" - monitors: MonitorsConfig = MonitorsConfig() + monitors: TrainMonitorsConfig = TrainMonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" run: RunConfig = Field(default_factory=RunConfig) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 544f65d1266..6bb965cb4db 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -5,7 +5,7 @@ from pydantic import Field, field_validator, model_validator -from prime_rl.configs.monitors import TrainerMonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig from prime_rl.configs.shared import ( BaseModelConfig, ControlConfig, @@ -705,7 +705,7 @@ class TrainerConfig(BaseTrainerConfig): rollout_transport: TransportConfig = ZMQTransportConfig() """Transport used to ship rollouts from orchestrator to trainer.""" - monitors: TrainerMonitorsConfig = TrainerMonitorsConfig() + monitors: MonitorsConfig = MonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``).""" output_dir: Path = Field(default_factory=default_output_dir) diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index dea87611034..7dc43b88a07 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -5,6 +5,7 @@ from pathlib import Path from prime_rl.configs.eval import SFTOnlineEvalConfig +from prime_rl.configs.monitors import EvalMonitorsConfig, TrainMonitorsConfig from prime_rl.configs.orchestrator import OnlineEvalSourceConfig from prime_rl.configs.sft import SFTConfig from prime_rl.configs.shared import LogConfig @@ -70,6 +71,12 @@ def resolve_resume_step(config: SFTConfig) -> int | None: return resolve_latest_ckpt_step(get_ckpt_dir(get_ckpt_base(config))) +def build_online_eval_monitors(monitors: TrainMonitorsConfig) -> EvalMonitorsConfig: + """The online-eval process's monitors: the trainer's, so it logs to the trainer's + W&B and Prime runs.""" + return EvalMonitorsConfig(wandb=monitors.wandb, file=monitors.file, prime=monitors.prime) + + def build_online_eval_config(config: SFTConfig) -> SFTOnlineEvalConfig: """The online-eval process's config: the resolved ``[eval]`` block with the run-level fields filled from the SFT config. The launcher spawns the env servers itself; the @@ -90,7 +97,7 @@ def build_online_eval_config(config: SFTConfig) -> SFTOnlineEvalConfig: # by directory separation. output_dir=config.run_dir, log=LogConfig(level=config.log.level, json_logging=config.log.json_logging), - monitors=config.monitors, + monitors=build_online_eval_monitors(config.monitors), ) return SFTOnlineEvalConfig(**{**eval_config.model_dump(exclude=set(run_fields)), **run_fields}) diff --git a/src/prime_rl/eval/runner.py b/src/prime_rl/eval/runner.py index e2d536cfb18..1102a4e4c37 100644 --- a/src/prime_rl/eval/runner.py +++ b/src/prime_rl/eval/runner.py @@ -253,7 +253,7 @@ async def land(self, episode: vf.Episode, pending: set[str]) -> None: """One episode of the epoch, arrived or restored: through the monitors and into its env's batch, which is finalized once the epoch is complete.""" step = eval_work(episode).step - await monitors.log([episode], step, "eval") + await monitors.log([episode], step, "eval", "all") eval_batch = self.eval_sink.add(episode) if eval_batch is not None: await self.finalize_eval_batch(eval_batch) @@ -261,12 +261,13 @@ async def land(self, episode: vf.Episode, pending: set[str]) -> None: async def finalize_eval_batch(self, batch: EvalBatch) -> None: """Persist + log one completed eval epoch through the monitors, mirroring the - orchestrator: the effective subset's annotations plus the ``eval/{env}/...`` metric dict.""" + orchestrator: effective episodes plus the ``eval/{env}/...`` metric dict.""" if not batch.episodes and not batch.failures and not batch.cancelled: get_logger().warning(f"Eval @ step={batch.step} env={batch.env_name}: no attempts returned, skipping log") return if batch.episodes.effective: + await monitors.log(batch.episodes.effective.vf_episodes, batch.step, "eval", "effective") await monitors.log_annotations(stamp_batch(batch.episodes.effective.vf_episodes, batch.step)) await monitors.log_eval_epoch(batch.env_name, batch.step, batch.episodes.vf_episodes) diff --git a/src/prime_rl/monitors/__init__.py b/src/prime_rl/monitors/__init__.py index 83c90092d9d..4f0b1dbcb1c 100644 --- a/src/prime_rl/monitors/__init__.py +++ b/src/prime_rl/monitors/__init__.py @@ -13,8 +13,12 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, overload -from prime_rl.configs.monitors import FileMonitorConfig, PrimeMonitorConfig, WandbMonitorConfig -from prime_rl.monitors.base import Kind, Monitor, MonitorContext +from prime_rl.configs.monitors import ( + FileMonitorConfig, + PrimeMonitorConfig, + WandbMonitorConfig, +) +from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.monitors.file import FileMonitor from prime_rl.monitors.prime import PrimeMonitor from prime_rl.monitors.wandb import WandbMonitor @@ -26,7 +30,6 @@ __all__ = [ "Monitor", - "MonitorContext", "WandbMonitor", "PrimeMonitor", "FileMonitor", @@ -63,24 +66,35 @@ async def setup( if rank != 0: return - context = MonitorContext( - output_dir=output_dir, - producer=producer, - run_config=run_config, - train_env_names=train_env_names or [], - eval_env_names=eval_env_names or [], - overview_flavor=overview_flavor, - ) - configured = [(PrimeMonitor, prime), (WandbMonitor, wandb), (FileMonitor, file)] - monitors = [cls(config, context) for cls, config in configured if config is not None] + monitors: list[tuple[str, Monitor, dict[str, Any]]] = [] + if prime is not None: + monitors.append( + ("prime", PrimeMonitor(prime), dict(config=run_config, output_dir=output_dir, producer=producer)) + ) + if wandb is not None: + monitors.append( + ( + "wandb", + WandbMonitor(wandb), + dict( + output_dir=output_dir, + config=run_config, + train_env_names=train_env_names, + eval_env_names=eval_env_names, + overview_flavor=overview_flavor, + ), + ) + ) + if file is not None: + monitors.append(("file", FileMonitor(file), dict(output_dir=output_dir, producer=producer))) if not monitors: return - get_logger().info(f"Initializing monitors ({', '.join(type(monitor).__name__ for monitor in monitors)})") + get_logger().info(f"Initializing monitors ({', '.join(name for name, _, _ in monitors)})") t0 = time.perf_counter() - for monitor in monitors: - get_logger().info(f"Initializing {type(monitor).__name__} ({monitor.config})") - await monitor.init() + for name, monitor, init_kwargs in monitors: + get_logger().info(f"Initializing {name} monitor ({monitor.config})") + await monitor.init(**init_kwargs) MONITORS.append(monitor) get_logger().debug(f"Initialized monitors in {format_time(time.perf_counter() - t0)}") @@ -91,48 +105,69 @@ def get(monitor_cls: type[Monitor]) -> Monitor | None: return next((monitor for monitor in MONITORS if isinstance(monitor, monitor_cls)), None) -async def _fan_out(hook: str, *args: Any, **kwargs: Any) -> None: - for monitor in MONITORS: - try: - await getattr(monitor, hook)(*args, **kwargs) - except Exception as e: - get_logger().warning(f"Failed to {hook} on {type(monitor).__name__}: {e}") - - @overload async def log(data: dict[str, Any], step: int | None) -> None: ... @overload -async def log(data: vf.Episode | list[vf.Episode], step: int, kind: Kind) -> None: ... +async def log(data: vf.Episode | list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: ... -async def log(data: dict[str, Any] | vf.Episode | list[vf.Episode], step: int | None, kind: Kind = "train") -> None: - """Log to all registered monitors: a dict of scalar metrics, or episodes as they - arrive.""" - await _fan_out("log", data, step=step, kind=kind) +async def log( + data: dict[str, Any] | vf.Episode | list[vf.Episode], + step: int | None, + kind: Kind = "train", + subset: Subset = "effective", +) -> None: + """Log to all registered monitors: a dict of scalar metrics, or episodes + with their cohort coordinates (train/eval x all/effective).""" + for monitor in MONITORS: + try: + await monitor.log(data, step=step, kind=kind, subset=subset) + except Exception as e: + get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") async def log_annotations(updates: list[dict[str, Any]]) -> None: """Log trace updates to all registered monitors.""" - await _fan_out("log_annotations", updates) + for monitor in MONITORS: + try: + await monitor.log_annotations(updates) + except Exception as e: + get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") async def log_eval_plan(env_name: str, step: int, expected: int) -> None: """Log an eval epoch's expected episode count to all registered monitors.""" - await _fan_out("log_eval_plan", env_name, step, expected) + for monitor in MONITORS: + try: + await monitor.log_eval_plan(env_name, step, expected) + except Exception as e: + get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") async def log_live(events: list[dict[str, Any]]) -> None: """Log live-trace events to all registered monitors.""" - await _fan_out("log_live", events) + for monitor in MONITORS: + try: + await monitor.log_live(events) + except Exception as e: + get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") async def log_eval_epoch(env_name: str, step: int, episodes: list[vf.Episode]) -> None: """Log one finished eval epoch to all registered monitors.""" - await _fan_out("log_eval_epoch", env_name, step, episodes) + for monitor in MONITORS: + try: + await monitor.log_eval_epoch(env_name, step, episodes) + except Exception as e: + get_logger().warning(f"Failed to log to {monitor.__class__.__name__}: {e}") async def finalize(paused: bool = False) -> None: """Finalize the run on all registered monitors. A paused run is not reported as completed.""" - await _fan_out("finalize", paused=paused) + for monitor in MONITORS: + try: + await monitor.finalize(paused=paused) + except Exception as e: + get_logger().warning(f"Failed to finalize {monitor.__class__.__name__}: {e}") diff --git a/src/prime_rl/monitors/base.py b/src/prime_rl/monitors/base.py index f9d88e9a30a..260cc0e56dd 100644 --- a/src/prime_rl/monitors/base.py +++ b/src/prime_rl/monitors/base.py @@ -1,8 +1,6 @@ from __future__ import annotations from abc import ABC, abstractmethod -from dataclasses import dataclass, field -from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, overload from prime_rl.utils.config import BaseConfig @@ -12,81 +10,66 @@ import verifiers.v1 as vf Kind = Literal["train", "eval"] - - -@dataclass(frozen=True) -class MonitorContext: - """What every monitor of a process knows about the run it logs.""" - - output_dir: Path - """The run directory.""" - producer: str | None = None - """The process logging: ``orchestrator``, ``trainer``, ``online-eval`` or ``eval``.""" - run_config: BaseConfig | None = None - """The process's resolved config.""" - train_env_names: list[str] = field(default_factory=list) - eval_env_names: list[str] = field(default_factory=list) - overview_flavor: Literal["rl", "sft", "eval"] = "rl" +Subset = Literal["all", "effective"] class Monitor(ABC): - """Base class for monitors. Every hook but ``log_metrics`` defaults to a no-op, so a - monitor implements only what it carries.""" + """Base class for monitors.""" - def __init__(self, config: BaseConfig, context: MonitorContext): + def __init__(self, config: BaseConfig): self.config = config - self.context = context self.logger = get_logger() - async def init(self) -> None: - """Open the monitor's run. Network or file I/O, before the first log.""" + async def init(self, **kwargs: Any) -> None: + """Initialize run. Overrides name their own kwargs.""" @overload async def log(self, data: dict[str, Any], step: int | None) -> None: ... @overload - async def log(self, data: vf.Episode | list[vf.Episode], step: int, kind: Kind) -> None: ... + async def log(self, data: vf.Episode | list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: ... async def log( self, data: dict[str, Any] | vf.Episode | list[vf.Episode], step: int | None, kind: Kind = "train", + subset: Subset = "effective", ) -> None: """Log scalar metrics, or episodes.""" if isinstance(data, dict): await self.log_metrics(data, step=step) - return - episodes = data if isinstance(data, list) else [data] - if episodes: + else: + episodes = data if isinstance(data, list) else [data] assert step is not None - await self.log_episodes(episodes, step=step, kind=kind) + await self.log_episodes(episodes, step=step, kind=kind, subset=subset) @abstractmethod async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: """Log scalar metrics. ``step=None`` logs a time-keyed row (e.g. inference metrics, which are sampled on wall time rather than the training step).""" - async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind) -> None: - """Log episodes as they arrive: every episode exactly once, errored and filtered - ones included. What the run learns about them later (batch membership, - advantages) arrives through ``log_annotations``.""" + @abstractmethod + async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: + """Log episodes.""" async def log_annotations(self, updates: list[dict[str, Any]]) -> None: - """Log trace updates — post-hoc facts about traces this run already logged.""" + """Log trace updates — post-hoc facts about traces this run already logged. + Monitors that only carry scalars ignore them.""" async def log_live(self, events: list[dict[str, Any]]) -> None: """Log the env servers' stream of in-flight traces: ``{"delta", "dispatch"}`` for - each delta as it arrived, ``{"done": trace_id}`` when a trace's episode finished.""" + each delta as it arrived, ``{"done": trace_id}`` when a trace's episode finished. + Monitors that only carry finished work ignore it.""" async def log_eval_plan(self, env_name: str, step: int, expected: int) -> None: """Log how many episodes the eval epoch of ``env_name`` at ``step`` will produce, - known once its tasks are counted.""" + known once its tasks are counted. Monitors that only carry results ignore it.""" async def log_eval_epoch(self, env_name: str, step: int, episodes: list[vf.Episode]) -> None: """Log one finished eval epoch: every episode ``env_name`` produced for ``step``, errored ones included. Fires once per epoch, after the episodes streamed through - ``log``.""" + ``log``. Monitors that carry episodes as they arrive ignore it.""" async def finalize(self, paused: bool = False) -> None: """Finalize run. A paused run stopped to resume later and did not complete.""" diff --git a/src/prime_rl/monitors/file/monitor.py b/src/prime_rl/monitors/file/monitor.py index 1631a34e9e7..03db7e219a7 100644 --- a/src/prime_rl/monitors/file/monitor.py +++ b/src/prime_rl/monitors/file/monitor.py @@ -12,7 +12,7 @@ from verifiers.v1.types import decode_ndarray from prime_rl.configs.monitors import FileMonitorConfig -from prime_rl.monitors.base import Kind, Monitor +from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.monitors.file.traces import get_annotations_dir, get_index_path, get_trace_stream from prime_rl.monitors.file.traces.chunks import ChunkedJsonl from prime_rl.monitors.file.traces.index import index_row @@ -47,9 +47,9 @@ class FileMonitor(Monitor): config: FileMonitorConfig file: TextIO - async def init(self) -> None: - output_dir = self.output_dir = self.context.output_dir - self.producer = self.context.producer + async def init(self, output_dir: Path, producer: str | None = None) -> None: + self.output_dir = output_dir + self.producer = producer index = get_index_path(get_trace_stream(output_dir)) # a relaunch appends to the stream it finds, so the numbering carries on self._logged = sum(1 for _ in index.open("rb")) if index.is_file() else 0 @@ -149,12 +149,14 @@ def _stream(self, directory: Path) -> tuple[ChunkedJsonl, BinaryIO]: self._streams[directory] = (stream, open(get_index_path(directory), "ab")) return self._streams[directory] - async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind) -> None: + async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: """Append each episode to the trace stream as it completes — every episode is serialized exactly once, in arrival order, whatever its kind, so an in-progress run can be tailed. Episode-level failures are preserved even when no trace was - produced. What the shipped cohort learns arrives as annotations that readers fold - back onto the stream.""" + produced. The shipped cohort writes no second copy: what it learns arrives as + annotations that readers fold back onto the stream.""" + if subset == "effective": + return def write() -> None: stream, index = self._stream(get_trace_stream(self.output_dir)) diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 62f3538b946..1f3c91e6c0d 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -13,7 +13,7 @@ from prime_rl.configs.eval import EvalConfig, SFTOnlineEvalConfig from prime_rl.configs.monitors import PrimeMonitorConfig from prime_rl.configs.sft import SFTConfig -from prime_rl.monitors.base import Kind, Monitor +from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.utils.config import BaseConfig from prime_rl.utils.links import ATTEMPT_VAR, PRIME_LINKS, publish_link, register_links from prime_rl.utils.pathing import get_platform_run_path @@ -30,15 +30,12 @@ # default (300 s) is sized for eval sample batches; a crashed training process should # not linger that long, and a clean finish rarely has more than the last step queued. FINISH_TIMEOUT = 60.0 -# How long an SFT online-eval process waits for the trainer to register the run: the -# trainer starts with it, but registers only after its imports and config setup. +# How long an SFT online-eval process waits for the trainer to register the run. TRAIN_RUN_TIMEOUT_S = 1200.0 -TRAIN_RUN_POLL_S = 1.0 -OFF_HINT = "turn the monitor off with --no-monitors.prime" def write_platform_record(output_dir: Path, record: dict[str, Any]) -> None: - """Leave the run's Prime identity in the run directory for the dashboard's + """Leave the run's platform identity in the run directory for the dashboard's "View on Prime" link.""" write_record(get_platform_run_path(output_dir), record) @@ -47,16 +44,27 @@ def read_platform_record(output_dir: Path) -> dict[str, Any] | None: return read_record(get_platform_run_path(output_dir)) -def _merge_platform_record(output_dir: Path, merge: Callable[[dict[str, Any]], dict[str, Any]]) -> None: - """Locked read-modify-write of the platform record. Best-effort, see ``update_record``.""" +def _merge_platform_record( + output_dir: Path, + merge: Callable[[dict[str, Any]], dict[str, Any]], +) -> None: + """Locked read-modify-write of the platform record: the trainer and the + online-eval process merge into the SAME run.json. Best-effort, see + ``update_record``.""" update_record(get_platform_run_path(output_dir), merge) def update_platform_record(output_dir: Path, update: dict[str, Any]) -> None: - """Overlay ``update`` on the platform record, keeping the keys it doesn't set.""" + """Best-effort locked read-modify-write of the platform record: `update` + overlays the current record (preserving keys it doesn't set). Use for + every writer that shares the file — a naive overwrite drops concurrent + writers' keys.""" _merge_platform_record(output_dir, lambda record: {**record, **update}) +OFF_HINT = "turn the monitor off with --no-monitors.prime" + + def prime_mode() -> str: """$PRIME_RUNS_MODE when set, else online. Like the W&B monitor, a configured monitor without credentials fails the run instead of leaving it untracked.""" @@ -76,42 +84,6 @@ def _base_url() -> str | None: return base.rstrip("/").removesuffix("/rft") if base else None -def _registration(name: str | None, config: BaseConfig | None) -> dict[str, Any]: - """``pr.init`` arguments that register a new training run for ``config``.""" - if isinstance(config, SFTConfig): - # SFT trains on a dataset, not on environments, and has no rollouts. - return dict( - name=name, - model=config.model.name, - environments=[], - training=pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.data.batch_size, - seq_len=config.data.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, - ), - config=config.model_dump(exclude_none=True, mode="json"), - ) - if config is not None: - return dict( - name=name, - model=config.model.name, - environments=[env.env_id for env in config.train.source], - training=pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.batch_size, - rollouts_per_example=max( - (source.group_size for source in config.train.source), default=config.train.group_size - ), - seq_len=config.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, - ), - config=config.model_dump(exclude_none=True, mode="json"), - ) - # The RFT API requires a base model. - return {"name": name, "model": "unknown"} - - def wait_for_train_run(output_dir: Path, timeout: float = TRAIN_RUN_TIMEOUT_S) -> str: """The id of the training run this attempt's trainer registered, from the platform record it leaves in the shared run dir.""" @@ -126,104 +98,131 @@ def wait_for_train_run(output_dir: Path, timeout: float = TRAIN_RUN_TIMEOUT_S) - return record["id"] if time.monotonic() >= deadline: raise RuntimeError(f"The trainer registered no Prime run within {timeout:.0f}s; {OFF_HINT}") - time.sleep(TRAIN_RUN_POLL_S) + time.sleep(1.0) class PrimeMonitor(Monitor): - """Logs a run to Prime through ``prime_runs``. - - A training run's processes (the rl orchestrator, the sft trainer and its online-eval - process) log to one Prime training run: metrics, and every episode as it arrives. - The process with its own config registers the run, or attaches to the one a - hosted launch created ($RUN_ID); an SFT online-eval process attaches to its - trainer's. ``$PRL_RUN_FINISHER`` names the process that finishes a shared run. - - A standalone eval opens one Prime evaluation per env and epoch when the epoch - starts (``log_eval_plan``), streams its episodes into it, and finishes it with the - epoch's aggregates (``log_eval_epoch``). - - Delivery is best-effort: the SDK's uploader thread owns retries and backpressure, - and its atexit hook reports a process that exits without finishing. An abrupt kill - skips the hook, and the platform owns the terminal state of attached runs. + """Logs metrics and episodes to the Prime platform through ``prime_runs``. + + A training run's processes log to one Prime training run: the rl orchestrator or + sft trainer registers it, and an sft online-eval process attaches to its trainer's. + ``$PRL_RUN_FINISHER`` names the process that finishes a shared run; the others + only flush. A standalone eval instead opens one evaluation per env and epoch when + the epoch starts (``log_eval_plan``), streams its episodes into it, and finishes + it with the epoch's aggregates (``log_eval_epoch``). + + The run handle owns the RFT lifecycle (register or attach, finalize), the + per-step metrics POSTs, the episode uploads to Prime Traces, and the + terminal status. Delivery is BEST-EFFORT: the SDK's atexit hook + reports a process that exits without finalizing, but an abrupt kill + (SIGKILL, OOM, node loss) skips the hook and a lost finalize is not + retried here — the platform owns the terminal state of attached + ($RUN_ID) and managed runs. `FINISH_TIMEOUT` bounds the queued-upload + drain, not the SDK's finalize HTTP retries. + + ``init``/``finish`` do network I/O and run in worker threads; the log + calls are queue puts onto the SDK's uploader thread, which owns retries + and backpressure, so they never stall the loop. """ config: PrimeMonitorConfig - async def init(self) -> None: - config = self.context.run_config - self.mode = prime_mode() + async def init(self, config: BaseConfig | None, output_dir: Path, producer: str | None = None) -> None: + self.output_dir = output_dir + self.producer = producer + self.mode = mode = prime_mode() self.run: pr.Run | None = None # (env, step) -> its open evaluation; None once opening it failed, so the epoch's - # episodes do not retry Prime on every arrival - self.evaluations: dict[tuple[str, int], pr.Run | None] = {} + # episodes do not retry the platform on every arrival + self.runs: dict[tuple[str, int], pr.Run | None] = {} if isinstance(config, EvalConfig): await self._init_evaluations(config) return online_eval = isinstance(config, SFTOnlineEvalConfig) - self.run = await asyncio.to_thread(self._open_run, config, online_eval) - if not self.run.url: - self.logger.info("Prime run disabled") + init_kwargs: dict[str, Any] + if run_id := os.getenv("RUN_ID"): + # A managed launch pre-created the platform run and injected its id - + # attach instead of registering a duplicate. The platform owns the + # attached run's failure/completion marking; a clean finish() is + # delivered best-effort and does not by itself mark it completed. + init_kwargs = {"id": run_id} + elif online_eval: + # The trainer registered the run; its online-eval process logs to it. + init_kwargs = {"id": await asyncio.to_thread(wait_for_train_run, output_dir)} if mode == "online" else {} + elif isinstance(config, SFTConfig): + # SFT trains on a dataset, not on environments, and has no rollouts - the + # platform run is registered without both. + init_kwargs = dict( + name=self.config.name, + model=config.model.name, + environments=[], + training=pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.data.batch_size, + seq_len=config.data.seq_len, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + ), + config=config.model_dump(exclude_none=True, mode="json"), + ) + elif config is not None: + init_kwargs = dict( + name=self.config.name, + model=config.model.name, + environments=[env.env_id for env in config.train.source], + training=pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.batch_size, + rollouts_per_example=max( + (source.group_size for source in config.train.source), default=config.train.group_size + ), + seq_len=config.seq_len, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + ), + config=config.model_dump(exclude_none=True, mode="json"), + ) else: - attached = " (attached)" if self.run.attached else "" - self.logger.info(f"Logging metrics and episodes to Prime run {self.run.id} ({self.run.url}){attached}") - if not online_eval: - # Off the event loop: the flock retry must not stall it. - await asyncio.to_thread(self._record_run) + # The RFT API requires a base model; "unknown" is what the + # pre-SDK register sent for a config-less init. + init_kwargs = {"name": self.config.name, "model": "unknown"} - def _open_run(self, config: BaseConfig | None, online_eval: bool) -> pr.Run: - """Register the training run, or attach to it. Blocking: runs in a worker thread.""" - run_id = os.getenv("RUN_ID") - if online_eval and not run_id and self.mode == "online": - run_id = wait_for_train_run(self.context.output_dir) - kwargs = {"id": run_id} if run_id else _registration(self.config.name, config) try: - return pr.init(kind="train", mode=self.mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT, **kwargs) + self.run = await asyncio.to_thread( + pr.init, kind="train", mode=mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT, **init_kwargs + ) except (pr.ConfigurationError, pr.APIError) as e: # Registering fails without a team, or for a team outside the external-runs # allowlist. An attached run's error is the launcher's to explain. - if run_id: + if "id" in init_kwargs: raise raise RuntimeError( f"Failed to register the Prime run ({type(e).__name__}: {e}). Select a team with " f"`prime switch` (or set PRIME_TEAM_ID), or {OFF_HINT}" ) from e + if self.run.url: + attached = " (attached via $RUN_ID)" if self.run.attached else "" + self.logger.info(f"Logging metrics and episodes to Prime run {self.run.id} ({self.run.url}){attached}") + else: + self.logger.info("Prime run disabled") + if not online_eval: + # Off the event loop: the flock retry must not stall it. + await asyncio.to_thread(self._record_run) def _record_run(self) -> None: - output_dir = self.context.output_dir - register_links(output_dir, PRIME_LINKS, 1 if self.run.url else 0) + register_links(self.output_dir, PRIME_LINKS, 1 if self.run.url else 0) if self.run.url: + # Stamped with the attempt: the online-eval process attaches to this run. update_platform_record( - output_dir, + self.output_dir, {"kind": "train", "id": self.run.id, "url": self.run.url, "attempt": os.getenv(ATTEMPT_VAR)}, ) - publish_link(output_dir, PRIME_LINKS, "run", self.run.url) - - async def _init_evaluations(self, config: EvalConfig) -> None: - self.model = config.model - self.sources = {source.resolved_name: source for source in config.source} - self.evaluation_id = os.getenv(EVAL_ID_VAR) - if self.evaluation_id and len(self.sources) != 1: - raise ValueError( - f"${EVAL_ID_VAR} names one Prime evaluation, so the run needs " - f"exactly one eval source (got {len(self.sources)})" - ) - online = self.mode == "online" - self.logger.info("Streaming eval epochs to Prime" if online else "Prime evaluations disabled") - - def record() -> None: - if online: - write_platform_record(self.context.output_dir, {"kind": "eval", "evaluations": {}}) - register_links(self.context.output_dir, PRIME_LINKS, len(self.sources) if online else 0) - - await asyncio.to_thread(record) + publish_link(self.output_dir, PRIME_LINKS, "run", self.run.url) async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: - """Onto the training run. A standalone eval's evaluations carry their own - aggregates through ``log_eval_epoch``.""" + # A standalone eval's evaluations carry their own aggregates (``log_eval_epoch``). if self.run is None: return - # The SDK also drops non-finite values, but silently; sanitize first so the - # dropped paths are named in the log. + # The SDK also drops non-finite values, but silently; sanitize first so + # the dropped paths are named in the log. metrics, dropped = sanitize(metrics) if dropped: self.logger.warning(f"Dropping {len(dropped)} non-finite metric value(s): {', '.join(dropped[:5])}") @@ -231,9 +230,13 @@ async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: # stamps `_timestamp` on every row, so step=None rows keep a time anchor. await asyncio.to_thread(self.run.log_metrics, metrics, step=step) - async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind) -> None: - """Every episode ships whole as it arrives, keyed by the SDK to its work type and - dispatch step off ``run.work``.""" + async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: + """Every episode ships whole as it arrives (the ``all`` subset), errored and + filtered ones included, online evals too. The SDK keys each to its dispatch + step and work type off ``run.work`` - the ``TrainRunInfo`` the dispatcher + stamps at emit time.""" + if subset != "all" or not episodes: + return if self.run is not None: # A queue put that can block briefly under backpressure - off the loop. await asyncio.to_thread(self.run.log_episodes, episodes) @@ -243,97 +246,135 @@ async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind) if episode.env.name is not None: by_env.setdefault(episode.env.name, []).append(episode) for env_name, batch in by_env.items(): - evaluation = await self._evaluation(env_name, step) - if evaluation is not None: - await asyncio.to_thread(evaluation.log_episodes, batch) + run = await self.run_for(env_name, step) + if run is not None: + await asyncio.to_thread(run.log_episodes, batch) # a queue put, off the loop + + async def finalize(self, paused: bool = False) -> None: + # An epoch the run did not finish leaves its evaluation open: close it as cancelled. + for (env_name, step), run in self.runs.items(): + if run is None or run.finished: + continue + try: + await asyncio.to_thread(run.finish, status=pr.RunStatus.CANCELLED, error="interrupted") + except Exception as e: + self.logger.warning(f"Failed to close the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") + if self.run is None: + return + if paused or os.getenv(RUN_FINISHER_VAR, self.producer) != self.producer: + # Drain only: a paused run continues on resume, and a shared run is finished + # by the process that outlives this one. A clean exit is not a crash. + await asyncio.to_thread(self.run.flush, FINISH_TIMEOUT) + atexit.unregister(self.run._atexit_hook) + return + # Drains queued uploads so the final step's metrics and episodes land, then + # finalizes (idempotent on the platform side); an attached run's failure + # marking stays with the launcher. + await asyncio.to_thread(self.run.finish) + + async def _init_evaluations(self, config: EvalConfig) -> None: + self.model: str = config.model + self.sources = {source.resolved_name: source for source in config.source} + self.run_id = os.getenv("PRL_RUN_ID") + self._lock = asyncio.Lock() + self.evaluation_id = os.getenv(EVAL_ID_VAR) + if self.evaluation_id and len(self.sources) != 1: + raise ValueError( + f"${EVAL_ID_VAR} names one platform evaluation, so the run needs " + f"exactly one eval source (got {len(self.sources)})" + ) + online = self.mode == "online" + self.logger.info("Streaming eval epochs to the Prime platform" if online else "Prime evaluations disabled") + + def record() -> None: + if online: + write_platform_record(self.output_dir, {"kind": "eval", "run_id": self.run_id, "evaluations": {}}) + register_links(self.output_dir, PRIME_LINKS, len(self.sources) if online else 0) + + # Off the event loop: the flock retry must not stall it. + await asyncio.to_thread(record) - def _open_evaluation(self, env_name: str, step: int, expected: int | None) -> pr.Run: - """Open the Prime evaluation of one epoch. Blocking: runs in a worker thread.""" - common = dict(kind="eval", mode=self.mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT) + def open(self, env_name: str, step: int, expected: int | None) -> pr.Run: + """Open the platform evaluation of one epoch. Blocking: runs in a worker thread.""" if self.evaluation_id: - # A hosted launch pre-created the evaluation and injected its id. - if self.evaluations: + # A hosted launch pre-created the platform evaluation and injected its id - + # attach instead of registering a duplicate. The backend owns its failure + # marking then; a clean finish still completes it. + if self.runs: raise RuntimeError(f"${EVAL_ID_VAR} holds one epoch, and it already took one") - return pr.init(id=self.evaluation_id, **common) + return pr.init( + kind="eval", + mode=self.mode, + base_url=_base_url(), + id=self.evaluation_id, + finish_timeout=FINISH_TIMEOUT, + ) source = self.sources[env_name] + name = self.config.name if len(self.sources) == 1 else f"{self.config.name}--{env_name}" return pr.init( - name=self.config.name if len(self.sources) == 1 else f"{self.config.name}--{env_name}", + kind="eval", + mode=self.mode, + base_url=_base_url(), + name=name, environments=[source.env.taskset.id], model=self.model, framework="prime-rl", + finish_timeout=FINISH_TIMEOUT, config={ "model": self.model, "step": step, - "run_id": os.getenv("PRL_RUN_ID"), + "run_id": self.run_id, "num_examples": expected // source.group_size if expected else None, "rollouts_per_example": source.group_size, }, - **common, ) - async def _evaluation(self, env_name: str, step: int, expected: int | None = None) -> pr.Run | None: + async def run_for(self, env_name: str, step: int, expected: int | None = None) -> pr.Run | None: """The epoch's evaluation, opened on first use.""" key = (env_name, step) - if key in self.evaluations: - return self.evaluations[key] - try: - evaluation = await asyncio.to_thread(self._open_evaluation, env_name, step, expected) - except Exception as e: - self.logger.warning(f"Failed to open the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") - evaluation = None - if key in self.evaluations: # another arrival opened it meanwhile - return self.evaluations[key] - self.evaluations[key] = evaluation - if evaluation is not None and evaluation.url: - attached = f" (attached via ${EVAL_ID_VAR})" if evaluation.attached else "" - self.logger.info(f"Streaming {env_name} (Step {step}) evaluation - {evaluation.url}{attached}") - await asyncio.to_thread(self._record_evaluation, env_name, step, evaluation) - return evaluation - - def _record_evaluation(self, env_name: str, step: int, evaluation: pr.Run | None) -> None: - url = evaluation.url if evaluation is not None else None - if url: + async with self._lock: + if key in self.runs: + return self.runs[key] + try: + run = await asyncio.to_thread(self.open, env_name, step, expected) + except Exception as e: + self.logger.warning(f"Failed to open the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") + self.runs[key] = None + await asyncio.to_thread(publish_link, self.output_dir, PRIME_LINKS, env_name, None) + return None + self.runs[key] = run + if run.url: + attached = f" (attached via ${EVAL_ID_VAR})" if run.attached else "" + self.logger.info(f"Streaming {env_name} (Step {step}) evaluation - {run.url}{attached}") + # Off the event loop: the flock retry must not stall it. + await asyncio.to_thread(self._record_evaluation, env_name, step, run) + return run + + def _record_evaluation(self, env_name: str, step: int, run: pr.Run) -> None: + if run.url: def update(record: dict[str, Any]) -> dict[str, Any]: - record.setdefault("evaluations", {})[env_name] = {"step": step, "id": evaluation.id, "url": url} + record.setdefault("kind", "eval") + record.setdefault("run_id", self.run_id) + record.setdefault("evaluations", {})[env_name] = {"step": step, "id": run.id, "url": run.url} return record - _merge_platform_record(self.context.output_dir, update) - publish_link(self.context.output_dir, PRIME_LINKS, env_name, url) + _merge_platform_record(self.output_dir, update) + publish_link(self.output_dir, PRIME_LINKS, env_name, run.url) async def log_eval_plan(self, env_name: str, step: int, expected: int) -> None: if self.run is None: - await self._evaluation(env_name, step, expected) + await self.run_for(env_name, step, expected) async def log_eval_epoch(self, env_name: str, step: int, episodes: list[vf.Episode]) -> None: if self.run is not None: return - evaluation = await self._evaluation(env_name, step) - if evaluation is None: + run = await self.run_for(env_name, step) + if run is None: return try: - await asyncio.to_thread(evaluation.finish, pr.metrics.from_episodes(episodes)) + await asyncio.to_thread(run.finish, pr.metrics.from_episodes(episodes)) except Exception as e: self.logger.warning(f"Failed to finish the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") return - self.logger.info(f"Uploaded {env_name} (Step {step}) evaluation - {evaluation.url}") - - async def finalize(self, paused: bool = False) -> None: - # An epoch the run did not finish leaves its evaluation open: close it as cancelled. - for (env_name, step), evaluation in self.evaluations.items(): - if evaluation is None or evaluation.finished: - continue - try: - await asyncio.to_thread(evaluation.finish, status=pr.RunStatus.CANCELLED, error="interrupted") - except Exception as e: - self.logger.warning(f"Failed to close the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") - if self.run is None: - return - producer = self.context.producer - if paused or os.getenv(RUN_FINISHER_VAR, producer) != producer: - # Drain only: a paused run continues on resume, and a shared run is finished - # by the process that outlives this one. A clean exit is not a crash. - await asyncio.to_thread(self.run.flush, FINISH_TIMEOUT) - atexit.unregister(self.run._atexit_hook) - return - await asyncio.to_thread(self.run.finish) + self.logger.info(f"Uploaded {env_name} (Step {step}) evaluation - {run.url}") diff --git a/src/prime_rl/monitors/wandb/monitor.py b/src/prime_rl/monitors/wandb/monitor.py index 11eb32be611..c85bf8c7baf 100644 --- a/src/prime_rl/monitors/wandb/monitor.py +++ b/src/prime_rl/monitors/wandb/monitor.py @@ -6,20 +6,21 @@ import sys import time from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import wandb from wandb.errors import CommError, UsageError from wandb.sdk.mailbox.mailbox_handle import ServerResponseError from prime_rl.configs.monitors import WandbMonitorConfig -from prime_rl.monitors.base import Monitor +from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.monitors.wandb.overview import ensure_overview_view +from prime_rl.utils.config import BaseConfig from prime_rl.utils.links import WANDB_LINKS, publish_link, register_links from prime_rl.utils.logger import format_time if TYPE_CHECKING: - pass + import verifiers.v1 as vf class WandbMonitor(Monitor): @@ -27,8 +28,14 @@ class WandbMonitor(Monitor): config: WandbMonitorConfig - async def init(self) -> None: - output_dir, config = self.context.output_dir, self.context.run_config + async def init( + self, + output_dir: Path, + config: BaseConfig | None = None, + train_env_names: list[str] | None = None, + eval_env_names: list[str] | None = None, + overview_flavor: Literal["rl", "sft", "eval"] = "rl", + ) -> None: # W&B reads the start command off sys.argv; the launcher passes the original # command to subprocesses via $WANDB_ARGS. wandb_args = os.environ.get("WANDB_ARGS") @@ -128,9 +135,9 @@ def init_wandb(max_retries: int): url = ensure_overview_view( self.wandb.entity, self.wandb.project, - train_envs=self.context.train_env_names, - eval_envs=self.context.eval_env_names, - flavor=self.context.overview_flavor, + train_envs=train_env_names or [], + eval_envs=eval_env_names or [], + flavor=overview_flavor, ) if url: self.logger.info(f"Created W&B overview view - {url}") @@ -159,6 +166,9 @@ async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: else: wandb.log({**metrics, "step": step, "_timestamp": time.time()}) + async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: + pass + async def finalize(self, paused: bool = False) -> None: self.logger.info(f"Finalizing W&B run {self.wandb.id}") t0 = time.perf_counter() diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 33698566c44..01fcfe4d36a 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -558,7 +558,7 @@ async def main_loop(self) -> None: kind = episode.run.work.type step = episode.run.work.step if kind == "eval" else self.progress.step stamp_arrival([episode], kind, step) - await monitors.log([episode], step, kind) + await monitors.log([episode], step, kind, "all") if kind == "eval": assert self.eval_sink is not None # eval rollouts only emitted when eval is configured @@ -679,6 +679,8 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: # records against each trace's arrival record - membership, advantages, the step # it shipped at - never a second episode copy. effective_episodes = await asyncio.to_thread(lambda: effective.vf_episodes) + # Monitors may retain the list for asynchronous upload after publication. + await monitors.log(effective_episodes.copy(), step, "train", "effective") annotations = await asyncio.to_thread(stamp_batch, effective_episodes, step) await monitors.log_annotations(annotations) @@ -909,9 +911,11 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: get_logger().warning(f"Eval @ step={batch.step} env={batch.env_name}: no attempts returned, skipping log") return - # Every episode was logged on arrival; the non-errored subset is annotated on epoch - # completion. + # The non-errored subset is logged on epoch completion (multiple eval envs share the + # step's trace file — each epoch appends its cohort once, and every record carries + # ``env_name``); the full returned cohort already streamed into ``all`` on arrival. if batch.episodes.effective: + await monitors.log(batch.episodes.effective.vf_episodes, batch.step, "eval", "effective") await monitors.log_annotations(stamp_batch(batch.episodes.effective.vf_episodes, batch.step)) policy_spans = [eval_work(episode).policy for episode in batch.episodes] if any(span is None for span in policy_spans): diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index 113a4b913a9..8d549dcc751 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -1,64 +1,56 @@ -"""The prime monitor's run wiring: SFT registration, online evals sharing the trainer's -run, and the platform record.""" +"""The SFT trainer's platform wiring: registration shape and monitor fan-out. + +``monitors.prime`` on ``SFTConfig`` must register the run with the dataset-batched +training fields (no environments, no rollouts) and flow the trainer's metrics and +finalize through the registered ``PrimeMonitor`` like the RL orchestrator does. +""" import asyncio -import errno -import fcntl -import os from types import SimpleNamespace import pytest -import prime_rl.monitors.prime as prime_module -import prime_rl.utils.record as record_module -from prime_rl.configs.monitors import PrimeMonitorConfig +from prime_rl import monitors +from prime_rl.configs.monitors import ( + EvalMonitorsConfig, + FileMonitorConfig, + PrimeMonitorConfig, + TrainMonitorsConfig, + WandbMonitorConfig, +) from prime_rl.configs.sft import SFTConfig -from prime_rl.monitors.base import MonitorContext -from prime_rl.monitors.prime import PrimeMonitor, read_platform_record, update_platform_record -from prime_rl.utils.links import ATTEMPT_VAR -from prime_rl.utils.pathing import get_platform_run_path -from prime_rl.utils.process import RUN_FINISHER_VAR -from prime_rl.utils.record import record_lock - - -class FakeRun(SimpleNamespace): - def __init__(self, **kwargs): - super().__init__(id=kwargs.get("id", "run-123"), url="https://prime/training/run-123", attached="id" in kwargs) - self.finished = False - self.flushed = False - self.kwargs = kwargs - - def _atexit_hook(self): - pass - - def finish(self, *args, **kwargs): - self.finished = True - - def flush(self, *args, **kwargs): - self.flushed = True - - def log_metrics(self, *args, **kwargs): - pass - - def log_episodes(self, *args, **kwargs): - pass +from prime_rl.monitors.prime import PrimeMonitor @pytest.fixture -def prime_init(monkeypatch, tmp_path): - """Replace the ``prime_runs`` SDK entry point with one that records its runs.""" - runs: list[FakeRun] = [] +def prime_init(monkeypatch): + """Replace the ``prime_runs`` SDK entry point with one that records its kwargs.""" + calls = [] def fake_init(**kwargs): - runs.append(FakeRun(**kwargs)) - return runs[-1] - - monkeypatch.setattr(prime_module.pr, "init", fake_init) - for var in ("RUN_ID", "PRIME_RUNS_MODE", RUN_FINISHER_VAR): - monkeypatch.delenv(var, raising=False) + calls.append(kwargs) + return SimpleNamespace( + id="run-123", + url="https://platform/training/run-123", + attached=False, + finished=False, + finish=lambda *args, **kwargs_: None, + log_metrics=lambda *args, **kwargs_: None, + log_episodes=lambda *args, **kwargs_: None, + ) + + monkeypatch.setattr("prime_rl.monitors.prime.pr.init", fake_init) + monkeypatch.delenv("RUN_ID", raising=False) + monkeypatch.delenv("PRIME_RUNS_MODE", raising=False) monkeypatch.setenv("PRIME_API_KEY", "test-key") - monkeypatch.setenv(ATTEMPT_VAR, str(tmp_path / "logs" / "attempt_1")) - return runs + return calls + + +@pytest.fixture +def clear_monitors(): + """``monitors.setup`` registers into a module-global list; empty it after each test.""" + yield + monitors.MONITORS.clear() def _sft_config(overrides: dict | None = None) -> SFTConfig: @@ -67,42 +59,158 @@ def _sft_config(overrides: dict | None = None) -> SFTConfig: "data": {"type": "fake", "batch_size": 8, "seq_len": 64}, "max_steps": 10, } - config.update(overrides or {}) + if overrides: + config.update(overrides) return SFTConfig.model_validate(config) -def _monitor(tmp_path, config, producer: str) -> PrimeMonitor: - return PrimeMonitor(PrimeMonitorConfig(name="my-run"), MonitorContext(tmp_path, producer, config)) +def test_sft_prime_monitor_name_inherits_run_name(): + config = _sft_config({"monitors": {"prime": {}}}) + + assert config.monitors.prime.name == config.run.name + + +def test_sft_prime_monitor_keeps_explicit_name(): + config = _sft_config({"monitors": {"prime": {"name": "my-experiment"}}}) + + assert config.monitors.prime.name == "my-experiment" def test_sft_registers_dataset_training_fields(prime_init, tmp_path): - config = _sft_config({"monitors": {"wandb": {"project": "sft-tests"}}}) - asyncio.run(_monitor(tmp_path, config, "trainer").init()) - - (run,) = prime_init - assert run.kwargs["kind"] == "train" - assert run.kwargs["environments"] == [] - training = run.kwargs["training"] - assert (training.max_steps, training.batch_size, training.seq_len) == (10, 8, 64) - assert training.rollouts_per_example is None - assert training.wandb_project == "sft-tests" - assert read_platform_record(tmp_path)["id"] == "run-123" - - -def test_online_evals_share_the_trainers_run(prime_init, tmp_path, monkeypatch): - """The online-eval process attaches to the run the trainer registered, and only it - finishes the run: it drains its final evals after the trainer exits.""" + config = _sft_config({"monitors": {"prime": {}}}) + + asyncio.run(PrimeMonitor(config.monitors.prime).init(config=config, output_dir=tmp_path)) + + (call,) = prime_init + assert call["kind"] == "train" + assert call["name"] == config.run.name + assert call["model"] == "PrimeIntellect/Qwen3-0.6B" + # SFT trains on a dataset: no environments, no rollouts per example + assert call["environments"] == [] + assert call["training"].max_steps == 10 + assert call["training"].batch_size == 8 + assert call["training"].seq_len == 64 + assert call["training"].rollouts_per_example is None + assert call["config"]["model"]["name"] == "PrimeIntellect/Qwen3-0.6B" + + +def test_sft_registers_wandb_project(prime_init, tmp_path): + config = _sft_config( + { + "monitors": { + "wandb": {"project": "sft-tests"}, + "prime": {}, + } + } + ) + + asyncio.run(PrimeMonitor(config.monitors.prime).init(config=config, output_dir=tmp_path)) + + assert prime_init[0]["training"].wandb_project == "sft-tests" + + +def test_online_eval_monitors_keep_the_trainers_prime_monitor(): + """The online-eval process logs to the trainer's Prime run, under the same config.""" + from prime_rl.entrypoints.sft import build_online_eval_monitors + + monitors = build_online_eval_monitors( + TrainMonitorsConfig( + wandb=WandbMonitorConfig(project="sft-tests"), + file=FileMonitorConfig(), + prime=PrimeMonitorConfig(name="my-run"), + ) + ) + + assert isinstance(monitors, EvalMonitorsConfig) + assert monitors.prime == PrimeMonitorConfig(name="my-run") + assert monitors.wandb.project == "sft-tests" + assert build_online_eval_monitors(TrainMonitorsConfig(prime=None)).prime is None + + +def test_sft_setup_registers_prime_monitor(prime_init, tmp_path, clear_monitors): + config = _sft_config({"monitors": {"prime": {}}}) + + asyncio.run( + monitors.setup( + prime=config.monitors.prime, + output_dir=tmp_path, + run_config=config, + ) + ) + + assert monitors.get(PrimeMonitor) is not None + + # The trainer's per-step metrics fan out to the platform run, and finalize + # closes it out (drains the SDK's upload queue). + asyncio.run(monitors.log({"loss/mean": 1.0, "step": 1}, step=1)) + asyncio.run(monitors.log({"val/loss": 0.5, "step": 1}, step=1)) + asyncio.run(monitors.finalize()) + + +def test_online_eval_config_keeps_prime_monitor(): + """SFTOnlineEvalConfig.monitors must be EvalMonitorsConfig so the + trainer's prime monitor survives the eval.json dump/re-parse boundary + (base MonitorsConfig drops/forbids prime).""" + from prime_rl.configs.eval import SFTOnlineEvalConfig + from prime_rl.configs.monitors import EvalMonitorsConfig + from prime_rl.entrypoints.sft import build_online_eval_monitors + + field = SFTOnlineEvalConfig.model_fields["monitors"] + assert field.annotation is EvalMonitorsConfig + + train_monitors = TrainMonitorsConfig.model_validate({"prime": {"name": "r1"}}) + eval_monitors = build_online_eval_monitors(train_monitors) + assert isinstance(eval_monitors.prime, PrimeMonitorConfig) + assert eval_monitors.prime.name == "r1" + + # The online-eval process passes the prime monitor to monitors.setup + # (src/prime_rl/eval/online.py) — same wiring as the standalone eval + # entrypoint. Read the source text (importing the module pulls torch, + # which this torch-free test env deliberately lacks). + from pathlib import Path + + import prime_rl + + online_src = Path(prime_rl.__path__[0], "eval", "online.py").read_text() + assert "prime=config.monitors.prime" in online_src + + +def test_online_evals_share_the_trainers_run(tmp_path, monkeypatch): + """The online-eval process shares the trainer's run dir and attaches to the run the + trainer registered. Only it finishes the run: it drains its final evals after the + trainer exits.""" from prime_rl.entrypoints.sft import build_online_eval_config + from prime_rl.utils.links import ATTEMPT_VAR + from prime_rl.utils.process import RUN_FINISHER_VAR - config = _sft_config({"eval": {"interval": 5, "source": [{"env": {"taskset": {"id": "harbor"}}}]}, "inference": {}}) - eval_config = build_online_eval_config(config) - assert eval_config.monitors.prime == config.monitors.prime + runs = [] + + def fake_init(**kwargs): + run = SimpleNamespace(id=kwargs.get("id", "run-123"), url="https://x/run-123", attached="id" in kwargs) + run.kwargs, run.finished, run.flushed = kwargs, False, False + run._atexit_hook = lambda: None + run.finish = lambda *a, **k: setattr(run, "finished", True) + run.flush = lambda *a, **k: setattr(run, "flushed", True) + runs.append(run) + return run + + monkeypatch.setattr("prime_rl.monitors.prime.pr.init", fake_init) + for var in ("RUN_ID", "PRIME_RUNS_MODE"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("PRIME_API_KEY", "test-key") + monkeypatch.setenv(ATTEMPT_VAR, str(tmp_path / "logs" / "attempt_1")) monkeypatch.setenv(RUN_FINISHER_VAR, "online-eval") - trainer = _monitor(tmp_path, config, "trainer") - asyncio.run(trainer.init()) - online_eval = _monitor(tmp_path, eval_config, "online-eval") - asyncio.run(online_eval.init()) + config = _sft_config( + {"eval": {"interval": 50, "source": [{"name": "rev", "env": {"taskset": {"id": "harbor"}}}]}, "inference": {}} + ) + eval_config = build_online_eval_config(config) + # Same run dir: dashboard visibility for the eval artifacts. + assert eval_config.output_dir == config.run_dir + + trainer, online_eval = PrimeMonitor(config.monitors.prime), PrimeMonitor(eval_config.monitors.prime) + asyncio.run(trainer.init(config=config, output_dir=tmp_path, producer="trainer")) + asyncio.run(online_eval.init(config=eval_config, output_dir=tmp_path, producer="online-eval")) assert online_eval.run.kwargs["id"] == trainer.run.id asyncio.run(trainer.finalize()) @@ -111,28 +219,94 @@ def test_online_evals_share_the_trainers_run(prime_init, tmp_path, monkeypatch): assert online_eval.run.finished -def test_online_eval_ignores_a_previous_attempts_run(prime_init, tmp_path, monkeypatch): - update_platform_record(tmp_path, {"kind": "train", "id": "old", "attempt": "attempt_0"}) - with pytest.raises(RuntimeError, match="registered no Prime run"): - prime_module.wait_for_train_run(tmp_path, timeout=0.0) - - def test_platform_record_lock_keeps_a_stable_inode(tmp_path): - """The record lock never unlinks its lock file: a waiting writer could hold the - unlinked inode's lock while a later writer locks a fresh file.""" + """The record lock must NEVER unlink its lock file: a writer waiting on + the flock can still hold the unlinked inode's lock while a later writer + creates and locks a fresh file - two simultaneous "exclusive" locks and + a lost merge. The lock file must persist with the SAME inode across + acquisitions, and the locked update helper must merge two sequential + writers into one record. (Interprocess contention is exercised by the + two-process update_platform_record probe; here the writers are + sequential.)""" + import os + + from prime_rl.monitors.prime import update_platform_record + from prime_rl.utils.pathing import get_platform_run_path + from prime_rl.utils.record import record_lock + lock_path = get_platform_run_path(tmp_path).with_suffix(".json.lock") inodes = [] for _ in range(2): with record_lock(get_platform_run_path(tmp_path)): + # The lock file exists DURING the critical section, always. + assert lock_path.is_file() inodes.append(os.stat(lock_path).st_ino) - assert lock_path.is_file() and len(set(inodes)) == 1 - update_platform_record(tmp_path, {"kind": "train", "id": "run-1"}) - update_platform_record(tmp_path, {"url": "https://x/run-1"}) - assert read_platform_record(tmp_path) == {"kind": "train", "id": "run-1", "url": "https://x/run-1"} + # The lock file SURVIVES release, and every acquisition locked the same inode. + assert lock_path.is_file() + assert len(set(inodes)) == 1 + + # Two sequential locked updates merge instead of clobbering. + update_platform_record(tmp_path, {"kind": "train", "id": "run-1", "url": "https://x/run-1"}) + update_platform_record(tmp_path, {"evaluations": {"rev": {"step": 5, "id": "ev-1", "url": "https://x/ev-1"}}}) + from prime_rl.monitors.prime import read_platform_record + + record = read_platform_record(tmp_path) + assert record["kind"] == "train" + assert record["id"] == "run-1" + assert record["evaluations"]["rev"]["id"] == "ev-1" + + +def test_record_persistence_runs_off_the_event_loop(prime_init, tmp_path, monkeypatch): + """The record merge (bounded flock retry included) must run in a worker + thread via asyncio.to_thread, not on the event loop thread: contention + for the full lock bound would otherwise freeze the loop.""" + import threading + from pathlib import Path + + import prime_rl.monitors.prime as prime_module + + seen_threads = [] + real_merge = prime_module._merge_platform_record + + def spy(output_dir, merge): + seen_threads.append(threading.get_ident()) + return real_merge(output_dir, merge) + + monkeypatch.setattr(prime_module, "_merge_platform_record", spy) + + config = _sft_config({"monitors": {"prime": {}}}) + asyncio.run(PrimeMonitor(config.monitors.prime).init(config=config, output_dir=tmp_path)) + + assert seen_threads, "the record merge never ran" + assert threading.get_ident() not in seen_threads, "record merge ran on the caller thread" + + # The async write sites (train run, evaluation) go through to_thread; only the sync + # helpers stay callable directly. + import re + + import prime_rl + + src = Path(prime_rl.__path__[0], "monitors", "prime.py").read_text() + assert len(re.findall(r"asyncio\.to_thread\(\s*self\._record", src)) == 2 + + from prime_rl.monitors.prime import read_platform_record + + record = read_platform_record(tmp_path) + assert record["kind"] == "train" + assert record["id"] == "run-123" def test_platform_record_lock_acquisition_is_bounded(tmp_path, monkeypatch): + """flock(LOCK_EX) must never block the event-loop thread indefinitely: + while another process holds the lock, acquisition retries + non-blockingly and raises TimeoutError past the bound instead of + hanging.""" + import fcntl + + import prime_rl.utils.record as record_module + from prime_rl.utils.pathing import get_platform_run_path + monkeypatch.setattr(record_module, "RECORD_LOCK_TIMEOUT", 0.2) path = get_platform_run_path(tmp_path) lock_path = path.with_suffix(".json.lock") @@ -140,16 +314,23 @@ def test_platform_record_lock_acquisition_is_bounded(tmp_path, monkeypatch): with open(lock_path, "a+b") as held: fcntl.flock(held.fileno(), fcntl.LOCK_EX) with pytest.raises(TimeoutError, match="still held"): - with record_lock(path): + with record_module.record_lock(path): pass def test_platform_record_persistence_is_best_effort(tmp_path, monkeypatch): - """A storage failure on the record (e.g. ENOLCK) is skipped, never fatal.""" + """A storage failure on the OPTIONAL dashboard record (e.g. ENOLCK on a + filesystem without flock support) must not abort training: the write is + skipped with a warning, and no record is left behind.""" + import errno + + import prime_rl.monitors.prime as prime_module + import prime_rl.utils.record as record_module def enolck(fd, op): raise OSError(errno.ENOLCK, "No locks available") monkeypatch.setattr(record_module.fcntl, "flock", enolck) - update_platform_record(tmp_path, {"kind": "train", "id": "run-1"}) - assert read_platform_record(tmp_path) is None + # Must not raise: the platform record is best-effort. + prime_module.update_platform_record(tmp_path, {"kind": "train", "id": "run-1", "url": "https://x/run-1"}) + assert prime_module.read_platform_record(tmp_path) is None diff --git a/tests/unit/utils/test_links.py b/tests/unit/utils/test_links.py index d58043b27e8..f4914fb0d7d 100644 --- a/tests/unit/utils/test_links.py +++ b/tests/unit/utils/test_links.py @@ -6,7 +6,6 @@ import prime_rl.monitors.prime as prime_module from prime_rl.configs.monitors import PrimeMonitorConfig -from prime_rl.monitors.base import MonitorContext from prime_rl.monitors.prime import PrimeMonitor from prime_rl.utils.links import ATTEMPT_VAR, publish_link, register_links, wait_for_links @@ -20,16 +19,16 @@ def attempt(tmp_path, monkeypatch): def test_waits_for_every_registered_link(tmp_path, attempt): register_links(tmp_path, "prime", 0) - register_links(tmp_path, "other", 2) + register_links(tmp_path, "prime/eval", 2) register_links(tmp_path, "wandb", 1) publish_link(tmp_path, "wandb", "run", "https://wandb.ai/run") - publish_link(tmp_path, "other", "gsm8k", "https://prime/ev-1") + publish_link(tmp_path, "prime/eval", "gsm8k", "https://prime/ev-1") # A page that failed to open publishes None, which resolves the wait without a link. - publish_link(tmp_path, "other", "aime", None) + publish_link(tmp_path, "prime/eval", "aime", None) start = time.monotonic() - assert wait_for_links(tmp_path, attempt, ["prime", "other", "wandb"], lambda: False, timeout=5.0) == { + assert wait_for_links(tmp_path, attempt, ["prime", "prime/eval", "wandb"], lambda: False, timeout=5.0) == { "prime": {}, - "other": {"gsm8k": "https://prime/ev-1"}, + "prime/eval": {"gsm8k": "https://prime/ev-1"}, "wandb": {"run": "https://wandb.ai/run"}, } assert time.monotonic() - start < 1.0 @@ -49,4 +48,4 @@ def test_prime_monitor_without_credentials_fails_informatively(tmp_path, monkeyp monkeypatch.delenv("RUN_ID", raising=False) monkeypatch.setattr(prime_module, "PrimeConfig", lambda: SimpleNamespace(api_key=None)) with pytest.raises(RuntimeError, match=r"prime login.*--no-monitors\.prime"): - asyncio.run(PrimeMonitor(PrimeMonitorConfig(), MonitorContext(tmp_path)).init()) + asyncio.run(PrimeMonitor(PrimeMonitorConfig()).init(config=None, output_dir=tmp_path)) From 1962bf262c5c4cd3e7de31c82dc1d311bc7d3ac2 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sun, 25 Oct 2026 18:15:51 +0000 Subject: [PATCH 03/12] register the prime run in the rl and sft launchers The launcher registers the run before it starts any component and makes the run's id the run identity: $PRL_RUN_ID (and with it the W&B run id) and $RUN_ID, which every component's prime monitor attaches to - the same shape as a hosted launch, whose platform pre-creates the run. The sft online-eval process no longer polls run.json for the trainer's run, and the record drops its attempt stamp. An attached run leaves its failure to the process that registered it, so a launch that fails marks the run failed. Components started without a launcher still register their own run. Co-Authored-By: Claude Opus 5.5 --- docs/training.md | 2 +- src/prime_rl/entrypoints/rl.py | 12 +- src/prime_rl/entrypoints/sft.py | 11 +- src/prime_rl/monitors/prime.py | 159 +++++++++++---------- tests/unit/train/sft/test_prime_monitor.py | 76 +++++++--- 5 files changed, 154 insertions(+), 106 deletions(-) diff --git a/docs/training.md b/docs/training.md index fac3923a8d7..4ba26dac97b 100644 --- a/docs/training.md +++ b/docs/training.md @@ -459,7 +459,7 @@ Name the Prime run in TOML: name = "my-experiment" ``` -On `sft` the monitor registers the run with the dataset-batched training fields (`max_steps`, `batch_size`, `seq_len`) and streams the trainer's per-step metrics and validation losses. Online evals log to the same run, like an `rl` run's: the online-eval process attaches to the trainer's run, uploads its eval metrics and every eval episode, and finishes the run after its final evals. +The `rl` and `sft` launchers register the run before they start the components, and every component logs to it. The run's id is also the run identity (`PRL_RUN_ID`) and the W&B run id. On `sft` the run carries the dataset-batched training fields (`max_steps`, `batch_size`, `seq_len`) and the trainer's per-step metrics and validation losses. Online evals log to the same run, like an `rl` run's: the online-eval process uploads its eval metrics and every eval episode, and finishes the run after its final evals. A launch that fails marks the run failed. The platform run needs `PRIME_API_KEY` (`prime login` or the env var) and a team (`PRIME_TEAM_ID`, or the team selected with `prime login`) enabled for external runs. Like the W&B monitor, a run without a key, or one whose registration fails (no team, a team outside the allowlist), fails at launch: authenticate, select a team with `prime switch`, or pass `--no-monitors.prime`. `PRIME_RUNS_MODE=disabled` keeps the monitor configured but opens no platform run; `RUN_ID=` attaches to an external run a launcher already created instead of registering a new one. Currently internal-only. diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index a2c3be8f90a..ba19902ef5a 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -480,7 +480,6 @@ def rl_slurm(config: RLConfig): def rl(config: RLConfig): - # TODO: fetch the run id from the Prime SDK once runs are registered there. ckpt_output_dir = config.ckpt.output_dir if config.ckpt else None prepare_run(config, ckpt_output_dir) @@ -511,10 +510,13 @@ def rl(config: RLConfig): pre_download_model(config.trainer.model.name, skip_weights=config.trainer.model.debug.random_init) - if config.slurm is not None: - rl_slurm(config) - else: - rl_local(config) + from prime_rl.monitors.prime import launch_run + + with launch_run(config.orchestrator, None if config.dry_run else config.orchestrator.monitors.prime): + if config.slurm is not None: + rl_slurm(config) + else: + rl_local(config) def main(): diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 7dc43b88a07..dd61a6c06f5 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -479,10 +479,13 @@ def sft(config: SFTConfig): if config.val is not None: pre_download_data(config.val.data, config.env_vars) - if config.slurm is not None: - sft_slurm(config) - else: - sft_local(config) + from prime_rl.monitors.prime import launch_run + + with launch_run(config, None if config.dry_run else config.monitors.prime): + if config.slurm is not None: + sft_slurm(config) + else: + sft_local(config) def main(): diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 1f3c91e6c0d..9df0c35555a 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -3,7 +3,8 @@ import asyncio import atexit import os -import time +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path from typing import TYPE_CHECKING, Any, Callable @@ -12,10 +13,11 @@ from prime_rl.configs.eval import EvalConfig, SFTOnlineEvalConfig from prime_rl.configs.monitors import PrimeMonitorConfig +from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.configs.sft import SFTConfig from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.utils.config import BaseConfig -from prime_rl.utils.links import ATTEMPT_VAR, PRIME_LINKS, publish_link, register_links +from prime_rl.utils.links import PRIME_LINKS, publish_link, register_links from prime_rl.utils.pathing import get_platform_run_path from prime_rl.utils.process import RUN_FINISHER_VAR from prime_rl.utils.record import read_record, update_record, write_record @@ -30,8 +32,6 @@ # default (300 s) is sized for eval sample batches; a crashed training process should # not linger that long, and a clean finish rarely has more than the last step queued. FINISH_TIMEOUT = 60.0 -# How long an SFT online-eval process waits for the trainer to register the run. -TRAIN_RUN_TIMEOUT_S = 1200.0 def write_platform_record(output_dir: Path, record: dict[str, Any]) -> None: @@ -84,28 +84,78 @@ def _base_url() -> str | None: return base.rstrip("/").removesuffix("/rft") if base else None -def wait_for_train_run(output_dir: Path, timeout: float = TRAIN_RUN_TIMEOUT_S) -> str: - """The id of the training run this attempt's trainer registered, from the platform - record it leaves in the shared run dir.""" - attempt = os.getenv(ATTEMPT_VAR) - deadline = time.monotonic() + timeout - while True: - try: - record = read_platform_record(output_dir) or {} - except (OSError, ValueError): - record = {} - if record.get("kind") == "train" and record.get("attempt") == attempt and record.get("id"): - return record["id"] - if time.monotonic() >= deadline: - raise RuntimeError(f"The trainer registered no Prime run within {timeout:.0f}s; {OFF_HINT}") - time.sleep(1.0) +def registration(config: OrchestratorConfig | SFTConfig, name: str | None) -> dict[str, Any]: + """The ``pr.init`` arguments that register ``config``'s training run.""" + if isinstance(config, SFTConfig): + # SFT trains on a dataset, not on environments, and has no rollouts - the + # platform run is registered without both. + environments = [] + training = pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.data.batch_size, + seq_len=config.data.seq_len, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + ) + else: + environments = [env.env_id for env in config.train.source] + training = pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.batch_size, + rollouts_per_example=max( + (source.group_size for source in config.train.source), default=config.train.group_size + ), + seq_len=config.seq_len, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + ) + return dict( + name=name, + model=config.model.name, + environments=environments, + training=training, + config=config.model_dump(exclude_none=True, mode="json"), + ) + + +def register(init_kwargs: dict[str, Any], mode: str) -> pr.Run: + try: + return pr.init(kind="train", mode=mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT, **init_kwargs) + except (pr.ConfigurationError, pr.APIError) as e: + # Registering fails without a team, or for a team outside the external-runs + # allowlist. An attached run's error is the launcher's to explain. + if "id" in init_kwargs: + raise + raise RuntimeError( + f"Failed to register the Prime run ({type(e).__name__}: {e}). Select a team with " + f"`prime switch` (or set PRIME_TEAM_ID), or {OFF_HINT}" + ) from e + + +@contextmanager +def launch_run(config: OrchestratorConfig | SFTConfig, prime: PrimeMonitorConfig | None) -> Iterator[None]: + """Register the launch's Prime run around starting its components, which attach to + it: its id becomes ``$PRL_RUN_ID`` and ``$RUN_ID``. A launch that fails marks the + run failed - the SDK leaves an attached run's failure to the process that + registered it. A run a managed launch pre-created (``$RUN_ID``) is the platform's.""" + if prime is None or os.getenv("RUN_ID") or prime_mode() != "online": + yield + return + run = register(registration(config, prime.name), "online") + os.environ["PRL_RUN_ID"] = os.environ["RUN_ID"] = run.id + try: + yield + except BaseException as e: + run.fail(e, timeout=FINISH_TIMEOUT) + raise + # The finisher component completes the run; this process only exits. + atexit.unregister(run._atexit_hook) class PrimeMonitor(Monitor): """Logs metrics and episodes to the Prime platform through ``prime_runs``. - A training run's processes log to one Prime training run: the rl orchestrator or - sft trainer registers it, and an sft online-eval process attaches to its trainer's. + A training run's processes log to one Prime training run: the ``rl`` or ``sft`` + launcher registers it (``launch_run``) and its components attach through ``$RUN_ID``. + A component started without a launcher registers its own. ``$PRL_RUN_FINISHER`` names the process that finishes a shared run; the others only flush. A standalone eval instead opens one evaluation per env and epoch when the epoch starts (``log_eval_plan``), streams its episodes into it, and finishes @@ -139,65 +189,24 @@ async def init(self, config: BaseConfig | None, output_dir: Path, producer: str await self._init_evaluations(config) return online_eval = isinstance(config, SFTOnlineEvalConfig) - init_kwargs: dict[str, Any] if run_id := os.getenv("RUN_ID"): - # A managed launch pre-created the platform run and injected its id - - # attach instead of registering a duplicate. The platform owns the - # attached run's failure/completion marking; a clean finish() is - # delivered best-effort and does not by itself mark it completed. + # The launcher registered the run (or a managed launch pre-created it) - + # attach instead of registering a duplicate. The registering process owns + # the attached run's failure marking; a clean finish() completes it. init_kwargs = {"id": run_id} elif online_eval: - # The trainer registered the run; its online-eval process logs to it. - init_kwargs = {"id": await asyncio.to_thread(wait_for_train_run, output_dir)} if mode == "online" else {} - elif isinstance(config, SFTConfig): - # SFT trains on a dataset, not on environments, and has no rollouts - the - # platform run is registered without both. - init_kwargs = dict( - name=self.config.name, - model=config.model.name, - environments=[], - training=pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.data.batch_size, - seq_len=config.data.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, - ), - config=config.model_dump(exclude_none=True, mode="json"), - ) + if mode == "online": + raise RuntimeError( + f"An online-eval process logs to its trainer's run: launch it with `sft`, or {OFF_HINT}" + ) + init_kwargs = {} elif config is not None: - init_kwargs = dict( - name=self.config.name, - model=config.model.name, - environments=[env.env_id for env in config.train.source], - training=pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.batch_size, - rollouts_per_example=max( - (source.group_size for source in config.train.source), default=config.train.group_size - ), - seq_len=config.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, - ), - config=config.model_dump(exclude_none=True, mode="json"), - ) + init_kwargs = registration(config, self.config.name) else: # The RFT API requires a base model; "unknown" is what the # pre-SDK register sent for a config-less init. init_kwargs = {"name": self.config.name, "model": "unknown"} - - try: - self.run = await asyncio.to_thread( - pr.init, kind="train", mode=mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT, **init_kwargs - ) - except (pr.ConfigurationError, pr.APIError) as e: - # Registering fails without a team, or for a team outside the external-runs - # allowlist. An attached run's error is the launcher's to explain. - if "id" in init_kwargs: - raise - raise RuntimeError( - f"Failed to register the Prime run ({type(e).__name__}: {e}). Select a team with " - f"`prime switch` (or set PRIME_TEAM_ID), or {OFF_HINT}" - ) from e + self.run = await asyncio.to_thread(register, init_kwargs, mode) if self.run.url: attached = " (attached via $RUN_ID)" if self.run.attached else "" self.logger.info(f"Logging metrics and episodes to Prime run {self.run.id} ({self.run.url}){attached}") @@ -210,11 +219,7 @@ async def init(self, config: BaseConfig | None, output_dir: Path, producer: str def _record_run(self) -> None: register_links(self.output_dir, PRIME_LINKS, 1 if self.run.url else 0) if self.run.url: - # Stamped with the attempt: the online-eval process attaches to this run. - update_platform_record( - self.output_dir, - {"kind": "train", "id": self.run.id, "url": self.run.url, "attempt": os.getenv(ATTEMPT_VAR)}, - ) + update_platform_record(self.output_dir, {"kind": "train", "id": self.run.id, "url": self.run.url}) publish_link(self.output_dir, PRIME_LINKS, "run", self.run.url) async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index 8d549dcc751..8f43a952955 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -6,6 +6,7 @@ """ import asyncio +import os from types import SimpleNamespace import pytest @@ -175,22 +176,18 @@ def test_online_eval_config_keeps_prime_monitor(): assert "prime=config.monitors.prime" in online_src -def test_online_evals_share_the_trainers_run(tmp_path, monkeypatch): - """The online-eval process shares the trainer's run dir and attaches to the run the - trainer registered. Only it finishes the run: it drains its final evals after the - trainer exits.""" - from prime_rl.entrypoints.sft import build_online_eval_config - from prime_rl.utils.links import ATTEMPT_VAR - from prime_rl.utils.process import RUN_FINISHER_VAR - +@pytest.fixture +def sdk_runs(monkeypatch): + """Fake ``pr.init`` runs that record their kwargs and how they ended.""" runs = [] def fake_init(**kwargs): run = SimpleNamespace(id=kwargs.get("id", "run-123"), url="https://x/run-123", attached="id" in kwargs) - run.kwargs, run.finished, run.flushed = kwargs, False, False + run.kwargs, run.finished, run.flushed, run.failed = kwargs, False, False, None run._atexit_hook = lambda: None run.finish = lambda *a, **k: setattr(run, "finished", True) run.flush = lambda *a, **k: setattr(run, "flushed", True) + run.fail = lambda error, **k: setattr(run, "failed", error) runs.append(run) return run @@ -198,9 +195,19 @@ def fake_init(**kwargs): for var in ("RUN_ID", "PRIME_RUNS_MODE"): monkeypatch.delenv(var, raising=False) monkeypatch.setenv("PRIME_API_KEY", "test-key") - monkeypatch.setenv(ATTEMPT_VAR, str(tmp_path / "logs" / "attempt_1")) - monkeypatch.setenv(RUN_FINISHER_VAR, "online-eval") + monkeypatch.setenv("PRL_RUN_ID", "local-uuid") + return runs + +def test_online_evals_share_the_launchers_run(sdk_runs, tmp_path, monkeypatch): + """The sft launcher registers the run and makes its id the run identity; the trainer + and the online-eval process attach to it. Only the online-eval process finishes the + run: it drains its final evals after the trainer exits.""" + from prime_rl.entrypoints.sft import build_online_eval_config + from prime_rl.monitors.prime import launch_run + from prime_rl.utils.process import RUN_FINISHER_VAR + + monkeypatch.setenv(RUN_FINISHER_VAR, "online-eval") config = _sft_config( {"eval": {"interval": 50, "source": [{"name": "rev", "env": {"taskset": {"id": "harbor"}}}]}, "inference": {}} ) @@ -209,14 +216,45 @@ def fake_init(**kwargs): assert eval_config.output_dir == config.run_dir trainer, online_eval = PrimeMonitor(config.monitors.prime), PrimeMonitor(eval_config.monitors.prime) - asyncio.run(trainer.init(config=config, output_dir=tmp_path, producer="trainer")) - asyncio.run(online_eval.init(config=eval_config, output_dir=tmp_path, producer="online-eval")) - assert online_eval.run.kwargs["id"] == trainer.run.id - - asyncio.run(trainer.finalize()) - assert trainer.run.flushed and not trainer.run.finished - asyncio.run(online_eval.finalize()) - assert online_eval.run.finished + with launch_run(config, config.monitors.prime): + launched = sdk_runs[0] + assert launched.kwargs["model"] == "PrimeIntellect/Qwen3-0.6B" and not launched.attached + assert os.environ["PRL_RUN_ID"] == os.environ["RUN_ID"] == launched.id + asyncio.run(trainer.init(config=config, output_dir=tmp_path, producer="trainer")) + asyncio.run(online_eval.init(config=eval_config, output_dir=tmp_path, producer="online-eval")) + assert trainer.run.attached and online_eval.run.attached + assert trainer.run.id == online_eval.run.id == launched.id + + asyncio.run(trainer.finalize()) + assert trainer.run.flushed and not trainer.run.finished + asyncio.run(online_eval.finalize()) + assert online_eval.run.finished + assert launched.failed is None + + +def test_a_failed_launch_fails_the_run(sdk_runs): + """A component's crash only flushes its attached run; the launcher, which registered + the run, marks it failed.""" + from prime_rl.monitors.prime import launch_run + + config = _sft_config() + error = RuntimeError("trainer exited with code 1") + with pytest.raises(RuntimeError), launch_run(config, config.monitors.prime): + raise error + assert sdk_runs[0].failed is error + + +@pytest.mark.parametrize("env", [{"RUN_ID": "hosted-1"}, {"PRIME_RUNS_MODE": "disabled"}]) +def test_launch_registers_no_run_it_does_not_own(sdk_runs, monkeypatch, env): + """A managed launch pre-created the run, or the monitor opens no platform run.""" + from prime_rl.monitors.prime import launch_run + + for var, value in env.items(): + monkeypatch.setenv(var, value) + config = _sft_config() + with launch_run(config, config.monitors.prime): + pass + assert sdk_runs == [] and os.environ["PRL_RUN_ID"] == "local-uuid" def test_platform_record_lock_keeps_a_stable_inode(tmp_path): From c4ffe6b5c7914b2303610952221bc926229e7d7b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sun, 25 Oct 2026 18:40:24 +0000 Subject: [PATCH 04/12] name the failed component on the prime run ProcessGroup.wait chains the component's error to its SystemExit, so the launcher marks the run failed with that error instead of "SystemExit: 1". The exit code stays 1. Co-Authored-By: Claude Opus 5.5 --- src/prime_rl/monitors/prime.py | 3 ++- src/prime_rl/utils/process.py | 2 +- tests/unit/train/sft/test_prime_monitor.py | 6 +++--- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 9df0c35555a..f7f6a5f2b2e 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -144,7 +144,8 @@ def launch_run(config: OrchestratorConfig | SFTConfig, prime: PrimeMonitorConfig try: yield except BaseException as e: - run.fail(e, timeout=FINISH_TIMEOUT) + # A failed component exits the launcher with a SystemExit chained to its error. + run.fail(e.__cause__ or e, timeout=FINISH_TIMEOUT) raise # The finisher component completes the run; this process only exits. atexit.unregister(run._atexit_hook) diff --git a/src/prime_rl/utils/process.py b/src/prime_rl/utils/process.py index bebbe9b4274..80854a0ba4d 100644 --- a/src/prime_rl/utils/process.py +++ b/src/prime_rl/utils/process.py @@ -214,7 +214,7 @@ def wait(self, *names: str) -> None: if self.error_queue: get_logger().error(f"Error: {self.error_queue[0]}") get_logger().error("Terminating all processes...") - sys.exit(1) + raise SystemExit(1) from self.error_queue[0] if not pending: return pending[0].wait(timeout=1) diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index 8f43a952955..4dd09a418e7 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -238,9 +238,9 @@ def test_a_failed_launch_fails_the_run(sdk_runs): from prime_rl.monitors.prime import launch_run config = _sft_config() - error = RuntimeError("trainer exited with code 1") - with pytest.raises(RuntimeError), launch_run(config, config.monitors.prime): - raise error + error = RuntimeError("Trainer failed with exit code -9") + with pytest.raises(SystemExit), launch_run(config, config.monitors.prime): + raise SystemExit(1) from error assert sdk_runs[0].failed is error From 9ab592044e2383f1ba782031fae6998643eecb5c Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 28 Oct 2026 23:35:52 +0000 Subject: [PATCH 05/12] collapse the monitors configs into one MonitorsConfig With one prime monitor config, TrainMonitorsConfig and EvalMonitorsConfig were the same type. The online-eval process now takes the trainer's MonitorsConfig as is. The RL trainer gains a monitors.prime field that it does not set up yet. Also export RUN_ID in the multi-node SFT script, which renders PRL_RUN_ID explicitly instead of inheriting the submitter environment. Co-Authored-By: Claude Opus 5.5 --- .../src/prime_rl/configs/eval.py | 13 ++--- .../src/prime_rl/configs/monitors.py | 14 ++--- .../src/prime_rl/configs/orchestrator.py | 4 +- .../src/prime_rl/configs/sft.py | 4 +- src/prime_rl/entrypoints/sft.py | 10 +--- .../templates/multi_node_sft.sbatch.j2 | 4 ++ tests/unit/train/sft/test_prime_monitor.py | 53 ------------------- 7 files changed, 18 insertions(+), 84 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/eval.py b/packages/prime-rl-configs/src/prime_rl/configs/eval.py index 3269993002f..12fdc92637c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/eval.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/eval.py @@ -4,7 +4,7 @@ import verifiers.v1 as vf from pydantic import AliasChoices, Field, model_validator -from prime_rl.configs.monitors import EvalMonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig from prime_rl.configs.orchestrator import ConcurrencyConfig, EvalSourcesConfig, ScheduledEvalConfig from prime_rl.configs.shared import ( ClientConfig, @@ -108,7 +108,7 @@ class EvalConfig(ServedEvalConfig): log: LogConfig = LogConfig() - monitors: EvalMonitorsConfig = EvalMonitorsConfig() + monitors: MonitorsConfig = MonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" @property @@ -172,12 +172,9 @@ class SFTOnlineEvalConfig(ScheduledEvalConfig, ServedEvalConfig): log: LogConfig = LogConfig() - monitors: EvalMonitorsConfig = EvalMonitorsConfig() - """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``). - - ``EvalMonitorsConfig``, not the base ``MonitorsConfig``: the online-eval - process logs to the trainer's Prime run, and the base type would drop (or - forbid) the prime block at the eval.json dump/re-parse boundary.""" + monitors: MonitorsConfig = MonitorsConfig() + """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``): the + trainer's, so the process logs to the trainer's W&B and Prime runs.""" @model_validator(mode="after") def auto_setup_broadcasts_dir(self): diff --git a/packages/prime-rl-configs/src/prime_rl/configs/monitors.py b/packages/prime-rl-configs/src/prime_rl/configs/monitors.py index cc3cde5f9f0..822bf0bb5ef 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/monitors.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/monitors.py @@ -55,15 +55,7 @@ class MonitorsConfig(BaseConfig): file: FileMonitorConfig | None = FileMonitorConfig() """Log metrics and episode traces to the run's output directory. On by default; disable with ``--no-monitors.file``.""" - -class TrainMonitorsConfig(MonitorsConfig): - prime: PrimeMonitorConfig | None = PrimeMonitorConfig() - """Log metrics and episodes to the Prime Intellect platform. On by default and - needs a Prime API key; disable with ``--no-monitors.prime``.""" - - -class EvalMonitorsConfig(MonitorsConfig): prime: PrimeMonitorConfig | None = PrimeMonitorConfig() - """Upload every eval source's epoch as an evaluation on the Prime Intellect platform (an - SFT online-eval process logs to the trainer's run instead). On by default and needs a - Prime API key; disable with ``--no-monitors.prime``.""" + """Log metrics and episodes to the Prime Intellect platform: a training run's processes + log to one platform run, and a standalone eval uploads each source's epoch as an + evaluation. On by default and needs a Prime API key; disable with ``--no-monitors.prime``.""" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index c4634ecb3b3..7ee47e3b5ce 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -14,7 +14,7 @@ GRPOAlgoConfig, ) from prime_rl.configs.diff import LIVE_INSIDE -from prime_rl.configs.monitors import TrainMonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig from prime_rl.configs.shared import ( BaseModelConfig, ControlConfig, @@ -500,7 +500,7 @@ class OrchestratorConfig(BaseConfig): env_vars: EnvVars = Field({}, frozen=True) """Extra environment variables for the orchestrator process(es). Merged on top of the launcher defaults.""" - monitors: TrainMonitorsConfig = Field(TrainMonitorsConfig(), frozen=True) + monitors: MonitorsConfig = Field(MonitorsConfig(), frozen=True) """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" collect_inference_metrics: bool = Field(True, frozen=True) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index a9f25597ea2..1d8fdbcdc51 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -11,7 +11,7 @@ from prime_rl.configs.eval import SFTOnlineEvalConfig from prime_rl.configs.inference import InferenceConfig from prime_rl.configs.inference import WeightBroadcastConfig as InferenceWeightBroadcastConfig -from prime_rl.configs.monitors import TrainMonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig from prime_rl.configs.shared import ( BaseMultiNodeDeploymentConfig, BaseSingleNodeDeploymentConfig, @@ -217,7 +217,7 @@ class SFTConfig(BaseTrainerConfig): resume: ResumeConfig | None = None """Resume the run from a checkpoint (point at it with the previous run's ``run.name``). Without ``[ckpt]`` the run loads the checkpoint but saves no new ones. If None, does not resume.""" - monitors: TrainMonitorsConfig = TrainMonitorsConfig() + monitors: MonitorsConfig = MonitorsConfig() """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" run: RunConfig = Field(default_factory=RunConfig) diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index dd61a6c06f5..8e62325fec6 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -5,7 +5,6 @@ from pathlib import Path from prime_rl.configs.eval import SFTOnlineEvalConfig -from prime_rl.configs.monitors import EvalMonitorsConfig, TrainMonitorsConfig from prime_rl.configs.orchestrator import OnlineEvalSourceConfig from prime_rl.configs.sft import SFTConfig from prime_rl.configs.shared import LogConfig @@ -71,12 +70,6 @@ def resolve_resume_step(config: SFTConfig) -> int | None: return resolve_latest_ckpt_step(get_ckpt_dir(get_ckpt_base(config))) -def build_online_eval_monitors(monitors: TrainMonitorsConfig) -> EvalMonitorsConfig: - """The online-eval process's monitors: the trainer's, so it logs to the trainer's - W&B and Prime runs.""" - return EvalMonitorsConfig(wandb=monitors.wandb, file=monitors.file, prime=monitors.prime) - - def build_online_eval_config(config: SFTConfig) -> SFTOnlineEvalConfig: """The online-eval process's config: the resolved ``[eval]`` block with the run-level fields filled from the SFT config. The launcher spawns the env servers itself; the @@ -97,7 +90,7 @@ def build_online_eval_config(config: SFTConfig) -> SFTOnlineEvalConfig: # by directory separation. output_dir=config.run_dir, log=LogConfig(level=config.log.level, json_logging=config.log.json_logging), - monitors=build_online_eval_monitors(config.monitors), + monitors=config.monitors, ) return SFTOnlineEvalConfig(**{**eval_config.model_dump(exclude=set(run_fields)), **run_fields}) @@ -220,6 +213,7 @@ def write_slurm_script( gpus_per_node=config.deployment.gpus_per_node, ranks_filter=",".join(map(str, config.log.ranks_filter)), prl_run_id=prl_run_id, + prime_run_id=os.environ.get("RUN_ID"), wandb_shared=wandb_shared, run_name=config.run.name, online_eval=online_eval, diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index 3579a093b61..3f05ca54fb4 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -65,6 +65,10 @@ mkdir -p "$LAUNCHER_DIR" rm -f "$TRAINER_DONE"{% if online_eval %} "$EVAL_DONE"{% endif %} export PRL_RUN_ID={{ prl_run_id }} +{% if prime_run_id %} +# The Prime run the launcher registered: every process attaches to it. +export RUN_ID={{ prime_run_id }} +{% endif %} export PRL_RUN_NAME={{ run_name }} {% if wandb_shared %} # Trainer and online-eval log to one shared W&B run. The trainer creates the run, diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index 4dd09a418e7..98e0731e39a 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -12,13 +12,6 @@ import pytest from prime_rl import monitors -from prime_rl.configs.monitors import ( - EvalMonitorsConfig, - FileMonitorConfig, - PrimeMonitorConfig, - TrainMonitorsConfig, - WandbMonitorConfig, -) from prime_rl.configs.sft import SFTConfig from prime_rl.monitors.prime import PrimeMonitor @@ -110,24 +103,6 @@ def test_sft_registers_wandb_project(prime_init, tmp_path): assert prime_init[0]["training"].wandb_project == "sft-tests" -def test_online_eval_monitors_keep_the_trainers_prime_monitor(): - """The online-eval process logs to the trainer's Prime run, under the same config.""" - from prime_rl.entrypoints.sft import build_online_eval_monitors - - monitors = build_online_eval_monitors( - TrainMonitorsConfig( - wandb=WandbMonitorConfig(project="sft-tests"), - file=FileMonitorConfig(), - prime=PrimeMonitorConfig(name="my-run"), - ) - ) - - assert isinstance(monitors, EvalMonitorsConfig) - assert monitors.prime == PrimeMonitorConfig(name="my-run") - assert monitors.wandb.project == "sft-tests" - assert build_online_eval_monitors(TrainMonitorsConfig(prime=None)).prime is None - - def test_sft_setup_registers_prime_monitor(prime_init, tmp_path, clear_monitors): config = _sft_config({"monitors": {"prime": {}}}) @@ -148,34 +123,6 @@ def test_sft_setup_registers_prime_monitor(prime_init, tmp_path, clear_monitors) asyncio.run(monitors.finalize()) -def test_online_eval_config_keeps_prime_monitor(): - """SFTOnlineEvalConfig.monitors must be EvalMonitorsConfig so the - trainer's prime monitor survives the eval.json dump/re-parse boundary - (base MonitorsConfig drops/forbids prime).""" - from prime_rl.configs.eval import SFTOnlineEvalConfig - from prime_rl.configs.monitors import EvalMonitorsConfig - from prime_rl.entrypoints.sft import build_online_eval_monitors - - field = SFTOnlineEvalConfig.model_fields["monitors"] - assert field.annotation is EvalMonitorsConfig - - train_monitors = TrainMonitorsConfig.model_validate({"prime": {"name": "r1"}}) - eval_monitors = build_online_eval_monitors(train_monitors) - assert isinstance(eval_monitors.prime, PrimeMonitorConfig) - assert eval_monitors.prime.name == "r1" - - # The online-eval process passes the prime monitor to monitors.setup - # (src/prime_rl/eval/online.py) — same wiring as the standalone eval - # entrypoint. Read the source text (importing the module pulls torch, - # which this torch-free test env deliberately lacks). - from pathlib import Path - - import prime_rl - - online_src = Path(prime_rl.__path__[0], "eval", "online.py").read_text() - assert "prime=config.monitors.prime" in online_src - - @pytest.fixture def sdk_runs(monkeypatch): """Fake ``pr.init`` runs that record their kwargs and how they ended.""" From a33bdec44d3490bd37cdfd45e3824ad586320bdc Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 28 Oct 2026 23:42:09 +0000 Subject: [PATCH 06/12] drop the prime monitor's redundant helpers - Key the links record by the monitor's config field ("prime", "wandb"); run_links takes the MonitorsConfig. - Write the platform record through update_record instead of four wrappers. - Read $PRIME_RUNS_MODE inline; register turns the SDK's missing-key error into the login hint. - Fold registration() into register(). Co-Authored-By: Claude Opus 5.5 --- src/prime_rl/entrypoints/dashboard.py | 26 +-- src/prime_rl/entrypoints/eval.py | 4 +- src/prime_rl/entrypoints/rl.py | 5 +- src/prime_rl/entrypoints/sft.py | 5 +- src/prime_rl/monitors/prime.py | 175 ++++++++------------- src/prime_rl/monitors/wandb/monitor.py | 6 +- src/prime_rl/utils/links.py | 3 - tests/unit/train/sft/test_prime_monitor.py | 34 ++-- tests/unit/utils/test_links.py | 4 +- 9 files changed, 107 insertions(+), 155 deletions(-) diff --git a/src/prime_rl/entrypoints/dashboard.py b/src/prime_rl/entrypoints/dashboard.py index 7292dc503a6..1428816f532 100644 --- a/src/prime_rl/entrypoints/dashboard.py +++ b/src/prime_rl/entrypoints/dashboard.py @@ -26,12 +26,14 @@ import urllib.request from contextlib import contextmanager from pathlib import Path -from typing import Callable +from typing import TYPE_CHECKING, Callable -from prime_rl.utils.links import PRIME_LINKS, WANDB_LINKS from prime_rl.utils.logger import get_logger from prime_rl.utils.pathing import CACHE_DIR +if TYPE_CHECKING: + from prime_rl.configs.monitors import MonitorsConfig + STATE_DIR = CACHE_DIR / "dashboard" DAEMON_FILE = STATE_DIR / "daemon.json" DIRS_FILE = STATE_DIR / "dirs.json" @@ -118,7 +120,7 @@ def ensure_dashboard(output_dir: Path, logger) -> str | None: return None -LINK_NAMES = {PRIME_LINKS: "Prime", WANDB_LINKS: "W&B"} +LINK_NAMES = {"prime": "Prime", "wandb": "W&B"} def log_run_links(dashboard_url: str | None, links: Callable[[], dict[str, dict[str, str]]] | None = None) -> None: @@ -145,21 +147,19 @@ def show(name: str, url: str, suffix: str) -> None: def run_links( - run_dir: Path, attempt_log_dir: Path, monitors: dict[str, object | None], stop: Callable[[], bool] + run_dir: Path, attempt_log_dir: Path, monitors: "MonitorsConfig", stop: Callable[[], bool] ) -> Callable[[], dict[str, dict[str, str]]] | None: - """The ``links`` argument of ``log_run_links``: waits for every configured monitor - in ``monitors`` (links-record key -> its config, None when off) that can open a - page. A monitor turned off by its mode variable (``$PRIME_RUNS_MODE``, an offline + """The ``links`` argument of ``log_run_links``: waits for every monitor in + ``monitors`` that is on and can open a page, keyed by its config field. A monitor turned off by its mode variable (``$PRIME_RUNS_MODE``, an offline or disabled ``$WANDB_MODE``) opens none; one without credentials fails its component, which ends the wait through ``stop``.""" from prime_rl.utils.links import wait_for_links - def opens_pages(key: str) -> bool: - if key == WANDB_LINKS: - return os.environ.get("WANDB_MODE") not in ("offline", "disabled") - return os.environ.get("PRIME_RUNS_MODE", "online") == "online" - - waited = [key for key, config in monitors.items() if config is not None and opens_pages(key)] + opens_pages = { + "prime": monitors.prime is not None and (os.environ.get("PRIME_RUNS_MODE") or "online") == "online", + "wandb": monitors.wandb is not None and os.environ.get("WANDB_MODE") not in ("offline", "disabled"), + } + waited = [key for key, opens in opens_pages.items() if opens] if not waited: return None return functools.partial(wait_for_links, run_dir, attempt_log_dir, waited, stop) diff --git a/src/prime_rl/entrypoints/eval.py b/src/prime_rl/entrypoints/eval.py index 9dc3aa68bb7..0101eb525f0 100644 --- a/src/prime_rl/entrypoints/eval.py +++ b/src/prime_rl/entrypoints/eval.py @@ -121,7 +121,6 @@ def main(): sys.argv = [sys.argv[0], *argv] from prime_rl.entrypoints.dashboard import ensure_dashboard, log_run_links, run_links - from prime_rl.utils.links import PRIME_LINKS, WANDB_LINKS from prime_rl.utils.logger import setup_logger from prime_rl.utils.pathing import ( format_config_message, @@ -195,8 +194,7 @@ def sigterm_handler(signum, frame): # Each source's Prime evaluation opens once the eval counts its tasks, so the # banner waits for the links next to the running eval. eval_done = threading.Event() - monitors = {PRIME_LINKS: config.monitors.prime, WANDB_LINKS: config.monitors.wandb} - links = run_links(config.run_dir, log_dir, monitors, eval_done.is_set) + links = run_links(config.run_dir, log_dir, config.monitors, eval_done.is_set) banner = threading.Thread(target=log_run_links, args=(dashboard_url, links), daemon=True) banner.start() try: diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 789c5e0bb03..b83b0a5b31e 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -12,7 +12,6 @@ from prime_rl.entrypoints.dashboard import ensure_dashboard, log_run_links, run_links from prime_rl.entrypoints.inference import mooncake_offloads, vllm_overrides_fragment, write_pd_configs from prime_rl.utils.config import cli, dump_resolved_config, dump_rl_components -from prime_rl.utils.links import PRIME_LINKS, WANDB_LINKS from prime_rl.utils.logger import get_logger, setup_logger from prime_rl.utils.pathing import ( clean_future_steps, @@ -254,8 +253,8 @@ def rl_local(config: RLConfig): ) logger.success("Launcher complete") - monitors = {PRIME_LINKS: config.orchestrator.monitors.prime, WANDB_LINKS: config.orchestrator.monitors.wandb} - log_run_links(dashboard_url, run_links(config.run_dir, log_dir, monitors, lambda: bool(processes.error_queue))) + links = run_links(config.run_dir, log_dir, config.orchestrator.monitors, lambda: bool(processes.error_queue)) + log_run_links(dashboard_url, links) # Trainer and orchestrator completion is the successful stop condition. processes.wait("trainer", "orchestrator") diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 8e62325fec6..234f39e35a8 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -10,7 +10,6 @@ from prime_rl.configs.shared import LogConfig from prime_rl.entrypoints.dashboard import ensure_dashboard, log_run_links, run_links from prime_rl.utils.config import cli, dump_resolved_config, find_package_resource -from prime_rl.utils.links import PRIME_LINKS, WANDB_LINKS from prime_rl.utils.logger import setup_logger from prime_rl.utils.pathing import ( clean_future_steps, @@ -427,8 +426,8 @@ def sft_local(config: SFTConfig): ) logger.success("Launcher complete") - monitors = {PRIME_LINKS: config.monitors.prime, WANDB_LINKS: config.monitors.wandb} - log_run_links(dashboard_url, run_links(config.run_dir, log_dir, monitors, lambda: bool(processes.error_queue))) + links = run_links(config.run_dir, log_dir, config.monitors, lambda: bool(processes.error_queue)) + log_run_links(dashboard_url, links) # Wait for the trainer (and the online-eval process, which drains its final evals # after the trainer's last checkpoint) while surfacing any process failure. diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index f7f6a5f2b2e..192b9d44eb6 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -6,10 +6,9 @@ from collections.abc import Iterator from contextlib import contextmanager from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable +from typing import TYPE_CHECKING, Any import prime_runs as pr -from prime_sandboxes import Config as PrimeConfig from prime_rl.configs.eval import EvalConfig, SFTOnlineEvalConfig from prime_rl.configs.monitors import PrimeMonitorConfig @@ -17,10 +16,10 @@ from prime_rl.configs.sft import SFTConfig from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.utils.config import BaseConfig -from prime_rl.utils.links import PRIME_LINKS, publish_link, register_links +from prime_rl.utils.links import publish_link, register_links from prime_rl.utils.pathing import get_platform_run_path from prime_rl.utils.process import RUN_FINISHER_VAR -from prime_rl.utils.record import read_record, update_record, write_record +from prime_rl.utils.record import update_record from prime_rl.utils.utils import sanitize if TYPE_CHECKING: @@ -34,49 +33,9 @@ FINISH_TIMEOUT = 60.0 -def write_platform_record(output_dir: Path, record: dict[str, Any]) -> None: - """Leave the run's platform identity in the run directory for the dashboard's - "View on Prime" link.""" - write_record(get_platform_run_path(output_dir), record) - - -def read_platform_record(output_dir: Path) -> dict[str, Any] | None: - return read_record(get_platform_run_path(output_dir)) - - -def _merge_platform_record( - output_dir: Path, - merge: Callable[[dict[str, Any]], dict[str, Any]], -) -> None: - """Locked read-modify-write of the platform record: the trainer and the - online-eval process merge into the SAME run.json. Best-effort, see - ``update_record``.""" - update_record(get_platform_run_path(output_dir), merge) - - -def update_platform_record(output_dir: Path, update: dict[str, Any]) -> None: - """Best-effort locked read-modify-write of the platform record: `update` - overlays the current record (preserving keys it doesn't set). Use for - every writer that shares the file — a naive overwrite drops concurrent - writers' keys.""" - _merge_platform_record(output_dir, lambda record: {**record, **update}) - - OFF_HINT = "turn the monitor off with --no-monitors.prime" -def prime_mode() -> str: - """$PRIME_RUNS_MODE when set, else online. Like the W&B monitor, a configured - monitor without credentials fails the run instead of leaving it untracked.""" - if mode := os.getenv(pr.MODE_ENV): - return mode - if not (os.getenv("PRIME_API_KEY") or PrimeConfig().api_key): - raise RuntimeError( - f"The prime monitor found no Prime API key. Authenticate with `prime login` (or set PRIME_API_KEY), or {OFF_HINT}" - ) - return "online" - - def _base_url() -> str | None: """$PRIME_API_BASE historically points at the RFT API root (``.../api/v1/rft``); the SDK takes the platform base URL. Unset means the SDK resolves it.""" @@ -84,46 +43,56 @@ def _base_url() -> str | None: return base.rstrip("/").removesuffix("/rft") if base else None -def registration(config: OrchestratorConfig | SFTConfig, name: str | None) -> dict[str, Any]: - """The ``pr.init`` arguments that register ``config``'s training run.""" - if isinstance(config, SFTConfig): - # SFT trains on a dataset, not on environments, and has no rollouts - the - # platform run is registered without both. - environments = [] - training = pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.data.batch_size, - seq_len=config.data.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, - ) +def register( + config: OrchestratorConfig | SFTConfig | None, name: str | None, mode: str, id: str | None = None +) -> pr.Run: + """Attach to the platform run ``id``, or register ``config``'s training run.""" + init_kwargs: dict[str, Any] + if id is not None: + init_kwargs = {"id": id} + elif config is None: + # The RFT API requires a base model. + init_kwargs = {"name": name, "model": "unknown"} else: - environments = [env.env_id for env in config.train.source] - training = pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.batch_size, - rollouts_per_example=max( - (source.group_size for source in config.train.source), default=config.train.group_size - ), - seq_len=config.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + wandb_project = config.monitors.wandb.project if config.monitors.wandb else None + if isinstance(config, SFTConfig): + # SFT trains on a dataset, not on environments, and has no rollouts. + environments = [] + training = pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.data.batch_size, + seq_len=config.data.seq_len, + wandb_project=wandb_project, + ) + else: + environments = [env.env_id for env in config.train.source] + training = pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.batch_size, + rollouts_per_example=max( + (source.group_size for source in config.train.source), default=config.train.group_size + ), + seq_len=config.seq_len, + wandb_project=wandb_project, + ) + init_kwargs = dict( + name=name, + model=config.model.name, + environments=environments, + training=training, + config=config.model_dump(exclude_none=True, mode="json"), ) - return dict( - name=name, - model=config.model.name, - environments=environments, - training=training, - config=config.model_dump(exclude_none=True, mode="json"), - ) - - -def register(init_kwargs: dict[str, Any], mode: str) -> pr.Run: try: return pr.init(kind="train", mode=mode, base_url=_base_url(), finish_timeout=FINISH_TIMEOUT, **init_kwargs) except (pr.ConfigurationError, pr.APIError) as e: - # Registering fails without a team, or for a team outside the external-runs - # allowlist. An attached run's error is the launcher's to explain. - if "id" in init_kwargs: + # An attached run's error is the launcher's to explain. + if id is not None: raise + if "needs an API key" in str(e): + raise RuntimeError( + f"The prime monitor found no Prime API key. Authenticate with `prime login` (or set PRIME_API_KEY), or {OFF_HINT}" + ) from e + # Registering fails without a team, or for a team outside the external-runs allowlist. raise RuntimeError( f"Failed to register the Prime run ({type(e).__name__}: {e}). Select a team with " f"`prime switch` (or set PRIME_TEAM_ID), or {OFF_HINT}" @@ -136,10 +105,10 @@ def launch_run(config: OrchestratorConfig | SFTConfig, prime: PrimeMonitorConfig it: its id becomes ``$PRL_RUN_ID`` and ``$RUN_ID``. A launch that fails marks the run failed - the SDK leaves an attached run's failure to the process that registered it. A run a managed launch pre-created (``$RUN_ID``) is the platform's.""" - if prime is None or os.getenv("RUN_ID") or prime_mode() != "online": + if prime is None or os.getenv("RUN_ID") or (os.getenv(pr.MODE_ENV) or "online") != "online": yield return - run = register(registration(config, prime.name), "online") + run = register(config, prime.name, "online") os.environ["PRL_RUN_ID"] = os.environ["RUN_ID"] = run.id try: yield @@ -181,7 +150,7 @@ class PrimeMonitor(Monitor): async def init(self, config: BaseConfig | None, output_dir: Path, producer: str | None = None) -> None: self.output_dir = output_dir self.producer = producer - self.mode = mode = prime_mode() + self.mode = mode = os.getenv(pr.MODE_ENV) or "online" self.run: pr.Run | None = None # (env, step) -> its open evaluation; None once opening it failed, so the epoch's # episodes do not retry the platform on every arrival @@ -190,24 +159,14 @@ async def init(self, config: BaseConfig | None, output_dir: Path, producer: str await self._init_evaluations(config) return online_eval = isinstance(config, SFTOnlineEvalConfig) - if run_id := os.getenv("RUN_ID"): - # The launcher registered the run (or a managed launch pre-created it) - - # attach instead of registering a duplicate. The registering process owns - # the attached run's failure marking; a clean finish() completes it. - init_kwargs = {"id": run_id} - elif online_eval: - if mode == "online": - raise RuntimeError( - f"An online-eval process logs to its trainer's run: launch it with `sft`, or {OFF_HINT}" - ) - init_kwargs = {} - elif config is not None: - init_kwargs = registration(config, self.config.name) - else: - # The RFT API requires a base model; "unknown" is what the - # pre-SDK register sent for a config-less init. - init_kwargs = {"name": self.config.name, "model": "unknown"} - self.run = await asyncio.to_thread(register, init_kwargs, mode) + run_id = os.getenv("RUN_ID") + if online_eval and run_id is None and mode == "online": + raise RuntimeError(f"An online-eval process logs to its trainer's run: launch it with `sft`, or {OFF_HINT}") + # The launcher registered the run (or a managed launch pre-created it) - attach + # instead of registering a duplicate. The registering process owns the attached + # run's failure marking; a clean finish() completes it. + train_config = None if online_eval else config + self.run = await asyncio.to_thread(register, train_config, self.config.name, mode, run_id) if self.run.url: attached = " (attached via $RUN_ID)" if self.run.attached else "" self.logger.info(f"Logging metrics and episodes to Prime run {self.run.id} ({self.run.url}){attached}") @@ -218,10 +177,11 @@ async def init(self, config: BaseConfig | None, output_dir: Path, producer: str await asyncio.to_thread(self._record_run) def _record_run(self) -> None: - register_links(self.output_dir, PRIME_LINKS, 1 if self.run.url else 0) + register_links(self.output_dir, "prime", 1 if self.run.url else 0) if self.run.url: - update_platform_record(self.output_dir, {"kind": "train", "id": self.run.id, "url": self.run.url}) - publish_link(self.output_dir, PRIME_LINKS, "run", self.run.url) + record = {"kind": "train", "id": self.run.id, "url": self.run.url} + update_record(get_platform_run_path(self.output_dir), lambda _: record) + publish_link(self.output_dir, "prime", "run", self.run.url) async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: # A standalone eval's evaluations carry their own aggregates (``log_eval_epoch``). @@ -294,8 +254,9 @@ async def _init_evaluations(self, config: EvalConfig) -> None: def record() -> None: if online: - write_platform_record(self.output_dir, {"kind": "eval", "run_id": self.run_id, "evaluations": {}}) - register_links(self.output_dir, PRIME_LINKS, len(self.sources) if online else 0) + record = {"kind": "eval", "run_id": self.run_id, "evaluations": {}} + update_record(get_platform_run_path(self.output_dir), lambda _: record) + register_links(self.output_dir, "prime", len(self.sources) if online else 0) # Off the event loop: the flock retry must not stall it. await asyncio.to_thread(record) @@ -346,7 +307,7 @@ async def run_for(self, env_name: str, step: int, expected: int | None = None) - except Exception as e: self.logger.warning(f"Failed to open the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") self.runs[key] = None - await asyncio.to_thread(publish_link, self.output_dir, PRIME_LINKS, env_name, None) + await asyncio.to_thread(publish_link, self.output_dir, "prime", env_name, None) return None self.runs[key] = run if run.url: @@ -365,8 +326,8 @@ def update(record: dict[str, Any]) -> dict[str, Any]: record.setdefault("evaluations", {})[env_name] = {"step": step, "id": run.id, "url": run.url} return record - _merge_platform_record(self.output_dir, update) - publish_link(self.output_dir, PRIME_LINKS, env_name, run.url) + update_record(get_platform_run_path(self.output_dir), update) + publish_link(self.output_dir, "prime", env_name, run.url) async def log_eval_plan(self, env_name: str, step: int, expected: int) -> None: if self.run is None: diff --git a/src/prime_rl/monitors/wandb/monitor.py b/src/prime_rl/monitors/wandb/monitor.py index c85bf8c7baf..138858f908d 100644 --- a/src/prime_rl/monitors/wandb/monitor.py +++ b/src/prime_rl/monitors/wandb/monitor.py @@ -16,7 +16,7 @@ from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.monitors.wandb.overview import ensure_overview_view from prime_rl.utils.config import BaseConfig -from prime_rl.utils.links import WANDB_LINKS, publish_link, register_links +from prime_rl.utils.links import publish_link, register_links from prime_rl.utils.logger import format_time if TYPE_CHECKING: @@ -149,9 +149,9 @@ def init_wandb(max_retries: int): @staticmethod def _record(output_dir: Path, url: str | None) -> None: - register_links(output_dir, WANDB_LINKS, 1 if url else 0) + register_links(output_dir, "wandb", 1 if url else 0) if url: - publish_link(output_dir, WANDB_LINKS, "run", url) + publish_link(output_dir, "wandb", "run", url) async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: # every log carries the monitor's own wall-time stamp diff --git a/src/prime_rl/utils/links.py b/src/prime_rl/utils/links.py index f2304d8f910..9f46a7ef8cb 100644 --- a/src/prime_rl/utils/links.py +++ b/src/prime_rl/utils/links.py @@ -16,9 +16,6 @@ # The launcher exports it to every component: entries stamped with another attempt are # a previous launch's. ATTEMPT_VAR = "PRL_ATTEMPT_LOG_DIR" -# Monitor keys in the links record. -PRIME_LINKS = "prime" -WANDB_LINKS = "wandb" LINK_TIMEOUT_S = 300.0 LINK_POLL_S = 0.5 diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index 98e0731e39a..98b857c7ab2 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -211,13 +211,11 @@ def test_platform_record_lock_keeps_a_stable_inode(tmp_path): a lost merge. The lock file must persist with the SAME inode across acquisitions, and the locked update helper must merge two sequential writers into one record. (Interprocess contention is exercised by the - two-process update_platform_record probe; here the writers are - sequential.)""" + two-process update_record probe; here the writers are sequential.)""" import os - from prime_rl.monitors.prime import update_platform_record from prime_rl.utils.pathing import get_platform_run_path - from prime_rl.utils.record import record_lock + from prime_rl.utils.record import read_record, record_lock, update_record lock_path = get_platform_run_path(tmp_path).with_suffix(".json.lock") inodes = [] @@ -232,11 +230,11 @@ def test_platform_record_lock_keeps_a_stable_inode(tmp_path): assert len(set(inodes)) == 1 # Two sequential locked updates merge instead of clobbering. - update_platform_record(tmp_path, {"kind": "train", "id": "run-1", "url": "https://x/run-1"}) - update_platform_record(tmp_path, {"evaluations": {"rev": {"step": 5, "id": "ev-1", "url": "https://x/ev-1"}}}) - from prime_rl.monitors.prime import read_platform_record + path = get_platform_run_path(tmp_path) + update_record(path, lambda record: {**record, "kind": "train", "id": "run-1", "url": "https://x/run-1"}) + update_record(path, lambda record: {**record, "evaluations": {"rev": {"step": 5, "id": "ev-1"}}}) - record = read_platform_record(tmp_path) + record = read_record(path) assert record["kind"] == "train" assert record["id"] == "run-1" assert record["evaluations"]["rev"]["id"] == "ev-1" @@ -252,13 +250,13 @@ def test_record_persistence_runs_off_the_event_loop(prime_init, tmp_path, monkey import prime_rl.monitors.prime as prime_module seen_threads = [] - real_merge = prime_module._merge_platform_record + real_update = prime_module.update_record - def spy(output_dir, merge): + def spy(path, merge): seen_threads.append(threading.get_ident()) - return real_merge(output_dir, merge) + return real_update(path, merge) - monkeypatch.setattr(prime_module, "_merge_platform_record", spy) + monkeypatch.setattr(prime_module, "update_record", spy) config = _sft_config({"monitors": {"prime": {}}}) asyncio.run(PrimeMonitor(config.monitors.prime).init(config=config, output_dir=tmp_path)) @@ -275,9 +273,10 @@ def spy(output_dir, merge): src = Path(prime_rl.__path__[0], "monitors", "prime.py").read_text() assert len(re.findall(r"asyncio\.to_thread\(\s*self\._record", src)) == 2 - from prime_rl.monitors.prime import read_platform_record + from prime_rl.utils.pathing import get_platform_run_path + from prime_rl.utils.record import read_record - record = read_platform_record(tmp_path) + record = read_record(get_platform_run_path(tmp_path)) assert record["kind"] == "train" assert record["id"] == "run-123" @@ -309,13 +308,14 @@ def test_platform_record_persistence_is_best_effort(tmp_path, monkeypatch): skipped with a warning, and no record is left behind.""" import errno - import prime_rl.monitors.prime as prime_module import prime_rl.utils.record as record_module + from prime_rl.utils.pathing import get_platform_run_path def enolck(fd, op): raise OSError(errno.ENOLCK, "No locks available") monkeypatch.setattr(record_module.fcntl, "flock", enolck) + path = get_platform_run_path(tmp_path) # Must not raise: the platform record is best-effort. - prime_module.update_platform_record(tmp_path, {"kind": "train", "id": "run-1", "url": "https://x/run-1"}) - assert prime_module.read_platform_record(tmp_path) is None + record_module.update_record(path, lambda _: {"kind": "train", "id": "run-1", "url": "https://x/run-1"}) + assert record_module.read_record(path) is None diff --git a/tests/unit/utils/test_links.py b/tests/unit/utils/test_links.py index f4914fb0d7d..1546b19f3c8 100644 --- a/tests/unit/utils/test_links.py +++ b/tests/unit/utils/test_links.py @@ -1,10 +1,8 @@ import asyncio import time -from types import SimpleNamespace import pytest -import prime_rl.monitors.prime as prime_module from prime_rl.configs.monitors import PrimeMonitorConfig from prime_rl.monitors.prime import PrimeMonitor from prime_rl.utils.links import ATTEMPT_VAR, publish_link, register_links, wait_for_links @@ -46,6 +44,6 @@ def test_prime_monitor_without_credentials_fails_informatively(tmp_path, monkeyp monkeypatch.delenv("PRIME_API_KEY", raising=False) monkeypatch.delenv("PRIME_RUNS_MODE", raising=False) monkeypatch.delenv("RUN_ID", raising=False) - monkeypatch.setattr(prime_module, "PrimeConfig", lambda: SimpleNamespace(api_key=None)) + monkeypatch.setenv("HOME", str(tmp_path)) with pytest.raises(RuntimeError, match=r"prime login.*--no-monitors\.prime"): asyncio.run(PrimeMonitor(PrimeMonitorConfig()).init(config=None, output_dir=tmp_path)) From f492450de111b44c1675cd2d28c1f31eb9e8f449 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 29 Oct 2026 00:18:21 +0000 Subject: [PATCH 07/12] name the orchestrator's monitor producer orchestrator The producer stamps each metrics row and names the orchestrator's annotation stream, so it now matches the component name used everywhere else (orchestrator.log, PRL_RUN_FINISHER, W&B labels). Co-Authored-By: Claude Opus 5.5 --- skills/training/monitor-run/SKILL.md | 2 +- src/prime_rl/dashboard/static/app.js | 4 ++-- src/prime_rl/monitors/file/traces/__init__.py | 2 +- src/prime_rl/orchestrator/orchestrator.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 3029543f36d..a3928a0e0d0 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -217,7 +217,7 @@ stream records, newest winning. wc -l {run_dir}/monitors/file/traces/stream.index.jsonl zstd -dcf {run_dir}/monitors/file/traces/stream/* | jq '.traces[].rewards' zstd -dcf {run_dir}/monitors/file/traces/stream/* | jq 'select(.ok | not) | {id, env: .env.id, errors}' -jq '{trace_id, info}' {run_dir}/monitors/file/traces/annotations/orch.index.jsonl +jq '{trace_id, info}' {run_dir}/monitors/file/traces/annotations/orchestrator.index.jsonl ``` The batches consumed by the trainer are shipped over ZMQ by default, so nothing binary is written. With `rollout_transport.type = "filesystem"` they land at `{run_dir}/batches/step_{n}/rank_.bin` (one packed micro-batch file per trainer DP rank). diff --git a/src/prime_rl/dashboard/static/app.js b/src/prime_rl/dashboard/static/app.js index 115b8d562b4..3ba126a5d85 100644 --- a/src/prime_rl/dashboard/static/app.js +++ b/src/prime_rl/dashboard/static/app.js @@ -469,9 +469,9 @@ function rowProducer(row, meta) { if (meta?.type === "sft") return "trainer"; for (const key of Object.keys(row)) { if (TRAINER_KEY_RE.test(key)) return "trainer"; - if (ORCH_KEY_RE.test(key)) return "orch"; + if (ORCH_KEY_RE.test(key)) return "orchestrator"; } - return Object.keys(row).some((k) => k.startsWith("progress/")) ? "orch" : "trainer"; + return Object.keys(row).some((k) => k.startsWith("progress/")) ? "orchestrator" : "trainer"; } /* store = {byKey, timeKeys, timeZero} — the primary run's is state.metrics, diff --git a/src/prime_rl/monitors/file/traces/__init__.py b/src/prime_rl/monitors/file/traces/__init__.py index 0c0b52b56a6..73b46bbd94a 100644 --- a/src/prime_rl/monitors/file/traces/__init__.py +++ b/src/prime_rl/monitors/file/traces/__init__.py @@ -20,7 +20,7 @@ def get_trace_stream(output_dir: Path) -> Path: def get_annotations_dir(output_dir: Path) -> Path: - """One stream per producer of trace updates: orch ship-time facts, trainer streams.""" + """One stream per producer of trace updates: orchestrator ship-time facts, trainer streams.""" return get_trace_dir(output_dir) / "annotations" diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index a0b364f3136..50f72d21472 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -226,7 +226,7 @@ async def setup(self) -> None: ) await monitors.setup( - producer="orch", + producer="orchestrator", wandb=config.monitors.wandb, prime=config.monitors.prime, file=config.monitors.file, From 6058999ae39558e8919b927551f5144d962ffce7 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 29 Oct 2026 22:05:43 +0000 Subject: [PATCH 08/12] fail a standalone eval at launch without a prime key Evaluations open when their epoch starts, so a missing key only showed as a warning per env and the eval finished with nothing on Prime. Also name the dashboard's orchestrator key pattern in full. Co-Authored-By: Claude Opus 5.5 --- src/prime_rl/dashboard/static/app.js | 4 ++-- src/prime_rl/monitors/prime.py | 11 ++++++++--- tests/unit/utils/test_links.py | 12 ++++++++++++ 3 files changed, 22 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/dashboard/static/app.js b/src/prime_rl/dashboard/static/app.js index 9412fae8146..eb8c3cf7ba7 100644 --- a/src/prime_rl/dashboard/static/app.js +++ b/src/prime_rl/dashboard/static/app.js @@ -461,7 +461,7 @@ const INFERENCE_PANELS = [ ]; const TRAINER_KEY_RE = /^(perf|optim|loss|entropy|system|mismatch_kl|kl_ent_ratio|is_masked|masked_|unmasked_|max_vio|routing_|ref_kl|val)[/_]?/; -const ORCH_KEY_RE = /^(train|batch|staleness|curriculum|eval)\//; +const ORCHESTRATOR_KEY_RE = /^(train|batch|staleness|curriculum|eval)\//; function rowProducer(row, meta) { if (typeof row.producer === "string") return row.producer; // stamped by FileMonitor @@ -469,7 +469,7 @@ function rowProducer(row, meta) { if (meta?.type === "sft") return "trainer"; for (const key of Object.keys(row)) { if (TRAINER_KEY_RE.test(key)) return "trainer"; - if (ORCH_KEY_RE.test(key)) return "orchestrator"; + if (ORCHESTRATOR_KEY_RE.test(key)) return "orchestrator"; } return Object.keys(row).some((k) => k.startsWith("progress/")) ? "orchestrator" : "trainer"; } diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 192b9d44eb6..65c7e80e3de 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any import prime_runs as pr +from prime_runs.config import Config from prime_rl.configs.eval import EvalConfig, SFTOnlineEvalConfig from prime_rl.configs.monitors import PrimeMonitorConfig @@ -34,6 +35,9 @@ OFF_HINT = "turn the monitor off with --no-monitors.prime" +NO_KEY = ( + f"The prime monitor found no Prime API key. Authenticate with `prime login` (or set PRIME_API_KEY), or {OFF_HINT}" +) def _base_url() -> str | None: @@ -89,9 +93,7 @@ def register( if id is not None: raise if "needs an API key" in str(e): - raise RuntimeError( - f"The prime monitor found no Prime API key. Authenticate with `prime login` (or set PRIME_API_KEY), or {OFF_HINT}" - ) from e + raise RuntimeError(NO_KEY) from e # Registering fails without a team, or for a team outside the external-runs allowlist. raise RuntimeError( f"Failed to register the Prime run ({type(e).__name__}: {e}). Select a team with " @@ -250,6 +252,9 @@ async def _init_evaluations(self, config: EvalConfig) -> None: f"exactly one eval source (got {len(self.sources)})" ) online = self.mode == "online" + # Evaluations open lazily, when their epoch starts; fail at launch instead. + if online and not Config().api_key: + raise RuntimeError(NO_KEY) self.logger.info("Streaming eval epochs to the Prime platform" if online else "Prime evaluations disabled") def record() -> None: diff --git a/tests/unit/utils/test_links.py b/tests/unit/utils/test_links.py index 1546b19f3c8..cef8227ff71 100644 --- a/tests/unit/utils/test_links.py +++ b/tests/unit/utils/test_links.py @@ -1,5 +1,6 @@ import asyncio import time +from types import SimpleNamespace import pytest @@ -47,3 +48,14 @@ def test_prime_monitor_without_credentials_fails_informatively(tmp_path, monkeyp monkeypatch.setenv("HOME", str(tmp_path)) with pytest.raises(RuntimeError, match=r"prime login.*--no-monitors\.prime"): asyncio.run(PrimeMonitor(PrimeMonitorConfig()).init(config=None, output_dir=tmp_path)) + + +def test_eval_prime_monitor_without_credentials_fails_at_launch(tmp_path, monkeypatch): + monkeypatch.delenv("PRIME_API_KEY", raising=False) + monkeypatch.delenv("PRIME_RUNS_MODE", raising=False) + monkeypatch.setenv("HOME", str(tmp_path)) + monitor = PrimeMonitor(PrimeMonitorConfig()) + monitor.output_dir, monitor.mode, monitor.runs = tmp_path, "online", {} + eval_config = SimpleNamespace(model="m", source=[]) + with pytest.raises(RuntimeError, match=r"prime login.*--no-monitors\.prime"): + asyncio.run(monitor._init_evaluations(eval_config)) From 00b293836bbb5b779c3db7d82b7efab63ee3ebda Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 29 Oct 2026 22:44:28 +0000 Subject: [PATCH 09/12] let the slurm job own its prime run The submitter registered the run and exited after sbatch, and the SDK never fails an attached run, so a failed SLURM job left its run RUNNING. The supervising process now registers the run: the local launcher, the in-job launcher on a single node, or a run owner that wraps a multi-node job's step and fails the run when the step fails or SLURM sends SIGTERM. Co-Authored-By: Claude Opus 5.5 --- docs/training.md | 2 +- src/prime_rl/entrypoints/rl.py | 10 +++--- src/prime_rl/entrypoints/run_owner.py | 31 +++++++++++++++++++ src/prime_rl/entrypoints/sft.py | 11 ++++--- src/prime_rl/monitors/prime.py | 4 ++- .../templates/multi_node_rl.sbatch.j2 | 8 +++-- .../templates/multi_node_sft.sbatch.j2 | 13 ++++---- tests/unit/train/sft/test_prime_monitor.py | 25 +++++++++++++++ 8 files changed, 85 insertions(+), 19 deletions(-) create mode 100644 src/prime_rl/entrypoints/run_owner.py diff --git a/docs/training.md b/docs/training.md index 2812616866a..b4605fdb659 100644 --- a/docs/training.md +++ b/docs/training.md @@ -491,7 +491,7 @@ Name the Prime run in TOML: name = "my-experiment" ``` -The `rl` and `sft` launchers register the run before they start the components, and every component logs to it. The run's id is also the run identity (`PRL_RUN_ID`) and the W&B run id. On `sft` the run carries the dataset-batched training fields (`max_steps`, `batch_size`, `seq_len`) and the trainer's per-step metrics and validation losses. Online evals log to the same run, like an `rl` run's: the online-eval process uploads its eval metrics and every eval episode, and finishes the run after its final evals. A launch that fails marks the run failed. +The process that supervises the launch registers the run before the components start, and every component logs to it. That is the `rl` or `sft` launcher, also inside a single-node SLURM job, or the run owner that wraps a multi-node SLURM job's step. The run's id is also the run identity (`PRL_RUN_ID`) and the W&B run id. On `sft` the run carries the dataset-batched training fields (`max_steps`, `batch_size`, `seq_len`) and the trainer's per-step metrics and validation losses. Online evals log to the same run, like an `rl` run's: the online-eval process uploads its eval metrics and every eval episode, and finishes the run after its final evals. A launch that fails marks the run failed, and so does a SLURM job that SLURM cancels or stops at its time limit. The platform run needs `PRIME_API_KEY` (`prime login` or the env var) and a team (`PRIME_TEAM_ID`, or the team selected with `prime login`) enabled for external runs. Like the W&B monitor, a run without a key, or one whose registration fails (no team, a team outside the allowlist), fails at launch: authenticate, select a team with `prime switch`, or pass `--no-monitors.prime`. `PRIME_RUNS_MODE=disabled` keeps the monitor configured but opens no platform run; `RUN_ID=` attaches to an external run a launcher already created instead of registering a new one. Currently internal-only. diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 99954761b1a..ab09b309fd3 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -591,11 +591,13 @@ def rl(config: RLConfig): from prime_rl.monitors.prime import launch_run + # On SLURM the job registers the run: the in-job `rl` (single node) or the + # batch script's run owner (multi node), so that a live process owns its status. + if config.slurm is not None: + rl_slurm(config) + return with launch_run(config.orchestrator, None if config.dry_run else config.orchestrator.monitors.prime): - if config.slurm is not None: - rl_slurm(config) - else: - rl_local(config) + rl_local(config) def main(): diff --git a/src/prime_rl/entrypoints/run_owner.py b/src/prime_rl/entrypoints/run_owner.py new file mode 100644 index 00000000000..43358cb8cfc --- /dev/null +++ b/src/prime_rl/entrypoints/run_owner.py @@ -0,0 +1,31 @@ +import signal +import subprocess +import sys +from pathlib import Path +from typing import Any + +from prime_rl.configs.orchestrator import OrchestratorConfig +from prime_rl.configs.sft import SFTConfig +from prime_rl.monitors.prime import launch_run + + +def main() -> None: + """``python -m prime_rl.entrypoints.run_owner {rl,sft} ``: run a multi-node + SLURM job's command as the owner of its Prime run (``launch_run``). The command's + non-zero exit, or SLURM's SIGTERM on cancel or time limit, fails the run.""" + kind, path, *command = sys.argv[1:] + config_cls = OrchestratorConfig if kind == "rl" else SFTConfig + config = config_cls.model_validate_json(Path(path).read_text()) + + def terminate(signum: int, frame: Any) -> None: + raise SystemExit(128 + signum) from RuntimeError("SLURM terminated the job") + + signal.signal(signal.SIGTERM, terminate) + with launch_run(config, config.monitors.prime): + code = subprocess.run(command).returncode + if code != 0: + raise SystemExit(code) from RuntimeError(f"The SLURM job step exited with code {code}") + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 234f39e35a8..4d5adad67c9 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -212,7 +212,6 @@ def write_slurm_script( gpus_per_node=config.deployment.gpus_per_node, ranks_filter=",".join(map(str, config.log.ranks_filter)), prl_run_id=prl_run_id, - prime_run_id=os.environ.get("RUN_ID"), wandb_shared=wandb_shared, run_name=config.run.name, online_eval=online_eval, @@ -474,11 +473,13 @@ def sft(config: SFTConfig): from prime_rl.monitors.prime import launch_run + # On SLURM the job registers the run: the in-job `sft` (single node) or the + # batch script's run owner (multi node), so that a live process owns its status. + if config.slurm is not None: + sft_slurm(config) + return with launch_run(config, None if config.dry_run else config.monitors.prime): - if config.slurm is not None: - sft_slurm(config) - else: - sft_local(config) + sft_local(config) def main(): diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 65c7e80e3de..82e3bbe3ad2 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -106,7 +106,9 @@ def launch_run(config: OrchestratorConfig | SFTConfig, prime: PrimeMonitorConfig """Register the launch's Prime run around starting its components, which attach to it: its id becomes ``$PRL_RUN_ID`` and ``$RUN_ID``. A launch that fails marks the run failed - the SDK leaves an attached run's failure to the process that - registered it. A run a managed launch pre-created (``$RUN_ID``) is the platform's.""" + registered it, so that process supervises the launch: the local ``rl`` or ``sft`` + launcher, or a multi-node SLURM job's ``run_owner``. A run a managed launch pre-created + (``$RUN_ID``) is the platform's.""" if prime is None or os.getenv("RUN_ID") or (os.getenv(pr.MODE_ENV) or "online") != "online": yield return diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 8bea2667936..98d64f178dd 100644 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -80,7 +80,6 @@ rm -f "$TRAINER_DONE"{% if num_infer_nodes > 0 %} "$ORCHESTRATOR_DONE"{% endif % export PRL_RUN_ID=${PRL_RUN_ID:-$(python3 -c "import uuid; print(uuid.uuid4().hex)")} export PRL_RUN_NAME={{ run_name }} -export WANDB_RUN_ID=$PRL_RUN_ID export WANDB_SHARED_MODE=1 # Networking @@ -233,11 +232,16 @@ CLEANUP_SH # quoted delimiter passes the script through verbatim, so apostrophes and quotes # (including in the shared _launch_*.sh.j2 includes) need no escaping. Nothing in # the body reads stdin, which would otherwise eat the rest of the script. -srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' +# +# The run owner registers the Prime run around the step and fails it when the +# step fails; every component attaches through the $RUN_ID it exports. +python -m prime_rl.entrypoints.run_owner rl $CONFIG_DIR/orchestrator.json srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' # pipefail ensures `cmd | tee log` propagates `cmd` failures, and the # wait/cleanup block at the end of this script surfaces background # process crashes (router, orchestrator) as task failures. set -o pipefail + # The run owner set $PRL_RUN_ID to the Prime run's id. + export WANDB_RUN_ID=$PRL_RUN_ID # Source environment cd $PROJECT_DIR diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index 3f05ca54fb4..0b276bbe9f4 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -65,15 +65,10 @@ mkdir -p "$LAUNCHER_DIR" rm -f "$TRAINER_DONE"{% if online_eval %} "$EVAL_DONE"{% endif %} export PRL_RUN_ID={{ prl_run_id }} -{% if prime_run_id %} -# The Prime run the launcher registered: every process attaches to it. -export RUN_ID={{ prime_run_id }} -{% endif %} export PRL_RUN_NAME={{ run_name }} {% if wandb_shared %} # Trainer and online-eval log to one shared W&B run. The trainer creates the run, # and the online-eval process finalizes it after the final epoch. -export WANDB_RUN_ID=$PRL_RUN_ID export WANDB_SHARED_MODE=1 export WANDB_SHARED_PRIMARY=trainer export WANDB_SHARED_FINISHER=online-eval @@ -150,8 +145,14 @@ CLEANUP_SH # Run trainer and online-eval roles in one SLURM step. A non-zero exit on any # node terminates the other roles. Clean trainer exits let the final eval drain. -srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' +# The run owner registers the Prime run around the step and fails it when the +# step fails; every role attaches through the $RUN_ID it exports. +python -m prime_rl.entrypoints.run_owner sft $TRAIN_CONFIG_PATH srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' set -o pipefail +{%- if wandb_shared %} + # The run owner set $PRL_RUN_ID to the Prime run's id. + export WANDB_RUN_ID=$PRL_RUN_ID +{%- endif %} cd $PROJECT_DIR [ -f .env ] && source .env diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index 98b857c7ab2..0b53f72f983 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -191,6 +191,31 @@ def test_a_failed_launch_fails_the_run(sdk_runs): assert sdk_runs[0].failed is error +@pytest.mark.parametrize("command,code", [(["true"], 0), (["false"], 1)]) +def test_slurm_run_owner_fails_the_run_with_its_step(sdk_runs, tmp_path, monkeypatch, command, code): + """A multi-node SLURM job's run owner registers the run around the job step, exports + its id to the step, and marks the run failed when the step fails.""" + from prime_rl.entrypoints import run_owner + from prime_rl.entrypoints.sft import write_config + + monkeypatch.setenv("RUN_ID", "") + monkeypatch.setattr(run_owner.signal, "signal", lambda *args: None) + path = tmp_path / "sft.json" + write_config(_sft_config(), path) + monkeypatch.setattr(run_owner.sys, "argv", ["prime", "sft", str(path), *command]) + + if code: + with pytest.raises(SystemExit) as exit_info: + run_owner.main() + assert exit_info.value.code == code + else: + run_owner.main() + + (run,) = sdk_runs + assert not run.attached and os.environ["RUN_ID"] == os.environ["PRL_RUN_ID"] == run.id + assert (run.failed is not None) == bool(code) + + @pytest.mark.parametrize("env", [{"RUN_ID": "hosted-1"}, {"PRIME_RUNS_MODE": "disabled"}]) def test_launch_registers_no_run_it_does_not_own(sdk_runs, monkeypatch, env): """A managed launch pre-created the run, or the monitor opens no platform run.""" From d4a13d7448e567debdb897237d4f79bcf7b19c25 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 29 Oct 2026 22:58:31 +0000 Subject: [PATCH 10/12] register the whole launch config with the prime run An rl launch registered only its orchestrator config, so the run lost the trainer, inference and deployment sections. The launcher now registers the full RL config. Multi-node SLURM jobs also write the full rl.json and sft.json for the run owner; the SFT trainer runs its own trainer.json, as the RL components already do. Co-Authored-By: Claude Opus 5.5 --- .../src/prime_rl/utils/config.py | 2 +- src/prime_rl/entrypoints/rl.py | 11 +++-- src/prime_rl/entrypoints/run_owner.py | 8 ++-- src/prime_rl/entrypoints/sft.py | 42 +++++++++++++------ src/prime_rl/monitors/prime.py | 28 +++++++------ .../templates/multi_node_rl.sbatch.j2 | 2 +- .../templates/multi_node_sft.sbatch.j2 | 2 +- 7 files changed, 60 insertions(+), 35 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/utils/config.py b/packages/prime-rl-configs/src/prime_rl/utils/config.py index a6e9ac04cfb..d40a25e6fe5 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/config.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/config.py @@ -15,7 +15,7 @@ def default_output_dir() -> Path: return Path(os.environ.get("PRL_OUTPUT_DIR", "outputs")) -def dump_resolved_config(config: BaseModel, exclude: set[str] | None = None) -> dict: +def dump_resolved_config(config: BaseModel, exclude: set[str] | dict[str, Any] | None = None) -> dict: """Dump a resolved config for a machine-written JSON artifact. Resolved configs are written as JSON, not TOML: JSON keeps nulls, so explicit None diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index ab09b309fd3..d57a30e6689 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -5,6 +5,7 @@ import sys import time from pathlib import Path +from typing import Any from urllib.parse import urlparse from prime_rl.configs.algorithm import FrozenModelConfig @@ -67,7 +68,7 @@ def env_server_names(config: RLConfig, split: str) -> list[str]: return [source.resolved_name for source_split, source in env_servers(config) if source_split == split] -def write_config(config: RLConfig, output_dir: Path, exclude: set[str] | None = None) -> None: +def write_config(config: RLConfig, output_dir: Path, exclude: set[str] | dict[str, Any] | None = None) -> None: """Write resolved config to disk, excluding launcher-only fields.""" output_dir.mkdir(parents=True, exist_ok=True) with open(output_dir / RL_CONFIG, "w") as f: @@ -521,8 +522,12 @@ def rl_slurm(config: RLConfig): env_names={"train": train_env_names, "eval": eval_env_names}, ) else: + # The job's run owner registers the Prime run with the whole config; nothing runs it. + # It keeps [slurm], which multi-node needs, so it drops the copy in [inference]. + write_config(config, config_dir, exclude={"dry_run": True, "clean": True, "inference": {"slurm"}}) write_subconfigs(config, config_dir) - logger.info(f"Configs:\n{format_config_message(config_dir, 'rl', rl_config_components(config, config_dir))}") + components = [("RL", config_dir / RL_CONFIG), *rl_config_components(config, config_dir)] + logger.info(f"Configs:\n{format_config_message(config_dir, 'rl', components)}") train_env_names = env_server_names(config, "train") eval_env_names = env_server_names(config, "eval") @@ -596,7 +601,7 @@ def rl(config: RLConfig): if config.slurm is not None: rl_slurm(config) return - with launch_run(config.orchestrator, None if config.dry_run else config.orchestrator.monitors.prime): + with launch_run(config, None if config.dry_run else config.orchestrator.monitors.prime): rl_local(config) diff --git a/src/prime_rl/entrypoints/run_owner.py b/src/prime_rl/entrypoints/run_owner.py index 43358cb8cfc..0bd9911308d 100644 --- a/src/prime_rl/entrypoints/run_owner.py +++ b/src/prime_rl/entrypoints/run_owner.py @@ -4,7 +4,7 @@ from pathlib import Path from typing import Any -from prime_rl.configs.orchestrator import OrchestratorConfig +from prime_rl.configs.rl import RLConfig from prime_rl.configs.sft import SFTConfig from prime_rl.monitors.prime import launch_run @@ -14,14 +14,14 @@ def main() -> None: SLURM job's command as the owner of its Prime run (``launch_run``). The command's non-zero exit, or SLURM's SIGTERM on cancel or time limit, fails the run.""" kind, path, *command = sys.argv[1:] - config_cls = OrchestratorConfig if kind == "rl" else SFTConfig - config = config_cls.model_validate_json(Path(path).read_text()) + config = (RLConfig if kind == "rl" else SFTConfig).model_validate_json(Path(path).read_text()) + prime = config.orchestrator.monitors.prime if isinstance(config, RLConfig) else config.monitors.prime def terminate(signum: int, frame: Any) -> None: raise SystemExit(128 + signum) from RuntimeError("SLURM terminated the job") signal.signal(signal.SIGTERM, terminate) - with launch_run(config, config.monitors.prime): + with launch_run(config, prime): code = subprocess.run(command).returncode if code != 0: raise SystemExit(code) from RuntimeError(f"The SLURM job step exited with code {code}") diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 4d5adad67c9..08c80007a83 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -38,6 +38,7 @@ ) SFT_CONFIG = "sft.json" +TRAINER_CONFIG = "trainer.json" SFT_SBATCH = "sft.sbatch" INFERENCE_CONFIG = "inference.json" @@ -98,6 +99,8 @@ def sft_config_components(config: SFTConfig, config_dir: Path) -> list[tuple[str """The resolved configs a launch leaves in ``config_dir``: the SFT config and, with online evals, the inference, online-eval and env-server configs.""" components: list[tuple[str, Path | str]] = [("SFT", config_dir / SFT_CONFIG)] + if config.deployment.type == "multi_node": + components.append(("Trainer", config_dir / TRAINER_CONFIG)) if config.eval is not None: if config.inference is not None: components.append(("Inference", config_dir / INFERENCE_CONFIG)) @@ -139,9 +142,16 @@ def write_eval_subconfigs(config: SFTConfig, config_dir: Path, strip_router: boo def write_slurm_script( - config: SFTConfig, config_path: Path, log_dir: Path, script_path: Path, prl_run_id: str, wandb_shared: bool + config: SFTConfig, + config_path: Path, + run_config_path: Path, + log_dir: Path, + script_path: Path, + prl_run_id: str, + wandb_shared: bool, ) -> None: - """Write the SLURM script to disk.""" + """Write the SLURM script to disk. The trainer runs ``config_path``; a multi-node + job's run owner registers the Prime run from ``run_config_path``.""" from jinja2 import Environment, FileSystemLoader assert config.slurm is not None @@ -200,6 +210,7 @@ def write_slurm_script( script = template.render( **config.slurm.template_vars, config_path=config_path, + run_config_path=run_config_path, config_dir=config_path.parent, log_dir=log_dir, output_dir=config.run_dir, @@ -240,16 +251,19 @@ def sft_slurm(config: SFTConfig): config_dir, log_dir = prepare_attempt_dirs(config.run_dir) write_launch_artifacts(config_dir, "sft") config_path = config_dir / SFT_CONFIG - exclude = ( - {"deployment", "slurm", "dry_run", "clean"} - if config.deployment.type == "multi_node" - else {"slurm", "dry_run", "clean"} - ) - if online_eval: - # The trainer job only needs [eval] for the weight-broadcast cadence; the - # inference pool runs on its dedicated nodes in the same allocation. - exclude = exclude | {"inference"} - write_config(config, config_path, exclude=exclude) + if config.deployment.type == "multi_node": + # The job's run owner registers the Prime run with the whole config; nothing runs it. + write_config(config, config_path, exclude={"dry_run", "clean"}) + trainer_config_path = config_dir / TRAINER_CONFIG + exclude = {"deployment", "slurm", "dry_run", "clean"} + if online_eval: + # The trainer job only needs [eval] for the weight-broadcast cadence; the + # inference pool runs on its dedicated nodes in the same allocation. + exclude = exclude | {"inference"} + write_config(config, trainer_config_path, exclude=exclude) + else: + write_config(config, config_path, exclude={"slurm", "dry_run", "clean"}) + trainer_config_path = config_path # Trainer and online-eval processes log to a single shared W&B run. wandb_shared = online_eval and config.monitors.wandb is not None @@ -259,7 +273,9 @@ def sft_slurm(config: SFTConfig): write_eval_subconfigs(config, config_dir, strip_router=True) logger.info(f"Configs:\n{format_config_message(config_dir, 'sft', sft_config_components(config, config_dir))}") script_path = launcher_dir / SFT_SBATCH - write_slurm_script(config, config_path, log_dir, script_path, os.environ["PRL_RUN_ID"], wandb_shared) + write_slurm_script( + config, trainer_config_path, config_path, log_dir, script_path, os.environ["PRL_RUN_ID"], wandb_shared + ) logger.info(f"Wrote SLURM script to {script_path}") num_nodes = config.deployment.num_train_nodes if config.deployment.type == "multi_node" else 1 diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 82e3bbe3ad2..12c1e45de86 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -14,6 +14,7 @@ from prime_rl.configs.eval import EvalConfig, SFTOnlineEvalConfig from prime_rl.configs.monitors import PrimeMonitorConfig from prime_rl.configs.orchestrator import OrchestratorConfig +from prime_rl.configs.rl import RLConfig from prime_rl.configs.sft import SFTConfig from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.utils.config import BaseConfig @@ -48,9 +49,10 @@ def _base_url() -> str | None: def register( - config: OrchestratorConfig | SFTConfig | None, name: str | None, mode: str, id: str | None = None + config: RLConfig | OrchestratorConfig | SFTConfig | None, name: str | None, mode: str, id: str | None = None ) -> pr.Run: - """Attach to the platform run ``id``, or register ``config``'s training run.""" + """Attach to the platform run ``id``, or register ``config``'s training run with the + whole config: an ``rl`` launch's covers the trainer and inference too.""" init_kwargs: dict[str, Any] if id is not None: init_kwargs = {"id": id} @@ -58,30 +60,32 @@ def register( # The RFT API requires a base model. init_kwargs = {"name": name, "model": "unknown"} else: - wandb_project = config.monitors.wandb.project if config.monitors.wandb else None if isinstance(config, SFTConfig): + model = config.model.name # SFT trains on a dataset, not on environments, and has no rollouts. environments = [] training = pr.TrainingSpec( max_steps=config.max_steps or 0, batch_size=config.data.batch_size, seq_len=config.data.seq_len, - wandb_project=wandb_project, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, ) else: - environments = [env.env_id for env in config.train.source] + orchestrator = config.orchestrator if isinstance(config, RLConfig) else config + model = orchestrator.model.name + environments = [env.env_id for env in orchestrator.train.source] training = pr.TrainingSpec( - max_steps=config.max_steps or 0, - batch_size=config.batch_size, + max_steps=orchestrator.max_steps or 0, + batch_size=orchestrator.batch_size, rollouts_per_example=max( - (source.group_size for source in config.train.source), default=config.train.group_size + (source.group_size for source in orchestrator.train.source), default=orchestrator.train.group_size ), - seq_len=config.seq_len, - wandb_project=wandb_project, + seq_len=orchestrator.seq_len, + wandb_project=orchestrator.monitors.wandb.project if orchestrator.monitors.wandb else None, ) init_kwargs = dict( name=name, - model=config.model.name, + model=model, environments=environments, training=training, config=config.model_dump(exclude_none=True, mode="json"), @@ -102,7 +106,7 @@ def register( @contextmanager -def launch_run(config: OrchestratorConfig | SFTConfig, prime: PrimeMonitorConfig | None) -> Iterator[None]: +def launch_run(config: RLConfig | SFTConfig, prime: PrimeMonitorConfig | None) -> Iterator[None]: """Register the launch's Prime run around starting its components, which attach to it: its id becomes ``$PRL_RUN_ID`` and ``$RUN_ID``. A launch that fails marks the run failed - the SDK leaves an attached run's failure to the process that diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 98d64f178dd..59f36da23ca 100644 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -235,7 +235,7 @@ CLEANUP_SH # # The run owner registers the Prime run around the step and fails it when the # step fails; every component attaches through the $RUN_ID it exports. -python -m prime_rl.entrypoints.run_owner rl $CONFIG_DIR/orchestrator.json srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' +python -m prime_rl.entrypoints.run_owner rl $CONFIG_DIR/rl.json srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' # pipefail ensures `cmd | tee log` propagates `cmd` failures, and the # wait/cleanup block at the end of this script surfaces background # process crashes (router, orchestrator) as task failures. diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index 0b276bbe9f4..ada10594f60 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -147,7 +147,7 @@ CLEANUP_SH # node terminates the other roles. Clean trainer exits let the final eval drain. # The run owner registers the Prime run around the step and fails it when the # step fails; every role attaches through the $RUN_ID it exports. -python -m prime_rl.entrypoints.run_owner sft $TRAIN_CONFIG_PATH srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' +python -m prime_rl.entrypoints.run_owner sft {{ run_config_path }} srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' set -o pipefail {%- if wandb_shared %} # The run owner set $PRL_RUN_ID to the Prime run's id. From 7e0c55fc313db82c95f6f80e0a24aa3e3b0bdefd Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 29 Oct 2026 23:25:56 +0000 Subject: [PATCH 11/12] let the run's owner finish it The process that registers a Prime run now also finishes it, once every component exited: the rl or sft launcher, or the multi-node SLURM run owner. Attached components only flush, so no component can finish the run while another still logs to it, and PRL_RUN_FINISHER goes away. The run owner forwards SLURM's SIGTERM to srun and waits for the step instead of killing it. An RL config can no longer turn a monitor on in one component only. Co-Authored-By: Claude Opus 5.5 --- docs/training.md | 2 +- .../src/prime_rl/configs/rl.py | 10 +++++++ src/prime_rl/entrypoints/run_owner.py | 27 +++++++++++++----- src/prime_rl/entrypoints/sft.py | 5 ---- src/prime_rl/monitors/prime.py | 28 ++++++++----------- .../templates/multi_node_sft.sbatch.j2 | 2 -- src/prime_rl/utils/process.py | 4 --- tests/unit/test_configs.py | 12 ++++++++ tests/unit/train/sft/test_prime_monitor.py | 19 ++++++------- 9 files changed, 64 insertions(+), 45 deletions(-) diff --git a/docs/training.md b/docs/training.md index b4605fdb659..a6f954d18fe 100644 --- a/docs/training.md +++ b/docs/training.md @@ -491,7 +491,7 @@ Name the Prime run in TOML: name = "my-experiment" ``` -The process that supervises the launch registers the run before the components start, and every component logs to it. That is the `rl` or `sft` launcher, also inside a single-node SLURM job, or the run owner that wraps a multi-node SLURM job's step. The run's id is also the run identity (`PRL_RUN_ID`) and the W&B run id. On `sft` the run carries the dataset-batched training fields (`max_steps`, `batch_size`, `seq_len`) and the trainer's per-step metrics and validation losses. Online evals log to the same run, like an `rl` run's: the online-eval process uploads its eval metrics and every eval episode, and finishes the run after its final evals. A launch that fails marks the run failed, and so does a SLURM job that SLURM cancels or stops at its time limit. +The process that supervises the launch registers the run before the components start, and every component logs to it. That is the `rl` or `sft` launcher, also inside a single-node SLURM job, or the run owner that wraps a multi-node SLURM job's step. The run's id is also the run identity (`PRL_RUN_ID`) and the W&B run id. On `sft` the run carries the dataset-batched training fields (`max_steps`, `batch_size`, `seq_len`) and the trainer's per-step metrics and validation losses. Online evals log to the same run, like an `rl` run's: the online-eval process uploads its eval metrics and every eval episode. The components only flush their uploads; the process that registered the run finishes it once every component exited. A launch that fails marks the run failed, and so does a SLURM job that SLURM cancels or stops at its time limit. The platform run needs `PRIME_API_KEY` (`prime login` or the env var) and a team (`PRIME_TEAM_ID`, or the team selected with `prime login`) enabled for external runs. Like the W&B monitor, a run without a key, or one whose registration fails (no team, a team outside the allowlist), fails at launch: authenticate, select a team with `prime switch`, or pass `--no-monitors.prime`. `PRIME_RUNS_MODE=disabled` keeps the monitor configured but opens no platform run; `RUN_ID=` attaches to an external run a launcher already created instead of registering a new one. Currently internal-only. 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 757d41fadfe..97b446cefb6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -342,6 +342,16 @@ def require_equal(name: str, values: dict[str, Any]) -> None: # The orchestrator queries the inference server it is paired with; without one, it # must match the trainer whose weights the external server receives. + # The components log to one run per monitor, so a monitor is on in all or none. + for name in ("wandb", "file"): + require_equal( + f"monitors.{name}", + { + f"{component}.monitors.{name}": getattr(config.monitors, name) is not None + for component, config in (("trainer", trainer), ("orchestrator", orchestrator)) + }, + ) + if inference is not None: require_equal( "deployment.gpus_per_node", diff --git a/src/prime_rl/entrypoints/run_owner.py b/src/prime_rl/entrypoints/run_owner.py index 0bd9911308d..5ac5e36be5e 100644 --- a/src/prime_rl/entrypoints/run_owner.py +++ b/src/prime_rl/entrypoints/run_owner.py @@ -10,21 +10,34 @@ def main() -> None: - """``python -m prime_rl.entrypoints.run_owner {rl,sft} ``: run a multi-node - SLURM job's command as the owner of its Prime run (``launch_run``). The command's - non-zero exit, or SLURM's SIGTERM on cancel or time limit, fails the run.""" + """``python -m prime_rl.entrypoints.run_owner {rl,sft} ``: run a + multi-node SLURM job's step as the owner of its Prime run (``launch_run``). The step's + clean exit finishes the run; its non-zero exit, or SLURM's SIGTERM on cancel or time + limit, fails it.""" kind, path, *command = sys.argv[1:] config = (RLConfig if kind == "rl" else SFTConfig).model_validate_json(Path(path).read_text()) prime = config.orchestrator.monitors.prime if isinstance(config, RLConfig) else config.monitors.prime + step = subprocess.Popen(command) + + def forward(signum: int, frame: Any) -> None: + step.send_signal(signum) def terminate(signum: int, frame: Any) -> None: + # srun relays the signal to the step's tasks, which tear down on their own + # (e.g. finish an in-flight checkpoint); a later SIGTERM only reaches srun. + forward(signum, frame) + signal.signal(signal.SIGTERM, forward) raise SystemExit(128 + signum) from RuntimeError("SLURM terminated the job") signal.signal(signal.SIGTERM, terminate) - with launch_run(config, prime): - code = subprocess.run(command).returncode - if code != 0: - raise SystemExit(code) from RuntimeError(f"The SLURM job step exited with code {code}") + try: + with launch_run(config, prime): + code = step.wait() + if code != 0: + raise SystemExit(code) from RuntimeError(f"The SLURM job step exited with code {code}") + finally: + # The batch script must outlive the step: SLURM ends the job's steps with it. + step.wait() if __name__ == "__main__": diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 08c80007a83..8daaddcc7c7 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -30,7 +30,6 @@ DEFAULT_COMMON_ENV_VARS, DEFAULT_INFERENCE_ENV_VARS, DEFAULT_TRAINER_ENV_VARS, - RUN_FINISHER_VAR, ProcessGroup, partition_gpus, set_proc_title, @@ -358,8 +357,6 @@ def sft_local(config: SFTConfig): "WANDB_PROGRAM": "uv run sft", "WANDB_ARGS": json.dumps(sys.argv), } - # The online-eval process outlives the trainer, so it finishes the shared Prime run. - shared_run_env = {RUN_FINISHER_VAR: "online-eval"} if config.eval is not None else {} with ProcessGroup() as processes: # Optionally, start the inference server for online evals @@ -405,7 +402,6 @@ def sft_local(config: SFTConfig): "PRL_ATTEMPT_LOG_DIR": str(log_dir), "PRL_LOG_DIR": str(log_dir), **wandb_shared_env, - **shared_run_env, "WANDB_SHARED_LABEL": "online-eval", }, log_path=log_dir / "eval.log", @@ -419,7 +415,6 @@ def sft_local(config: SFTConfig): **DEFAULT_TRAINER_ENV_VARS, **config.env_vars, **wandb_shared_env, - **shared_run_env, **launcher_attempt_env(config_dir, log_dir), } if config.eval is not None: diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 12c1e45de86..8fcfa0c683b 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -20,7 +20,6 @@ from prime_rl.utils.config import BaseConfig from prime_rl.utils.links import publish_link, register_links from prime_rl.utils.pathing import get_platform_run_path -from prime_rl.utils.process import RUN_FINISHER_VAR from prime_rl.utils.record import update_record from prime_rl.utils.utils import sanitize @@ -109,10 +108,10 @@ def register( def launch_run(config: RLConfig | SFTConfig, prime: PrimeMonitorConfig | None) -> Iterator[None]: """Register the launch's Prime run around starting its components, which attach to it: its id becomes ``$PRL_RUN_ID`` and ``$RUN_ID``. A launch that fails marks the - run failed - the SDK leaves an attached run's failure to the process that - registered it, so that process supervises the launch: the local ``rl`` or ``sft`` - launcher, or a multi-node SLURM job's ``run_owner``. A run a managed launch pre-created - (``$RUN_ID``) is the platform's.""" + run failed, and one that ends cleanly finishes it once every component exited - the + components only flush their attached run. So the process that registers the run + supervises the launch: the local ``rl`` or ``sft`` launcher, or a multi-node SLURM + job's ``run_owner``. A run a managed launch pre-created (``$RUN_ID``) is the platform's.""" if prime is None or os.getenv("RUN_ID") or (os.getenv(pr.MODE_ENV) or "online") != "online": yield return @@ -124,8 +123,7 @@ def launch_run(config: RLConfig | SFTConfig, prime: PrimeMonitorConfig | None) - # A failed component exits the launcher with a SystemExit chained to its error. run.fail(e.__cause__ or e, timeout=FINISH_TIMEOUT) raise - # The finisher component completes the run; this process only exits. - atexit.unregister(run._atexit_hook) + run.finish(timeout=FINISH_TIMEOUT) class PrimeMonitor(Monitor): @@ -133,9 +131,8 @@ class PrimeMonitor(Monitor): A training run's processes log to one Prime training run: the ``rl`` or ``sft`` launcher registers it (``launch_run``) and its components attach through ``$RUN_ID``. - A component started without a launcher registers its own. - ``$PRL_RUN_FINISHER`` names the process that finishes a shared run; the others - only flush. A standalone eval instead opens one evaluation per env and epoch when + The process that registered a run finishes it; an attached process only flushes, and + a component started without a launcher registers and finishes its own. A standalone eval instead opens one evaluation per env and epoch when the epoch starts (``log_eval_plan``), streams its episodes into it, and finishes it with the epoch's aggregates (``log_eval_epoch``). @@ -235,15 +232,14 @@ async def finalize(self, paused: bool = False) -> None: self.logger.warning(f"Failed to close the {env_name} (Step {step}) evaluation: {type(e).__name__}: {e}") if self.run is None: return - if paused or os.getenv(RUN_FINISHER_VAR, self.producer) != self.producer: - # Drain only: a paused run continues on resume, and a shared run is finished - # by the process that outlives this one. A clean exit is not a crash. + if paused or self.run.attached: + # Drain only: a paused run continues on resume, and an attached run is + # finished by the process that registered it, once every component exited. + # A clean exit is not a crash. await asyncio.to_thread(self.run.flush, FINISH_TIMEOUT) atexit.unregister(self.run._atexit_hook) return - # Drains queued uploads so the final step's metrics and episodes land, then - # finalizes (idempotent on the platform side); an attached run's failure - # marking stays with the launcher. + # Drains queued uploads so the final step's metrics and episodes land, then finalizes. await asyncio.to_thread(self.run.finish) async def _init_evaluations(self, config: EvalConfig) -> None: diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index ada10594f60..bb9b588009e 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -58,8 +58,6 @@ ln -sfn inference/node_0.log $LOG_DIR/inference.log export TRAINER_DONE=$LAUNCHER_DIR/.sft_trainer_done_$SLURM_JOB_ID {% if online_eval -%} export EVAL_DONE=$LAUNCHER_DIR/.sft_eval_done_$SLURM_JOB_ID -# The online-eval process outlives the trainer, so it finishes the shared Prime run. -export PRL_RUN_FINISHER=online-eval {%- endif %} mkdir -p "$LAUNCHER_DIR" rm -f "$TRAINER_DONE"{% if online_eval %} "$EVAL_DONE"{% endif %} diff --git a/src/prime_rl/utils/process.py b/src/prime_rl/utils/process.py index 80854a0ba4d..464251215b4 100644 --- a/src/prime_rl/utils/process.py +++ b/src/prime_rl/utils/process.py @@ -26,10 +26,6 @@ "TRITON_CACHE_DIR": f"/tmp/triton-{os.getuid()}", } -# The producer that finishes a run several processes log to, e.g. the SFT online-eval -# process, which drains its final evals after the trainer exits. -RUN_FINISHER_VAR = "PRL_RUN_FINISHER" - DEFAULT_TRAINER_ENV_VARS: dict[str, str] = { "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", # NCCL flight recorder: on a collective timeout every rank dumps its last 2000 collectives diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index ad46d07b9e3..e5f56287e70 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -1059,6 +1059,18 @@ def test_shared_monitor_disable_and_prime_propagate(): assert config.orchestrator.monitors.prime.name == "shared-prime" +def test_monitor_on_in_one_component_only_is_rejected(): + """The components log to one run per monitor, so a monitor is on in all or none.""" + with pytest.raises(ValueError, match=r"monitors.wandb must match across components"): + RLConfig.model_validate( + { + "model": {"name": "Qwen/Qwen3-0.6B"}, + "trainer": {"monitors": {"wandb": {"project": "p"}}}, + "orchestrator": {"renderer": {"name": "default"}}, + } + ) + + def test_empty_shared_ckpt_block_does_not_conflict_with_subconfig_ckpt(): """An empty shared [ckpt] block is a presence-only signal, not a field setting — it should not conflict with a non-empty [trainer.ckpt].""" diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index 0b53f72f983..fc3f49ccc58 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -146,15 +146,13 @@ def fake_init(**kwargs): return runs -def test_online_evals_share_the_launchers_run(sdk_runs, tmp_path, monkeypatch): +def test_online_evals_share_the_launchers_run(sdk_runs, tmp_path): """The sft launcher registers the run and makes its id the run identity; the trainer - and the online-eval process attach to it. Only the online-eval process finishes the - run: it drains its final evals after the trainer exits.""" + and the online-eval process attach to it and only flush. The launcher finishes the + run once both exited.""" from prime_rl.entrypoints.sft import build_online_eval_config from prime_rl.monitors.prime import launch_run - from prime_rl.utils.process import RUN_FINISHER_VAR - monkeypatch.setenv(RUN_FINISHER_VAR, "online-eval") config = _sft_config( {"eval": {"interval": 50, "source": [{"name": "rev", "env": {"taskset": {"id": "harbor"}}}]}, "inference": {}} ) @@ -173,10 +171,10 @@ def test_online_evals_share_the_launchers_run(sdk_runs, tmp_path, monkeypatch): assert trainer.run.id == online_eval.run.id == launched.id asyncio.run(trainer.finalize()) - assert trainer.run.flushed and not trainer.run.finished asyncio.run(online_eval.finalize()) - assert online_eval.run.finished - assert launched.failed is None + assert trainer.run.flushed and online_eval.run.flushed + assert not (trainer.run.finished or online_eval.run.finished or launched.finished) + assert launched.finished and launched.failed is None def test_a_failed_launch_fails_the_run(sdk_runs): @@ -194,7 +192,8 @@ def test_a_failed_launch_fails_the_run(sdk_runs): @pytest.mark.parametrize("command,code", [(["true"], 0), (["false"], 1)]) def test_slurm_run_owner_fails_the_run_with_its_step(sdk_runs, tmp_path, monkeypatch, command, code): """A multi-node SLURM job's run owner registers the run around the job step, exports - its id to the step, and marks the run failed when the step fails.""" + its id to the step, and finishes the run when the step succeeds or fails it when the + step fails.""" from prime_rl.entrypoints import run_owner from prime_rl.entrypoints.sft import write_config @@ -213,7 +212,7 @@ def test_slurm_run_owner_fails_the_run_with_its_step(sdk_runs, tmp_path, monkeyp (run,) = sdk_runs assert not run.attached and os.environ["RUN_ID"] == os.environ["PRL_RUN_ID"] == run.id - assert (run.failed is not None) == bool(code) + assert (run.failed is not None) == bool(code) and run.finished == (not code) @pytest.mark.parametrize("env", [{"RUN_ID": "hosted-1"}, {"PRIME_RUNS_MODE": "disabled"}]) From 298960f5a9289ae0908842d3f657ec5bd6d625f8 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 29 Oct 2026 23:32:14 +0000 Subject: [PATCH 12/12] start the slurm job step inside the run owner's run The step started before launch_run exported the run's id, so its components had no $RUN_ID to attach to. Co-Authored-By: Claude Opus 5.5 --- src/prime_rl/entrypoints/run_owner.py | 10 +++++++--- tests/unit/train/sft/test_prime_monitor.py | 2 +- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/prime_rl/entrypoints/run_owner.py b/src/prime_rl/entrypoints/run_owner.py index 5ac5e36be5e..86222852f60 100644 --- a/src/prime_rl/entrypoints/run_owner.py +++ b/src/prime_rl/entrypoints/run_owner.py @@ -17,10 +17,11 @@ def main() -> None: kind, path, *command = sys.argv[1:] config = (RLConfig if kind == "rl" else SFTConfig).model_validate_json(Path(path).read_text()) prime = config.orchestrator.monitors.prime if isinstance(config, RLConfig) else config.monitors.prime - step = subprocess.Popen(command) + step: subprocess.Popen | None = None def forward(signum: int, frame: Any) -> None: - step.send_signal(signum) + if step is not None: + step.send_signal(signum) def terminate(signum: int, frame: Any) -> None: # srun relays the signal to the step's tasks, which tear down on their own @@ -32,12 +33,15 @@ def terminate(signum: int, frame: Any) -> None: signal.signal(signal.SIGTERM, terminate) try: with launch_run(config, prime): + # Started inside launch_run, so the step inherits the run's $RUN_ID. + step = subprocess.Popen(command) code = step.wait() if code != 0: raise SystemExit(code) from RuntimeError(f"The SLURM job step exited with code {code}") finally: # The batch script must outlive the step: SLURM ends the job's steps with it. - step.wait() + if step is not None: + step.wait() if __name__ == "__main__": diff --git a/tests/unit/train/sft/test_prime_monitor.py b/tests/unit/train/sft/test_prime_monitor.py index fc3f49ccc58..0e5a7be2666 100644 --- a/tests/unit/train/sft/test_prime_monitor.py +++ b/tests/unit/train/sft/test_prime_monitor.py @@ -189,7 +189,7 @@ def test_a_failed_launch_fails_the_run(sdk_runs): assert sdk_runs[0].failed is error -@pytest.mark.parametrize("command,code", [(["true"], 0), (["false"], 1)]) +@pytest.mark.parametrize("command,code", [(["bash", "-c", '[ "$RUN_ID" = run-123 ]'], 0), (["false"], 1)]) def test_slurm_run_owner_fails_the_run_with_its_step(sdk_runs, tmp_path, monkeypatch, command, code): """A multi-node SLURM job's run owner registers the run around the job step, exports its id to the step, and finishes the run when the step succeeds or fails it when the