V1 (verl-project/verl main) plugin shims for the Arctic verl integration

arctic_platform/client/config.py

@@ -90,6 +90,31 @@ def resolve_pat(self) -> str | None:
         """The PAT for host/PAT auth: explicit `pat`, else the `pat_env_var` value."""
         return self.pat if self.pat is not None else os.environ.get(self.pat_env_var)
 
+    @classmethod
+    def from_env(cls, **overrides: Any) -> Self:
+        """Build a ``CortexConfig`` from ``ARCTIC_CORTEX_*`` env vars.
+
+        The one call-site framework adapters (SkyRL shim / verl adapter) use to
+        flip to Cortex from a shell-level env, so the env-var contract lives in
+        one place instead of being reinvented per integration. Explicit
+        ``overrides`` win.
+        """
+        env: dict[str, Any] = {}
+        for key, field in (
+            ("base_url", "ARCTIC_CORTEX_BASE_URL"),
+            ("host", "ARCTIC_CORTEX_HOST"),
+            ("pat_env_var", "ARCTIC_CORTEX_PAT_ENV_VAR"),
+            ("database", "ARCTIC_CORTEX_DATABASE"),
+            ("endpoint", "ARCTIC_CORTEX_ENDPOINT"),
+        ):
+            v = os.environ.get(field)
+            if v:
+                env[key] = v
+        schema = os.environ.get("ARCTIC_CORTEX_SCHEMA")
+        if schema:
+            env["schema"] = schema
+        return cls(**{**env, **overrides})
+
     @model_validator(mode="after")
     def _check(self) -> Self:
         if not (self.base_url or self.host):

arctic_platform/client/transports/cortex.py

@@ -64,6 +64,11 @@
 # forward-backward carries the large gradient frame; a mid-stream chunk-group
 # desync (GS restart) is recoverable only by re-posting the whole group.
 _GROUP_RESTART_OPS = {"forward-backward"}
+# Cortex sub-jobs are always awake; the client's colocated wake/sleep lifecycle
+# is a no-op here. Short-circuiting in the transport means the shim doesn't have
+# to wrap every wake/sleep call individually — including the ones ``sync_weights``
+# invokes internally.
+_NOOP_OPS = {"wake-inference", "sleep-inference", "wake-training", "sleep-training"}
 _CHUNK_GROUP_RESTART_REQUIRED = "chunk_group_restart_required"
 _CHUNK_GROUP_ERROR_CODES = {_CHUNK_GROUP_RESTART_REQUIRED, "chunk_group_conflict", "chunk_group_missing_chunks"}
 _JOB_TERMINAL = ("failed", "done", "cancelled", "canceled")
@@ -222,12 +227,16 @@ def shutdown(self) -> None:
 
     # ── deliver one op: submit + poll to completion ──────────────────────────
     def call(self, request: Request) -> dict:
+        if request.op in _NOOP_OPS:
+            return {}
         result = self._poll(self._submit(request))
         # generate returns token ids as DSSST1 tensors; on-prem returns plain
         # lists, so match that contract.
         return _to_python(result) if request.op == "generate" else result
 
     async def acall(self, request: Request) -> dict:
+        if request.op in _NOOP_OPS:
+            return {}
         result = await self._apoll(await self._asubmit(request))
         return _to_python(result) if request.op == "generate" else result
 
@@ -359,11 +368,25 @@ def _wait_running(self) -> None:
         deadline = time.monotonic() + self.poll_timeout
         delay = self.poll_interval
         while time.monotonic() < deadline:
-            state = _short(self._job().get("status"))
+            job = self._job()
+            state = _short(job.get("status"))
             if state == "running":
                 return
             if state in _JOB_TERMINAL:
-                raise RuntimeError(f"cortex job {self.job_id} reached terminal state '{state}'")
+                # Surface the server-side reason and any per-sub-job status so
+                # callers can distinguish rate limits, allowlist rejections,
+                # capacity exhaustion, and genuine sub-job crashes.
+                reason = job.get("reason") or "(no reason)"
+                sub_states = ", ".join(
+                    f"{_short(sj.get('job_type', ''), 'job_type_')}={_short(sj.get('status'))}"
+                    for sj in (job.get("sub_jobs") or [])
+                )
+                detail = f" reason={reason!r}"
+                if sub_states:
+                    detail += f" sub_jobs=[{sub_states}]"
+                raise RuntimeError(
+                    f"cortex job {self.job_id} reached terminal state '{state}';{detail}"
+                )
             time.sleep(delay)
             delay = _next_delay(delay)
         raise TimeoutError(f"cortex job {self.job_id} did not become running within {self.poll_timeout}s")

