feat(rl): Cortex-serverless dispatch via arctic_platform.rl (zero integration change)

arctic_platform/client/client.py

@@ -29,7 +29,16 @@
 from arctic_platform.client.transport import Transport
 
 
-def make_transport(config: ArcticRLClientConfig) -> Transport:
+def make_transport(config: ArcticRLClientConfig, *, server_state: Any = None) -> Transport:
+    """Build the transport for `config`, optionally reattaching to an existing
+    on-prem Ray state actor.
+
+    `server_state` is only meaningful for `backend=onprem, comm_protocol=ray`:
+    verl's forwarder worker and SkyRL's `@ray.remote skyrl_entrypoint` both
+    rebuild the client in a different process from the driver and pass the
+    driver's state actor back in so the second process reattaches rather than
+    racing a fresh set of Ray actors. Every other transport ignores it.
+    """
     if config.backend == "cortex":
         from arctic_platform.client.transports.cortex import CortexTransport
 
@@ -38,18 +47,78 @@ def make_transport(config: ArcticRLClientConfig) -> Transport:
     from arctic_platform.client.transports.onprem_ray import RayTransport
 
     if config.backend == "onprem" and config.comm_protocol == "ray":
-        return RayTransport(config)
+        return RayTransport(config, server_state=server_state)
     return HttpTransport(config)  # onprem (HTTP)
 
 
+def _flatten_metrics(result: Any) -> Any:
+    """Bring ``result["metrics"]`` keys to the top level for legacy readers.
+
+    SkyRL reads ``client.step().get("grad_norm")`` and verl reads
+    ``client.fwd_bwd()["avg_loss"]`` at the top level, but the on-prem server
+    nests scalars under ``metrics`` and Cortex returns loss-only. Both callers
+    predate the unified schema. Flattening once here means the transports can
+    keep returning their native shape and callers see a superset dict:
+
+        {"loss": ..., "avg_loss": ..., "grad_norm": ..., "metrics": {...}, ...}
+
+    Existing top-level keys always win over ``metrics`` keys of the same name.
+    """
+    if not isinstance(result, dict):
+        return result
+    metrics = result.get("metrics")
+    if not isinstance(metrics, dict):
+        return result
+    merged = {**metrics, **result}
+    if "avg_loss" not in merged and "loss" in merged:
+        merged["avg_loss"] = merged["loss"]
+    return merged
+
+
 class ArcticRLClient:
-    def __init__(self, config: ArcticRLClientConfig) -> None:
+    def __init__(self, config: ArcticRLClientConfig, *, server_state: Any = None) -> None:
         self.config = config
-        self.transport = make_transport(config)
+        # `server_state` is forwarded to `make_transport` so the on-prem Ray
+        # transport can reattach to the driver's state actor in a forwarder
+        # worker / Ray-remote entrypoint. See `make_transport` for details.
+        self.transport = make_transport(config, server_state=server_state)
         self.jobs = self.transport.initialize()
 
+    # ── legacy job-id attributes ─────────────────────────────────────────
+    # SkyRL's entrypoint reads pre_client.training_job_id / sampling_job_id /
+    # log_prob_job_id after initialize(). Preserve that surface as attributes
+    # backed by JobHandles rather than forcing the integration to reach into
+    # `client.jobs.*`.
+    @property
+    def training_job_id(self) -> Any:
+        return self.jobs.training
+
+    @property
+    def sampling_job_id(self) -> Any:
+        return self.jobs.sampling
+
+    @property
+    def log_prob_job_id(self) -> Any:
+        return self.jobs.log_prob
+
+    def get_server_state(self) -> Any:
+        """Expose the transport's server state to callers that reconnect Ray.
+
+        The SkyRL entrypoint reads this after creating the driver-side client
+        and forwards it to the Ray worker so the worker can reattach to the
+        same on-prem Ray server. Non-Ray transports return ``None``.
+        """
+        getter = getattr(self.transport, "get_server_state", None)
+        return getter() if callable(getter) else None
+
     # ── training ─────────────────────────────────────────────────────────
-    def fwd_bwd(self, batch: dict, processing: dict | None = None, router_replay: Any = None) -> dict:
+    def fwd_bwd(
+        self,
+        batch: dict,
+        processing: dict | None = None,
+        router_replay: Any = None,
+        **legacy_kwargs: Any,
+    ) -> dict:
         # NOTE: the call *signature* is unified across backends, but `batch`'s
         # *content* is not: Cortex expects an RPC-style
         # {"args": [...], "kwargs": {...}} that the server tokenizes, while on-prem
