prime-runs: RftBackend for training runs

packages/prime-runs/README.md

@@ -1,6 +1,6 @@
 # Prime Runs SDK
 
-Track eval runs on the Prime Intellect platform.
+Track eval and training runs on the Prime Intellect platform.
 
 ```bash
 pip install prime-runs
@@ -105,6 +105,68 @@ disables the run with a warning — it never silently writes somewhere else.
   block — is `cancelled`; one that stopped without saying — an exit that never
   reached `finish()` — is `crashed`.
 
+## Training runs
+
+```python
+import prime_runs as pr
+
+run = pr.init(
+    kind="train",
+    name="qwen3-8b-gsm8k-rl",
+    model="Qwen/Qwen3-8B",                  # the base model
+    environments=["primeintellect/gsm8k"],  # hub ids, passed through as given
+    training=pr.TrainingSpec(max_steps=1000, batch_size=64, rollouts_per_example=8),
+    config=train_config.model_dump(),       # or the path to the launched TOML
+    team_id="team_...",                     # external runs are created for a team
+)
+
+for step, (episodes, metrics) in enumerate(training_loop):
+    run.log_episodes(episodes)              # carry run.work.step (verifiers TrainRunInfo)
+    run.log_metrics(metrics, step=step)
+
+run.finish()
+```
+
+A training run lives on `/api/v1/rft/external-runs` — what prime-rl's monitor
+did by hand. Where it differs from an eval run:
+
+- **A team is required.** The platform enables external runs per team. A team
+  outside the allowlist gets a `ForbiddenError` from `init()` with the
+  platform's reason; the producer decides whether that means "run locally".
+- **Metrics.** `log_metrics(values, step=...)` streams one row per call on its
+  own uploader, so a slow sample upload never holds a step's metrics back.
+  `_timestamp` is stamped on every row; non-finite values are dropped. The
+  endpoint allows 60 rows a minute per token. A row whose response was lost
+  is re-sent: the platform keeps one row per step, so nothing doubles.
+- **Samples.** Every record reaches Prime Traces. The training viewer's sample
+  table is a Parquet object per upload, keyed by the training step an episode
+  was *dispatched* at — `run.work.step` on verifiers' `TrainRunInfo`, which
+  prime-rl stamps on every dispatched episode — and uploaded every 10th step
+  (prime-rl's cadence) for training-work episodes. Objects are additive: a
+  step logged in several `log_episodes` calls (an off-policy episode landing
+  in a later batch) gets one object per call, each numbering its `sample_id`s
+  in its own range, and the viewer shows their union. So "step N" in the
+  sample viewer means dispatched at N, while the metrics logged at N describe
+  the batch trained at N. Encoding needs pyarrow: `pip install
+  'prime-runs[train]'`. Without it the table is skipped with a warning and
+  traces still flow.
+- **Attach.** `init(kind="train", id=os.environ["RUN_ID"])` joins a run a
+  launcher created: nothing is registered, the platform keeps the run's
+  failure marking, and a clean `finish()` still completes it. The id must be
+  an external run's (one created through `POST /rft/external-runs`): the
+  platform's monitoring endpoints answer 400 for a hosted run.
+- **Status.** The RFT vocabulary is `completed | failed`. `cancelled` and
+  `crashed` are reported as `failed` with the reason in `error_message`; a
+  clean `finish()` is an idempotent finalize, safe to replay.
+- **Outages.** A sink that strikes out on transient failures (three in a row,
+  each already retried) is paused for five minutes and then tried again, not
+  retired for the run: a platform deploy costs a multi-day run a window of
+  rows, never the rest of its curves. Records that arrive while a sink is
+  paused are counted in `run.failed_records` for that sink.
+- **No config update.** `config=` is registered once as `run_config` and there
+  is no summary document, so `finish(summary=...)` is not sent for training
+  runs — log final numbers with `log_metrics`.
+
 ## From async code
 
 `log_traces()` / `log_episodes()` are a queue put, not a request, so they are
@@ -144,8 +206,8 @@ package is imported — this is a leaf package, because the `prime` CLI depends
 
 ## Status
 
-Eval runs only. Training runs will arrive with a backend over
-`/api/v1/rft/external-runs`, designed against prime-rl's actual needs.
+Eval runs over `/api/v1/evaluations`, training runs over
+`/api/v1/rft/external-runs` (see "Training runs").
 
 The evaluations API this version targets has no producer-facing way to mark an
 evaluation **failed** or **cancelled**: the SDK records the terminal state under

packages/prime-runs/pyproject.toml

@@ -43,6 +43,11 @@ Documentation = "https://github.com/PrimeIntellect-ai/prime/tree/main/packages/p
 Repository = "https://github.com/PrimeIntellect-ai/prime.git"
 
 [project.optional-dependencies]
+# The training sample table is Parquet. Optional: an eval producer never needs
+# it, and prime-rl already carries pyarrow.
+train = [
+    "pyarrow>=15.0.0",
+]
 dev = [
     "pytest>=7.0.0",
     "pytest-asyncio>=0.23.0",

packages/prime-runs/src/prime_runs/__init__.py

@@ -1,6 +1,6 @@
 """Prime Intellect Runs SDK.
 
-Track eval runs on the Prime platform::
+Track eval and training runs on the Prime platform::
 
     import prime_runs as pr
 
@@ -31,7 +31,14 @@
     TransportError,
     UnauthorizedError,
 )
-from .models import CONFIG_SOURCE_KEY, ConfigSource, EnvironmentRef, RunStatus
+from .models import (
+    CONFIG_SOURCE_KEY,
+    ConfigSource,
+    EnvironmentRef,
+    RunKind,
+    RunStatus,
+    TrainingSpec,
+)
 from .run import MODE_ENV, Run, init
 
 __version__ = "0.1.0"
@@ -40,6 +47,8 @@
     "init",
     "Run",
     "RunStatus",
+    "RunKind",
+    "TrainingSpec",
     "ConfigSource",
     "CONFIG_SOURCE_KEY",
     "EnvironmentRef",

packages/prime-runs/src/prime_runs/backend.py

@@ -20,7 +20,7 @@
 
 from ._http import PlatformClient
 from .exceptions import APIError, ConfigurationError, EnvironmentResolutionError
-from .models import EnvironmentRef, RunHandle, RunSpec, RunStatus
+from .models import EnvironmentRef, RunHandle, RunSpec, RunStatus, TrainingSpec
 
 logger = logging.getLogger(__name__)
 
@@ -267,3 +267,186 @@ def finalize(self, run_id: str, **kwargs: Any) -> None:
 
     def close(self) -> None:
         return None
+
+
+# ------------------------------------------------------------------- rft
+
+#: The RFT API's status vocabulary is ``completed | failed``; the SDK's finer
+#: distinctions travel in ``error_message``.
+RFT_TERMINAL_STATUS = {
+    RunStatus.COMPLETED: "completed",
+    RunStatus.FAILED: "failed",
+    RunStatus.CANCELLED: "failed",
+    RunStatus.CRASHED: "failed",
+}
+
+
+class RftBackend:
+    """Lifecycle for external training runs over ``/api/v1/rft/*``.
+
+    What prime-rl's ``TrainRun`` did, in the SDK: register a run (or attach to
+    one a managed launch created), close it out. External runs are created
+    *for a team* and the platform enables them per team, so ``team_id`` is
+    required and a team outside the allowlist gets a 403 from :meth:`create`;
+    the producer decides whether that means "run locally" (the way verifiers
+    treats a missing key) or "stop".
+
+    The API has no config/summary update — ``run_config`` travels with
+    :meth:`create` and :meth:`update` is a no-op — and its status vocabulary
+    is ``completed | failed``: ``cancelled`` and ``crashed`` are reported as
+    ``failed`` with the reason in ``error_message``.
+    """
+
+    def __init__(
+        self,
+        client: PlatformClient,
+        *,
+        frontend_url: str,
+        team_id: Optional[str] = None,
+    ) -> None:
+        self._client = client
+        self._frontend_url = frontend_url.rstrip("/")
+        self._team_id = team_id
+
+    # ------------------------------------------------------------------ create
+
+    def create(self, spec: RunSpec) -> RunHandle:
+        if not spec.model:
+            raise ConfigurationError("A training run needs model= — the base model being trained.")
+        team_id = spec.team_id or self._team_id
+        if not team_id:
+            raise ConfigurationError(
+                "A training run needs a team: pass team_id= or set PRIME_TEAM_ID. "
+                "External training runs are created for a team, and the platform "
+                "enables them per team."
+            )
+        training = spec.training or TrainingSpec()
+        payload: Dict[str, Any] = {
+            "base_model": spec.model,
+            "max_steps": int(training.max_steps),
+            # Environment ids are passed through: training environments are
+            # named by hub id, and the RFT API does its own resolution.
+            "environments": [_training_environment(ref) for ref in spec.environments],
+            "team_id": team_id,
+        }
+        _set_if(payload, "name", spec.name)
+        _set_if(payload, "batch_size", training.batch_size)
+        _set_if(payload, "rollouts_per_example", training.rollouts_per_example)
+        _set_if(payload, "seq_len", training.seq_len)
+        _set_if(payload, "run_config", spec.config or None)
+        _set_if(payload, "wandb_project", training.wandb_project)
+        _set_if(payload, "wandb_entity", training.wandb_entity)
+        _set_if(payload, "wandb_run_name", training.wandb_run_name)
+
+        # Not replayable: a retry after a lost response would register a
+        # second run and only the second would be tracked.
+        response = self._client.post("/rft/external-runs", json_body=payload)
+        run = response.get("run") or {}
+        run_id = run.get("id")
+        if not run_id:
+            raise APIError(f"POST /rft/external-runs returned no run id (keys: {sorted(response)})")
+        return RunHandle(
+            id=str(run_id),
+            name=str(run.get("name") or spec.name or "") or None,
+            url=self.url_for(str(run_id)),
+        )
+
+    def attach(self, run_id: str) -> RunHandle:
+        """A handle for a run the platform already created — a launcher that
+        injected its id. No request: the first write proves access, and it
+        must be an *external* run (``POST /rft/external-runs``): the v1
+        monitoring endpoints answer 400 for a hosted run's id."""
+        return RunHandle(id=run_id, url=self.url_for(run_id))
+
+    def url_for(self, run_id: str) -> str:
+        return f"{self._frontend_url}/dashboard/training/{run_id}"
+
+    # ------------------------------------------------------------------ update
+
+    def update(
+        self,
+        run_id: str,
+        *,
+        config: Optional[Dict[str, Any]] = None,
+        summary: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """No-op: an RFT run takes its config at creation and has no summary
+        document. Run-level outputs go through ``log_metrics``."""
+        if config or summary:
+            logger.debug("Run %s: the RFT API has no config/summary update; nothing sent", run_id)
+
+    # ---------------------------------------------------------------- finalize
+
+    def finalize(
+        self,
+        run_id: str,
+        *,
+        status: RunStatus,
+        summary: Optional[Dict[str, Any]] = None,
+        error: Optional[str] = None,
+        config: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        if status is RunStatus.COMPLETED:
+            # Compare-and-set on the server, and "already finalized" is a
+            # success, so a lost response is safe to replay.
+            try:
+                self._client.post(
+                    "/rft/finalize",
+                    json_body={"run_id": run_id, "exit_code": 0},
+                    idempotent=True,
+                )
+            except APIError as exc:
+                # Preserve TrainRun's fallback: finalization may be unavailable
+                # even though the status endpoint can still close the run out.
+                logger.warning(
+                    "Run %s could not be finalized (%s); falling back to a status update",
+                    run_id,
+                    exc,
+                )
+                self._set_terminal_status(run_id, status=status)
+            return
+
+        message = error or status.value
+        if status is not RunStatus.FAILED:
+            message = f"{status.value}: {message}"
+        self._set_terminal_status(run_id, status=status, error_message=message)
+
+    def _set_terminal_status(
+        self,
+        run_id: str,
+        *,
+        status: RunStatus,
+        error_message: Optional[str] = None,
+    ) -> None:
+        payload: Dict[str, Any] = {"status": RFT_TERMINAL_STATUS[status]}
+        _set_if(payload, "error_message", error_message)
+        try:
+            self._client.put(
+                f"/rft/external-runs/{run_id}/status",
+                json_body=payload,
+            )
+        except APIError as exc:
+            if exc.status_code == 409:
+                # Already closed out (a managed launcher, or an earlier attempt
+                # whose response was lost). The outcome is recorded either way.
+                logger.info(
+                    "Run %s is already closed on the platform; not marking it %s",
+                    run_id,
+                    status.value,
+                )
+                return
+            raise
+
+    def log_metrics(self, run_id: str, metrics: Dict[str, Any]) -> None:
+        """One step's metrics, synchronously. The run handle batches these
+        through its uploader; this is the direct call."""
+        self._client.post("/rft/metrics", json_body={"run_id": run_id, "metrics": metrics})
+
+    def close(self) -> None:
+        self._client.close()
+
+
+def _training_environment(ref: EnvironmentRef) -> Dict[str, Any]:
+    entry: Dict[str, Any] = {"id": ref.id or ref.slug or ref.name}
+    _set_if(entry, "version_id", ref.version_id)
+    return entry

packages/prime-runs/src/prime_runs/models.py

@@ -18,8 +18,14 @@
 
 OnError = Literal["warn", "raise"]
 
-RUN_KIND = "eval"
-"""Stamped as ``run.type`` on records and sent as upload provenance."""
+RunKind = Literal["eval", "train"]
+"""What a run is: a standalone evaluation (``/api/v1/evaluations``) or an
+external training run (``/api/v1/rft/external-runs``). Stamped as ``run.type``
+on records and sent as upload provenance — the traces service's vocabulary,
+which is also what verifiers' ``EvalRunInfo`` / ``TrainRunInfo`` carry."""
+
+RUN_KIND: RunKind = "eval"
+"""The default kind."""
 
 
 class RunStatus(str, Enum):
@@ -191,6 +197,26 @@ def __post_init__(self) -> None:
             )
 
 
+@dataclass
+class TrainingSpec:
+    """What a training run is registered with, beyond what every run has
+    (``model`` is the base model, ``environments`` the training environments,
+    ``config`` the run config). All of it is for display on the dashboard.
+    ``max_steps`` is required by the platform; ``0`` means unknown."""
+
+    max_steps: int = 0
+    batch_size: Optional[int] = None
+    rollouts_per_example: Optional[int] = None
+    seq_len: Optional[int] = None
+    wandb_project: Optional[str] = None
+    wandb_entity: Optional[str] = None
+    wandb_run_name: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        if self.max_steps < 0:
+            raise ValueError("max_steps must be non-negative")
+
+
 @dataclass
 class RunSpec:
     """Everything a backend needs to open a run: ``init()``'s arguments after
@@ -205,6 +231,8 @@ class RunSpec:
     tags: List[str] = field(default_factory=list)
     team_id: Optional[str] = None
     config: Dict[str, Any] = field(default_factory=dict)
+    kind: RunKind = RUN_KIND
+    training: Optional[TrainingSpec] = None
 
 
 @dataclass