feat: one prime monitor, sft online evals on the training run

docs/eval.md

@@ -130,7 +130,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
 

docs/training.md

@@ -490,7 +490,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.
+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=<id>` attaches to an external run a launcher already created instead of registering a new one. Currently internal-only.
 

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

@@ -175,11 +175,9 @@ class SFTOnlineEvalConfig(ScheduledEvalConfig, ServedEvalConfig):
     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."""
+    ``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):

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

@@ -43,14 +43,9 @@ 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."""
-
-
-class PrimeEvalMonitorConfig(BaseConfig):
-    name: str | None = None
-    """Evaluation name shown on the platform. Inherits ``run.name`` when unset."""
+    """Run or evaluation name shown on Prime. Inherits ``run.name`` when unset."""
 
 
 class MonitorsConfig(BaseConfig):
@@ -62,12 +57,13 @@ class MonitorsConfig(BaseConfig):
 
 
 class TrainMonitorsConfig(MonitorsConfig):
-    prime: PrimeTrainMonitorConfig | None = PrimeTrainMonitorConfig()
+    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: 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``."""
+    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``."""

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, TrainSamplingConfig
 from prime_rl.configs.shared import (
     BaseMultiNodeDeploymentConfig,
@@ -80,7 +80,7 @@ class SharedMonitorsConfig(BaseConfig):
     file: FileMonitorConfig | None = None
     """Shared local JSONL metric sink. If set, enables ``<output_dir>/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."""
 
 

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:

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()

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():

src/prime_rl/entrypoints/sft.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 from prime_rl.configs.eval import SFTOnlineEvalConfig
-from prime_rl.configs.monitors import EvalMonitorsConfig, PrimeEvalMonitorConfig, TrainMonitorsConfig
+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
@@ -32,6 +32,7 @@
     DEFAULT_COMMON_ENV_VARS,
     DEFAULT_INFERENCE_ENV_VARS,
     DEFAULT_TRAINER_ENV_VARS,
+    RUN_FINISHER_VAR,
     ProcessGroup,
     partition_gpus,
     set_proc_title,
@@ -71,13 +72,9 @@ def resolve_resume_step(config: SFTConfig) -> int | None:
 
 
 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,
-    )
+    """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:
@@ -353,6 +350,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 +397,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 +411,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:
@@ -478,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():

src/prime_rl/monitors/__init__.py

@@ -15,13 +15,12 @@
 
 from prime_rl.configs.monitors import (
     FileMonitorConfig,
-    PrimeEvalMonitorConfig,
-    PrimeTrainMonitorConfig,
+    PrimeMonitorConfig,
     WandbMonitorConfig,
 )
 from prime_rl.monitors.base import Kind, Monitor, Subset
 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
@@ -32,8 +31,7 @@
 __all__ = [
     "Monitor",
     "WandbMonitor",
-    "PrimeTrainMonitor",
-    "PrimeEvalMonitor",
+    "PrimeMonitor",
     "FileMonitor",
     "setup",
     "get",
@@ -47,7 +45,7 @@
 
 async def setup(
     wandb: WandbMonitorConfig | None = None,
-    prime: PrimeTrainMonitorConfig | PrimeEvalMonitorConfig | None = None,
+    prime: PrimeMonitorConfig | None = None,
     file: FileMonitorConfig | None = None,
     *,
     output_dir: Path,
@@ -69,10 +67,10 @@ async def setup(
         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 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(
             (