arctic_platform/integrations/_cortex_shared.py

@@ -0,0 +1,54 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""Cortex ``forward-backward`` payload reshape, shared by the SkyRL shim and
+the verl adapter (both construct ``{batch, meta, processing}``; Cortex takes
+``{args, kwargs, context, processing}``)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def to_cortex_fwd_bwd_payload(
+    batch: dict,
+    *,
+    dp_size: int,
+    processing: dict | None = None,
+) -> dict:
+    """Reshape ``{batch, meta, processing}`` -> Cortex ``{args, kwargs, context, processing}``.
+
+    Omits ``old_log_probs_shifted``; the server-side GRPO loss restores
+    ``π_old = π_new`` via ``logprobs.detach()``. Correct for single-epoch
+    on-policy GRPO only — recipes with ``ppo_epochs > 1`` or KL-to-reference
+    need the real rollout-time snapshot and should fail before reaching here.
+    """
+    import torch
+
+    b = batch["batch"] if isinstance(batch, dict) and "batch" in batch else batch
+    meta = batch.get("meta", {}) if isinstance(batch, dict) else {}
+    processing = processing or (batch.get("processing") if isinstance(batch, dict) else None) or {}
+
+    input_ids = b["input_ids"]
+    attention_mask = b.get("attention_mask")
+    prompt_len = int(meta.get("prompt_len", 0))
+    if not torch.is_tensor(input_ids):
+        raise TypeError(f"cortex fwd_bwd: input_ids must be a tensor, got {type(input_ids).__name__}")
+
+    total = int(input_ids.shape[-1])
+    response_len = max(total - prompt_len, 1)
+
+    args: list[Any] = [input_ids]
+    kwargs: dict[str, Any] = {
+        "response_length": response_len,
+        "prompt_length": prompt_len,
+    }
+    if attention_mask is not None:
+        kwargs["attention_mask"] = attention_mask
+    context: dict[str, Any] = {
+        "advantages": b.get("advantages"),
+        "response_mask": b.get("response_mask"),
+        "dp_size": int(dp_size),
+    }
+    context = {k: v for k, v in context.items() if v is not None}
+
+    return {"args": args, "kwargs": kwargs, "context": context, "processing": processing}

arctic_platform/integrations/skyrl/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""SkyRL integration helpers for the Cortex backend."""
+
+from __future__ import annotations
+
+
+def install_cortex_driver_shims() -> None:
+    """Patch ``skyrl.train.utils.utils.peer_access_supported`` to ``False``.
+
+    SkyRL's ``prepare_runtime_environment`` probes ``cudaCanAccessPeer`` by
+    requesting a ``{"CPU":1,"GPU":2}`` Ray placement group, which hangs a
+    CPU-only Cortex driver. No-op unless SkyRL is importable.
+    """
+    try:
+        from skyrl.train.utils import utils as _skyrl_utils
+    except Exception:
+        return
+    if getattr(_skyrl_utils, "_cortex_shimmed", False):
+        return
+    _skyrl_utils.peer_access_supported = lambda max_num_gpus_per_node=1: False
+    _skyrl_utils._cortex_shimmed = True
+
+
+__all__ = ["install_cortex_driver_shims"]

arctic_platform/integrations/skyrl/__main__.py

@@ -0,0 +1,23 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""SkyRL launcher that installs the Cortex driver shims first.
+
+    python -m arctic_platform.integrations.skyrl <hydra args>
+"""
+
+from __future__ import annotations
+
+import runpy
+import sys
+
+from arctic_platform.integrations.skyrl import install_cortex_driver_shims
+
+install_cortex_driver_shims()
+
+# Forward argv to SkyRL's own entrypoint. This intentionally re-uses whatever
+# SkyRL treats as its main module so we don't fork its argument surface.
+if __name__ == "__main__":
+    # Drop our own ``-m arctic_platform.integrations.skyrl`` from sys.argv so
+    # SkyRL sees the same CLI it would from ``python -m skyrl.train.entrypoints.main_base``.
+    sys.argv = [sys.argv[0]] + sys.argv[1:]
+    runpy.run_module("skyrl.train.entrypoints.main_base", run_name="__main__", alter_sys=True)

