diff --git a/integrations/arctic-rl/README.md b/integrations/arctic-rl/README.md
new file mode 100644
index 0000000000..1acb646e0f
--- /dev/null
+++ b/integrations/arctic-rl/README.md
@@ -0,0 +1,54 @@
+# Arctic RL backend for PRIME-RL
+
+Opt-in integration that delegates training (forward/backward/optimizer/weight-sync)
+and rollout generation to a remote Arctic RL server. Lives entirely under
+`integrations/arctic-rl/` so core PRIME-RL has zero Arctic-specific code.
+
+## Enabling
+
+Add this to any PRIME-RL `rl.toml`:
+
+```toml
+[trainer]
+backend = "arctic_rl"
+
+[arctic]
+backend = "remote"
+url = "http://your-arctic-server:7000"
+```
+
+The PRIME-RL launcher (`prime_rl.entrypoints.rl:main`) peeks at
+`trainer.backend` before strict config parse and dynamically dispatches to
+`arctic_rl.entrypoint:main` when it's set to `"arctic_rl"`. Otherwise the
+native PRIME-RL path runs unchanged and `arctic_rl` is never imported.
+
+## Install
+
+Install this integration alongside PRIME-RL:
+
+```
+uv pip install -e ./integrations/arctic-rl
+```
+
+Plus the Arctic RL client (private until release — see Arctic RL release
+notes for the install command).
+
+## What gets replaced
+
+In Arctic mode the launcher swaps two of PRIME-RL's three subprocesses:
+
+| Slot | Native PRIME-RL | Arctic mode |
+|---|---|---|
+| Trainer | `torchrun` + FSDP2 | `arctic-trainer` (single CPU process, HTTP client) |
+| Inference | vLLM server | `arctic-shim` (FastAPI OpenAI-compat proxy) |
+| Orchestrator | verifiers + rollouts | unchanged |
+
+The orchestrator points at the shim's `base_url` (rewritten by the
+launcher); rollout files / STABLE markers / weight-sync semantics are
+all preserved.
+
+## Out of scope (rejected by validator)
+
+- Multi-node deployment
+- LoRA, multi-run, `teacher_inference`, `[inference]` (mutually exclusive)
+- Multimodal
diff --git a/integrations/arctic-rl/arctic_rl/__init__.py b/integrations/arctic-rl/arctic_rl/__init__.py
new file mode 100644
index 0000000000..30ce1d8c76
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/__init__.py
@@ -0,0 +1,4 @@
+"""PRIME-RL ↔ Arctic RL backend integration.
+
+Activated by ``trainer.backend = "arctic_rl"`` in ``rl.toml``.
+"""
diff --git a/integrations/arctic-rl/arctic_rl/_trainer_entrypoint.py b/integrations/arctic-rl/arctic_rl/_trainer_entrypoint.py
new file mode 100644
index 0000000000..988065ade5
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/_trainer_entrypoint.py
@@ -0,0 +1,52 @@
+"""arctic-trainer entrypoint.
+
+Mirror of prime_rl.trainer.rl.train:main, but single-process (no torchrun)
+and running the Arctic HTTP adapter instead of a local FSDP2 model.
+
+The launcher writes trainer.toml and arctic.toml side-by-side. We load the
+trainer config via the standard `cli(TrainerConfig)` mechanism and pick up
+arctic.toml via the ARCTIC_CONFIG_TOML env var (populated by the launcher
+because pydantic-config's cli() takes one positional @-path arg at a time).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import tomli
+from loguru import logger
+
+from arctic_rl.config import ArcticConfig
+from arctic_rl.trainer import ArcticTrainerAdapter
+from prime_rl.configs.trainer import TrainerConfig
+from prime_rl.utils.config import cli
+from prime_rl.utils.process import set_proc_title
+
+
+def _load_arctic_config() -> ArcticConfig:
+ toml_path = os.environ.get("ARCTIC_CONFIG_TOML")
+ if not toml_path:
+ raise RuntimeError(
+ "ARCTIC_CONFIG_TOML env var not set. The arctic-trainer is meant to be "
+ "launched by rl_arctic_local which sets this variable pointing to arctic.toml."
+ )
+ path = Path(toml_path)
+ if not path.exists():
+ raise RuntimeError(f"ARCTIC_CONFIG_TOML points at non-existent path: {path}")
+ with open(path, "rb") as f:
+ data = tomli.load(f)
+ return ArcticConfig(**data)
+
+
+def main():
+ set_proc_title("ArcticTrainer")
+ trainer_cfg: TrainerConfig = cli(TrainerConfig)
+ arctic_cfg = _load_arctic_config()
+ logger.info("ArcticTrainer starting (backend={})", arctic_cfg.backend)
+ ArcticTrainerAdapter(trainer_cfg, arctic_cfg).run()
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/integrations/arctic-rl/arctic_rl/client.py b/integrations/arctic-rl/arctic_rl/client.py
new file mode 100644
index 0000000000..4e46d17951
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/client.py
@@ -0,0 +1,68 @@
+"""Build the ArcticRLClient that the trainer uses to talk to the Arctic RL server."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from loguru import logger
+
+from arctic_rl.config import ArcticConfig
+
+if TYPE_CHECKING:
+ from prime_rl.configs.trainer import TrainerConfig
+
+
+def _build_training_config(trainer_cfg: TrainerConfig, arctic_cfg: ArcticConfig) -> dict:
+ """Build the dict the Arctic RL server expects at /initialize time."""
+ optim = trainer_cfg.optim
+ return {
+ "dtype": "bfloat16",
+ "gradient_checkpointing": True,
+ "max_seq_len": trainer_cfg.model.seq_len,
+ "n_gpus": arctic_cfg.training_gpus,
+ "optimizer": {
+ "lr": getattr(optim, "lr", 1e-5),
+ "weight_decay": getattr(optim, "weight_decay", 0.0),
+ "beta1": getattr(optim, "betas", (0.9, 0.999))[0] if hasattr(optim, "betas") else 0.9,
+ "beta2": getattr(optim, "betas", (0.9, 0.999))[1] if hasattr(optim, "betas") else 0.999,
+ "eps": getattr(optim, "eps", 1e-8),
+ "gradient_clipping": 1.0,
+ "lr_scheduler_type": "constant",
+ "warmup_steps_proportion": 0.0,
+ },
+ }
+
+
+def _build_vllm_config(arctic_cfg: ArcticConfig, trainer_cfg: TrainerConfig) -> dict:
+ vllm_config = dict(arctic_cfg.vllm_config or {})
+ vllm_config.setdefault("max_model_len", trainer_cfg.model.seq_len)
+ vllm_config.setdefault("tensor_parallel_size", arctic_cfg.sampling_tensor_parallel_size)
+ return vllm_config
+
+
+def build_arctic_client(arctic_cfg: ArcticConfig, trainer_cfg: TrainerConfig):
+ """Build and return an ArcticRLClient. Blocks until all jobs are RUNNING."""
+ # Lazy-imported so the integration package doesn't pull arctic_training
+ # at module-load time.
+ from arctic_training.arctic_rl.client import ArcticRLClient
+ from arctic_training.arctic_rl.config import ArcticRLClientConfig
+
+ model_name = trainer_cfg.model.name
+ client_config = ArcticRLClientConfig(
+ backend="local",
+ model_name=model_name,
+ training_config=_build_training_config(trainer_cfg, arctic_cfg),
+ vllm_config=_build_vllm_config(arctic_cfg, trainer_cfg),
+ training_gpus=arctic_cfg.training_gpus,
+ sampling_gpus=arctic_cfg.sampling_tensor_parallel_size,
+ log_prob_gpus=arctic_cfg.log_prob_gpus,
+ )
+ logger.info("Initializing ArcticRLClient (model={})", model_name)
+ client = ArcticRLClient(client_config)
+ logger.info(
+ "ArcticRLClient ready: training={} sampling={} log_prob={}",
+ client.training_job_id,
+ client.sampling_job_id,
+ getattr(client, "log_prob_job_id", None),
+ )
+ return client
diff --git a/integrations/arctic-rl/arctic_rl/config.py b/integrations/arctic-rl/arctic_rl/config.py
new file mode 100644
index 0000000000..0e1506ef1a
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/config.py
@@ -0,0 +1,85 @@
+"""Arctic adapter config.
+
+`ArcticConfig` is attached to `RLConfig.arctic`. When `backend` is None, the
+entire adapter is inactive and Prime-RL runs natively.
+"""
+
+from typing import Annotated, Any, Literal
+
+from pydantic import Field
+
+from prime_rl.utils.config import BaseConfig
+
+
+class ArcticConfig(BaseConfig):
+ """Arctic RL integration config.
+
+ When `backend = "local"`, the launcher dispatches to the Arctic
+ entrypoint, which spawns the Arctic RL server in-process via
+ `ArcticRLClient` and replaces prime-rl's native trainer + vLLM with
+ HTTP calls to it.
+ """
+
+ backend: Annotated[
+ Literal["local"] | None,
+ Field(
+ description=(
+ "Arctic backend mode. 'local' spawns the Arctic RL server in-process. "
+ "None (default) disables Arctic and uses the native Prime-RL path."
+ )
+ ),
+ ] = None
+
+ training_gpus: Annotated[
+ int,
+ Field(description="GPUs requested for the training job.", gt=0),
+ ] = 2
+ sampling_tensor_parallel_size: Annotated[
+ int,
+ Field(description="Tensor-parallel size for the sampling engine (forwarded as vLLM tensor_parallel_size).", gt=0),
+ ] = 1
+ log_prob_gpus: Annotated[
+ int,
+ Field(description="GPUs for the log-prob job. 0 disables it."),
+ ] = 0
+
+ vllm_config: Annotated[
+ dict[str, Any] | None,
+ Field(
+ description=(
+ "Optional vLLM overrides forwarded to Arctic RL sampling/log-prob engines. "
+ "max_model_len defaults to trainer.model.seq_len when omitted."
+ )
+ ),
+ ] = None
+
+ enable_thinking: Annotated[
+ bool,
+ Field(description="Pass enable_thinking to tokenizer.apply_chat_template."),
+ ] = False
+
+ use_cispo_loss: Annotated[
+ bool,
+ Field(
+ description=(
+ "Use CISPO (Clipped IS-weight Policy Optimization) instead of vanilla "
+ "PPO-CLIP. CISPO clips the importance-sampling ratio with a stop-gradient, "
+ "so every token including clipped ones contributes a non-zero gradient. "
+ "Recommended for async-pipeline off-policy training (ScaleRL §3.2). "
+ "Uses asymmetric clips eps=0.2 / eps_higher=0.28 per the paper."
+ )
+ ),
+ ] = False
+
+ loss_agg_mode: Annotated[
+ str,
+ Field(
+ description=(
+ "Loss aggregation mode forwarded to the server's grpo_loss. "
+ "'token-mean' (default): average over all unmasked tokens. "
+ "'prompt-mean': average per rollout first, then across rollouts — "
+ "requires prompt_group_ids (automatically built when example_ids are tracked). "
+ "Other valid values: 'seq-mean-token-sum', 'seq-mean-token-sum-norm', 'seq-mean-token-mean'."
+ )
+ ),
+ ] = "token-mean"
diff --git a/integrations/arctic-rl/arctic_rl/context.py b/integrations/arctic-rl/arctic_rl/context.py
new file mode 100644
index 0000000000..20135d75a7
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/context.py
@@ -0,0 +1,118 @@
+"""Translate PRIME-RL packed microbatches to Arctic RL's batch format.
+
+PRIME-RL packs multiple rollouts into ``[1, T]`` microbatches; Arctic RL
+expects ``[B, max_S]`` with one rollout per row. We unpack, apply a
+per-rollout left-shift (so labels match Arctic's shifted-index convention)
+and zero the loss mask at the wrap-around position.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from arctic_rl.unpack import iter_rollout_slices, unpack_packed_microbatch
+from prime_rl.trainer.rl.data import TensorMicroBatch
+
+# Pad values mirror PRIME-RL's pad_micro_batch so the padded batch is
+# numerically indistinguishable from a natively-padded one.
+_PAD_VALUES: dict[str, float | int | bool] = {
+ "input_ids": 1,
+ "position_ids": 0,
+ "old_log_probs_shifted": 0.0,
+ "advantages": 0.0,
+ "teacher_log_probs_shifted": 0.0,
+ "loss_mask": False,
+}
+
+
+def _extract_and_roll_rollout(mb: TensorMicroBatch, start: int, end: int) -> dict:
+ """Extract one rollout's slice from a packed mb, applying per-rollout `torch.roll(-1)`.
+
+ Returns 1-D tensors (no batch dim). The roll is scoped to this rollout so
+ cross-rollout contamination at boundary positions is impossible.
+ """
+ out: dict = {
+ "input_ids": mb["input_ids"][0, start:end].clone(),
+ "position_ids": mb["position_ids"][0, start:end].clone(),
+ "old_log_probs_shifted": torch.roll(mb["inference_logprobs"][0, start:end], shifts=-1, dims=-1),
+ "advantages": torch.roll(mb["advantages"][0, start:end], shifts=-1, dims=-1),
+ }
+ loss_mask = torch.roll(mb["loss_mask"][0, start:end], shifts=-1, dims=-1).clone()
+ # The last position wraps around within this rollout; not a valid target.
+ loss_mask[-1] = False
+ out["loss_mask"] = loss_mask
+
+ teacher = mb.get("teacher_logprobs")
+ if teacher is not None:
+ out["teacher_log_probs_shifted"] = torch.roll(teacher[0, start:end], shifts=-1, dims=-1)
+ return out
+
+
+def microbatches_to_arctic_context(mbs: list[TensorMicroBatch]) -> dict:
+ """Consolidate all rollouts from a list of packed microbatches into one `[B, max_S]` batch.
+
+ Output:
+ A dict of `[B_total, max_S]` tensors where `B_total = sum of rollouts
+ across all mbs` (trailing pad per mb is dropped) and `max_S` is the
+ longest rollout across the batch. Includes a real `attention_mask`.
+ """
+ assert mbs, "Expected at least one microbatch"
+
+ rollouts: list[dict] = []
+ raw_example_ids: list[int] = []
+ for mb in mbs:
+ slices = iter_rollout_slices(mb)
+ mb_ids = mb.get("example_ids") or []
+ for i, (start, end) in enumerate(slices):
+ rollouts.append(_extract_and_roll_rollout(mb, start, end))
+ if i < len(mb_ids):
+ raw_example_ids.append(mb_ids[i])
+
+ b = len(rollouts)
+ max_s = max(r["input_ids"].shape[0] for r in rollouts)
+ ref = rollouts[0]
+
+ out: dict = {}
+ for key, ref_tensor in ref.items():
+ pad_val = _PAD_VALUES.get(key, 0)
+ tensor = torch.full((b, max_s), pad_val, dtype=ref_tensor.dtype, device=ref_tensor.device)
+ for i, r in enumerate(rollouts):
+ length = r[key].shape[0]
+ tensor[i, :length] = r[key]
+ out[key] = tensor
+
+ attention_mask = torch.zeros((b, max_s), dtype=ref["input_ids"].dtype, device=ref["input_ids"].device)
+ for i, r in enumerate(rollouts):
+ attention_mask[i, : r["input_ids"].shape[0]] = 1
+ out["attention_mask"] = attention_mask
+
+ # Build prompt_group_ids for prompt-mean loss aggregation. Maps each rollout row to its
+ # example group so agg_loss("prompt-mean") averages token losses per rollout before
+ # averaging across rollouts. Only present when all rollouts carried an example_id.
+ if len(raw_example_ids) == b:
+ out["prompt_group_ids"] = torch.tensor(raw_example_ids, dtype=torch.long)
+
+ return out
+
+
+def microbatch_to_arctic_context(mb: TensorMicroBatch) -> dict:
+ """Legacy single-microbatch path kept for unit tests.
+
+ Prefer `microbatches_to_arctic_context([mb, ...])` in the training loop
+ so that a single-rollout bin doesn't trip the server's DP-chunk assertion.
+ """
+ loss_mask = torch.roll(mb["loss_mask"], shifts=-1, dims=-1).clone()
+ loss_mask[..., -1] = False
+
+ rolled: dict = {
+ "input_ids": mb["input_ids"],
+ "position_ids": mb["position_ids"],
+ "old_log_probs_shifted": torch.roll(mb["inference_logprobs"], shifts=-1, dims=-1),
+ "advantages": torch.roll(mb["advantages"], shifts=-1, dims=-1),
+ "loss_mask": loss_mask,
+ }
+ teacher = mb.get("teacher_logprobs")
+ if teacher is not None:
+ rolled["teacher_log_probs_shifted"] = torch.roll(teacher, shifts=-1, dims=-1)
+
+ return unpack_packed_microbatch(rolled)
diff --git a/integrations/arctic-rl/arctic_rl/entrypoint.py b/integrations/arctic-rl/arctic_rl/entrypoint.py
new file mode 100644
index 0000000000..9cf16547ec
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/entrypoint.py
@@ -0,0 +1,278 @@
+"""Arctic RL backend entrypoint.
+
+Selected via `[trainer] backend = "arctic_rl"` in `rl.toml`. Dispatched
+generically by `prime_rl.entrypoints.rl:main`, which dynamically imports
+`arctic_rl.entrypoint:main` — no Arctic-specific code lives in core
+prime-rl.
+
+The integration parses its own extended config (RLConfig + an `arctic`
+block), then spawns:
+ - `arctic-trainer` (single-process HTTP client to the Arctic RL server)
+ - `arctic-shim` (FastAPI OAI-compat proxy to Arctic RL `/generate`)
+ - the prime-rl orchestrator (unchanged), pointed at the shim's base_url.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import signal
+import sys
+import time
+from pathlib import Path
+from subprocess import Popen
+from threading import Event, Thread
+from typing import Annotated
+
+import tomli_w
+from pydantic import Field, model_validator
+
+from prime_rl.configs.rl import RLConfig
+from prime_rl.utils.config import cli
+from prime_rl.utils.logger import setup_logger
+from prime_rl.utils.pathing import (
+
+ get_log_dir,
+
+)
+from prime_rl.utils.process import cleanup_processes, cleanup_threads, monitor_process, set_proc_title
+
+from arctic_rl.config import ArcticConfig
+
+ARCTIC_TOML = "arctic.toml"
+TRAINER_TOML = "trainer.toml"
+ORCHESTRATOR_TOML = "orchestrator.toml"
+
+
+class ArcticRLConfig(RLConfig):
+ """Extended RLConfig with the integration-specific `arctic` block.
+
+ Only loaded when the launcher dispatches to `arctic_rl.entrypoint:main`
+ (i.e. when `trainer.backend = "arctic_rl"`). Core prime-rl never sees
+ this schema.
+ """
+
+ arctic: Annotated[
+ ArcticConfig,
+ Field(description="Arctic RL backend config. Required when trainer.backend = 'arctic_rl'."),
+ ]
+
+ @model_validator(mode="after")
+ def _validate_arctic_exclusivity(self) -> "ArcticRLConfig":
+ # `[inference]` is a no-op in Arctic mode — Arctic owns sampling via
+ # its own server, the native vLLM block is silently ignored. We don't
+ # reject it so a native recipe can be flipped to Arctic with just an
+ # overlay snippet (no edits to the original file).
+ if self.inference is not None:
+ self.inference = None
+ if getattr(self, "teacher_inference", None) is not None:
+ raise ValueError("Arctic mode does not support teacher_inference.")
+ if self.deployment.type == "multi_node":
+ raise ValueError("Arctic is single-node only in the initial integration.")
+ if self.trainer.model.lora is not None:
+ raise ValueError("Arctic does not support LoRA.")
+ if self.trainer.max_concurrent_runs > 1:
+ raise ValueError("Arctic does not support multi-run.")
+ return self
+
+
+def write_arctic_subconfigs(config: ArcticRLConfig, output_dir: Path) -> None:
+ """Write trainer/orchestrator/arctic subconfigs to disk."""
+ output_dir.mkdir(parents=True, exist_ok=True)
+ with open(output_dir / TRAINER_TOML, "wb") as f:
+ tomli_w.dump(config.trainer.model_dump(exclude_none=True, mode="json"), f)
+ with open(output_dir / ORCHESTRATOR_TOML, "wb") as f:
+ tomli_w.dump(config.orchestrator.model_dump(exclude_none=True, mode="json"), f)
+ with open(output_dir / ARCTIC_TOML, "wb") as f:
+ tomli_w.dump(config.arctic.model_dump(exclude_none=True, mode="json"), f)
+
+
+def _run_arctic(config: ArcticRLConfig) -> None:
+ """Launch arctic-trainer + orchestrator (no shim).
+
+ The orchestrator's client is rewired to load
+ ``arctic_rl.verifiers_backend.ArcticClient`` via verifiers'
+ ``client_type="custom"`` extension; the ArcticClient calls the Arctic
+ RL server directly in-process, replacing the FastAPI shim subprocess
+ we used in the earlier version of this integration.
+ """
+ arctic_cfg = config.arctic
+ logger = setup_logger(
+ config.log.level or os.environ.get("PRIME_LOG_LEVEL", "info"),
+ json_logging=config.log.json_logging,
+ )
+
+ config_dir = config.output_dir / "configs"
+ write_arctic_subconfigs(config, config_dir)
+ arctic_toml_path = config_dir / ARCTIC_TOML
+ logger.info(f"Wrote subconfigs (including arctic.toml) to {config_dir}")
+
+ if config.dry_run:
+ logger.success("Dry run complete (arctic mode).")
+ return
+
+ log_dir = get_log_dir(config.output_dir)
+ log_dir.mkdir(parents=True, exist_ok=True)
+
+ processes: list[Popen] = []
+ monitor_threads: list[Thread] = []
+ error_queue: list[Exception] = []
+ stop_events: dict[str, Event] = {}
+
+ def sigterm_handler(signum, frame):
+ logger.warning("Received SIGTERM, terminating arctic processes...")
+ cleanup_threads(monitor_threads)
+ cleanup_processes(processes)
+ sys.exit(1)
+
+ signal.signal(signal.SIGTERM, sigterm_handler)
+
+ try:
+ # 1. Start arctic-trainer. Trainer signals readiness by writing
+ # reconnect.json; clear any stale copy first.
+ (config_dir / "reconnect.json").unlink(missing_ok=True)
+
+ trainer_log_path = log_dir / "arctic_trainer.log"
+ trainer_cmd = ["arctic-trainer", "@", (config_dir / TRAINER_TOML).as_posix()]
+ logger.info("Starting arctic-trainer (%s)", " ".join(trainer_cmd))
+ trainer_env = {
+ **os.environ,
+ "ARCTIC_CONFIG_TOML": arctic_toml_path.as_posix(),
+ "PYTHONUNBUFFERED": "1",
+ "LOGURU_FORCE_COLORS": "1",
+ }
+ # Ray 2.40+ auto-detects uv from env vars and re-launches its raylet
+ # workers via `uv run python`, which lands them in a freshly-created
+ # empty venv without ray. Strip uv-specific vars and disable Ray's uv
+ # hook so workers run under sys.executable (the prime-rl venv python).
+ for _k in (
+ "VIRTUAL_ENV",
+ "UV_PROJECT_ENVIRONMENT",
+ "UV_ACTIVE",
+ "UV_PROJECT",
+ "UV_RUN_RECURSION_DEPTH",
+ ):
+ trainer_env.pop(_k, None)
+ trainer_env["RAY_ENABLE_UV_RUN_RUNTIME_ENV"] = "0"
+ with open(trainer_log_path, "w") as f:
+ trainer_process = Popen(trainer_cmd, env=trainer_env, stdout=f, stderr=f)
+ processes.append(trainer_process)
+
+ # Wait for reconnect.json (no timeout; trainer-side init can take 10+ min).
+ reconnect_path = config_dir / "reconnect.json"
+ logger.info("Waiting for arctic-trainer to initialize Arctic RL jobs…")
+ while not reconnect_path.exists():
+ if trainer_process.poll() is not None:
+ raise RuntimeError(
+ f"arctic-trainer exited (code {trainer_process.returncode}) "
+ f"before writing reconnect.json. See {trainer_log_path}."
+ )
+ time.sleep(1)
+ logger.success("arctic-trainer ready (reconnect.json written)")
+
+ stop_event = Event()
+ stop_events["arctic_trainer"] = stop_event
+ t = Thread(
+ target=monitor_process,
+ args=(trainer_process, stop_event, error_queue, "arctic_trainer"),
+ daemon=True,
+ )
+ t.start()
+ monitor_threads.append(t)
+
+ # 2. Rewrite orchestrator config to use the in-process ArcticClient
+ # via verifiers' custom client_type extension. Arctic RL connection
+ # info is passed via env var read by the ArcticClient at first call.
+ config.orchestrator.student.client.client_type = "custom"
+ config.orchestrator.student.client.class_path = (
+ "arctic_rl.verifiers_backend.ArcticClient"
+ )
+ # The Arctic RL coordinator (spawned in-process by ArcticRLClient with
+ # backend="local") listens on localhost:7000.
+ arctic_url = "http://localhost:7000"
+ config.orchestrator.student.client.base_url = [arctic_url]
+ config.orchestrator.student.client.skip_model_check = True
+ write_arctic_subconfigs(config, config_dir)
+
+ orch_env = {
+ **os.environ,
+ "LOGURU_FORCE_COLORS": "1",
+ "WANDB_PROGRAM": "uv run rl (arctic)",
+ "WANDB_ARGS": json.dumps(sys.argv),
+ "ARCTIC_RL_RECONNECT_CONFIG": reconnect_path.as_posix(),
+ "PRIME_RL_DIRECT_ARCTIC_RECONNECT_CONFIG": reconnect_path.as_posix(),
+ "PRIME_RL_ARCTIC_TOKENIZER_NAME": config.trainer.model.name,
+ "PRIME_RL_ARCTIC_ENABLE_THINKING": "1" if arctic_cfg.enable_thinking else "0",
+ }
+
+ # 3. Start orchestrator.
+ orch_cmd = ["orchestrator", "@", (config_dir / ORCHESTRATOR_TOML).as_posix()]
+ logger.info("Starting orchestrator")
+ with open(log_dir / "orchestrator.log", "w") as f:
+ orch_process = Popen(orch_cmd, stdout=f, stderr=f, env=orch_env)
+ processes.append(orch_process)
+
+ stop_event = Event()
+ stop_events["orchestrator"] = stop_event
+ t = Thread(
+ target=monitor_process,
+ args=(orch_process, stop_event, error_queue, "orchestrator"),
+ daemon=True,
+ )
+ t.start()
+ monitor_threads.append(t)
+
+ logger.success("Arctic startup complete. Tailing trainer log...")
+ tail_process = Popen(f"tail -F '{trainer_log_path}'", shell=True)
+ processes.append(tail_process)
+
+ while not (stop_events["orchestrator"].is_set() and stop_events["arctic_trainer"].is_set()):
+ if error_queue:
+ logger.error(f"Error: {error_queue[0]}")
+ logger.error("Terminating all arctic processes...")
+ cleanup_threads(monitor_threads)
+ cleanup_processes(processes)
+ sys.exit(1)
+ time.sleep(1)
+
+ if orch_process.returncode != 0:
+ logger.error(f"Orchestrator failed with exit code {orch_process.returncode}")
+ cleanup_threads(monitor_threads)
+ cleanup_processes(processes)
+ sys.exit(1)
+ if trainer_process.returncode != 0:
+ logger.error(f"Arctic trainer failed with exit code {trainer_process.returncode}")
+ cleanup_threads(monitor_threads)
+ cleanup_processes(processes)
+ sys.exit(1)
+
+ logger.success("Arctic RL training finished.")
+ cleanup_threads(monitor_threads)
+ cleanup_processes(processes)
+
+ except KeyboardInterrupt:
+ logger.warning("Received interrupt, terminating arctic processes...")
+ cleanup_threads(monitor_threads)
+ cleanup_processes(processes)
+ sys.exit(1)
+ except Exception:
+ cleanup_threads(monitor_threads)
+ cleanup_processes(processes)
+ raise
+
+
+
+def main() -> None:
+ """Console-script entrypoint for the `arctic_rl` backend.
+
+ Dispatched generically by `prime_rl.entrypoints.rl:main` when
+ `trainer.backend = "arctic_rl"`. Parses the integration's extended
+ schema (`ArcticRLConfig`) and runs the launcher.
+ """
+ set_proc_title("Launcher (arctic_rl)")
+ config = cli(ArcticRLConfig)
+ _run_arctic(config)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/integrations/arctic-rl/arctic_rl/loss_config.py b/integrations/arctic-rl/arctic_rl/loss_config.py
new file mode 100644
index 0000000000..9936385a87
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/loss_config.py
@@ -0,0 +1,41 @@
+"""Build the config dict sent to Arctic RL's grpo_loss on each /fwd-bwd call."""
+
+from __future__ import annotations
+
+from loguru import logger
+
+_CISPO_LOGGED = False
+
+
+def build_grpo_loss_config(
+ current_version: int,
+ eps_clip: float = 0.2,
+ use_cispo: bool = False,
+ loss_agg_mode: str = "token-mean",
+) -> dict:
+ """Build the ``processing.config`` dict for Arctic RL's grpo_loss."""
+ global _CISPO_LOGGED
+ if use_cispo:
+ if not _CISPO_LOGGED:
+ logger.info(f"LOSS CONFIG: CISPO (eps=0.2/0.28, agg={loss_agg_mode})")
+ _CISPO_LOGGED = True
+ return {
+ "use_cispo_loss": True,
+ "eps_clip": eps_clip,
+ "eps_clip_higher": 0.28,
+ "loss_agg_mode": loss_agg_mode,
+ "prox_logp_method": "recompute",
+ "importance_sampling_level": "token",
+ "current_version": current_version,
+ }
+
+ if not _CISPO_LOGGED:
+ logger.info(f"LOSS CONFIG: vanilla PPO (agg={loss_agg_mode})")
+ _CISPO_LOGGED = True
+ return {
+ "eps_clip": eps_clip,
+ "loss_agg_mode": loss_agg_mode,
+ "prox_logp_method": "recompute",
+ "importance_sampling_level": "token",
+ "current_version": current_version,
+ }
diff --git a/integrations/arctic-rl/arctic_rl/oai_translation.py b/integrations/arctic-rl/arctic_rl/oai_translation.py
new file mode 100644
index 0000000000..5dc947eee5
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/oai_translation.py
@@ -0,0 +1,227 @@
+"""Pure OpenAI-chat <-> Arctic /generate translation functions."""
+
+from __future__ import annotations
+
+import json
+import re
+import time
+import uuid
+from typing import Any
+
+TOOL_CALL_RE = re.compile(r"\s*(.*?)\s*", re.DOTALL)
+
+
+def oai_sampling_params(body: dict) -> dict:
+ """Extract Arctic sampling params from an OpenAI chat-completions body."""
+ params: dict[str, Any] = {}
+
+ if (temperature := body.get("temperature")) is not None:
+ params["temperature"] = float(temperature)
+ if (top_p := body.get("top_p")) is not None:
+ params["top_p"] = float(top_p)
+ if (top_k := body.get("top_k")) is not None:
+ params["top_k"] = int(top_k)
+ if (max_tokens := body.get("max_tokens") or body.get("max_completion_tokens")) is not None:
+ params["max_tokens"] = int(max_tokens)
+
+ nested = body.get("extra_body") or {}
+ if (stop := body.get("stop", nested.get("stop"))) is not None:
+ params["stop"] = stop
+ if (n := body.get("n")) is not None:
+ params["n"] = int(n)
+ if (seed := body.get("seed")) is not None:
+ params["seed"] = int(seed)
+
+ if body.get("logprobs") is True and (top_logprobs := body.get("top_logprobs")) is not None:
+ params["logprobs"] = int(top_logprobs)
+ elif body.get("logprobs") is True:
+ params["logprobs"] = 1
+
+ for key in (
+ "min_tokens",
+ "repetition_penalty",
+ "include_stop_str_in_output",
+ "skip_special_tokens",
+ "return_sampled_logprobs_only",
+ ):
+ if key in body:
+ params[key] = body[key]
+ elif key in nested:
+ params[key] = nested[key]
+
+ return params
+
+
+def extract_routing_metadata(body: dict) -> tuple[str | None, bool]:
+ """Pull the optional Arctic routing annotation from an OpenAI request body."""
+ nested = body.get("extra_body") or {}
+ routing_key = body.get("routing_key")
+ if routing_key is None:
+ routing_key = nested.get("routing_key")
+ if routing_key is not None and not isinstance(routing_key, str):
+ routing_key = str(routing_key)
+ strict = bool(body.get("routing_strict") or nested.get("routing_strict"))
+ return routing_key, strict
+
+
+def _serialize_tool_arguments(arguments: Any) -> str | None:
+ if arguments is None:
+ return "{}"
+ if isinstance(arguments, str):
+ return arguments
+ if isinstance(arguments, dict):
+ return json.dumps(arguments)
+ return None
+
+
+def _tool_call_from_payload(payload: dict[str, Any], index: int) -> dict[str, Any] | None:
+ function = payload.get("function")
+ if isinstance(function, dict):
+ name = function.get("name")
+ arguments = function.get("arguments")
+ else:
+ name = payload.get("name")
+ arguments = payload.get("arguments")
+
+ if not isinstance(name, str) or not name:
+ return None
+
+ serialized_arguments = _serialize_tool_arguments(arguments)
+ if serialized_arguments is None:
+ return None
+
+ tool_call_id = payload.get("id")
+ if not isinstance(tool_call_id, str) or not tool_call_id:
+ tool_call_id = f"call_{index}"
+
+ return {
+ "id": tool_call_id,
+ "type": "function",
+ "function": {
+ "name": name,
+ "arguments": serialized_arguments,
+ },
+ }
+
+
+def parse_tool_calls(text: str) -> list[dict[str, Any]] | None:
+ matches = TOOL_CALL_RE.findall(text)
+ if not matches:
+ return None
+
+ tool_calls: list[dict[str, Any]] = []
+ for raw_payload in matches:
+ try:
+ parsed_payload = json.loads(raw_payload)
+ except json.JSONDecodeError:
+ return None
+
+ payloads = parsed_payload if isinstance(parsed_payload, list) else [parsed_payload]
+ for payload in payloads:
+ if not isinstance(payload, dict):
+ return None
+ tool_call = _tool_call_from_payload(payload, len(tool_calls))
+ if tool_call is None:
+ return None
+ tool_calls.append(tool_call)
+
+ return tool_calls or None
+
+
+def _tool_content_from_text(text: str) -> str | None:
+ content = TOOL_CALL_RE.sub("", text).strip()
+ return content or None
+
+
+def _logprob_value(position: Any, token_id: int) -> float:
+ if isinstance(position, (int, float)):
+ return float(position)
+ if not isinstance(position, dict) or not position:
+ raise ValueError(f"Missing logprob for sampled token_id={token_id}")
+
+ if "logprob" in position:
+ return float(position["logprob"])
+
+ entry = position.get(token_id)
+ if entry is None:
+ entry = position.get(str(token_id))
+ if entry is None:
+ raise ValueError(f"Missing logprob for sampled token_id={token_id}")
+
+ if isinstance(entry, dict):
+ if "logprob" not in entry:
+ raise ValueError(f"Missing logprob for sampled token_id={token_id}")
+ return float(entry["logprob"])
+ return float(entry)
+
+
+def _oai_logprobs_content(token_ids: list[int], logprobs: Any) -> list[dict[str, Any]]:
+ if not isinstance(logprobs, list):
+ return []
+ return [
+ {
+ "token": "",
+ "bytes": [],
+ "logprob": _logprob_value(logprobs[index] if index < len(logprobs) else None, token_id),
+ "top_logprobs": [],
+ }
+ for index, token_id in enumerate(token_ids)
+ ]
+
+
+def _arctic_result_to_oai_choice(result: dict, index: int) -> dict:
+ """Convert one Arctic /generate result to one OpenAI chat choice."""
+ text = result.get("text", "")
+ tool_calls = parse_tool_calls(text) if isinstance(text, str) else None
+ message: dict[str, Any] = {
+ "role": "assistant",
+ "content": text,
+ }
+ finish_reason = result.get("finish_reason", "stop")
+ if tool_calls is not None:
+ message = {
+ "role": "assistant",
+ "content": _tool_content_from_text(text),
+ "tool_calls": tool_calls,
+ }
+ finish_reason = "tool_calls"
+
+ choice: dict[str, Any] = {
+ "index": index,
+ "message": message,
+ "finish_reason": finish_reason,
+ }
+ token_ids = result.get("token_ids")
+ if isinstance(token_ids, list):
+ choice["token_ids"] = token_ids
+ if (logprobs := result.get("logprobs")) is not None:
+ choice["logprobs"] = {"content": _oai_logprobs_content(token_ids, logprobs)}
+ return choice
+
+
+def _arctic_results_to_oai_response(
+ results: list[dict],
+ model: str,
+ request_id: str | None = None,
+ prompt_token_ids: list[int] | None = None,
+) -> dict:
+ """Shape a full OpenAI chat.completion response from Arctic results."""
+ completion_tokens = sum(
+ len(token_ids) for result in results if isinstance((token_ids := result.get("token_ids")), list)
+ )
+ prompt_tokens = len(prompt_token_ids or [])
+ response = {
+ "id": request_id or f"chatcmpl-{uuid.uuid4().hex}",
+ "object": "chat.completion",
+ "created": int(time.time()),
+ "model": model,
+ "choices": [_arctic_result_to_oai_choice(result, index) for index, result in enumerate(results)],
+ "usage": {
+ "prompt_tokens": prompt_tokens,
+ "completion_tokens": completion_tokens,
+ "total_tokens": prompt_tokens + completion_tokens,
+ },
+ }
+ if prompt_token_ids is not None:
+ response["prompt_token_ids"] = prompt_token_ids
+ return response
diff --git a/integrations/arctic-rl/arctic_rl/trainer.py b/integrations/arctic-rl/arctic_rl/trainer.py
new file mode 100644
index 0000000000..3492db9ab1
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/trainer.py
@@ -0,0 +1,230 @@
+"""Arctic trainer adapter.
+
+Replaces PRIME-RL's torchrun+FSDP2 trainer with a single CPU process that
+issues ``/fwd-bwd``, ``/step``, and ``/sync-weights`` over HTTP to the
+Arctic RL server. Reuses PRIME-RL's DataLoader and writes the STABLE
+marker after each weight sync so the orchestrator's polling loop is
+unchanged.
+"""
+
+from __future__ import annotations
+
+import os
+import socket
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from loguru import logger
+
+from arctic_rl.client import build_arctic_client
+from arctic_rl.config import ArcticConfig
+from arctic_rl.context import microbatches_to_arctic_context
+from arctic_rl.loss_config import build_grpo_loss_config
+from prime_rl.configs.trainer import TrainerConfig
+from prime_rl.trainer.runs import get_multi_run_manager
+from prime_rl.trainer.scheduler import setup_scheduler
+
+
+def _get_free_port() -> int:
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+ s.bind(("127.0.0.1", 0))
+ return s.getsockname()[1]
+
+
+def _init_single_process_dist() -> None:
+ """world_size=1 gloo group so PRIME-RL's DataLoader can call get_world()."""
+ if dist.is_initialized():
+ return
+ os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+ os.environ.setdefault("MASTER_PORT", str(_get_free_port()))
+ os.environ.setdefault("RANK", "0")
+ os.environ.setdefault("WORLD_SIZE", "1")
+ os.environ.setdefault("LOCAL_RANK", "0")
+ dist.init_process_group(backend="gloo", world_size=1, rank=0)
+
+
+def _write_stable_marker(broadcast_dir: Path, step: int) -> None:
+ """Write the zero-byte STABLE file the orchestrator polls before each step."""
+ stable_path = broadcast_dir / f"step_{step}" / "STABLE"
+ stable_path.parent.mkdir(parents=True, exist_ok=True)
+ stable_path.touch()
+
+
+def _concat_microbatches(mbs: list[dict]) -> dict:
+ """Concatenate the DataLoader's per-DP microbatches into one consolidated batch."""
+ if len(mbs) == 1:
+ return dict(mbs[0])
+
+ out: dict = {}
+ for key, value in mbs[0].items():
+ if value is None:
+ out[key] = None
+ elif isinstance(value, torch.Tensor):
+ out[key] = torch.cat([mb[key] for mb in mbs], dim=0)
+ else:
+ out[key] = value
+ return out
+
+
+class ArcticTrainerAdapter:
+ """Single-process trainer that delegates fwd/bwd/optimizer to Arctic RL.
+
+ Reuses PRIME-RL's DataLoader and scheduler unchanged; writes the
+ STABLE marker after each ``sync_weights`` so the orchestrator's
+ polling loop is unaffected.
+ """
+
+ def __init__(self, trainer_cfg: TrainerConfig, arctic_cfg: ArcticConfig):
+ self.trainer_cfg = trainer_cfg
+ self.arctic_cfg = arctic_cfg
+
+ _init_single_process_dist()
+
+ from transformers import AutoTokenizer
+
+ from prime_rl.trainer.rl.data import DataLoader, FakeDataLoader
+ from prime_rl.trainer.runs import Progress, setup_multi_run_manager
+
+ self.progress = Progress()
+
+ # MultiRunManager is a singleton DataLoader depends on.
+ setup_multi_run_manager(
+ output_dir=trainer_cfg.output_dir,
+ max_runs=trainer_cfg.max_concurrent_runs,
+ device=torch.device("cpu"),
+ )
+
+ if trainer_cfg.data.fake is not None:
+ self.loader = FakeDataLoader(
+ config=trainer_cfg.data.fake,
+ seq_len=trainer_cfg.data.fake.seq_len if hasattr(trainer_cfg.data.fake, "seq_len") else 512,
+ dp_world_size=1,
+ )
+ self._using_fake = True
+ else:
+ tokenizer = AutoTokenizer.from_pretrained(trainer_cfg.tokenizer.name)
+ self.loader = DataLoader(
+ output_dir=trainer_cfg.output_dir,
+ start_step=self.progress.step,
+ dp_world_size=1,
+ seq_len=trainer_cfg.data.seq_len if hasattr(trainer_cfg.data, "seq_len") else 2048,
+ pad_to_multiple_of=trainer_cfg.data.pad_to_multiple_of
+ if hasattr(trainer_cfg.data, "pad_to_multiple_of")
+ else 64,
+ tokenizer=tokenizer,
+ config=trainer_cfg.rollout_transport,
+ )
+ self._using_fake = False
+
+ self.client = build_arctic_client(arctic_cfg, trainer_cfg)
+
+ # Dummy single-parameter optimizer used only to drive the LR schedule.
+ # The actual optimizer lives server-side; we pass the computed LR via
+ # adam_params on each /step call, which overrides the server's own LR.
+ _dummy = torch.nn.Parameter(torch.zeros(1))
+ _optim = torch.optim.AdamW([_dummy], lr=trainer_cfg.optim.lr)
+ self._lr_scheduler = setup_scheduler(
+ _optim, trainer_cfg.scheduler, trainer_cfg.max_steps or 1, trainer_cfg.optim.lr
+ )
+ self._optim = _optim
+
+ # /run_default/broadcasts matches PRIME-RL's native
+ # broadcast directory layout.
+ self.broadcast_dir = Path(trainer_cfg.output_dir) / "run_default" / "broadcasts"
+
+ # ArcticRLClientConfig marks job ID fields as Field(exclude=True),
+ # so model_dump_json() drops them — build the payload manually.
+ reconnect_path = Path(trainer_cfg.output_dir) / "configs" / "reconnect.json"
+ reconnect_path.parent.mkdir(parents=True, exist_ok=True)
+ rc = self.client.reconnect_config()
+ import json as _json
+
+ reconnect_path.write_text(
+ _json.dumps(
+ {
+ "host": rc.host,
+ "port": rc.port,
+ "backend": rc.backend,
+ "model_name": rc.model_name,
+ "training_job_id": rc.training_job_id,
+ "sampling_job_id": rc.sampling_job_id,
+ "log_prob_job_id": rc.log_prob_job_id,
+ }
+ )
+ )
+ logger.info("Wrote reconnect config → {}", reconnect_path)
+
+ def run(self) -> None:
+ max_steps = self.trainer_cfg.max_steps or 100
+ while self.progress.step < max_steps:
+ step = self.progress.step
+ logger.info("Step {} — waiting for batch", step)
+ self.loader.wait_for_batch()
+ mbs = self.loader.get_batch()
+
+ processing = {
+ "loss_fn": "arctic_training.arctic_rl.processors.grpo_loss",
+ "config": build_grpo_loss_config(
+ current_version=step,
+ use_cispo=self.arctic_cfg.use_cispo_loss,
+ loss_agg_mode=self.arctic_cfg.loss_agg_mode,
+ ),
+ "post": ["compute_logprobs"],
+ }
+
+ # Send all rollouts in one consolidated [B_total, max_S] batch
+ # so the server's torch.chunk(dim=0, world_size) splits cleanly.
+ # The server repacks per DP shard before the model forward, so
+ # activation memory matches a per-microbatch packed call.
+ kwargs = microbatches_to_arctic_context(mbs)
+ logger.info(
+ "Step {} — /fwd-bwd (B={}, S={}, across {} source microbatch(es))",
+ step,
+ kwargs["input_ids"].shape[0],
+ kwargs["input_ids"].shape[1],
+ len(mbs),
+ )
+ # Write the scheduled LR before fwd-bwd so the orchestrator can
+ # read it during rollout generation for this same step.
+ self._lr_scheduler.step()
+ current_lr = self._optim.param_groups[0]["lr"]
+ (self.broadcast_dir.parent / "last_lr").write_text(str(current_lr))
+
+ result = self.client.fwd_bwd({"args": (), "kwargs": kwargs}, processing=processing)
+ avg_loss = result.get("avg_loss") or result.get("loss") or float("nan")
+ logger.info("Step {} — avg_loss={:.4f}", step, avg_loss)
+
+ logger.info("Step {} — /step", step)
+ # Older ArcticRLClient.step() takes no kwargs; newer accepts learning_rate.
+ import inspect as _inspect
+ try:
+ if "learning_rate" in _inspect.signature(self.client.step).parameters:
+ step_result = self.client.step(learning_rate=current_lr) or {}
+ else:
+ step_result = self.client.step() or {}
+ except (TypeError, ValueError):
+ step_result = self.client.step() or {}
+
+ if step > 0:
+ logger.info("Step {} — /sync-weights", step)
+ self.client.sync_weights()
+ _write_stable_marker(self.broadcast_dir, step)
+
+ # Clear ready_to_update for the next step. Native PRIME-RL does
+ # this via FileSystemWeightBroadcast.broadcast_weights() which we
+ # skip (Arctic owns the weights).
+ if not self._using_fake:
+ mgr = get_multi_run_manager()
+ for idx in mgr.used_idxs:
+ mgr.ready_to_update[idx] = False
+
+ self.progress.step += 1
+ logger.info(
+ "Step {} done — last_lr={} grad_norm={}",
+ step,
+ step_result.get("last_lr"),
+ step_result.get("grad_norm"),
+ )
+
+ logger.success("Training finished after {} steps", self.progress.step)
diff --git a/integrations/arctic-rl/arctic_rl/unpack.py b/integrations/arctic-rl/arctic_rl/unpack.py
new file mode 100644
index 0000000000..8a8666674b
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/unpack.py
@@ -0,0 +1,108 @@
+"""Unpack PRIME-RL packed ``[1, T]`` microbatches to ``[B, S]`` padded form.
+
+Required because Arctic RL's per-replica chunk along dim 0 fails when
+``world_size > 1`` and the input batch is ``[1, T]``.
+"""
+
+from __future__ import annotations
+
+import torch
+
+# Pad values mirror PRIME-RL's pad_micro_batch. loss_mask=False is the
+# authoritative guard at padded positions.
+_PAD_VALUES: dict[str, float | int | bool] = {
+ "input_ids": 1,
+ "position_ids": 0,
+ "advantages": 0.0,
+ "inference_logprobs": 0.0,
+ "old_log_probs_shifted": 0.0,
+ "teacher_logprobs": 0.0,
+ "teacher_log_probs_shifted": 0.0,
+ "loss_mask": False,
+ "temperatures": 1.0,
+}
+
+
+def _detect_rollout_starts(position_ids: torch.Tensor) -> list[int]:
+ """Indices where a new rollout begins in a packed ``[1, T]`` tensor.
+
+ A boundary is any position where position_ids decreases (or doesn't
+ increment by 1), plus position 0.
+ """
+ assert position_ids.dim() == 2 and position_ids.shape[0] == 1, (
+ f"Expected position_ids of shape [1, T], got {tuple(position_ids.shape)}"
+ )
+ pos_flat = position_ids.squeeze(0)
+ t = pos_flat.shape[0]
+ is_start = torch.zeros(t, dtype=torch.bool, device=pos_flat.device)
+ is_start[0] = True
+ is_start[1:] = pos_flat[1:] - pos_flat[:-1] != 1
+ return is_start.nonzero(as_tuple=True)[0].tolist()
+
+
+def iter_rollout_slices(mb: dict) -> list[tuple[int, int]]:
+ """Per-rollout ``(start, end)`` slices, dropping the trailing pad segment.
+
+ Guarded so a fully-masked real rollout in a non-trailing position
+ isn't silently discarded.
+ """
+ position_ids = mb["position_ids"]
+ loss_mask = mb["loss_mask"]
+
+ t = position_ids.shape[1]
+ starts = _detect_rollout_starts(position_ids)
+ ends = starts[1:] + [t]
+ slices = list(zip(starts, ends))
+
+ loss_mask_flat = loss_mask.squeeze(0)
+ if len(slices) > 1 and not loss_mask_flat[slices[-1][0] : slices[-1][1]].any():
+ slices = slices[:-1]
+
+ assert len(slices) >= 1, "Expected at least one rollout after dropping trailing pad"
+ return slices
+
+
+def unpack_packed_microbatch(rolled: dict) -> dict:
+ """Convert a ``[1, T]`` packed microbatch to ``[B, S]`` padded form.
+
+ Caller must have already applied any per-token alignment shift.
+ Adds an ``attention_mask`` to the output.
+ """
+ position_ids = rolled["position_ids"]
+ loss_mask = rolled["loss_mask"]
+
+ t = position_ids.shape[1]
+ starts = _detect_rollout_starts(position_ids)
+ n_segs = len(starts)
+ ends = starts[1:] + [t]
+ lengths = [e - s for s, e in zip(starts, ends)]
+
+ loss_mask_flat = loss_mask.squeeze(0)
+ if n_segs > 1 and not loss_mask_flat[starts[-1] : ends[-1]].any():
+ starts = starts[:-1]
+ ends = ends[:-1]
+ lengths = lengths[:-1]
+
+ assert len(lengths) >= 1, "Expected at least one rollout after dropping trailing pad"
+
+ b = len(lengths)
+ s = max(lengths)
+
+ out: dict = {}
+ for key, value in rolled.items():
+ if not isinstance(value, torch.Tensor) or value.dim() != 2 or value.shape[0] != 1:
+ out[key] = value
+ continue
+ pad_val = _PAD_VALUES.get(key, 0)
+ flat = value.squeeze(0)
+ padded = torch.full((b, s), pad_val, dtype=value.dtype, device=value.device)
+ for i, (seg_start, seg_len) in enumerate(zip(starts, lengths)):
+ padded[i, :seg_len] = flat[seg_start : seg_start + seg_len]
+ out[key] = padded
+
+ attention_mask = torch.zeros((b, s), dtype=position_ids.dtype, device=position_ids.device)
+ for i, seg_len in enumerate(lengths):
+ attention_mask[i, :seg_len] = 1
+ out["attention_mask"] = attention_mask
+
+ return out
diff --git a/integrations/arctic-rl/arctic_rl/verifiers_backend.py b/integrations/arctic-rl/arctic_rl/verifiers_backend.py
new file mode 100644
index 0000000000..e777754cad
--- /dev/null
+++ b/integrations/arctic-rl/arctic_rl/verifiers_backend.py
@@ -0,0 +1,780 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import threading
+import time
+import uuid
+from collections.abc import Mapping
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Any, cast
+
+import requests
+from openai.types.chat import ChatCompletion
+
+from arctic_rl.oai_translation import (
+ _arctic_results_to_oai_response,
+ extract_routing_metadata,
+ oai_sampling_params,
+)
+from loguru import logger as _logger
+
+def get_logger():
+ return _logger
+
+_DISABLED_ENV_VALUES = {"0", "false", "no", "off"}
+_DEFAULT_MAX_BATCH = 64
+_DEFAULT_MAX_INFLIGHT_BATCHES = 12
+_DEFAULT_FLUSH_INTERVAL_S = 0.05
+_DEFAULT_MAX_QUEUE = 4096
+_DEFAULT_LOG_EVERY_N = 256
+_DEFAULT_MAX_RETRIES = 2
+_DEFAULT_RETRY_BASE_DELAY_S = 0.2
+_TOKENIZER: Any | None = None
+_TOKENIZER_ENABLE_THINKING = False
+
+
+class _ArcticStreamUnavailable(RuntimeError):
+ pass
+
+
+def _env_flag(name: str, default: bool = False) -> bool:
+ value = os.environ.get(name)
+ if value is None:
+ return default
+ return value.lower() not in _DISABLED_ENV_VALUES
+
+
+def _env_int(name: str, default: int, *, minimum: int = 1) -> int:
+ value = os.environ.get(name)
+ if value is None or value == "":
+ return default
+ try:
+ parsed = int(value)
+ except ValueError:
+ get_logger().warning("{}={} is not an integer; using {}", name, value, default)
+ return default
+ return max(minimum, parsed)
+
+
+def _ensure_verifiers_tokenizer() -> Any:
+ global _TOKENIZER, _TOKENIZER_ENABLE_THINKING
+ if _TOKENIZER is not None:
+ return _TOKENIZER
+
+ tokenizer_name = os.environ.get("PRIME_RL_ARCTIC_TOKENIZER_NAME")
+ if not tokenizer_name:
+ raise RuntimeError("Arctic verifiers tokenizer is not registered and PRIME_RL_ARCTIC_TOKENIZER_NAME is not set")
+
+ from transformers import AutoTokenizer
+
+ trust_remote_code_env = os.environ.get("PRIME_RL_ARCTIC_TOKENIZER_TRUST_REMOTE_CODE")
+ trust_remote_code = (
+ None if trust_remote_code_env is None else trust_remote_code_env.lower() not in _DISABLED_ENV_VALUES
+ )
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=trust_remote_code)
+ if chat_template := os.environ.get("PRIME_RL_ARCTIC_TOKENIZER_CHAT_TEMPLATE"):
+ template_path = Path(chat_template)
+ tokenizer.chat_template = template_path.read_text() if template_path.is_file() else chat_template
+ tokenizer.pad_token_id = tokenizer.eos_token_id
+
+ _TOKENIZER = tokenizer
+ _TOKENIZER_ENABLE_THINKING = _env_flag("PRIME_RL_ARCTIC_ENABLE_THINKING", default=False)
+ get_logger().info(
+ "Loaded Arctic verifiers tokenizer from env (name={}, enable_thinking={})",
+ tokenizer_name,
+ _TOKENIZER_ENABLE_THINKING,
+ )
+ return _TOKENIZER
+
+
+ guard = getattr(native_response.choices[0], _REASONING_GUARD_KEY, None)
+ if isinstance(guard, dict):
+ return guard
+ return None
+
+
+ tokens = getattr(response.message, "tokens", None)
+ if tokens is None:
+ return response
+
+ completion_ids = list(tokens.completion_ids)
+ completion_mask = [int(value) for value in tokens.completion_mask]
+ completion_logprobs = list(tokens.completion_logprobs)
+ if len(completion_mask) != len(completion_ids) or len(completion_logprobs) != len(completion_ids):
+ raise RuntimeError(
+ "reasoning guard received misaligned completion token fields "
+ f"(ids={len(completion_ids)}, mask={len(completion_mask)}, logprobs={len(completion_logprobs)})"
+ )
+
+ for span in guard.get("masked_token_spans") or []:
+ if not isinstance(span, (list, tuple)) or len(span) != 2:
+ continue
+ start, end = int(span[0]), int(span[1])
+ start = max(0, min(start, len(completion_mask)))
+ end = max(start, min(end, len(completion_mask)))
+ for idx in range(start, end):
+ completion_mask[idx] = 0
+ completion_logprobs[idx] = 0.0
+
+ tokens.completion_mask = completion_mask
+ tokens.completion_logprobs = completion_logprobs
+ setattr(tokens, _REASONING_GUARD_KEY, guard)
+ response.message.tokens = tokens
+ setattr(response.message, _REASONING_GUARD_KEY, guard)
+ return response
+
+
+def _jsonable(value: Any) -> Any:
+ if hasattr(value, "model_dump"):
+ try:
+ value = value.model_dump(exclude_none=True)
+ except TypeError:
+ value = value.model_dump()
+ if isinstance(value, Mapping):
+ return {str(key): _jsonable(item) for key, item in value.items() if item is not None}
+ if isinstance(value, list):
+ return [_jsonable(item) for item in value]
+ return value
+
+
+def _coerce_token_ids(tokens: Any) -> list[int]:
+ if isinstance(tokens, Mapping):
+ tokens = tokens.get("input_ids")
+ elif hasattr(tokens, "data") and isinstance(tokens.data, Mapping):
+ tokens = tokens.data.get("input_ids")
+ elif not isinstance(tokens, list):
+ try:
+ tokens = tokens["input_ids"]
+ except (KeyError, TypeError, AttributeError):
+ pass
+ if hasattr(tokens, "tolist"):
+ tokens = tokens.tolist()
+ if isinstance(tokens, list) and len(tokens) == 1 and isinstance(tokens[0], list):
+ tokens = tokens[0]
+ if not isinstance(tokens, list):
+ raise TypeError(f"Expected tokenizer output to be a token-id list, got {type(tokens).__name__}")
+ return [int(token_id) for token_id in tokens]
+
+
+def _local_tokenize(
+ *,
+ messages: str | list[Any],
+ tools: list[Any] | None,
+ extra_kwargs: dict[str, Any] | None = None,
+) -> list[int]:
+ tokenizer = _ensure_verifiers_tokenizer()
+
+ extra_kwargs = dict(extra_kwargs or {})
+ if isinstance(messages, str):
+ return _coerce_token_ids(
+ tokenizer.encode(
+ messages,
+ add_special_tokens=bool(extra_kwargs.pop("add_special_tokens", False)),
+ )
+ )
+
+ template_kwargs: dict[str, Any] = {
+ "tokenize": True,
+ "add_generation_prompt": bool(extra_kwargs.pop("add_generation_prompt", True)),
+ "enable_thinking": extra_kwargs.pop("enable_thinking", _TOKENIZER_ENABLE_THINKING),
+ }
+ template_kwargs.update(extra_kwargs)
+ if tools is not None:
+ template_kwargs["tools"] = _jsonable(tools)
+ return _coerce_token_ids(tokenizer.apply_chat_template(_jsonable(messages), **template_kwargs))
+
+
+class _ArcticGenerateBatcher:
+ def __init__(self, reconnect_config: Path):
+ reconnect = json.loads(reconnect_config.read_text())
+ host = reconnect["host"]
+ port = int(reconnect["port"])
+ self.base_url = f"http://{host}:{port}"
+ self.sampling_job_id = int(reconnect["sampling_job_id"])
+ self.max_batch = max(1, int(os.environ.get("PRIME_RL_ARCTIC_MAX_BATCH", str(_DEFAULT_MAX_BATCH))))
+ self.max_inflight_batches = max(
+ 1, int(os.environ.get("PRIME_RL_ARCTIC_MAX_INFLIGHT_BATCHES", str(_DEFAULT_MAX_INFLIGHT_BATCHES)))
+ )
+ self.flush_interval_s = max(
+ 0.0, float(os.environ.get("PRIME_RL_ARCTIC_FLUSH_INTERVAL_S", str(_DEFAULT_FLUSH_INTERVAL_S)))
+ )
+ self.max_queue = max(1, int(os.environ.get("PRIME_RL_ARCTIC_MAX_QUEUE", str(_DEFAULT_MAX_QUEUE))))
+ self.log_every_n = max(0, int(os.environ.get("PRIME_RL_ARCTIC_LOG_EVERY_N", str(_DEFAULT_LOG_EVERY_N))))
+ self.max_retries = max(0, int(os.environ.get("PRIME_RL_ARCTIC_MAX_RETRIES", str(_DEFAULT_MAX_RETRIES))))
+ self.retry_base_delay_s = max(
+ 0.0,
+ float(os.environ.get("PRIME_RL_ARCTIC_RETRY_BASE_DELAY_S", str(_DEFAULT_RETRY_BASE_DELAY_S))),
+ )
+ self.split_on_retryable_error = os.environ.get("PRIME_RL_ARCTIC_SPLIT_ON_5XX", "1") != "0"
+ self.stream_results = os.environ.get("PRIME_RL_ARCTIC_STREAM", "1").lower() not in _DISABLED_ENV_VALUES
+ self.stream_fallback_logged = False
+ self.queue: list[dict[str, Any]] = []
+ self.condition = asyncio.Condition()
+ self.semaphore = asyncio.Semaphore(self.max_inflight_batches)
+ self.tasks: set[asyncio.Task] = set()
+ self.loop_task: asyncio.Task | None = None
+ self.thread_local = threading.local()
+ self.executor = ThreadPoolExecutor(
+ max_workers=self.max_inflight_batches,
+ thread_name_prefix="prime-rl-arctic-generate",
+ )
+ self.completed = 0
+ self.failed = 0
+ self.backend_calls = 0
+ self.backend_prompts = 0
+ self.backend_elapsed_total = 0.0
+ self.backend_elapsed_max = 0.0
+ self.queue_wait_total = 0.0
+ self.queue_wait_max = 0.0
+ self.backend_failures = 0
+ self.backend_retries = 0
+ self.backend_splits = 0
+ get_logger().info(
+ "Arctic generate batcher ready "
+ f"(url={self.base_url} job_id={self.sampling_job_id} max_batch={self.max_batch} "
+ f"max_inflight_batches={self.max_inflight_batches} flush_s={self.flush_interval_s:.3f} "
+ f"max_retries={self.max_retries} split_on_5xx={self.split_on_retryable_error} "
+ f"stream={self.stream_results} backend=ArcticHTTP)"
+ )
+
+ def ensure_started(self) -> None:
+ if self.loop_task is None:
+ self.loop_task = asyncio.create_task(self._batch_loop(), name="arctic-generate-batch-loop")
+
+ async def generate(
+ self,
+ *,
+ prompt_ids: list[int],
+ sampling_params: dict[str, Any],
+ routing_key: str | None,
+ strict: bool,
+ ) -> dict[str, Any]:
+ self.ensure_started()
+ future = asyncio.get_running_loop().create_future()
+ async with self.condition:
+ if len(self.queue) >= self.max_queue:
+ raise RuntimeError(f"Arctic generate queue exceeded max_queue={self.max_queue}")
+ self.queue.append(
+ {
+ "prompt_ids": prompt_ids,
+ "sampling_params": sampling_params,
+ "routing_key": routing_key,
+ "strict": strict,
+ "future": future,
+ "enqueued_at": time.perf_counter(),
+ }
+ )
+ self.condition.notify()
+ return await future
+
+ async def _batch_loop(self) -> None:
+ while True:
+ await self.semaphore.acquire()
+ try:
+ batch = await self._pop_batch()
+ except BaseException:
+ self.semaphore.release()
+ raise
+ task = asyncio.create_task(self._run_batch(batch))
+ self.tasks.add(task)
+ task.add_done_callback(self.tasks.discard)
+
+ async def _pop_batch(self) -> list[dict[str, Any]]:
+ async with self.condition:
+ while not self.queue:
+ await self.condition.wait()
+ strict = bool(self.queue[0]["strict"])
+ deadline = asyncio.get_running_loop().time() + self.flush_interval_s
+ while self._matching_queue_size(strict) < self.max_batch:
+ remaining = deadline - asyncio.get_running_loop().time()
+ if remaining <= 0:
+ break
+ try:
+ await asyncio.wait_for(self.condition.wait(), timeout=remaining)
+ except asyncio.TimeoutError:
+ break
+
+ selected: list[dict[str, Any]] = []
+ remaining_items: list[dict[str, Any]] = []
+ for item in self.queue:
+ if len(selected) < self.max_batch and bool(item["strict"]) == strict:
+ selected.append(item)
+ else:
+ remaining_items.append(item)
+ self.queue = remaining_items
+ return selected
+
+ def _matching_queue_size(self, strict: bool) -> int:
+ return sum(1 for item in self.queue if bool(item["strict"]) == strict)
+
+ def _thread_session(self) -> requests.Session:
+ session = getattr(self.thread_local, "session", None)
+ if session is None:
+ session = requests.Session()
+ self.thread_local.session = session
+ return session
+
+ def _generate_http(
+ self,
+ *,
+ prompts: list[list[int]],
+ sampling_params: list[dict[str, Any]],
+ routing_key: list[str | None],
+ strict: bool,
+ ) -> list[dict[str, Any]]:
+ payload: dict[str, Any] = {
+ "prompts": prompts,
+ "sampling_params": sampling_params,
+ }
+ if any(key is not None for key in routing_key):
+ payload["routing_key"] = routing_key
+ if strict:
+ payload["strict"] = True
+ response = self._thread_session().post(
+ f"{self.base_url}/generate",
+ params={"job_id": self.sampling_job_id},
+ json=payload,
+ )
+ response.raise_for_status()
+ return response.json()["results"]
+
+ def _resolve_stream_result(self, item: dict[str, Any], result: dict[str, Any]) -> None:
+ if not item["future"].done():
+ item["future"].set_result(result)
+ self.completed += 1
+
+ def _resolve_stream_error(self, item: dict[str, Any], exc: Exception) -> None:
+ if not item["future"].done():
+ item["future"].set_exception(exc)
+ self.failed += 1
+
+ def _event_index(self, event: dict[str, Any], batch_size: int) -> int:
+ try:
+ index = int(event["index"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise RuntimeError(f"Arctic stream event missing valid index: {event!r}") from exc
+ if index < 0 or index >= batch_size:
+ raise RuntimeError(f"Arctic stream event index {index} outside batch size {batch_size}")
+ return index
+
+ def _generate_http_stream(
+ self,
+ *,
+ batch: list[dict[str, Any]],
+ prompts: list[list[int]],
+ sampling_params: list[dict[str, Any]],
+ routing_key: list[str | None],
+ strict: bool,
+ loop: asyncio.AbstractEventLoop,
+ ) -> None:
+ payload: dict[str, Any] = {
+ "prompts": prompts,
+ "sampling_params": sampling_params,
+ }
+ if any(key is not None for key in routing_key):
+ payload["routing_key"] = routing_key
+ if strict:
+ payload["strict"] = True
+
+ response = self._thread_session().post(
+ f"{self.base_url}/generate-stream",
+ params={"job_id": self.sampling_job_id},
+ json=payload,
+ stream=True,
+ )
+ if response.status_code == 404:
+ response.close()
+ raise _ArcticStreamUnavailable("/generate-stream endpoint is unavailable")
+ try:
+ response.raise_for_status()
+ except Exception:
+ response.close()
+ raise
+
+ seen: set[int] = set()
+ done = False
+ try:
+ for raw_line in response.iter_lines(decode_unicode=True):
+ if not raw_line:
+ continue
+ event = json.loads(raw_line)
+ event_type = event.get("type")
+ if event_type == "result":
+ index = self._event_index(event, len(batch))
+ seen.add(index)
+ loop.call_soon_threadsafe(
+ self._resolve_stream_result,
+ batch[index],
+ event["result"],
+ )
+ elif event_type == "error":
+ index = self._event_index(event, len(batch))
+ seen.add(index)
+ loop.call_soon_threadsafe(
+ self._resolve_stream_error,
+ batch[index],
+ RuntimeError(str(event.get("error", "Arctic stream prompt failed"))),
+ )
+ elif event_type == "stream_error":
+ raise RuntimeError(str(event.get("error", "Arctic stream failed")))
+ elif event_type == "done":
+ done = True
+ break
+ else:
+ raise RuntimeError(f"Unknown Arctic stream event type: {event!r}")
+ finally:
+ response.close()
+
+ if not done:
+ raise RuntimeError("Arctic stream ended before a done event")
+ missing = sorted(set(range(len(batch))) - seen)
+ if missing:
+ raise RuntimeError(f"Arctic stream ended without results for indexes {missing[:16]}")
+
+ async def _run_batch(self, batch: list[dict[str, Any]]) -> None:
+ now = time.perf_counter()
+ for item in batch:
+ queue_wait = now - float(item["enqueued_at"])
+ self.queue_wait_total += queue_wait
+ self.queue_wait_max = max(self.queue_wait_max, queue_wait)
+ try:
+ await self._resolve_batch(batch)
+ finally:
+ self.semaphore.release()
+ self._maybe_log_stats()
+
+ async def _resolve_batch(self, batch: list[dict[str, Any]], *, attempt: int = 0) -> None:
+ if self.stream_results:
+ await self._resolve_batch_stream(batch, attempt=attempt)
+ return
+
+ try:
+ results = await self._post_batch(batch)
+ if len(results) != len(batch):
+ raise RuntimeError(f"Arctic returned {len(results)} results for {len(batch)} prompts")
+ except Exception as exc:
+ if self._is_retryable(exc):
+ if self.split_on_retryable_error and len(batch) > 1:
+ midpoint = max(1, len(batch) // 2)
+ self.backend_splits += 1
+ get_logger().warning(
+ "arctic-generate batch failed; splitting batch size {} -> {}/{}: {}",
+ len(batch),
+ midpoint,
+ len(batch) - midpoint,
+ self._short_error(exc),
+ )
+ await self._resolve_batch(batch[:midpoint])
+ await self._resolve_batch(batch[midpoint:])
+ return
+ if attempt < self.max_retries:
+ self.backend_retries += 1
+ delay_s = self.retry_base_delay_s * (2**attempt)
+ if delay_s > 0:
+ await asyncio.sleep(delay_s)
+ get_logger().warning(
+ "arctic-generate batch failed; retrying batch size {} attempt {}/{}: {}",
+ len(batch),
+ attempt + 1,
+ self.max_retries,
+ self._short_error(exc),
+ )
+ await self._resolve_batch(batch, attempt=attempt + 1)
+ return
+
+ self.failed += len(batch)
+ for item in batch:
+ if not item["future"].done():
+ item["future"].set_exception(exc)
+ return
+
+ for item, result in zip(batch, results):
+ if not item["future"].done():
+ item["future"].set_result(result)
+ self.completed += len(batch)
+
+ async def _resolve_batch_stream(self, batch: list[dict[str, Any]], *, attempt: int = 0) -> None:
+ try:
+ await self._post_batch_stream(batch)
+ await asyncio.sleep(0)
+ unresolved = [item for item in batch if not item["future"].done()]
+ if unresolved:
+ raise RuntimeError(f"Arctic stream left {len(unresolved)} prompts unresolved")
+ return
+ except _ArcticStreamUnavailable:
+ self.stream_results = False
+ if not self.stream_fallback_logged:
+ self.stream_fallback_logged = True
+ get_logger().warning("/generate-stream unavailable; falling back to batch /generate")
+ unresolved = [item for item in batch if not item["future"].done()]
+ if unresolved:
+ await self._resolve_batch(unresolved, attempt=attempt)
+ return
+ except Exception as exc:
+ unresolved = [item for item in batch if not item["future"].done()]
+ if not unresolved:
+ return
+ if self._is_retryable(exc):
+ if self.split_on_retryable_error and len(unresolved) > 1:
+ midpoint = max(1, len(unresolved) // 2)
+ self.backend_splits += 1
+ get_logger().warning(
+ "arctic-generate stream failed; splitting unresolved batch size {} -> {}/{}: {}",
+ len(unresolved),
+ midpoint,
+ len(unresolved) - midpoint,
+ self._short_error(exc),
+ )
+ await self._resolve_batch_stream(unresolved[:midpoint])
+ await self._resolve_batch_stream(unresolved[midpoint:])
+ return
+ if attempt < self.max_retries:
+ self.backend_retries += 1
+ delay_s = self.retry_base_delay_s * (2**attempt)
+ if delay_s > 0:
+ await asyncio.sleep(delay_s)
+ get_logger().warning(
+ "arctic-generate stream failed; retrying unresolved batch size {} attempt {}/{}: {}",
+ len(unresolved),
+ attempt + 1,
+ self.max_retries,
+ self._short_error(exc),
+ )
+ await self._resolve_batch_stream(unresolved, attempt=attempt + 1)
+ return
+
+ self.failed += len(unresolved)
+ for item in unresolved:
+ if not item["future"].done():
+ item["future"].set_exception(exc)
+
+ async def _post_batch(self, batch: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ prompts = [item["prompt_ids"] for item in batch]
+ sampling_params = [item["sampling_params"] for item in batch]
+ routing_keys = [item["routing_key"] for item in batch]
+ strict = bool(batch[0]["strict"])
+ started = time.perf_counter()
+ try:
+ loop = asyncio.get_running_loop()
+ return await loop.run_in_executor(
+ self.executor,
+ lambda: self._generate_http(
+ prompts=prompts,
+ sampling_params=sampling_params,
+ routing_key=routing_keys,
+ strict=strict,
+ ),
+ )
+ except Exception:
+ self.backend_failures += 1
+ raise
+ finally:
+ elapsed = time.perf_counter() - started
+ self.backend_calls += 1
+ self.backend_prompts += len(batch)
+ self.backend_elapsed_total += elapsed
+ self.backend_elapsed_max = max(self.backend_elapsed_max, elapsed)
+
+ async def _post_batch_stream(self, batch: list[dict[str, Any]]) -> None:
+ prompts = [item["prompt_ids"] for item in batch]
+ sampling_params = [item["sampling_params"] for item in batch]
+ routing_keys = [item["routing_key"] for item in batch]
+ strict = bool(batch[0]["strict"])
+ started = time.perf_counter()
+ try:
+ loop = asyncio.get_running_loop()
+ await loop.run_in_executor(
+ self.executor,
+ lambda: self._generate_http_stream(
+ batch=batch,
+ prompts=prompts,
+ sampling_params=sampling_params,
+ routing_key=routing_keys,
+ strict=strict,
+ loop=loop,
+ ),
+ )
+ except Exception:
+ self.backend_failures += 1
+ raise
+ finally:
+ elapsed = time.perf_counter() - started
+ self.backend_calls += 1
+ self.backend_prompts += len(batch)
+ self.backend_elapsed_total += elapsed
+ self.backend_elapsed_max = max(self.backend_elapsed_max, elapsed)
+
+ def _is_retryable(self, exc: Exception) -> bool:
+ response = getattr(exc, "response", None)
+ status_code = getattr(response, "status_code", None)
+ if isinstance(status_code, int):
+ return status_code == 429 or status_code >= 500
+ return isinstance(
+ exc,
+ (
+ requests.exceptions.ConnectionError,
+ requests.exceptions.Timeout,
+ ),
+ )
+
+ def _short_error(self, exc: Exception) -> str:
+ response = getattr(exc, "response", None)
+ status_code = getattr(response, "status_code", None)
+ detail = getattr(response, "text", "") or ""
+ if isinstance(status_code, int):
+ return f"HTTP {status_code}: {detail.strip().replace(chr(10), ' ')[:300]}"
+ return repr(exc)[:300]
+
+ def _maybe_log_stats(self) -> None:
+ if self.log_every_n <= 0 or self.backend_prompts <= 0:
+ return
+ total = self.completed + self.failed
+ if total <= 0 or total % self.log_every_n != 0:
+ return
+ avg_backend_s = self.backend_elapsed_total / max(self.backend_calls, 1)
+ avg_queue_s = self.queue_wait_total / max(self.backend_prompts, 1)
+ get_logger().info(
+ "arctic-generate stats: "
+ f"completed={self.completed} failed={self.failed} queue_len={len(self.queue)} "
+ f"backend_calls={self.backend_calls} backend_prompts={self.backend_prompts} "
+ f"backend_failures={self.backend_failures} retries={self.backend_retries} splits={self.backend_splits} "
+ f"avg_backend_s={avg_backend_s:.2f} max_backend_s={self.backend_elapsed_max:.2f} "
+ f"avg_queue_s={avg_queue_s:.3f} max_queue_s={self.queue_wait_max:.3f}"
+ )
+
+
+_BATCHER: _ArcticGenerateBatcher | None = None
+_PATCHED = False
+
+
+def _get_batcher() -> _ArcticGenerateBatcher:
+ global _BATCHER
+ if _BATCHER is None:
+ reconnect = os.environ.get("PRIME_RL_DIRECT_ARCTIC_RECONNECT_CONFIG")
+ if not reconnect:
+ raise RuntimeError("PRIME_RL_DIRECT_ARCTIC_RECONNECT_CONFIG is required for Arctic generation")
+ _BATCHER = _ArcticGenerateBatcher(Path(reconnect))
+ return _BATCHER
+
+
+def _normalize_sampling_args(raw_sampling_args: Any) -> dict[str, Any]:
+ normalized = dict(raw_sampling_args)
+ if "max_tokens" in normalized:
+ normalized["max_completion_tokens"] = normalized.pop("max_tokens")
+ normalized["logprobs"] = True
+ extra_body = dict(return_token_ids=True, return_sampled_logprobs_only=True)
+ if "extra_body" in normalized:
+ normalized["extra_body"] = {
+ **normalized["extra_body"],
+ **extra_body,
+ }
+ else:
+ normalized["extra_body"] = extra_body
+ return {key: value for key, value in normalized.items() if value is not None}
+
+
+async def _prompt_ids_for_request(
+ client: Any,
+ state: dict[str, Any],
+ prompt: Any,
+ tools: list[Any] | None,
+ model: str,
+) -> list[int]:
+ if len(state["trajectory"]) == 0:
+ return await client.tokenize(messages=prompt, tools=tools, model=model)
+
+ prompt_ids = await client.get_prompt_ids(state, prompt, tools)
+ if prompt_ids is None:
+ get_logger().debug(
+ "Arctic token-id stitching failed for a multi-turn rollout; falling back to local full-prompt tokenization."
+ )
+ return await client.tokenize(messages=prompt, tools=tools, model=model)
+ return prompt_ids
+
+
+# ============================================================================
+# ArcticClient — verifiers.clients.Client subclass
+# Replaces the resolve-poc monkey-patch with a proper Client subclass that
+# integrators can plug in via verifiers' ClientConfig(client_type="custom",
+# class_path="arctic_rl.verifiers_backend.ArcticClient").
+# ============================================================================
+
+from openai.types.chat import ChatCompletion as _ChatCompletion
+from verifiers.clients.openai_chat_completions_token_client import (
+ OpenAIChatCompletionsTokenClient as _OpenAIChatCompletionsTokenClient,
+ _has_multimodal_content as _has_multimodal_content,
+)
+
+
+class ArcticClient(_OpenAIChatCompletionsTokenClient):
+ """Verifiers Client that calls an Arctic RL server directly (no shim).
+
+ Overrides three methods of ``OpenAIChatCompletionsTokenClient``:
+
+ - ``tokenize``: tokenize locally with the registered tokenizer.
+ - ``get_native_response``: build the request and route through the
+ shared :class:`_ArcticGenerateBatcher` instead of calling OpenAI.
+ - ``from_native_response``: apply the optional reasoning-guard mask.
+ """
+
+ async def tokenize(
+ self,
+ messages,
+ tools,
+ model,
+ extra_kwargs=None,
+ **kwargs,
+ ):
+ return _local_tokenize(messages=messages, tools=tools, extra_kwargs=extra_kwargs)
+
+ async def get_native_response(
+ self,
+ prompt,
+ model,
+ sampling_args,
+ tools=None,
+ **kwargs,
+ ):
+ from typing import cast as _cast
+
+ state = _cast(dict, kwargs.get("state"))
+ if state is None:
+ raise RuntimeError("ArcticClient requires verifiers state in get_native_response")
+
+ has_multimodal = _has_multimodal_content(prompt) or any(
+ _has_multimodal_content(step["prompt"]) for step in state["trajectory"]
+ )
+ if has_multimodal:
+ raise RuntimeError("ArcticClient does not support multimodal prompts")
+
+ normalized_sampling_args = _normalize_sampling_args(sampling_args)
+ prompt_ids = await _prompt_ids_for_request(self, state, prompt, tools, model)
+ extra_body = normalized_sampling_args.pop("extra_body", {})
+ body = dict(
+ model=model,
+ messages=prompt,
+ tools=tools,
+ tokens=prompt_ids,
+ **normalized_sampling_args,
+ **extra_body,
+ )
+ routing_key, strict = extract_routing_metadata(body)
+ result = await _get_batcher().generate(
+ prompt_ids=prompt_ids,
+ sampling_params=oai_sampling_params(body),
+ routing_key=routing_key,
+ strict=strict,
+ )
+ response = _arctic_results_to_oai_response(
+ [result],
+ model=model,
+ request_id=f"chatcmpl-{uuid.uuid4().hex}",
+ prompt_token_ids=prompt_ids,
+ )
+ return _ChatCompletion.model_validate(response)
+
+ async def from_native_response(self, response):
+ return await _OpenAIChatCompletionsTokenClient.from_native_response(self, response)
diff --git a/integrations/arctic-rl/examples/arctic_overlay.toml b/integrations/arctic-rl/examples/arctic_overlay.toml
new file mode 100644
index 0000000000..65c3e2d244
--- /dev/null
+++ b/integrations/arctic-rl/examples/arctic_overlay.toml
@@ -0,0 +1,22 @@
+# Arctic backend overlay snippet.
+#
+# Apply on top of any standard prime-rl recipe to flip dispatch to the
+# Arctic RL backend without editing the base recipe:
+#
+# uv run rl @ configs/gsm8k/rl.toml @ integrations/arctic-rl/examples/arctic_overlay.toml
+#
+# The base recipe's `[inference]` block is silently ignored in Arctic mode
+# (Arctic owns sampling). The launcher's `_peek_backend` reads `trainer.backend`
+# from the @-files, so this overlay must be passed via @.
+#
+# Restrictions still enforced: no LoRA, no multi-run, no multi-node, no
+# `[teacher_inference]`. Recipes using any of those need explicit edits.
+
+[trainer]
+backend = "arctic_rl"
+
+[arctic]
+backend = "local" # spawns the Arctic RL server in-process
+training_gpus = 1
+sampling_tensor_parallel_size = 1
+log_prob_gpus = 0 # disables the separate log-prob job
diff --git a/integrations/arctic-rl/examples/arctic_reverse_text/rl.toml b/integrations/arctic-rl/examples/arctic_reverse_text/rl.toml
new file mode 100644
index 0000000000..e89a6afd62
--- /dev/null
+++ b/integrations/arctic-rl/examples/arctic_reverse_text/rl.toml
@@ -0,0 +1,59 @@
+# Minimal Arctic-mode example: routes forward/backward and rollouts to a
+# remote Arctic RL server instead of running a local trainer + vLLM.
+#
+# How dispatch works:
+# trainer.backend = "arctic_rl" tells PRIME-RL's launcher to dispatch to
+# `arctic_rl.entrypoint:main` (this integration). When unset (or "native"),
+# PRIME-RL runs the standard torchrun+FSDP2 + vLLM path.
+#
+# Setup: install this integration package and the Arctic RL client; start
+# an Arctic RL server reachable at the URL below. See README.md.
+#
+# Then:
+# uv run rl @ integrations/arctic-rl/examples/arctic_reverse_text/rl.toml
+
+max_steps = 30
+seq_len = 2048
+
+[model]
+name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"
+
+[wandb]
+project = "primerl-arctic-reverse-text"
+name = "primerl-arctic-reverse-text"
+
+[trainer]
+backend = "arctic_rl"
+
+[arctic]
+backend = "remote"
+url = "http://localhost:7000"
+
+[orchestrator]
+batch_size = 128
+rollouts_per_example = 16
+
+[orchestrator.advantage]
+type = "default"
+
+[orchestrator.train.sampling]
+max_completion_tokens = 128
+
+[[orchestrator.train.env]]
+id = "reverse-text"
+
+[trainer.optim]
+lr = 3e-6
+
+[trainer.scheduler]
+type = "cosine"
+warmup_steps = 10
+
+[ckpt]
+
+# Model not in MODEL_RENDERER_MAP — opt into DefaultRenderer (apply_chat_template).
+[orchestrator.renderer]
+name = "default"
+
+# NOTE: no [inference] block — Arctic mode replaces the native inference server.
+# Validators reject the combination of [arctic] + [inference].
diff --git a/integrations/arctic-rl/pyproject.toml b/integrations/arctic-rl/pyproject.toml
new file mode 100644
index 0000000000..2fe38945a4
--- /dev/null
+++ b/integrations/arctic-rl/pyproject.toml
@@ -0,0 +1,32 @@
+[project]
+name = "prime-rl-arctic"
+version = "0.1.0"
+description = "Arctic RL backend integration for PRIME-RL"
+readme = "README.md"
+requires-python = ">=3.12"
+
+# Depends on prime-rl (installed separately) and the Arctic RL client
+# (`arctic-training`, private until release; install separately).
+dependencies = [
+ "fastapi>=0.115",
+ "uvicorn>=0.30",
+ "httpx>=0.27",
+ "tomli; python_version < '3.11'",
+ "tomli-w>=1.0",
+ "pydantic>=2",
+ "loguru",
+ "verifiers", # uses our fork via [tool.uv.sources] when installed
+]
+
+[tool.uv.sources]
+verifiers = { git = "https://github.com/sfc-gh-kganesan/verifiers.git", branch = "feat/custom-client-extension" }
+
+[project.scripts]
+arctic-trainer = "arctic_rl._trainer_entrypoint:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["arctic_rl"]
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index c0102a222e..e5c5b88336 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -155,6 +155,12 @@ class ClientConfig(BaseConfig):
router_url: str | None = None
"""vllm-router URL for load-aware inference routing. With elastic mode, inference requests go through the router while admin ops still hit discovered pods directly."""
+ client_type: str | None = None
+ """Override the verifiers client_type used for rollouts. None (default) uses the orchestrator's renderer/MITO selection. Set to ``"custom"`` to load a user-provided ``verifiers.Client`` subclass via ``class_path`` — used by external integrations (e.g. the Arctic RL backend) that plug in their own rollout client."""
+
+ class_path: str | None = None
+ """Dotted path to a ``verifiers.Client`` subclass. Only consulted when ``client_type="custom"``."""
+
@property
def is_elastic(self) -> bool:
"""Check if elastic mode is enabled."""
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 f4d37cd9d0..684b55e1c7 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -484,6 +484,9 @@ class TrainerExperimentalConfig(BaseConfig):
class TrainerConfig(BaseConfig):
+ backend: str = "native"
+ """Training backend. ``"native"`` (default) is the standard PRIME-RL torchrun+FSDP2 path; any other value names an installed integration package whose ``.entrypoint:main`` is dispatched to by the launcher."""
+
model: ModelConfig = ModelConfig()
tokenizer: TokenizerConfig = TokenizerConfig()
diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py
index 582d17116e..809861b53b 100644
--- a/src/prime_rl/entrypoints/rl.py
+++ b/src/prime_rl/entrypoints/rl.py
@@ -551,8 +551,48 @@ def rl(config: RLConfig):
rl_local(config)
+def _peek_backend(argv: list[str]) -> str:
+ """Read ``trainer.backend`` from ``@.toml`` and CLI overrides.
+
+ Run before strict config parse so integrations can extend ``RLConfig``
+ with their own fields without the core schema rejecting them.
+ """
+ import tomli
+
+ backend = "native"
+ i = 0
+ while i < len(argv):
+ token = argv[i]
+ if token == "@" and i + 1 < len(argv):
+ path, i = argv[i + 1], i + 2
+ elif token.startswith("@") and len(token) > 1:
+ path, i = token[1:], i + 1
+ else:
+ i += 1
+ continue
+ try:
+ with open(path, "rb") as f:
+ data = tomli.load(f)
+ tb = (data.get("trainer") or {}).get("backend")
+ if isinstance(tb, str):
+ backend = tb
+ except Exception:
+ pass
+ for j, arg in enumerate(argv):
+ if arg.startswith("--trainer.backend="):
+ backend = arg.split("=", 1)[1]
+ elif arg == "--trainer.backend" and j + 1 < len(argv):
+ backend = argv[j + 1]
+ return backend
+
+
def main():
set_proc_title("Launcher")
+ backend = _peek_backend(sys.argv[1:])
+ if backend != "native":
+ from importlib import import_module
+
+ return import_module(f"{backend}.entrypoint").main()
rl(cli(RLConfig))
diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py
index b24012d5fc..726d5a9102 100644
--- a/src/prime_rl/orchestrator/orchestrator.py
+++ b/src/prime_rl/orchestrator/orchestrator.py
@@ -886,6 +886,22 @@ async def setup_student_inference_pool(
client_config = config.student.client
model_name = config.student.model.name
+ if client_config.client_type is not None:
+ logger.info(f"Using custom rollout client_type={client_config.client_type!r}")
+ inference_pool = await setup_inference_pool(
+ client_config,
+ model_name=model_name,
+ train_client_type=client_config.client_type,
+ eval_client_type=client_config.client_type,
+ )
+ # Custom clients (e.g. Arctic RL) manage their own weight sync via
+ # the trainer-side backend; the orchestrator's HTTP-pause + filesystem
+ # broadcast flow doesn't apply.
+ async def _noop_update_weights(*args, **kwargs):
+ return None
+ inference_pool.update_weights = _noop_update_weights
+ return None, inference_pool
+
if config.use_renderer:
renderer = create_renderer(
tokenizer,
diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py
index beb41e8ab6..1b98a596a5 100644
--- a/src/prime_rl/utils/client.py
+++ b/src/prime_rl/utils/client.py
@@ -206,6 +206,7 @@ def setup_clients(
vf.ClientConfig(
client_idx=client_idx,
client_type=client_type,
+ class_path=client_config.class_path,
renderer=renderer_name,
renderer_model_name=renderer_model_name,
renderer_pool_size=renderer_pool_size,