@@ -60,12 +129,25 @@ def fwd_bwd(self, batch: dict, processing: dict | None = None, router_replay: An
         # contract (ideally Cortex's) so this frontend is truly backend-agnostic
         # and callers stop branching on backend. Until then `processing` is folded
         # into the body here, which is why callers can leave it inside `batch`.
+        # Legacy callers (verl `compute_log_prob`, SkyRL wire packer) pass
+        # `post_processors=[...]`, `context={...}` or `kwargs={...}` alongside
+        # `batch`; we fold them into the body so the transports can forward one
+        # uniform envelope.
         body = dict(batch)
         body.update({k: v for k, v in (("processing", processing), ("router_replay", router_replay)) if v is not None})
-        return self.transport.call(Request("fwd-bwd", self.jobs.require("training"), body, binary=True))
+        body.update({k: v for k, v in legacy_kwargs.items() if v is not None})
+        result = self.transport.call(Request("fwd-bwd", self.jobs.require("training"), body, binary=True))
+        return _flatten_metrics(result)
 
-    def fwd_no_grad(self, batch: dict) -> dict:
-        return self.transport.call(Request("fwd-no-grad", self.jobs.require("training"), dict(batch), binary=True))
+    def fwd_no_grad(self, batch: dict, **legacy_kwargs: Any) -> dict:
+        # SkyRL's arctic_generator / _ArcticDispatch calls this with
+        # `post_processors=["logprobs"]`; verl calls it with `context=...` and
+        # `processing=...`. Fold every extra kwarg into the body so transports
+        # receive one uniform envelope (see `CortexTransport._fwd_no_grad`).
+        body = dict(batch)
+        body.update({k: v for k, v in legacy_kwargs.items() if v is not None})
+        result = self.transport.call(Request("fwd-no-grad", self.jobs.require("training"), body, binary=True))
+        return _flatten_metrics(result)
 
     def step(self, learning_rate: float | None = None) -> dict:
         # NOTE: response *shapes* also diverge across backends -- on-prem returns
@@ -74,7 +156,8 @@ def step(self, learning_rate: float | None = None) -> dict:
         # via graceful key lookups today.
         # TODO(unify-backends): converge the server-side fwd_bwd/step *responses*
         # on ONE schema (loss + metrics) so callers never key on backend.
-        return self.transport.call(Request("step", self.jobs.require("training"), {"learning_rate": learning_rate}))
+        result = self.transport.call(Request("step", self.jobs.require("training"), {"learning_rate": learning_rate}))
+        return _flatten_metrics(result)
 
     def save_checkpoint(self, stage_info: dict | None = None, path: str | None = None) -> dict:
         # `path` is this call's destination; when None the server uses the job's
@@ -83,7 +166,12 @@ def save_checkpoint(self, stage_info: dict | None = None, path: str | None = Non
         return self.transport.call(Request("save-checkpoint", self.jobs.require("training"), body))
 
     # ── sampling / log-prob ──────────────────────────────────────────────
-    def generate(self, prompts: list, sampling_params: dict | None = None, routing_key: Any = None) -> list:
+    def generate(
+        self,
+        prompts: list,
+        sampling_params: dict | None = None,
+        routing_key: Any = None,
+    ) -> list:
         body = {"prompts": prompts, "sampling_params": sampling_params, "routing_key": routing_key}
         return self.transport.call(Request("generate", self.jobs.require("sampling"), body))["results"]
 
@@ -92,15 +180,78 @@ def log_probs(self, prompts: list, completions: list | None = None, top_k: int =
         return self.transport.call(Request("log-probs", self.jobs.require("log_prob"), body))
 
     # ── weight sync + cache ──────────────────────────────────────────────
-    def sync_weights(self) -> dict:
-        # sync-weights has no primary job id; both ids travel in the body.
+    def sync_weights(self, cuda_ipc: bool = False, low_memory: bool = False) -> dict:
+        # SkyRL's colocated path calls `sync_weights(cuda_ipc=True)`. On-prem
+        # honors the flag (same-GPU IPC handoff); Cortex has no colocation
+        # concept and silently ignores it. Both flags travel in the body so the
+        # transport decides what's meaningful.
         tid, sid = self.jobs.require("training"), self.jobs.require("sampling")
-        return self.transport.call(Request("sync-weights", None, {"training_job_id": tid, "sampling_job_id": sid}))
+        body = {
+            "training_job_id": tid,
+            "sampling_job_id": sid,
+            "cuda_ipc": cuda_ipc,
+            "low_memory": low_memory,
+        }
+        return self.transport.call(Request("sync-weights", None, body))
 
     def reset_prefix_cache(self, drain: bool = True, timeout_s: float = 60.0) -> dict:
         body = {"drain": drain, "timeout_s": timeout_s}
         return self.transport.call(Request("reset-prefix-cache", self.jobs.require("sampling"), body))
 
+    # ── colocation lifecycle (legacy SkyRL / verl surface) ───────────────
+    # SkyRL's `_ArcticDispatch.save_weights_for_sampler` and the on-prem
+    # colocated path call these unconditionally when `colocate=True`. On-prem
+    # implements each op server-side; Cortex has no colocation concept and
+    # returns an empty dict (its transport table registers them as no-ops).
+    # Exposing them here (rather than raising `AttributeError`) is what lets
+    # SkyRL / verl run against `backend=cortex` unmodified.
+    #
+    # verl passes `tags=...` on wake and `level=...` on sleep — those are
+    # engine hints (e.g. vLLM sleep level). We fold every extra kwarg into
+    # the body so on-prem can honor them and Cortex can safely ignore.
+    def wake_training(self, **kwargs: Any) -> dict:
+        return self._colo("wake-training", "training", body=kwargs)
+
+    def sleep_training(self, **kwargs: Any) -> dict:
+        return self._colo("sleep-training", "training", body=kwargs)
+
+    def wake_inference(self, **kwargs: Any) -> dict:
+        return self._colo("wake-inference", "sampling", body=kwargs)
+
+    def sleep_inference(self, **kwargs: Any) -> dict:
+        return self._colo("sleep-inference", "sampling", body=kwargs)
+
+    def wake_log_prob(self, **kwargs: Any) -> dict:
+        return self._colo("wake-log-prob", "log_prob", body=kwargs)
+
+    def sleep_log_prob(self, **kwargs: Any) -> dict:
+        return self._colo("sleep-log-prob", "log_prob", body=kwargs)
+
+    def empty_training_cache(self, **kwargs: Any) -> dict:
+        return self._colo("empty-training-cache", "training", body=kwargs)
+
+    def weight_norm(self, **kwargs: Any) -> dict:
+        return self._colo("weight-norm", "training", body=kwargs)
+
+    def save_weights(self, path: str | None = None) -> dict:
+        # Disk-based weight reload; deliberately optional (see UNIFICATION_NOTES).
+        body = {"path": path}
+        return self.transport.call(Request("save-weights", self.jobs.require("training"), body))
+
+    def _colo(self, op: str, job_type: str, *, body: dict[str, Any] | None = None) -> dict:
+        """Dispatch a colocation-lifecycle op if the target job exists.
+
+        The primary target is the requested job type; when it isn't set (e.g.
+        Cortex without a log-prob sub-job) the op is a no-op so callers don't
+        have to branch on backend. Extra kwargs from the caller (e.g. verl's
+        `tags`, `level`) travel in `body` so the transport can decide what's
+        meaningful.
+        """
+        job_id = getattr(self.jobs, job_type, None)
+        if job_id is None:
+            return {}
+        return self.transport.call(Request(op, job_id, dict(body or {})))
+
     # ── lifecycle ────────────────────────────────────────────────────────
     def reconnect_config(self) -> ArcticRLClientConfig:
         """A serializable config that reattaches to these jobs in another process."""
@@ -122,6 +273,15 @@ def __exit__(self, *exc: object) -> None:
         self.shutdown()
 
 
-def create_arctic_rl_client(config: ArcticRLClientConfig) -> ArcticRLClient:
-    """Factory matching the current OSS entrypoint shape."""
-    return ArcticRLClient(config)
+def create_arctic_rl_client(
+    config: ArcticRLClientConfig,
+    server_state: Any = None,
+) -> ArcticRLClient:
+    """Factory matching the current OSS entrypoint shape.
+
+    `server_state` is a positional shim for the legacy call site
+    ``create_arctic_rl_client(reconnect_config, rl_server_state)`` that the
+    verl adapter and SkyRL's `main_arctic_rl.py` still use. New callers should
+    keep it None (driver path) or thread it as a kwarg when reconnecting.
+    """
+    return ArcticRLClient(config, server_state=server_state)

arctic_platform/client/config.py

@@ -16,18 +16,54 @@
 
 from __future__ import annotations
 
+import warnings
 from typing import Any
 from typing import Literal
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
+from pydantic import model_validator
 
 JobId = int | str
 
+# Legacy backend labels accepted for compatibility with the ArcticTraining and
+# older arctic_platform clients that SkyRL / verl integrations were built
+# against. They collapse into the two canonical labels below at validation
+# time; downstream code only ever sees `onprem` or `cortex`.
+_BACKEND_ALIASES = {
+    "local": "onprem",
+    "dss-platform": "onprem",
+    "dss_platform": "onprem",
+    "neutrino": "cortex",
+}
+
+# Legacy field names accepted for compatibility. Mapped into the canonical
+# field before pydantic validation so downstream code only sees the canonical
+# name. Kept as (legacy_name, canonical_name) pairs.
+_FIELD_ALIASES = (
+    ("sample_gpus", "sampling_gpus"),
+    ("sampling_engine", "vllm_config"),  # legacy verl passes engine name; ignored below
+)
+
+# Legacy fields accepted-and-ignored. These exist on old ArcticTraining /
+# dss-client configs but have no analogue on the unified client. Silently
+# dropping them keeps SkyRL / verl integrations config-compatible.
+_LEGACY_IGNORED_FIELDS = frozenset({
+    "log_prob_engine",
+    "sampling_engine",
+    "reference_model",
+    "job_name",
+    "experiment_name",
+})
+
 
 class ArcticRLClientConfig(BaseModel):
-    model_config = ConfigDict(extra="forbid", validate_default=True)
+    # NOTE: extra="ignore" (was "forbid") so legacy fields from ArcticTraining /
+    # verl configs pass through without erroring. Legacy fields we care about
+    # are mapped into canonical fields by `_apply_legacy_aliases` below; the
+    # rest are dropped silently.
+    model_config = ConfigDict(extra="ignore", validate_default=True)
 
     backend: Literal["onprem", "cortex"] = Field("onprem", description="Deployment target.")
     comm_protocol: Literal["http", "ray"] = Field("http", description="onprem transport: HTTP or in-process Ray.")
@@ -62,6 +98,33 @@ class ArcticRLClientConfig(BaseModel):
             "that call and falls back to this dir when path is None."
         ),
     )
+    # Server-side init knobs threaded through `_init_payload` — these existed
+    # on the legacy `arctic_platform.rl.ArcticRLClientConfig` and are used by
+    # the on-prem DeepSpeed / Arctic-Inference servers. Kept typed as
+    # `dict | None` so verl / SkyRL can pass their yaml sub-blocks through
+    # without repeating the schema here.
+    log_prob_ds_config: dict[str, Any] | None = Field(
+        None, description="onprem: log-prob job's DeepSpeed engine config (dtype, batch, ...)."
+    )
+    ds_worker_config: dict[str, Any] | None = Field(
+        None,
+        description=(
+            "onprem: DeepSpeed worker config forwarded to `deepspeed_worker.py` "
+            "(use_liger, attn_implementation, zorro_train_*, logits_*)."
+        ),
+    )
+    arctic_inference_config: dict[str, Any] | None = Field(
+        None,
+        description=(
+            "onprem: Arctic-Inference rollout config (Forest Cascade Attention, "
+            "speculative decoding). Server keys on zorro_inference.enable / "
+            "speculative_decoding.model."
+        ),
+    )
+    full_determinism: bool = Field(
+        False,
+        description="onprem: enable full DeepSpeed determinism (reproducibility over throughput).",
+    )
 
     # cortex (SnowAPI)
     cortex_base_url: str | None = Field(None, description="Mock/direct GS URL; bypasses PAT auth.")
@@ -76,6 +139,44 @@ class ArcticRLClientConfig(BaseModel):
     sampling_job_id: JobId | None = None
     log_prob_job_id: JobId | None = None
 
+    @model_validator(mode="before")
+    @classmethod
+    def _apply_legacy_aliases(cls, data: Any) -> Any:
+        """Map legacy backend labels and field names to canonical ones.
+
+        Runs before field validation so SkyRL and verl configs authored against
+        the old ArcticTraining / dss-client contract validate unchanged:
+
+        - ``backend="local"`` / ``"dss-platform"`` -> ``"onprem"``
+        - ``backend="neutrino"``                    -> ``"cortex"``
+        - ``sample_gpus=`` alias for ``sampling_gpus=``
+        - ``log_prob_engine``/``sampling_engine`` are dropped (no analogue).
+        """
+        if not isinstance(data, dict):
+            return data
+        data = dict(data)
+        raw = data.get("backend")
+        if isinstance(raw, str) and raw in _BACKEND_ALIASES:
+            canonical = _BACKEND_ALIASES[raw]
+            warnings.warn(
+                f"ArcticRLClientConfig(backend={raw!r}) is a legacy alias for "
+                f"{canonical!r}; migrate to {canonical!r}.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            data["backend"] = canonical
+        for legacy, canonical in _FIELD_ALIASES:
+            if legacy in data and data.get(canonical) in (None, 0, "", {}):
+                if canonical in {name for name, _ in _FIELD_ALIASES}:
+                    continue
+                data[canonical] = data.pop(legacy)
+            elif legacy in data:
+                data.pop(legacy)
+        for name in list(data):
+            if name in _LEGACY_IGNORED_FIELDS:
+                data.pop(name)
+        return data
+
     def gpus_for(self, job_type: str) -> int:
         """GPU count allocated to a job type (0 == the job type is disabled)."""
         return getattr(self, f"{job_type}_gpus")