arctic_platform/integrations/verl/adapter.py

@@ -44,6 +44,7 @@
 from arctic_platform.client import SamplingConfig
 from arctic_platform.client import TrainingConfig
 from arctic_platform.client import create_arctic_rl_client
+from arctic_platform.integrations._cortex_shared import to_cortex_fwd_bwd_payload
 from arctic_platform.rl.ray_server import ArcticRLRayServerState
 
 _ARCTIC_METRIC_REDUCTION_FN = {
@@ -600,6 +601,16 @@ def _initialize_client(
             onprem_kwargs["port"] = 7000
             onprem_kwargs["launch_local_server"] = True
 
+        # verl's YAML has no backend discriminator, so ARCTIC_BACKEND=cortex is
+        # the one shell-level knob that flips this adapter to Cortex. Cortex
+        # settings hydrate from ARCTIC_CORTEX_* via CortexConfig.from_env.
+        if os.environ.get("ARCTIC_BACKEND", "").strip().lower() == "cortex":
+            from arctic_platform.client import CortexConfig
+
+            backend_config = CortexConfig.from_env()
+        else:
+            backend_config = OnPremConfig(**onprem_kwargs)
+
         # attn_implementation also lives on ds_worker_config; keep it there for
         # the DeepSpeed worker (matches recipe/rl-correctness).
         ds_worker_config = self._create_ds_worker_config()
@@ -612,7 +623,7 @@ def _initialize_client(
             training_gpus=n_training_gpus,
             sampling_gpus=n_sampling_gpus,
             log_prob_gpus=n_log_prob_gpus,
-            backend=OnPremConfig(**onprem_kwargs),
+            backend=backend_config,
             training=TrainingConfig(
                 full_determinism=self._backend_config.train.determinism.get("full", False),
                 checkpoint_path=self.config.trainer.default_local_dir,
@@ -662,11 +673,55 @@ async def generate(self, prompt_ids, sampling_params, routing_key=None) -> list:
     # backends are free to pick their own wire format; verl never calls
     # these directly.
 
+    def _is_cortex_backend(self) -> bool:
+        from arctic_platform.client import CortexConfig
+
+        return isinstance(self._client.config.backend, CortexConfig)
+
+    def _zero_logprob_response(self, payload: dict, *, caller: str) -> dict:
+        """Return zero log-probs / entropy shaped like a real fwd_no_grad.
+
+        Cortex has no ``/forward`` sub-job. Zero-filling is only safe when the
+        caller doesn't actually consume the values — single-epoch on-policy
+        GRPO with no KL-to-reference (server-side default
+        ``old_log_probs = logprobs.detach()`` restores π_old ≡ π_new). Recipes
+        with ``use_kl_loss`` / ``use_kl_in_reward`` on, or multi-epoch PPO,
+        need the real snapshot; fail loud instead of silently producing the
+        wrong training signal.
+        """
+        actor = self.config.actor_rollout_ref.actor
+        algo = getattr(self.config, "algorithm", None)
+        if getattr(actor, "use_kl_loss", False) or getattr(algo, "use_kl_in_reward", False):
+            raise NotImplementedError(
+                f"cortex backend: {caller} needs real ref log-probs "
+                "(actor.use_kl_loss / algorithm.use_kl_in_reward is on) but "
+                "Cortex has no /forward sub-job. Disable both, or use the "
+                "on-prem backend for KL-anchored training."
+            )
+        ppo_epochs = int(getattr(actor, "ppo_epochs", 1))
+        if ppo_epochs > 1:
+            raise NotImplementedError(
+                f"cortex backend: {caller} needs the rollout-time log-prob "
+                f"snapshot for ppo_epochs={ppo_epochs} (off-policy PPO). "
+                "Set actor.ppo_epochs=1, or use the on-prem backend."
+            )
+        # make_njt() slices from data["input_ids"], so match [B, T] with T ≥ 1.
+        b_data = payload.get("batch") if isinstance(payload, dict) else None
+        if not isinstance(b_data, dict):
+            b_data = payload if isinstance(payload, dict) else {}
+        ids = b_data.get("input_ids")
+        b = int(ids.shape[0]) if torch.is_tensor(ids) else 1
+        t = int(ids.shape[-1]) if torch.is_tensor(ids) else 1
+        z = torch.zeros((b, max(t, 1)), dtype=torch.float32)
+        return {"batch": {"log_probs": z, "entropy": z}}
+
     async def _send_compute_ref_log_prob(self, payload: dict):
         payload["processing"] = {
             "post": ["compute_entropy_and_logprobs"],
             "loss_fn": None,
         }
+        if self._is_cortex_backend():
+            return self._zero_logprob_response(payload, caller="compute_ref_log_prob")
         response = await self._client.fwd_no_grad(payload, reference_model=True)
         response["batch"]["log_probs"] = response["batch"].pop("logprobs")
         return response
@@ -676,6 +731,8 @@ async def _send_compute_log_prob(self, payload: dict):
             "post": ["compute_entropy_and_logprobs"],
             "loss_fn": None,
         }
+        if self._is_cortex_backend():
+            return self._zero_logprob_response(payload, caller="compute_log_prob")
         response = await self._client.fwd_no_grad(payload, reference_model=False)
         response["batch"]["log_probs"] = response["batch"].pop("logprobs")
         return response
@@ -700,6 +757,15 @@ def _left_pad(t: torch.Tensor, seq_len: int) -> torch.Tensor:
                 payload["batch"][name] = _left_pad(payload["batch"][name], seq_len)
         payload["batch"]["loss_mask"] = payload["batch"]["response_mask"]
 
+        if self._is_cortex_backend():
+            # Cortex takes {args, kwargs, context, processing}; on-prem takes
+            # {batch, meta, processing}. Reshape here, not on the wire.
+            payload = to_cortex_fwd_bwd_payload(
+                payload["batch"],
+                dp_size=int(self._client.config.training_gpus or 1),
+                processing=payload.get("processing"),
+            )
+
         fwd_bwd_response = await self._client.fwd_bwd(payload)
         step_response = await self._client.step()
         step_response["metrics"].update(**fwd_bwd_response["metrics"])

arctic_platform/integrations/verl/examples/README-cortex.md

@@ -0,0 +1,41 @@
+# verl × Arctic-Platform × Cortex-training
+
+Companion to [`run_gsm8k_grpo_arl.sh`](run_gsm8k_grpo_arl.sh): same verl
+`RemoteBackend` adapter and same GRPO recipe on Qwen3-0.6B / GSM8K.
+`ARCTIC_BACKEND=cortex` (read in
+[`adapter.py::_create_rl_client_config`](../adapter.py)) swaps the default
+`OnPremConfig` for `CortexConfig.from_env()`. The verl YAML stays backend-agnostic.
+
+## Install + env
+
+```bash
+pip install arctic-platform[cortex]
+export ARCTIC_BACKEND=cortex
+export ARCTIC_CORTEX_HOST=<account>.<region>.snowflakecomputing.com
+export ARCTIC_CORTEX_DATABASE=<db>
+export ARCTIC_CORTEX_SCHEMA=<schema>
+export CORTEX_PAT=<pat>
+```
+
+## Run
+
+```bash
+./run_gsm8k_grpo_cortex.sh
+```
+
+## Adapter changes for the Cortex path
+
+Two shape mismatches versus the on-prem wire format live in
+`arctic_platform/integrations/verl/adapter.py`:
+
+* `_send_compute_{ref_,}log_prob`: return zero-shaped `log_probs` / `entropy`;
+  Cortex has no `/forward` sub-job. Only correct for single-epoch on-policy
+  GRPO without KL — `use_kl_loss`, `use_kl_in_reward`, or `ppo_epochs > 1`
+  raise `NotImplementedError`.
+* `_send_update_actor`: reshape `{batch, meta, processing}` →
+  `{args, kwargs, context, processing}` via
+  [`to_cortex_fwd_bwd_payload`](../../_cortex_shared.py) (shared with the
+  SkyRL shim).
+
+Everything else (`generate`, `sync_weights`, `save_checkpoint`, wake/sleep)
+goes through `ArcticRLClient` unchanged.