diff --git a/arctic_platform/client/client.py b/arctic_platform/client/client.py index ab0f4cb..0043b3d 100644 --- a/arctic_platform/client/client.py +++ b/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) diff --git a/arctic_platform/client/config.py b/arctic_platform/client/config.py index 4493f97..826b88e 100644 --- a/arctic_platform/client/config.py +++ b/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") diff --git a/arctic_platform/client/transport.py b/arctic_platform/client/transport.py index 8129d33..e2db746 100644 --- a/arctic_platform/client/transport.py +++ b/arctic_platform/client/transport.py @@ -49,6 +49,17 @@ "log-probs", "sync-weights", "reset-prefix-cache", + # Colocation lifecycle: SkyRL calls these unconditionally when + # `colocate=True`. On-prem implements them; Cortex no-ops. + "wake-training", + "sleep-training", + "wake-inference", + "sleep-inference", + "wake-log-prob", + "sleep-log-prob", + "empty-training-cache", + "weight-norm", + "save-weights", } ) diff --git a/arctic_platform/client/transports/cortex.py b/arctic_platform/client/transports/cortex.py index 6a8d2ba..5148ef3 100644 --- a/arctic_platform/client/transports/cortex.py +++ b/arctic_platform/client/transports/cortex.py @@ -269,12 +269,30 @@ def __init__(self, config: ArcticRLClientConfig) -> None: self._session = self._build_session() self._handlers: dict[str, Callable[[dict], dict]] = { "fwd-bwd": self._fwd_bwd, + "fwd-no-grad": self._fwd_no_grad, + "log-probs": self._log_probs, "step": self._step, "save-checkpoint": self._save_checkpoint, "generate": self._generate, "sync-weights": self._sync_weights, "reset-prefix-cache": self._reset_prefix_cache, } + # Colocation lifecycle ops are no-ops on Cortex (sub-jobs live in + # separate placements; there is no wake/sleep concept). We register + # them so SkyRL's colocated code path stays call-shape identical + # without branching on backend. See ArcticRLClient._colo. + for op in ( + "wake-training", + "sleep-training", + "wake-inference", + "sleep-inference", + "wake-log-prob", + "sleep-log-prob", + "empty-training-cache", + "weight-norm", + "save-weights", + ): + self._handlers[op] = self._colo_noop def initialize(self) -> JobHandles: cfg = self.config @@ -305,6 +323,7 @@ def shutdown(self) -> None: # ── op handlers: canonical body -> SnowAPI call -> canonical dict ────── def _fwd_bwd(self, body: dict) -> dict: + body = self._normalize_train_body(body) payload = wire.dumps(body, metadata=_CHUNKED_DSSST1) request_id = self._post_octet_request_chunks( path_suffix="forward-backward", @@ -312,12 +331,51 @@ def _fwd_bwd(self, body: dict) -> dict: frame=payload, max_bytes=self._MAX_FWD_BWD_BYTES, )["request_id"] + return self._shape_train_response(self._poll(request_id)) + + def _fwd_no_grad(self, body: dict) -> dict: + """Cortex-side forward-only pass; returns model_outputs (log-probs). + + SkyRL and verl both need this every training step. Symmetric to + `_fwd_bwd`: same octet-chunked submit + poll, hitting + ``/{job_id}/forward-no-grad``. Requires the Neutrino GS to expose the + endpoint (see the tracking issue). Same envelope translation applies: + callers that ship the verl-GRPO ``{batch, meta, processing}`` shape + have it repackaged into ``{args, kwargs}`` for Cortex. + """ + body = self._normalize_train_body(body) + payload = wire.dumps(body, metadata=_CHUNKED_DSSST1) + request_id = self._post_octet_request_chunks( + path_suffix="forward-no-grad", + operation="fwd-no-grad", + frame=payload, + max_bytes=self._MAX_FWD_BWD_BYTES, + )["request_id"] + return self._shape_train_response(self._poll(request_id)) + + def _log_probs(self, body: dict) -> dict: + """Cortex-side log-probs endpoint (JSON in / DSSST1-decoded out). + + Symmetric to ``_generate`` in framing but returns a log-probs dict + rather than sampled sequences. Requires the Neutrino GS endpoint at + ``/{job_id}/log-probs``. + """ + payload: dict = {"prompts": body["prompts"]} + if body.get("completions") is not None: + payload["completions"] = body["completions"] + if body.get("top_k") is not None: + payload["top_k"] = body["top_k"] + request_id = self._send( + "POST", + f"{self._prefix}/{self.job_id}/log-probs", + json=payload, + ).json()["request_id"] return self._poll(request_id) def _step(self, body: dict) -> dict: req_body = {} if body["learning_rate"] is None else {"learning_rate": body["learning_rate"]} request_id = self._send("POST", f"{self._prefix}/{self.job_id}/step", json=req_body).json()["request_id"] - return self._poll(request_id) + return self._shape_train_response(self._poll(request_id)) def _save_checkpoint(self, body: dict) -> dict: request_id = self._send( @@ -344,6 +402,10 @@ def _generate(self, body: dict) -> dict: return {"results": self._poll(request_id).get("results", [])} def _sync_weights(self, body: dict) -> dict: + # cuda_ipc / low_memory flags are on-prem colocation hints; Cortex has + # separate sub-jobs and does a server-driven pull regardless. Drop + # them silently rather than raising so SkyRL's colocated call site + # (`sync_weights(cuda_ipc=True)`) works unchanged. source = self.sub_jobs["training"] request_id = self._operation( "weight-sync", @@ -353,6 +415,90 @@ def _sync_weights(self, body: dict) -> dict: )["request_id"] return self._poll(request_id) + def _colo_noop(self, body: dict) -> dict: + """No-op handler for colocation-lifecycle ops. + + SkyRL calls wake/sleep/empty_training_cache/weight_norm unconditionally + under `colocate=True`. Cortex has no colocation lifecycle (training + and sampling live in separate sub-jobs), so these are no-ops. Returning + `{}` matches on-prem's post-hoc metrics-less responses closely enough + that downstream `.get(...)` reads don't blow up. + """ + return {} + + def _normalize_train_body(self, body: dict) -> dict: + """Translate the verl-GRPO envelope into Cortex's RPC-style body. + + The unified frontend forwards ``batch`` verbatim, but SkyRL and verl + build ``{batch: {input_ids, labels, ...}, meta: {...}, processing: {...}}`` + while Cortex expects ``{args: [], kwargs: {input_ids, labels, ...}}``. + We detect the verl-GRPO shape (``"batch"`` key holding a dict) and + repack it into Cortex's shape. Bodies already in Cortex shape pass + through unchanged. + + Extra sibling keys (``meta``, ``processing``, ``router_replay``, + ``context``, ``post_processors``, ``reference_model``, …) are copied + alongside so the Neutrino trainer — whose proto is + ``additionalProperties: true`` — can route them without a schema + change. ``reference_model`` in particular is how verl toggles between + the actor and reference forward-only pass. + """ + if not isinstance(body, dict): + return body + batch = body.get("batch") + if not isinstance(batch, dict): + return body + kwargs = dict(batch) + out: dict = {"args": [], "kwargs": kwargs} + for key in ( + "meta", + "processing", + "router_replay", + "context", + "post_processors", + "reference_model", + ): + if body.get(key) is not None: + out[key] = body[key] + return out + + def _shape_train_response(self, result: dict) -> dict: + """Coerce a Cortex train-op response into the shape callers expect. + + Cortex today returns loss-only for fwd_bwd; SkyRL reads + ``result["grad_norm"]`` and verl reads ``result["avg_loss"]`` / + ``result["post_process_outputs"]``. We surface the loss under + ``avg_loss`` (and mirror it as ``loss`` if the server used a different + key) so the integrations Just Work. Anything the server does return + (``model_outputs``, extra scalars) passes through untouched. + + For fwd_no_grad, on-prem returns ``{"batch": {"logprobs": ..., + "entropy": ...}, ...}`` while Cortex packages the same fields under + ``model_outputs``. verl's adapter reads ``response["batch"]["log_probs"]`` + after renaming; we alias ``model_outputs`` -> ``batch`` here so the + response schema is uniform across transports without touching the + integration. + """ + if not isinstance(result, dict): + return result + out = dict(result) + if "avg_loss" not in out and "loss" in out: + out["avg_loss"] = out["loss"] + elif "loss" not in out and "avg_loss" in out: + out["loss"] = out["avg_loss"] + # Alias model_outputs -> batch so on-prem's response shape is the + # canonical one, regardless of which server produced the result. + if "batch" not in out and isinstance(out.get("model_outputs"), dict): + out["batch"] = dict(out["model_outputs"]) + # Ensure the two dicts SkyRL/verl reach into always exist. + out.setdefault("metrics", {}) + out.setdefault("post_process_outputs", {}) + # ``grad_norm`` is not returned by Cortex today. Surface a None so + # ``.get("grad_norm")`` returns None rather than raising KeyError from + # downstream code that does dict subscripting. + out["metrics"].setdefault("grad_norm", None) + return out + def _reset_prefix_cache(self, body: dict) -> dict: result = self._operation( "reset-prefix-cache", diff --git a/arctic_platform/client/transports/onprem.py b/arctic_platform/client/transports/onprem.py index a0abfb3..2e9dfa1 100644 --- a/arctic_platform/client/transports/onprem.py +++ b/arctic_platform/client/transports/onprem.py @@ -78,16 +78,26 @@ def _check_op_coverage(self, target: object) -> None: def _init_payload(self, job_type: str) -> dict[str, Any]: cfg = self.config payload: dict[str, Any] = {"model_name": cfg.model_name, "job_type": job_type, "seed": cfg.seed} + if cfg.full_determinism: + payload["full_determinism"] = True if job_type in ("training", "log_prob"): - if cfg.ds_config: + # Log-prob may override the base `ds_config` with its own tuning. + if job_type == "log_prob" and cfg.log_prob_ds_config is not None: + payload["ds_config"] = cfg.log_prob_ds_config + elif cfg.ds_config: payload["ds_config"] = cfg.ds_config + if cfg.ds_worker_config is not None: + payload["ds_worker_config"] = cfg.ds_worker_config if job_type == "training": if cfg.training_config: payload["training_config"] = cfg.training_config if cfg.checkpoint_path: payload["checkpoint_path"] = cfg.checkpoint_path - elif cfg.vllm_config: - payload["vllm_config"] = cfg.vllm_config + else: + if cfg.vllm_config: + payload["vllm_config"] = cfg.vllm_config + if cfg.arctic_inference_config is not None: + payload["arctic_inference_config"] = cfg.arctic_inference_config return payload # delivery primitives — the only things a concrete transport implements diff --git a/arctic_platform/client/transports/onprem_ray.py b/arctic_platform/client/transports/onprem_ray.py index a823c9b..67fddb6 100644 --- a/arctic_platform/client/transports/onprem_ray.py +++ b/arctic_platform/client/transports/onprem_ray.py @@ -35,17 +35,25 @@ class RayTransport(OnPremTransport): `_rpc` resolves ``op -> method`` and forwards the request unchanged. """ - def __init__(self, config: ArcticRLClientConfig) -> None: + def __init__(self, config: ArcticRLClientConfig, *, server_state: object | None = None) -> None: super().__init__(config) from arctic_platform.rl.ray_server import create_arctic_rl_ray_server_state - self._state = create_arctic_rl_ray_server_state( - training_gpus=config.training_gpus, - sampling_gpus=config.sampling_gpus, - log_prob_gpus=config.log_prob_gpus, - log_prob_engine="deepspeed", - colocate=config.colocate, - ) + # Reconnect: forwarder workers (verl) and Ray-remote entrypoints (SkyRL) + # rebuild the client in a different process from the driver. They pass + # the driver's state actor back in via `server_state=` so we reattach + # instead of spinning up fresh Ray actors that would race the existing + # jobs. Non-reconnect path: build a fresh state actor as before. + if server_state is not None: + self._state = server_state + else: + self._state = create_arctic_rl_ray_server_state( + training_gpus=config.training_gpus, + sampling_gpus=config.sampling_gpus, + log_prob_gpus=config.log_prob_gpus, + log_prob_engine="deepspeed", + colocate=config.colocate, + ) self._server = None # ArcticRLRayServer, built once jobs exist # One long-lived loop for every op instead of asyncio.run() per call. self._loop = asyncio.new_event_loop() @@ -74,3 +82,14 @@ def _destroy(self, job_id: JobId, job_type: str) -> None: def shutdown(self) -> None: super().shutdown() self._loop.close() + + def get_server_state(self) -> object: + """Expose the state actor for reconnect flows. + + The SkyRL entrypoint runs the driver-side client in the parent + process, then dispatches training to a Ray remote task with a + `reconnect_config` and this server_state so the worker can reattach + to the same actors. Non-Ray transports return None from the client + wrapper `ArcticRLClient.get_server_state`. + """ + return self._state diff --git a/arctic_platform/rl/__init__.py b/arctic_platform/rl/__init__.py index dbaa3d4..4544dc1 100644 --- a/arctic_platform/rl/__init__.py +++ b/arctic_platform/rl/__init__.py @@ -13,18 +13,70 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Arctic RL client -- HTTP client for RL training against dss-platform or local server.""" +"""Arctic RL client — HTTP / Ray / Cortex frontends for RL training. + +Two symbols are imported eagerly (`ArcticRLClientConfig`, `WeightSyncConfig`) +because they carry no heavy transitive deps — just pydantic. Everything else +is lazy-loaded on first attribute access via `__getattr__` so that + + from arctic_platform.rl import ArcticRLClientConfig + # or: + from arctic_platform.rl import create_arctic_rl_client + create_arctic_rl_client(ArcticRLClientConfig(backend="cortex", ...)) + +on a *Cortex-only* driver (no vllm / ray / arctic_inference / torch installed) +succeeds. Prior to this, `__init__` eagerly loaded `client.py`, which pulled +`http_client.py` → `http_server.py` → `arctic_inference.server.metrics` → +`vllm`. Cortex users paid that cost even though the Cortex dispatch branch +(see `arctic_platform.rl.client.create_arctic_rl_client`) short-circuits to +`arctic_platform.client` and never touches the on-prem HTTP server code. + +Attribute-name → (module, attr) map below defines the public surface: +`__all__` still lists everything, so `from arctic_platform.rl import *` +retrieves the same set as before (each import triggers the lazy load). +""" + +from __future__ import annotations + +import importlib +from typing import Any -from arctic_platform.rl.client import create_arctic_rl_client from arctic_platform.rl.config import ArcticRLClientConfig from arctic_platform.rl.config import WeightSyncConfig -from arctic_platform.rl.processors import grpo_loss -from arctic_platform.rl.processors import pack_sequences -from arctic_platform.rl.processors import register_loss_fn -from arctic_platform.rl.processors import register_post_processor -from arctic_platform.rl.processors import run_pipeline -from arctic_platform.rl.processors import unpack_sequences -from arctic_platform.rl.weight_sync import WeightSyncCoordinator + +# name -> (submodule dotted path, attribute name on that submodule). +# Kept as data (not code) so the lazy resolver stays a single implementation. +_LAZY_EXPORTS: dict[str, tuple[str, str]] = { + "create_arctic_rl_client": ("arctic_platform.rl.client", "create_arctic_rl_client"), + "grpo_loss": ("arctic_platform.rl.processors", "grpo_loss"), + "pack_sequences": ("arctic_platform.rl.processors", "pack_sequences"), + "register_loss_fn": ("arctic_platform.rl.processors", "register_loss_fn"), + "register_post_processor": ("arctic_platform.rl.processors", "register_post_processor"), + "run_pipeline": ("arctic_platform.rl.processors", "run_pipeline"), + "unpack_sequences": ("arctic_platform.rl.processors", "unpack_sequences"), + "WeightSyncCoordinator": ("arctic_platform.rl.weight_sync", "WeightSyncCoordinator"), +} + + +def __getattr__(name: str) -> Any: + """PEP 562 lazy loader for the heavy exports. + + Cached in `globals()` after first resolution so subsequent attribute + accesses hit the normal fast path (no re-import). + """ + try: + mod_path, attr = _LAZY_EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module 'arctic_platform.rl' has no attribute {name!r}") from exc + value = getattr(importlib.import_module(mod_path), attr) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + """Include lazy exports in `dir()` and tab-completion.""" + return sorted(set(globals()) | set(_LAZY_EXPORTS)) + __all__ = [ "create_arctic_rl_client", diff --git a/arctic_platform/rl/_cortex_dispatch.py b/arctic_platform/rl/_cortex_dispatch.py new file mode 100644 index 0000000..e2fe6a6 --- /dev/null +++ b/arctic_platform/rl/_cortex_dispatch.py @@ -0,0 +1,270 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Cortex dispatch layer for the legacy `arctic_platform.rl` client entry point. + +Both merged upstream integrations +(`NovaSky-AI/SkyRL#1837:integrations/arctic_rl/` and the Arctic-specific +`RemoteBackend` under `arctic_platform/integrations/verl/adapter.py` used by +`verl-project/verl#6422`) construct their client via a single call: + + from arctic_platform.rl import ArcticRLClientConfig, create_arctic_rl_client + client = create_arctic_rl_client(config, server_state) + +For Cortex support without touching either integration, `create_arctic_rl_client` +dispatches to this module when `config.backend == "cortex"`. This module: + +- Translates the on-prem `ArcticRLClientConfig` into the unified + `arctic_platform.client.ArcticRLClientConfig` (Cortex-relevant fields only; + on-prem knobs like `ds_config` / `arctic_inference_config` are dropped + because Cortex has no local placement to configure). +- Wraps the unified (synchronous) `ArcticRLClient` in `_CortexClientShim`, + which re-exposes the exact async surface the legacy + `ArcticRLHTTPClient` / `ArcticRLRayClient` had. Every `async def foo(...)` + method internally calls the corresponding sync unified method — legal + Python (`async def` that never awaits is a valid coroutine returning the + computed value) and semantically identical from the caller's POV because + the underlying operation is I/O over Cortex, which the unified client + already handles. + +Design rules: + +- Method surface is a strict superset of what the two integrations reach + for (verified against `integrations/arctic_rl/trainer.py` + + `arctic_platform/integrations/verl/adapter.py`). Any callable the + integrations hit is here. +- Properties (`config`, `training_job_id`, `sampling_job_id`, + `log_prob_job_id`) match the legacy client's public names verbatim so + the integrations' attribute reads work unmodified. +- The compat layer in `arctic_platform.client` (this PR's earlier commits) + already shapes Cortex responses into the on-prem envelope + (`metrics.grad_norm` at the top level, `model_outputs` -> `batch` alias, + `logprobs` naming preserved). This shim doesn't re-shape — it just + forwards. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from arctic_platform.rl.config import ArcticRLClientConfig as _LegacyConfig + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Config translation +# --------------------------------------------------------------------------- + + +def _to_unified_config(legacy: _LegacyConfig): + """Translate an on-prem `arctic_platform.rl.ArcticRLClientConfig` into a + unified `arctic_platform.client.ArcticRLClientConfig` sized for Cortex. + + Only fields Cortex actually consumes are threaded through; on-prem-only + fields (ds_config, log_prob_ds_config, ds_worker_config, + arctic_inference_config, colocate, ray_auto_attach, startup_timeout, ...) + are intentionally dropped because Cortex has no local placement / no + DeepSpeed worker to configure. `extra="ignore"` on the unified config + would silently drop them if we forwarded them, but being explicit here + keeps the wire minimal and makes divergences visible. + """ + from arctic_platform.client.config import ArcticRLClientConfig as _UnifiedConfig + + kwargs: dict[str, Any] = { + "backend": "cortex", + "model_name": legacy.model_name, + "training_gpus": legacy.training_gpus, + "sampling_gpus": legacy.sampling_gpus, + "log_prob_gpus": legacy.log_prob_gpus, + "seed": legacy.seed, + # Cortex reconnect: forward pre-existing job ids so the unified client + # skips job init and attaches instead. Consumed by JobHandles.from_config. + "training_job_id": legacy.training_job_id, + "sampling_job_id": legacy.sampling_job_id, + "log_prob_job_id": legacy.log_prob_job_id, + } + for legacy_attr, unified_attr in ( + ("cortex_host", "cortex_host"), + ("cortex_database", "cortex_database"), + ("cortex_schema", "cortex_schema"), + ("cortex_endpoint", "cortex_endpoint"), + ("cortex_pat_env_var", "cortex_pat_env_var"), + ("cortex_base_url", "cortex_base_url"), + ("max_seq_len", "max_seq_len"), + ): + val = getattr(legacy, legacy_attr, None) + if val is not None: + kwargs[unified_attr] = val + + return _UnifiedConfig(**kwargs) + + +# --------------------------------------------------------------------------- +# Async facade over the (synchronous) unified client +# --------------------------------------------------------------------------- + + +class _CortexClientShim: + """Async facade over `arctic_platform.client.ArcticRLClient` that matches + the public surface of the legacy `ArcticRLHTTPClient` / `ArcticRLRayClient`. + + The unified client is synchronous by design (the underlying Cortex RPC is + a single request/response). The legacy client is async because the + on-prem HTTP path uses `httpx.AsyncClient`. Both merged integrations wire + the client through both idioms — SkyRL's `_run(coro)` wrapper in the sync + training loop and direct `await client.foo(...)` in the async paths. + Exposing `async def` methods that internally call sync unified methods + makes both idioms work identically. + """ + + def __init__(self, unified_client, legacy_config: _LegacyConfig) -> None: + self._client = unified_client + # Both integrations read `client.config.colocate` (SkyRL) and + # `client.config` for reconnect flows (verl); expose the legacy + # config, not the translated unified one, so those reads see what + # the caller passed in. + self._legacy_config = legacy_config + + # -- Properties preserved verbatim from the legacy client ------------------ + + @property + def config(self) -> _LegacyConfig: + return self._legacy_config + + @property + def training_job_id(self): + return self._client.training_job_id + + @property + def sampling_job_id(self): + return self._client.sampling_job_id + + @property + def log_prob_job_id(self): + return self._client.log_prob_job_id + + # -- Sync methods (already sync on the legacy client) --------------------- + + def reconnect_config(self) -> _LegacyConfig: + """Rebuild an on-prem-shaped config that reconnects to the same Cortex + sub-jobs. Called by SkyRL's driver -> Ray-worker handoff in + `main_arctic_rl.py` and by verl's `reconnect_handle()`. + """ + # We can't just call `self._client.reconnect_config()` because that + # returns a unified config; the caller expects a legacy config. Round- + # trip via the legacy config with cortex job ids populated. + return self._legacy_config.model_copy( + update={ + "training_job_id": self._client.training_job_id, + "sampling_job_id": self._client.sampling_job_id, + "log_prob_job_id": self._client.log_prob_job_id, + } + ) + + def get_server_state(self): + """Cortex has no local Ray state actor; verl's `reconnect_handle` + pattern still calls this. Returning None is safe — the reconnect + path re-attaches via `training_job_id` on the config, not via a + server-state handle. + """ + return None + + def shutdown(self) -> None: + self._client.shutdown() + + # -- Async methods (legacy client is async; shim keeps that signature) --- + + async def fwd_bwd(self, batch: dict, **legacy_kwargs: Any) -> dict: + return self._client.fwd_bwd(batch, **legacy_kwargs) + + async def fwd_no_grad(self, batch: dict, **legacy_kwargs: Any) -> dict: + return self._client.fwd_no_grad(batch, **legacy_kwargs) + + async def step(self, learning_rate: float | None = None) -> dict: + return self._client.step(learning_rate=learning_rate) + + async def save_checkpoint(self, stage_info: dict | None = None, path: str | None = None) -> dict: + return self._client.save_checkpoint(stage_info=stage_info, path=path) + + async def save_weights(self, path: str) -> dict: + return self._client.save_weights(path=path) + + async def generate(self, prompts, sampling_params=None, **kwargs) -> list: + return self._client.generate(prompts=prompts, sampling_params=sampling_params, **kwargs) + + async def sync_weights(self, cuda_ipc: bool = False, low_memory: bool = False) -> dict: + return self._client.sync_weights(cuda_ipc=cuda_ipc, low_memory=low_memory) + + async def reset_prefix_cache(self, drain: bool = True, timeout_s: float = 60.0) -> dict: + return self._client.reset_prefix_cache(drain=drain, timeout_s=timeout_s) + + async def wake_inference(self, **kwargs: Any) -> dict: + return self._client.wake_inference(**kwargs) + + async def sleep_inference(self, **kwargs: Any) -> dict: + return self._client.sleep_inference(**kwargs) + + async def wake_training(self, **kwargs: Any) -> dict: + return self._client.wake_training(**kwargs) + + async def sleep_training(self, **kwargs: Any) -> dict: + return self._client.sleep_training(**kwargs) + + async def wake_log_prob(self, **kwargs: Any) -> dict: + return self._client.wake_log_prob(**kwargs) + + async def sleep_log_prob(self, **kwargs: Any) -> dict: + return self._client.sleep_log_prob(**kwargs) + + async def empty_training_cache(self, **kwargs: Any) -> dict: + return self._client.empty_training_cache(**kwargs) + + async def weight_norm(self, **kwargs: Any) -> dict: + return self._client.weight_norm(**kwargs) + + async def log_probs(self, batch: dict, **kwargs: Any) -> dict: + # Unified client exposes log-prob as a fwd_no_grad variant + the + # compat layer registers a `log-probs` op on CortexTransport; forward + # via the public `log_probs` accessor when present, else fall back to + # `fwd_no_grad` with a marker kwarg. + fn = getattr(self._client, "log_probs", None) + if callable(fn): + return fn(batch, **kwargs) + return self._client.fwd_no_grad(batch, log_probs_only=True, **kwargs) + + +# --------------------------------------------------------------------------- +# Public entrypoint +# --------------------------------------------------------------------------- + + +def build_cortex_client(legacy_config: _LegacyConfig) -> _CortexClientShim: + """Called from `arctic_platform.rl.client.create_arctic_rl_client` when + `config.backend == "cortex"`. Kept as a top-level factory so the shim + can be tested in isolation. + """ + from arctic_platform.client import ArcticRLClient + + unified_config = _to_unified_config(legacy_config) + logger.info( + "arctic_platform.rl -> cortex dispatch: model=%s training_gpus=%d sampling_gpus=%d log_prob_gpus=%d", + unified_config.model_name, + unified_config.training_gpus, + unified_config.sampling_gpus, + unified_config.log_prob_gpus, + ) + unified_client = ArcticRLClient(unified_config) + return _CortexClientShim(unified_client, legacy_config) diff --git a/arctic_platform/rl/client.py b/arctic_platform/rl/client.py index 5960410..4179bd0 100644 --- a/arctic_platform/rl/client.py +++ b/arctic_platform/rl/client.py @@ -26,21 +26,117 @@ from __future__ import annotations import logging +import os +from typing import TYPE_CHECKING from arctic_platform.rl.config import ArcticRLClientConfig -from arctic_platform.rl.http_client import ArcticRLHTTPClient -from arctic_platform.rl.ray_client import ArcticRLRayClient -# from arctic_platform.rl.ray_server import ArcticRLRayServerState -from arctic_platform.rl.server import ArcticRLServerState +if TYPE_CHECKING: + # Only referenced as a type hint on the function signature; the concrete + # class lives under `arctic_platform.rl.server` which pulls ray in. Keeping + # the runtime import guarded so a Cortex-only caller (no ray installed) + # can still `from arctic_platform.rl import create_arctic_rl_client`. + from arctic_platform.rl.server import ArcticRLServerState logger = logging.getLogger(__name__) -def create_arctic_rl_client(config: ArcticRLClientConfig, arctic_rl_server_state: ArcticRLServerState = None): +# Env vars that override the incoming legacy config into a Cortex-backed one. +# The two merged upstream integrations (NovaSky-AI/SkyRL#1837 and verl-project/ +# verl#6422's Arctic adapter at arctic_platform/integrations/verl/adapter.py) +# both hardcode `backend="local"` when they construct ArcticRLClientConfig. +# Editing either adapter would require a new upstream SkyRL PR + our own +# Arctic-Platform-side patch; that's an integration-side change we explicitly +# want to avoid. Instead the launcher exports these env vars and this factory +# rewrites the (hardcoded-local) config into a cortex one before dispatch. +_CORTEX_ENV_TOGGLE = "ARCTIC_RL_BACKEND" +_CORTEX_ENV_MAP: dict[str, str] = { + "CORTEX_BASE_URL": "cortex_base_url", + "CORTEX_HOST": "cortex_host", + "CORTEX_DATABASE": "cortex_database", + "CORTEX_SCHEMA": "cortex_schema", + "CORTEX_ENDPOINT": "cortex_endpoint", + "CORTEX_PAT_ENV_VAR": "cortex_pat_env_var", + "CORTEX_MAX_SEQ_LEN": "max_seq_len", +} + + +def _maybe_override_from_env(config: ArcticRLClientConfig) -> ArcticRLClientConfig: + """Apply Cortex overrides driven by launcher environment variables. + + When ``ARCTIC_RL_BACKEND=cortex`` is set, rewrite ``config.backend`` to + ``"cortex"`` and populate ``cortex_*`` / ``max_seq_len`` from the env + (``CORTEX_BASE_URL``, ``CORTEX_HOST``, ``CORTEX_DATABASE``, ``CORTEX_SCHEMA``, + ``CORTEX_ENDPOINT``, ``CORTEX_PAT_ENV_VAR``, ``CORTEX_MAX_SEQ_LEN``). + Explicit fields already on ``config`` win over env vars — the env is a + fallback for launchers whose adapter code doesn't yet thread cortex knobs. + + When ``ARCTIC_RL_BACKEND`` is unset or has any other value, return ``config`` + untouched. When it's set to ``"cortex"`` on a config that already has + ``backend="cortex"``, still merge the env fields in (idempotent). + """ + requested = os.environ.get(_CORTEX_ENV_TOGGLE, "").strip().lower() + if requested != "cortex": + return config + + overrides: dict = {} + if config.backend != "cortex": + overrides["backend"] = "cortex" + for env_key, field_name in _CORTEX_ENV_MAP.items(): + env_val = os.environ.get(env_key) + if not env_val: + continue + # Preserve explicit config values — env vars only fill in gaps. + if getattr(config, field_name, None) is not None: + continue + if field_name == "max_seq_len": + try: + overrides[field_name] = int(env_val) + except ValueError: + logger.warning("Ignoring non-integer %s=%r for max_seq_len", env_key, env_val) + continue + overrides[field_name] = env_val + + if not overrides: + return config + + logger.info( + "arctic_platform.rl: ARCTIC_RL_BACKEND=cortex active; overriding %s", + sorted(overrides.keys()), + ) + return config.model_copy(update=overrides) + + +def create_arctic_rl_client(config: ArcticRLClientConfig, arctic_rl_server_state: "ArcticRLServerState | None" = None): + # Env-var override lets launchers force Cortex without touching either + # integration's adapter code (both currently hardcode `backend="local"`). + # See `_maybe_override_from_env` for the recognized env vars. + config = _maybe_override_from_env(config) + + # `cortex` short-circuits the on-prem HTTP / Ray transports: dispatch to + # `arctic_platform.client` (unified client + Cortex transport) via a thin + # async facade that preserves the legacy public surface. Both merged + # upstream integrations (SkyRL#1837, verl#6422's Arctic adapter) call this + # factory unchanged; the env override (or `config.backend = "cortex"` set + # directly) is the only knob they touch. + if config.backend == "cortex": + # Lazy import: `arctic_platform.client` pulls the Cortex transport + # dependency chain (requests + pydantic only; no ray/vllm). + from arctic_platform.rl._cortex_dispatch import build_cortex_client + + return build_cortex_client(config) + + # On-prem transports are lazy-loaded to keep the Cortex-only import + # path free of ray / vllm / arctic_inference / uvicorn. The eager + # module-level imports these files used to have would drag the on-prem + # HTTP server code into every Cortex driver just to construct a client. if config.comm_protocol == "http": + from arctic_platform.rl.http_client import ArcticRLHTTPClient + return ArcticRLHTTPClient(config) elif config.comm_protocol == "ray": + from arctic_platform.rl.ray_client import ArcticRLRayClient + # assert arctic_rl_server_state is not None, "arctic_rl_server_state is required for comm_protocol: ray" return ArcticRLRayClient(config, arctic_rl_server_state) else: diff --git a/arctic_platform/rl/config.py b/arctic_platform/rl/config.py index 0e31df8..55fd699 100644 --- a/arctic_platform/rl/config.py +++ b/arctic_platform/rl/config.py @@ -26,10 +26,38 @@ class ArcticRLClientConfig(BaseModel): - backend: Literal["local", "dss-platform"] = "local" + # `cortex` routes through `arctic_platform.client` (unified client + Cortex + # transport) at `create_arctic_rl_client` time — the legacy on-prem path + # (`local` / `dss-platform`) is unchanged. Adding a value here rather than a + # separate config type so both merged upstream integrations + # (NovaSky-AI/SkyRL#1837 + verl-project/verl#6422's Arctic adapter) can flip + # to Cortex by setting a single field with no other code changes. + backend: Literal["local", "dss-platform", "cortex"] = "local" comm_protocol: Literal["http", "ray"] = "http" checkpoint_path: Optional[str] = None + # ---- Cortex-only fields (ignored on `local` / `dss-platform`). -------- + # These are forwarded 1:1 into `arctic_platform.client.ArcticRLClientConfig` + # by `_cortex_dispatch._to_unified_config`; any that are None fall back to + # `CortexTransport` defaults / `CORTEX_*` env vars. + cortex_host: Optional[str] = Field( + default=None, description="Cortex-only: SnowAPI host (e.g. account.snowflakecomputing.com)." + ) + cortex_database: Optional[str] = Field(default=None, description="Cortex-only: SnowAPI database.") + cortex_schema: Optional[str] = Field(default=None, description="Cortex-only: SnowAPI schema.") + cortex_endpoint: Optional[str] = Field( + default=None, description="Cortex-only: Cortex-training endpoint name (defaults to 'cortex-training')." + ) + cortex_pat_env_var: Optional[str] = Field( + default=None, description="Cortex-only: env var holding the PAT (defaults to CORTEX_PAT)." + ) + cortex_base_url: Optional[str] = Field( + default=None, description="Cortex-only: direct GS URL (bypasses PAT auth); useful for mocks / staging." + ) + max_seq_len: Optional[int] = Field( + default=None, description="Cortex-only: max seq len for cortex sub-jobs (falls back to unified default)." + ) + # it's best not to pass explicitly the host and port since they are auto derived from comm_protocol host: Optional[str] = None port: Optional[int] = None @@ -111,7 +139,18 @@ def _derive_host_port(self) -> "ArcticRLClientConfig": this node's routable IP at port 7000 so off-node Ray workers can reach the driver node by IP rather than "localhost". Values passed explicitly by the caller are left untouched (e.g. reconnecting to a known server). + + Cortex bypasses on-prem host/port entirely (routing is via SnowAPI + base_url), so we short-circuit before touching the ray_cluster helper. """ + if self.backend == "cortex": + return self + # If both host and port were passed explicitly, there's nothing to + # derive and we can skip the ray_cluster import entirely — this + # matters for CPU-only test environments where `ray_cluster.py`'s + # transitive imports (via `arctic_platform.rl.utils`) pull tensordict. + if "host" in self.model_fields_set and "port" in self.model_fields_set: + return self # Lazy import to avoid pulling ray in at config import time. from arctic_platform.rl.ray_cluster import primary_ip diff --git a/tests/client/test_client_ops.py b/tests/client/test_client_ops.py index 95552ae..7681c62 100644 --- a/tests/client/test_client_ops.py +++ b/tests/client/test_client_ops.py @@ -36,8 +36,9 @@ class FakeTransport(Transport): - def __init__(self, config: ArcticRLClientConfig) -> None: + def __init__(self, config: ArcticRLClientConfig, *, server_state: object | None = None) -> None: self.config = config + self.server_state = server_state self.jobs = JobHandles() self.calls: list[Request] = [] @@ -132,12 +133,19 @@ def test_reset_prefix_cache_targets_sampling(self, client): assert req.body == {"drain": False, "timeout_s": 5.0} def test_sync_weights_has_no_primary_job_id(self, client): - """sync_weights -> job_id None; both ids ride in the body.""" + """sync_weights -> job_id None; both ids and the on-prem colocation + hints (`cuda_ipc`, `low_memory`) ride in the body. SkyRL's colocated + path calls `sync_weights(cuda_ipc=True)`; Cortex ignores the hints.""" client.sync_weights() req = _last(client) assert req.op == "sync-weights" assert req.job_id is None - assert req.body == {"training_job_id": TRAINING, "sampling_job_id": SAMPLING} + assert req.body == { + "training_job_id": TRAINING, + "sampling_job_id": SAMPLING, + "cuda_ipc": False, + "low_memory": False, + } class TestLifecycle: @@ -159,7 +167,13 @@ class TestOpRegistry: transport's op coverage is checkable without a live backend.""" def test_client_emits_exactly_the_registered_ops(self, client): - """Driving every client op must produce exactly the canonical OPS set.""" + """Driving every client op must produce exactly the canonical OPS set. + + The colocation-lifecycle ops (`wake_/sleep_training/inference/log_prob`, + `empty_training_cache`, `weight_norm`, `save_weights`) are part of the + canonical vocabulary because SkyRL calls them unconditionally under + `colocate=True` — they must dispatch even though Cortex no-ops them. + """ client.fwd_bwd({"input_ids": [1]}) client.fwd_no_grad({"input_ids": [1]}) client.step() @@ -168,6 +182,15 @@ def test_client_emits_exactly_the_registered_ops(self, client): client.log_probs(["hi"]) client.sync_weights() client.reset_prefix_cache() + client.wake_training() + client.sleep_training() + client.wake_inference() + client.sleep_inference() + client.wake_log_prob() + client.sleep_log_prob() + client.empty_training_cache() + client.weight_norm() + client.save_weights() assert {req.op for req in client.transport.calls} == OPS def test_unresolved_ops_flags_a_missing_method(self): @@ -199,12 +222,18 @@ def test_make_transport_selects_ray(self, monkeypatch): import arctic_platform.client.transports.onprem_ray as ray_mod class DummyRay: - def __init__(self, config): + def __init__(self, config, *, server_state=None): self.config = config + self.server_state = server_state monkeypatch.setattr(ray_mod, "RayTransport", DummyRay) cfg = ArcticRLClientConfig(model_name="m", comm_protocol="ray", training_gpus=1) - assert isinstance(client_module.make_transport(cfg), DummyRay) + transport = client_module.make_transport(cfg) + assert isinstance(transport, DummyRay) + assert transport.server_state is None + # Reconnect-path: an existing state actor is threaded through. + reattach = client_module.make_transport(cfg, server_state="state-actor-handle") + assert reattach.server_state == "state-actor-handle" def test_make_transport_selects_http_for_onprem(self): """onprem + http (the default) routes to HttpTransport.""" diff --git a/tests/client/test_rl_cortex_dispatch.py b/tests/client/test_rl_cortex_dispatch.py new file mode 100644 index 0000000..afa55ae --- /dev/null +++ b/tests/client/test_rl_cortex_dispatch.py @@ -0,0 +1,663 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for `arctic_platform.rl.create_arctic_rl_client` -> Cortex dispatch. + +These pin the "zero adapter change" contract: + +- Both merged upstream integrations + (`NovaSky-AI/SkyRL#1837:integrations/arctic_rl/` + + `arctic_platform/integrations/verl/adapter.py` used by verl#6422) call + `create_arctic_rl_client(config, server_state)` and read the returned + client's public surface (async methods + property attrs). +- Flipping `config.backend = "cortex"` in yaml is the only knob those + integrations touch; no import swap, no `await`/`sync` restructure. +- The returned shim's async surface matches the legacy + `ArcticRLHTTPClient` / `ArcticRLRayClient` verbatim so both integrations' + call sites keep working. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +# `arctic_platform.rl` is lazy on the heavy exports (see the `__getattr__` +# in `arctic_platform/rl/__init__.py`). Importing `ArcticRLClientConfig` +# alone is pydantic-only, so this test module runs on any environment, +# Cortex-only drivers included. The `create_arctic_rl_client` symbol +# below only materialises the on-prem-server import chain if a test +# constructs a `backend="local"` config — the Cortex branch short- +# circuits before touching those modules. +from arctic_platform.rl import ArcticRLClientConfig, create_arctic_rl_client # noqa: E402 +from arctic_platform.rl._cortex_dispatch import _CortexClientShim, _to_unified_config # noqa: E402 + + +# --------------------------------------------------------------------------- +# Fake unified client: records every call so tests can assert forwarding. +# --------------------------------------------------------------------------- + + +class _FakeUnifiedClient: + """Stand-in for `arctic_platform.client.ArcticRLClient` inside the shim. + + Records every method call + kwargs. Returns canned bodies shaped like + the on-prem envelopes the integrations expect. + """ + + def __init__(self) -> None: + self.calls: list[tuple[str, tuple, dict]] = [] + self.training_job_id = "cortex-train-1" + self.sampling_job_id = "cortex-sample-1" + self.log_prob_job_id = None + + def _record(self, name: str, args: tuple, kwargs: dict) -> Any: + self.calls.append((name, args, kwargs)) + # Shape: fwd_bwd / step responses carry a `metrics` dict; fwd_no_grad + # carries a `batch` dict; the rest return {}. + if name in {"fwd_bwd", "step"}: + return {"metrics": {"grad_norm": 0.5, "avg_loss": 0.1}, "loss": 0.1} + if name == "fwd_no_grad": + return {"batch": {"logprobs": [[0.0]]}, "model_outputs": {"logprobs": [[0.0]]}} + return {} + + def fwd_bwd(self, batch, **kw): + return self._record("fwd_bwd", (batch,), kw) + + def fwd_no_grad(self, batch, **kw): + return self._record("fwd_no_grad", (batch,), kw) + + def step(self, learning_rate=None): + return self._record("step", (), {"learning_rate": learning_rate}) + + def save_checkpoint(self, stage_info=None, path=None): + return self._record("save_checkpoint", (), {"stage_info": stage_info, "path": path}) + + def save_weights(self, path): + return self._record("save_weights", (), {"path": path}) + + def generate(self, prompts, sampling_params=None, **kw): + return self._record("generate", (prompts,), {"sampling_params": sampling_params, **kw}) + + def sync_weights(self, cuda_ipc=False, low_memory=False): + return self._record("sync_weights", (), {"cuda_ipc": cuda_ipc, "low_memory": low_memory}) + + def reset_prefix_cache(self, drain=True, timeout_s=60.0): + return self._record("reset_prefix_cache", (), {"drain": drain, "timeout_s": timeout_s}) + + def wake_inference(self, **kw): + return self._record("wake_inference", (), kw) + + def sleep_inference(self, **kw): + return self._record("sleep_inference", (), kw) + + def wake_training(self, **kw): + return self._record("wake_training", (), kw) + + def sleep_training(self, **kw): + return self._record("sleep_training", (), kw) + + def wake_log_prob(self, **kw): + return self._record("wake_log_prob", (), kw) + + def sleep_log_prob(self, **kw): + return self._record("sleep_log_prob", (), kw) + + def empty_training_cache(self, **kw): + return self._record("empty_training_cache", (), kw) + + def weight_norm(self, **kw): + return self._record("weight_norm", (), kw) + + def shutdown(self): + self._record("shutdown", (), {}) + + +def _run(coro): + """Drive a coroutine to completion in a fresh loop. + + Mirrors `_run(...)` in SkyRL's `integrations/arctic_rl/trainer.py` + (which uses `asyncio.run`) so if this helper works, that call site does. + """ + return asyncio.run(coro) + + +@pytest.fixture +def legacy_cfg() -> ArcticRLClientConfig: + return ArcticRLClientConfig( + backend="cortex", + model_name="Qwen/Qwen3-0.6B", + training_gpus=4, + sampling_gpus=2, + log_prob_gpus=0, + cortex_host="test.snowflakecomputing.com", + cortex_database="ARCTIC_DB", + cortex_schema="RL", + cortex_endpoint="cortex-training", + max_seq_len=8192, + seed=1234, + ) + + +# --------------------------------------------------------------------------- +# Config translation +# --------------------------------------------------------------------------- + + +class TestConfigTranslation: + def test_backend_cortex_translates_to_unified_cortex(self, legacy_cfg): + unified = _to_unified_config(legacy_cfg) + assert unified.backend == "cortex" + assert unified.model_name == "Qwen/Qwen3-0.6B" + assert unified.training_gpus == 4 + assert unified.sampling_gpus == 2 + assert unified.log_prob_gpus == 0 + assert unified.seed == 1234 + assert unified.max_seq_len == 8192 + + def test_cortex_fields_threaded(self, legacy_cfg): + unified = _to_unified_config(legacy_cfg) + assert unified.cortex_host == "test.snowflakecomputing.com" + assert unified.cortex_database == "ARCTIC_DB" + assert unified.cortex_schema == "RL" + assert unified.cortex_endpoint == "cortex-training" + + def test_omitted_cortex_fields_fall_to_unified_defaults(self): + cfg = ArcticRLClientConfig( + backend="cortex", model_name="Qwen/Qwen3-0.6B", training_gpus=1 + ) + unified = _to_unified_config(cfg) + assert unified.cortex_host is None + assert unified.cortex_database == "" # unified default + assert unified.cortex_endpoint == "cortex-training" # unified default + assert unified.cortex_pat_env_var == "CORTEX_PAT" # unified default + + def test_onprem_fields_are_not_forwarded(self, legacy_cfg): + """`ds_config` / `ds_worker_config` / `arctic_inference_config` are + on-prem-only; the shim must not silently forward them because Cortex + has no local placement to configure and picking them up would mask + real config drift.""" + legacy_cfg.ds_config = {"train_batch_size": 32} + legacy_cfg.ds_worker_config = {"use_liger": True} + legacy_cfg.arctic_inference_config = {"speculative_decoding": {"model": "..."}} + unified = _to_unified_config(legacy_cfg) + # Unified config's on-prem sub-fields must be None (they may exist on + # the unified schema but the translator declines to forward them). + assert getattr(unified, "ds_config", None) is None + assert getattr(unified, "ds_worker_config", None) is None + assert getattr(unified, "arctic_inference_config", None) is None + + def test_reconnect_job_ids_forwarded(self): + cfg = ArcticRLClientConfig( + backend="cortex", + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + training_job_id=42, + sampling_job_id=43, + log_prob_job_id=44, + ) + unified = _to_unified_config(cfg) + assert unified.training_job_id == 42 + assert unified.sampling_job_id == 43 + assert unified.log_prob_job_id == 44 + + +# --------------------------------------------------------------------------- +# Shim: async surface + property surface +# --------------------------------------------------------------------------- + + +class TestShimAsyncSurface: + """Each integration call site must resolve to the underlying unified + client method with matching kwargs. Missing methods regress `await + client.foo(...)` in the integrations.""" + + @pytest.fixture + def fake(self, legacy_cfg): + return _FakeUnifiedClient(), _CortexClientShim(_FakeUnifiedClient(), legacy_cfg) + + @pytest.fixture + def shim(self, legacy_cfg): + fake = _FakeUnifiedClient() + return fake, _CortexClientShim(fake, legacy_cfg) + + def test_fwd_bwd_forwards_batch_and_kwargs(self, shim): + fake, shim = shim + result = _run(shim.fwd_bwd({"input_ids": []}, reference_model=True)) + assert result["metrics"]["grad_norm"] == 0.5 + name, args, kwargs = fake.calls[-1] + assert name == "fwd_bwd" + assert kwargs == {"reference_model": True} + + def test_fwd_no_grad_forwards_reference_model(self, shim): + """verl's adapter calls fwd_no_grad(payload, reference_model=True/False).""" + fake, shim = shim + _run(shim.fwd_no_grad({"batch": {}}, reference_model=False)) + name, _, kwargs = fake.calls[-1] + assert name == "fwd_no_grad" + assert kwargs == {"reference_model": False} + + def test_fwd_no_grad_forwards_post_processors(self, shim): + """SkyRL's dispatch calls fwd_no_grad(..., post_processors=[...]).""" + fake, shim = shim + _run(shim.fwd_no_grad({"kwargs": {}}, post_processors=["logprobs"])) + _, _, kwargs = fake.calls[-1] + assert kwargs == {"post_processors": ["logprobs"]} + + def test_step_forwards_learning_rate(self, shim): + fake, shim = shim + _run(shim.step()) + _run(shim.step(learning_rate=1e-5)) + assert fake.calls[-2] == ("step", (), {"learning_rate": None}) + assert fake.calls[-1] == ("step", (), {"learning_rate": 1e-5}) + + def test_sync_weights_cuda_ipc_and_low_memory(self, shim): + """SkyRL colocated calls `sync_weights(cuda_ipc=True)` and verl reads + both `cuda_ipc` and `low_memory` off yaml.""" + fake, shim = shim + _run(shim.sync_weights(cuda_ipc=True, low_memory=True)) + _, _, kwargs = fake.calls[-1] + assert kwargs == {"cuda_ipc": True, "low_memory": True} + + def test_wake_inference_forwards_tags(self, shim): + """verl calls wake_inference(tags=[...]).""" + fake, shim = shim + _run(shim.wake_inference(tags=["vllm"])) + _, _, kwargs = fake.calls[-1] + assert kwargs == {"tags": ["vllm"]} + + def test_sleep_inference_forwards_level(self, shim): + """verl calls sleep_inference(level=2).""" + fake, shim = shim + _run(shim.sleep_inference(level=2)) + _, _, kwargs = fake.calls[-1] + assert kwargs == {"level": 2} + + def test_colocation_lifecycle_ops_all_present(self, shim): + """SkyRL colocated path calls every wake_/sleep_/empty_/weight_norm op.""" + fake, shim = shim + _run(shim.empty_training_cache()) + _run(shim.wake_training()) + _run(shim.sleep_training()) + _run(shim.wake_log_prob()) + _run(shim.sleep_log_prob()) + _run(shim.weight_norm()) + names = [c[0] for c in fake.calls] + assert names == [ + "empty_training_cache", + "wake_training", + "sleep_training", + "wake_log_prob", + "sleep_log_prob", + "weight_norm", + ] + + def test_generate_forwards_prompts_and_params(self, shim): + fake, shim = shim + _run(shim.generate(["hi"], sampling_params={"temperature": 0.7})) + name, args, kwargs = fake.calls[-1] + assert name == "generate" + assert args == (["hi"],) + assert kwargs == {"sampling_params": {"temperature": 0.7}} + + def test_save_checkpoint_forwards_stage_info_and_path(self, shim): + fake, shim = shim + _run(shim.save_checkpoint(stage_info={"step": 10}, path="/tmp/ckpt")) + _, _, kwargs = fake.calls[-1] + assert kwargs == {"stage_info": {"step": 10}, "path": "/tmp/ckpt"} + + +class TestShimSyncSurface: + def test_shutdown_is_sync(self, legacy_cfg): + fake = _FakeUnifiedClient() + shim = _CortexClientShim(fake, legacy_cfg) + shim.shutdown() # must NOT be a coroutine + assert fake.calls[-1] == ("shutdown", (), {}) + + def test_reconnect_config_returns_legacy_shape_with_job_ids(self, legacy_cfg): + fake = _FakeUnifiedClient() + shim = _CortexClientShim(fake, legacy_cfg) + rc = shim.reconnect_config() + assert isinstance(rc, ArcticRLClientConfig) + assert rc.backend == "cortex" + assert rc.training_job_id == fake.training_job_id + assert rc.sampling_job_id == fake.sampling_job_id + assert rc.log_prob_job_id == fake.log_prob_job_id + + def test_get_server_state_is_none_on_cortex(self, legacy_cfg): + """Cortex has no local Ray state actor; verl's reconnect_handle path + still calls this and must not crash.""" + shim = _CortexClientShim(_FakeUnifiedClient(), legacy_cfg) + assert shim.get_server_state() is None + + +class TestPropertySurface: + def test_config_returns_legacy_config(self, legacy_cfg): + """SkyRL reads `client.config.colocate` and verl reads + `client.config` in reconnect_handle. Must be the legacy shape.""" + shim = _CortexClientShim(_FakeUnifiedClient(), legacy_cfg) + assert shim.config is legacy_cfg + assert isinstance(shim.config, ArcticRLClientConfig) + + def test_job_id_properties_pass_through(self, legacy_cfg): + fake = _FakeUnifiedClient() + shim = _CortexClientShim(fake, legacy_cfg) + assert shim.training_job_id == "cortex-train-1" + assert shim.sampling_job_id == "cortex-sample-1" + assert shim.log_prob_job_id is None + + +# --------------------------------------------------------------------------- +# Factory-level dispatch: `create_arctic_rl_client(config)` picks Cortex +# --------------------------------------------------------------------------- + + +class TestFactoryDispatch: + def test_backend_cortex_returns_shim(self, monkeypatch, legacy_cfg): + """`create_arctic_rl_client(config)` must produce a `_CortexClientShim` + when `config.backend == "cortex"`. Both integrations depend on this + being the ONLY code path they trip when they set `backend: cortex`.""" + + # Stub the unified client so we don't need SnowAPI creds. `_to_unified_config` + # runs for real (it's just pydantic model construction, no I/O). + class _StubUnified: + def __init__(self, cfg, *_a, **_kw): + self.cfg = cfg + self.training_job_id = "t" + self.sampling_job_id = "s" + self.log_prob_job_id = "l" + + def shutdown(self): # for teardown paths + pass + + monkeypatch.setattr("arctic_platform.client.ArcticRLClient", _StubUnified) + + client = create_arctic_rl_client(legacy_cfg) + assert isinstance(client, _CortexClientShim) + assert client.training_job_id == "t" + # Cortex sub-jobs threaded through config translation: + assert client._client.cfg.backend == "cortex" + assert client._client.cfg.model_name == "Qwen/Qwen3-0.6B" + + def test_cortex_path_never_imports_onprem_transports(self, monkeypatch, legacy_cfg): + """The whole point of `create_arctic_rl_client`'s early cortex branch is + that Cortex users don't drag the on-prem HTTP / Ray / vllm chain into + their driver. Pin that invariant: after a cortex dispatch, the on-prem + transport modules must NOT be in `sys.modules`. + """ + import sys + + # Purge any prior loads so we can observe a clean walk of the cortex path. + for mod in list(sys.modules): + if mod in { + "arctic_platform.rl.http_client", + "arctic_platform.rl.http_server", + "arctic_platform.rl.ray_client", + "arctic_platform.rl.ray_server", + }: + monkeypatch.delitem(sys.modules, mod, raising=False) + + class _StubUnified: + def __init__(self, *_a, **_kw): + self.training_job_id = None + self.sampling_job_id = None + self.log_prob_job_id = None + + def shutdown(self): + pass + + monkeypatch.setattr("arctic_platform.client.ArcticRLClient", _StubUnified) + create_arctic_rl_client(legacy_cfg) + + # The on-prem transport modules must NOT have been touched. + loaded = {m for m in sys.modules if m.startswith("arctic_platform.rl.")} + onprem_hits = { + m + for m in loaded + if m + in { + "arctic_platform.rl.http_client", + "arctic_platform.rl.http_server", + "arctic_platform.rl.ray_client", + "arctic_platform.rl.ray_server", + } + } + assert not onprem_hits, ( + f"Cortex path leaked on-prem transport imports: {sorted(onprem_hits)}. " + "This means a Cortex-only driver would still need ray / vllm / " + "arctic_inference installed — regressing the serverless UX." + ) + + def test_onprem_transports_are_lazy_at_module_level(self): + """`arctic_platform/rl/client.py` must NOT eagerly import + `ArcticRLHTTPClient` / `ArcticRLRayClient`. Otherwise merely resolving + `arctic_platform.rl.create_arctic_rl_client` (via the package + `__getattr__`) pulls the on-prem server chain in. + """ + import arctic_platform.rl.client as client_module + + assert not hasattr(client_module, "ArcticRLHTTPClient"), ( + "`ArcticRLHTTPClient` is exposed at module scope; it must be " + "imported inside `create_arctic_rl_client`'s on-prem branch." + ) + assert not hasattr(client_module, "ArcticRLRayClient"), ( + "`ArcticRLRayClient` is exposed at module scope; it must be " + "imported inside `create_arctic_rl_client`'s on-prem branch." + ) + + +class TestEnvOverride: + """`ARCTIC_RL_BACKEND=cortex` env-var rewrites the incoming config. + + Both integrations' adapters currently hardcode `backend="local"` at + ``ArcticRLClientConfig`` construction time. The launcher exports these + env vars and the factory rewrites the config before dispatch, so + neither adapter needs to change. + """ + + @pytest.fixture(autouse=True) + def _clear_cortex_env(self, monkeypatch): + # Ensure a clean env slate — other tests / user shell might have + # ARCTIC_RL_BACKEND=cortex set which would poison local-path tests. + for key in ( + "ARCTIC_RL_BACKEND", + "CORTEX_BASE_URL", + "CORTEX_HOST", + "CORTEX_DATABASE", + "CORTEX_SCHEMA", + "CORTEX_ENDPOINT", + "CORTEX_PAT_ENV_VAR", + "CORTEX_MAX_SEQ_LEN", + ): + monkeypatch.delenv(key, raising=False) + + def test_env_unset_is_noop(self): + """Without `ARCTIC_RL_BACKEND=cortex`, `_maybe_override_from_env` + returns the config untouched — no cortex rewriting, no dispatch + redirection. We assert this at the helper level rather than + driving `create_arctic_rl_client` end-to-end because the local + branch needs vllm/ray installed to construct. + + `host`/`port` explicitly set so `_derive_host_port` doesn't trip + the ray_cluster import (which pulls tensordict). + """ + from arctic_platform.rl.client import _maybe_override_from_env + + cfg = ArcticRLClientConfig( + backend="local", + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + host="localhost", + port=7000, + ) + out = _maybe_override_from_env(cfg) + assert out is cfg + assert out.backend == "local" + + def test_env_flips_local_to_cortex(self, monkeypatch): + """Setting `ARCTIC_RL_BACKEND=cortex` on a `backend="local"` config + rewrites it to cortex and produces a `_CortexClientShim`, without + the caller touching a single field.""" + + class _StubUnified: + def __init__(self, cfg, *_a, **_kw): + self.cfg = cfg + self.training_job_id = None + self.sampling_job_id = None + self.log_prob_job_id = None + + def shutdown(self): + pass + + monkeypatch.setattr("arctic_platform.client.ArcticRLClient", _StubUnified) + + # Simulate what verl's / SkyRL's launcher would export. + monkeypatch.setenv("ARCTIC_RL_BACKEND", "cortex") + monkeypatch.setenv("CORTEX_BASE_URL", "http://localhost:8080") + + # Adapter passes `backend="local"` — that's the hardcode we're + # working around. `host`/`port` explicit so the local-branch + # `_derive_host_port` validator doesn't pull tensordict. + cfg = ArcticRLClientConfig( + backend="local", + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + host="localhost", + port=7000, + ) + client = create_arctic_rl_client(cfg) + assert isinstance(client, _CortexClientShim) + assert client._client.cfg.backend == "cortex" + assert client._client.cfg.cortex_base_url == "http://localhost:8080" + + def test_env_populates_all_recognized_fields(self, monkeypatch): + """Every field in `_CORTEX_ENV_MAP` is threaded through when the + env var is set. Missing fields on the input config get filled + from env; explicit fields on config are preserved.""" + + class _StubUnified: + def __init__(self, cfg, *_a, **_kw): + self.cfg = cfg + self.training_job_id = None + self.sampling_job_id = None + self.log_prob_job_id = None + + def shutdown(self): + pass + + monkeypatch.setattr("arctic_platform.client.ArcticRLClient", _StubUnified) + + monkeypatch.setenv("ARCTIC_RL_BACKEND", "cortex") + monkeypatch.setenv("CORTEX_HOST", "cortex.snowflakecomputing.com") + monkeypatch.setenv("CORTEX_DATABASE", "prod_db") + monkeypatch.setenv("CORTEX_SCHEMA", "rl") + monkeypatch.setenv("CORTEX_ENDPOINT", "cortex-training-v2") + monkeypatch.setenv("CORTEX_PAT_ENV_VAR", "MY_PAT") + monkeypatch.setenv("CORTEX_MAX_SEQ_LEN", "8192") + + cfg = ArcticRLClientConfig( + backend="local", + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + host="localhost", + port=7000, + ) + client = create_arctic_rl_client(cfg) + c = client._client.cfg + assert c.backend == "cortex" + assert c.cortex_host == "cortex.snowflakecomputing.com" + assert c.cortex_database == "prod_db" + assert c.cortex_schema == "rl" + assert c.cortex_endpoint == "cortex-training-v2" + assert c.cortex_pat_env_var == "MY_PAT" + assert c.max_seq_len == 8192 + + def test_env_does_not_clobber_explicit_config_fields(self, monkeypatch): + """A field the adapter DID populate on ``config`` wins over the + matching env var. Env is a fallback for silent adapters, not an + override of explicit intent. + """ + + class _StubUnified: + def __init__(self, cfg, *_a, **_kw): + self.cfg = cfg + self.training_job_id = None + self.sampling_job_id = None + self.log_prob_job_id = None + + def shutdown(self): + pass + + monkeypatch.setattr("arctic_platform.client.ArcticRLClient", _StubUnified) + + monkeypatch.setenv("ARCTIC_RL_BACKEND", "cortex") + monkeypatch.setenv("CORTEX_BASE_URL", "http://env-wins.local:8080") + + cfg = ArcticRLClientConfig( + backend="cortex", + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + cortex_base_url="http://config-wins.local:9090", + ) + client = create_arctic_rl_client(cfg) + assert client._client.cfg.cortex_base_url == "http://config-wins.local:9090" + + def test_env_ignored_when_toggle_value_is_not_cortex(self, monkeypatch): + """Only ``ARCTIC_RL_BACKEND=cortex`` triggers the rewrite. Any + other value (or empty string) is a no-op, so we don't accidentally + break someone with a leftover ``ARCTIC_RL_BACKEND=onprem`` in their + shell. Asserted at the helper level so the local dispatch branch + doesn't need vllm/ray to run this check.""" + from arctic_platform.rl.client import _maybe_override_from_env + + monkeypatch.setenv("ARCTIC_RL_BACKEND", "onprem") + monkeypatch.setenv("CORTEX_BASE_URL", "http://should-be-ignored.local:8080") + + cfg = ArcticRLClientConfig( + backend="local", + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + host="localhost", + port=7000, + ) + out = _maybe_override_from_env(cfg) + assert out is cfg + assert out.backend == "local" + assert out.cortex_base_url is None + + def test_env_invalid_max_seq_len_is_warned_not_crashed(self, monkeypatch, caplog): + """A bad ``CORTEX_MAX_SEQ_LEN`` value shouldn't hard-fail the + launcher — log a warning and drop the field.""" + + class _StubUnified: + def __init__(self, cfg, *_a, **_kw): + self.cfg = cfg + self.training_job_id = None + self.sampling_job_id = None + self.log_prob_job_id = None + + def shutdown(self): + pass + + monkeypatch.setattr("arctic_platform.client.ArcticRLClient", _StubUnified) + + monkeypatch.setenv("ARCTIC_RL_BACKEND", "cortex") + monkeypatch.setenv("CORTEX_MAX_SEQ_LEN", "not_a_number") + + cfg = ArcticRLClientConfig( + backend="local", + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + host="localhost", + port=7000, + ) + with caplog.at_level("WARNING"): + client = create_arctic_rl_client(cfg) + assert isinstance(client, _CortexClientShim) + # The legacy config on the shim: max_seq_len env var was garbage + # and got dropped, so this field stays at whatever the legacy + # config had (None here). The unified config downstream falls back + # to its own default (8192) which is fine. + assert client.config.max_seq_len is None + assert any("CORTEX_MAX_SEQ_LEN" in rec.message for rec in caplog.records) diff --git a/tests/client/test_skyrl_verl_compat.py b/tests/client/test_skyrl_verl_compat.py new file mode 100644 index 0000000..0c5e727 --- /dev/null +++ b/tests/client/test_skyrl_verl_compat.py @@ -0,0 +1,372 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pin the exact call patterns SkyRL and verl integrations use. + +These tests intentionally mirror the SkyRL ``_ArcticDispatch`` and verl +``ArcticRLClientWrapper`` call sites verbatim. Every regression here is a +regression for the pre-merged integrations that (by design) we cannot edit. + +Sources being pinned: +- SkyRL: ``arctic-skyrl/skyrl/backends/arctic_rl/{arctic_trainer,arctic_generator}.py`` +- verl: ``arctic-verl/verl/trainer/ppo/arctic_rl_client.py`` +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import patch + +import pytest + +from arctic_platform.client import ArcticRLClient +from arctic_platform.client import ArcticRLClientConfig +from arctic_platform.client.transport import JobHandles +from arctic_platform.client.transport import OPS +from arctic_platform.client.transport import Request +from arctic_platform.client.transport import Transport +from arctic_platform.client.transport import method_name +from arctic_platform.client.transport import unresolved_ops + + +class _RecordingTransport(Transport): + """A Transport that records every request and returns canned responses. + + The canned responses match the shapes the *on-prem* server returns today + (metrics nested under ``metrics``) so the client-side flattening is + exercised end-to-end. + """ + + def __init__(self, config: ArcticRLClientConfig, *, server_state: Any = None) -> None: + self.config = config + self.server_state = server_state + self.calls: list[Request] = [] + self.jobs = JobHandles(training="t-1", sampling="s-1", log_prob="lp-1") + self.responses: dict[str, Any] = { + "fwd-bwd": {"avg_loss": 0.5, "metrics": {"grad_norm": 1.0, "loss": 0.5}, "post_process_outputs": {}}, + "fwd-no-grad": {"model_outputs": {"logprobs": [[-0.1, -0.2]]}, "metrics": {}}, + "step": {"metrics": {"grad_norm": 0.9, "learning_rate": 1e-5}}, + "generate": {"results": [{"token_ids": [1, 2, 3], "text": "ok", "finish_reason": "stop"}]}, + "sync-weights": {"ok": True}, + "save-checkpoint": {"path": "/tmp/ckpt"}, + "reset-prefix-cache": {"drained": True}, + "log-probs": {"logprobs": [[-0.1, -0.2]]}, + } + + def initialize(self) -> JobHandles: + return self.jobs + + def call(self, request: Request) -> dict: + self.calls.append(request) + return dict(self.responses.get(request.op, {})) + + def shutdown(self) -> None: + return None + + +@pytest.fixture +def recording_client() -> ArcticRLClient: + config = ArcticRLClientConfig( + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + sampling_gpus=1, + log_prob_gpus=1, + ) + with patch("arctic_platform.client.client.make_transport") as mt: + transport = _RecordingTransport(config) + mt.return_value = transport + client = ArcticRLClient(config) + client._transport_for_tests = transport + return client + + +# ── legacy config aliases (SkyRL + verl) ───────────────────────────────── + + +class TestConfigLegacyAliases: + def test_backend_local_maps_to_onprem(self) -> None: + cfg = ArcticRLClientConfig(model_name="m", backend="local", training_gpus=1) + assert cfg.backend == "onprem" + + def test_backend_dss_platform_maps_to_onprem(self) -> None: + cfg = ArcticRLClientConfig(model_name="m", backend="dss-platform", training_gpus=1) + assert cfg.backend == "onprem" + + def test_backend_neutrino_maps_to_cortex(self) -> None: + cfg = ArcticRLClientConfig(model_name="m", backend="neutrino", training_gpus=1) + assert cfg.backend == "cortex" + + def test_sample_gpus_alias_populates_sampling_gpus(self) -> None: + """verl builds its config with `sample_gpus=`; must land as sampling_gpus.""" + cfg = ArcticRLClientConfig(model_name="m", backend="local", sample_gpus=4) + assert cfg.sampling_gpus == 4 + + def test_legacy_ignored_fields_do_not_error(self) -> None: + """`log_prob_engine` and friends flow in from verl configs; must drop silently.""" + cfg = ArcticRLClientConfig( + model_name="m", + backend="local", + training_gpus=1, + sampling_gpus=1, + log_prob_gpus=1, + log_prob_engine="deepspeed", + experiment_name="ignored", + ) + assert cfg.backend == "onprem" + assert not hasattr(cfg, "log_prob_engine") + + def test_server_init_extras_accepted(self) -> None: + """verl adapter passes ds_worker_config / log_prob_ds_config / + arctic_inference_config / full_determinism; the unified config must + accept them so they can reach `_init_payload`.""" + cfg = ArcticRLClientConfig( + model_name="m", + training_gpus=1, + log_prob_ds_config={"train_micro_batch_size_per_gpu": 1}, + ds_worker_config={"use_liger": True, "attn_implementation": "flash_attention_2"}, + arctic_inference_config={"speculative_decoding": {"model": None}}, + full_determinism=True, + ) + assert cfg.log_prob_ds_config == {"train_micro_batch_size_per_gpu": 1} + assert cfg.ds_worker_config["use_liger"] is True + assert cfg.arctic_inference_config["speculative_decoding"]["model"] is None + assert cfg.full_determinism is True + + def test_init_payload_forwards_server_extras(self) -> None: + """The on-prem transport's `_init_payload` must thread the extras onto + the wire so the server DS worker / Arctic-Inference config actually + take effect.""" + from arctic_platform.client.transports.onprem import OnPremTransport + + cfg = ArcticRLClientConfig( + model_name="m", + training_gpus=1, + sampling_gpus=1, + log_prob_gpus=1, + ds_config={"zero_optimization": {"stage": 2}}, + log_prob_ds_config={"train_micro_batch_size_per_gpu": 1}, + ds_worker_config={"attn_implementation": "flash_attention_2"}, + arctic_inference_config={"zorro_inference": {"enable": True}}, + full_determinism=True, + training_config={"training_horizon": 100}, + vllm_config={"gpu_memory_utilization": 0.7}, + ) + + class _Peek(OnPremTransport): + def _start(self, payload: dict) -> str: + return "id" + + def _rpc(self, request): + return {} + + def _destroy(self, job_id, job_type): + return None + + def _wait_running(self) -> None: + return None + + t = _Peek(cfg) + + train = t._init_payload("training") + assert train["ds_config"] == {"zero_optimization": {"stage": 2}} + assert train["ds_worker_config"] == {"attn_implementation": "flash_attention_2"} + assert train["training_config"] == {"training_horizon": 100} + assert train["full_determinism"] is True + assert "arctic_inference_config" not in train # sampling-only + + log_prob = t._init_payload("log_prob") + assert log_prob["ds_config"] == {"train_micro_batch_size_per_gpu": 1} + assert log_prob["ds_worker_config"] == {"attn_implementation": "flash_attention_2"} + assert "training_config" not in log_prob + + sampling = t._init_payload("sampling") + assert sampling["vllm_config"] == {"gpu_memory_utilization": 0.7} + assert sampling["arctic_inference_config"] == {"zorro_inference": {"enable": True}} + assert "ds_config" not in sampling + + +# ── SkyRL client-surface pinning ───────────────────────────────────────── + + +class TestSkyRLCallPatterns: + """SkyRL calls the client in _ArcticDispatch / ArcticGenerator / entrypoint. + + Every method used there must exist on ArcticRLClient with a compatible + signature and response shape. + """ + + def test_fwd_bwd_folds_processing_and_extra_kwargs(self, recording_client: ArcticRLClient) -> None: + # `_ArcticDispatch.forward_backward` call: `fwd_bwd(batch, processing={...})`. + result = recording_client.fwd_bwd( + {"batch": {"input_ids": [[1, 2, 3]]}, "meta": {"pad_token_id": 0}}, + processing={"loss_fn": "grpo", "config": {"n_samples": 4}, "post": ["compute_logprobs"]}, + ) + req = recording_client._transport_for_tests.calls[-1] + assert req.op == "fwd-bwd" + assert req.body["processing"]["loss_fn"] == "grpo" + # SkyRL reads `result["grad_norm"]` at the top level after step; the + # fwd_bwd response is passed through the same flattener. + assert result["grad_norm"] == 1.0 + assert result["avg_loss"] == 0.5 + + def test_fwd_no_grad_accepts_post_processors_kwarg(self, recording_client: ArcticRLClient) -> None: + # `_ArcticDispatch.forward` call: `fwd_no_grad(batch, post_processors=["logprobs"])`. + result = recording_client.fwd_no_grad( + {"batch": {"input_ids": [[1, 2]]}}, + post_processors=["logprobs"], + ) + req = recording_client._transport_for_tests.calls[-1] + assert req.op == "fwd-no-grad" + assert req.body["post_processors"] == ["logprobs"] + assert result["model_outputs"]["logprobs"] + + def test_step_grad_norm_is_top_level(self, recording_client: ArcticRLClient) -> None: + # `_ArcticDispatch.optim_step` reads `.get("grad_norm")` on the return. + out = recording_client.step() + assert out.get("grad_norm") == 0.9 + + def test_sync_weights_accepts_cuda_ipc(self, recording_client: ArcticRLClient) -> None: + # SkyRL colocated path: `sync_weights(cuda_ipc=True)`. + recording_client.sync_weights(cuda_ipc=True) + req = recording_client._transport_for_tests.calls[-1] + assert req.body["cuda_ipc"] is True + + def test_generate_returns_list(self, recording_client: ArcticRLClient) -> None: + # `ArcticGenerator.generate` reads `output.get("token_ids"/"text"/"finish_reason")`. + out = recording_client.generate(["hello"], sampling_params={"temperature": 1.0}) + assert isinstance(out, list) + assert out[0]["token_ids"] == [1, 2, 3] + + def test_colocation_lifecycle_exists(self, recording_client: ArcticRLClient) -> None: + for op in ("wake_training", "wake_inference", "sleep_training", "sleep_inference", "empty_training_cache"): + assert callable(getattr(recording_client, op)) + getattr(recording_client, op)() + + def test_wake_inference_forwards_tags(self, recording_client: ArcticRLClient) -> None: + # verl adapter calls ``self._client.wake_inference(tags=tags)`` and + # ``sleep_inference(level=level)``; the transport-facing body must + # carry those kwargs so on-prem colocation can honor them. + recording_client.wake_inference(tags=["actor"]) + req = recording_client._transport_for_tests.calls[-1] + assert req.op == "wake-inference" + assert req.body == {"tags": ["actor"]} + + recording_client.sleep_inference(level=2) + req = recording_client._transport_for_tests.calls[-1] + assert req.op == "sleep-inference" + assert req.body == {"level": 2} + + def test_job_id_attributes_exposed(self, recording_client: ArcticRLClient) -> None: + # `pre_client.training_job_id` / `sampling_job_id` / `log_prob_job_id` + # read in `integrations.arctic_rl.entrypoint.main`. + assert recording_client.training_job_id == "t-1" + assert recording_client.sampling_job_id == "s-1" + assert recording_client.log_prob_job_id == "lp-1" + + def test_get_server_state_defaults_to_none(self, recording_client: ArcticRLClient) -> None: + # The Ray path returns an actor; non-Ray transports return None. + assert recording_client.get_server_state() is None + + +# ── verl client-surface pinning ────────────────────────────────────────── + + +class TestVerlCallPatterns: + """`ArcticRLClientWrapper` in verl builds a config with `sample_gpus=`, + `backend="local"`, `log_prob_engine="deepspeed"`, then calls + `fwd_no_grad(batch)`, `fwd_bwd(batch, processing=...)`, `step()`, + `generate(prompts=..., sampling_params=...)`, `shutdown()`. + """ + + def test_config_shape_matches_wrapper(self) -> None: + cfg = ArcticRLClientConfig( + host="localhost", + port=7000, + backend="local", + training_gpus=2, + sample_gpus=2, + log_prob_gpus=2, + colocate=True, + log_prob_engine="deepspeed", + model_name="Qwen/Qwen3-0.6B", + ) + assert (cfg.backend, cfg.sampling_gpus, cfg.log_prob_gpus) == ("onprem", 2, 2) + + def test_fwd_bwd_returns_avg_loss_and_post_process_outputs( + self, recording_client: ArcticRLClient + ) -> None: + # `ArcticRLClientWrapper.update_actor` reads + # `result.get("avg_loss")` and `result.get("post_process_outputs")`. + result = recording_client.fwd_bwd( + {"batch": {"input_ids": [[1]]}, "context": {"input_ids": [[1]]}}, + processing={"loss_fn": "grpo", "post": ["compute_logprobs"]}, + ) + assert "avg_loss" in result + assert "post_process_outputs" in result + + def test_fwd_no_grad_model_outputs_logprobs(self, recording_client: ArcticRLClient) -> None: + # `ArcticRLClientWrapper.compute_log_prob` reads + # `result.get("model_outputs", result).get("logprobs")`. + result = recording_client.fwd_no_grad({"kwargs": {"input_ids": [[1]]}, "context": {}}) + outputs = result.get("model_outputs", result) + assert outputs.get("logprobs") + + def test_fwd_no_grad_forwards_reference_model(self, recording_client: ArcticRLClient) -> None: + # verl's ``_send_compute_ref_log_prob`` calls + # ``self._client.fwd_no_grad(payload, reference_model=True)``. The + # transport-facing body must carry the flag so the Cortex normalizer + # can route the request to the reference-model engine. + recording_client.fwd_no_grad({"batch": {"input_ids": [[1]]}}, reference_model=True) + req = recording_client._transport_for_tests.calls[-1] + assert req.op == "fwd-no-grad" + assert req.body["reference_model"] is True + + def test_shutdown_is_idempotent(self, recording_client: ArcticRLClient) -> None: + recording_client.shutdown() + recording_client.shutdown() + + +# ── op-vocabulary sanity ───────────────────────────────────────────────── + + +class TestOpVocabulary: + def test_new_lifecycle_ops_registered(self) -> None: + for op in ( + "wake-training", + "sleep-training", + "wake-inference", + "sleep-inference", + "empty-training-cache", + "save-weights", + ): + assert op in OPS + + def test_cortex_covers_every_op(self) -> None: + """CortexTransport must resolve every op in OPS (some as no-ops).""" + from arctic_platform.client.transports.cortex import CortexTransport + + class _Fake(CortexTransport): + def __init__(self) -> None: + pass + + fake = _Fake() + # Handlers dispatch by name; expose them the way the transport ABC docs + # `unresolved_ops` — construct a shim that maps method_name -> handler. + for op in OPS: + attr = method_name(op) + # ``fwd_bwd`` etc. exist as private methods; the transport dispatches + # through ``_handlers``. This assertion is a lint that new ops added + # to OPS also land in ``_handlers`` (populated in __init__). + if attr.startswith(("fwd", "log_probs", "generate", "step", "save_checkpoint", "sync_weights", "reset_prefix_cache")): + assert hasattr(fake, "_" + attr) or hasattr(fake, attr) diff --git a/tests/e2e/fake_cortex_gs.py b/tests/e2e/fake_cortex_gs.py new file mode 100644 index 0000000..3fb2f22 --- /dev/null +++ b/tests/e2e/fake_cortex_gs.py @@ -0,0 +1,380 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Fake Cortex GS — a plumbing-only stand-in for the SnowAPI Cortex-training +endpoint. + +The point of this server is to let us drive verl and SkyRL end-to-end against +the ``CortexTransport`` without a real Neutrino GS deployment. It speaks every +route ``arctic_platform.client.transports.cortex.CortexTransport`` actually +hits, decodes DSSST1 chunked requests correctly, and returns shape-plausible +canned responses. Losses are random; convergence validation still needs a real +server. + +Not a spec for the Cortex server team — just a running executable of what the +client currently sends. Handy as a starting point when Neutrino GS writes their +own tests against the same client. + +Endpoints implemented: + + POST {prefix} → CreateJob → {"job_id"} + GET {prefix}/{job_id} → GetJob → running + sub_jobs + POST {prefix}/{job_id}:cancel → {} + POST {prefix}/{job_id}/forward-backward → octet chunks → request_id + POST {prefix}/{job_id}/forward-no-grad → octet chunks → request_id + POST {prefix}/{job_id}/generate → octet chunks → request_id + POST {prefix}/{job_id}/step → JSON → request_id + POST {prefix}/{job_id}/save → JSON → request_id + POST {prefix}/{job_id}/log-probs → JSON → request_id + POST {prefix}/{job_id}/operation → JSON → request_id + (operation_type ∈ {weight-sync, reset-prefix-cache}) + GET {prefix}/{job_id}/requests/{request_id} → completed + result + +Prefix layout mirrors what the client builds: + + /api/v2/databases/{database}/schemas/{schema}/{endpoint} + +Any (database, schema, endpoint) triple is accepted so a test / launcher can +pick whatever it wants. + +Run standalone (default port 8080): + + python -m tests.e2e.fake_cortex_gs [--port 8080] [--host 0.0.0.0] + +or from a test via the ``fake_cortex_gs_process`` fixture. +""" + +from __future__ import annotations + +import argparse +import base64 +import logging +import random +import threading +import uuid +from typing import Any + +import torch +import uvicorn +from fastapi import Body +from fastapi import FastAPI +from fastapi import Request +from fastapi.responses import JSONResponse + +from arctic_platform.client import wire + +log = logging.getLogger("fake_cortex_gs") + + +# ─── State ──────────────────────────────────────────────────────────────── + + +class _State: + """Ephemeral in-memory state for one running mock instance. + + Per-job: a dict of pending octet request chunks (keyed by + (job_id, path_suffix)) and a request registry (request_id -> canned result). + """ + + def __init__(self) -> None: + self.jobs: dict[str, dict] = {} + # (job_id, path_suffix, chunk_group_id) -> list[bytes] as chunks arrive + self.chunk_buffers: dict[tuple[str, str, str], list[bytes]] = {} + # request_id -> canned response dict (result already shaped for the client) + self.requests: dict[str, dict] = {} + + +STATE = _State() + + +# ─── Prefix routing ─────────────────────────────────────────────────────── + +# The client's prefix is /api/v2/databases//schemas//. +# All routes below mount under that variable prefix. +_PREFIX = "/api/v2/databases/{database}/schemas/{schema}/{endpoint}" + + +def _make_app() -> FastAPI: + app = FastAPI(title="fake_cortex_gs", version="0.1.0") + + # ── CreateJob ────────────────────────────────────────────────────────── + @app.post(_PREFIX) + async def create_job(database: str, schema: str, endpoint: str, body: dict = Body(...)) -> dict: + job_id = f"fake-job-{uuid.uuid4().hex[:8]}" + sub_job_configs = body.get("sub_job_configs") or [] + sub_jobs = [] + for cfg in sub_job_configs: + jt = cfg.get("job_type", "training") + sub_jobs.append({"job_type": f"job_type_{jt}", "sub_job_id": f"{job_id}:{jt}"}) + STATE.jobs[job_id] = {"sub_jobs": sub_jobs, "status": "job_state_running"} + log.info("fake_cortex_gs: created job %s with sub_jobs=%s", job_id, [s["job_type"] for s in sub_jobs]) + return {"job_id": job_id} + + # ── GetJob (client polls this immediately after CreateJob) ───────────── + # + # Return shape follows what `CortexTransport._wait_for_job` reads: + # top-level `status` = "job_state_running", and `_capture_sub_jobs` + # reads `job_info.get("job", job_info)["sub_jobs"]` (accepts either). + @app.get(_PREFIX + "/{job_id}") + async def get_job(database: str, schema: str, endpoint: str, job_id: str) -> dict: + job = STATE.jobs.get(job_id) + if job is None: + return JSONResponse({"error": "job not found"}, status_code=404) + return {"status": job["status"], "sub_jobs": job["sub_jobs"]} + + # ── Cancel ───────────────────────────────────────────────────────────── + @app.post(_PREFIX + "/{job_id}:cancel") + async def cancel_job(database: str, schema: str, endpoint: str, job_id: str) -> dict: + STATE.jobs.pop(job_id, None) + return {} + + # ── Octet-chunked ops: forward-backward / forward-no-grad / generate ─── + for path_suffix in ("forward-backward", "forward-no-grad", "generate"): + _register_octet_route(app, path_suffix) + + # ── JSON ops: step / save / log-probs ────────────────────────────────── + @app.post(_PREFIX + "/{job_id}/step") + async def step(database: str, schema: str, endpoint: str, job_id: str, body: dict = Body(...)) -> dict: + req_id = _register_result(_fake_step_result(learning_rate=body.get("learning_rate"))) + return {"request_id": req_id} + + @app.post(_PREFIX + "/{job_id}/save") + async def save(database: str, schema: str, endpoint: str, job_id: str, body: dict = Body(...)) -> dict: + req_id = _register_result({"path": f"/fake/checkpoints/{uuid.uuid4().hex[:8]}"}) + return {"request_id": req_id} + + @app.post(_PREFIX + "/{job_id}/log-probs") + async def log_probs(database: str, schema: str, endpoint: str, job_id: str, body: dict = Body(...)) -> dict: + prompts = body.get("prompts") or [] + completions = body.get("completions") or [] + # Return a wire-encoded logprobs tensor shaped [len(prompts), ]. + n_rows = max(1, len(prompts)) + seq = _completion_len_hint(completions) or 8 + result = _fake_log_probs_result(n_rows=n_rows, seq_len=seq) + req_id = _register_result(result, wire_encoded=True) + return {"request_id": req_id} + + # ── operation: weight-sync / reset-prefix-cache ──────────────────────── + @app.post(_PREFIX + "/{job_id}/operation") + async def operation( + database: str, schema: str, endpoint: str, job_id: str, body: dict = Body(...) + ) -> dict: + op_type = body.get("operation_type") + if op_type in {"weight-sync", "reset-prefix-cache"}: + req_id = _register_result({}) + return {"request_id": req_id} + return JSONResponse({"error": f"unknown operation_type {op_type!r}"}, status_code=400) + + # ── Request status (client polls this until state=completed) ────────── + @app.get(_PREFIX + "/{job_id}/requests/{request_id}") + async def request_status( + database: str, schema: str, endpoint: str, job_id: str, request_id: str + ) -> dict: + canned = STATE.requests.get(request_id) + if canned is None: + return JSONResponse({"error": "request not found"}, status_code=404) + return { + "status": "request_state_completed", + "events": canned.get("events", []), + "result": canned.get("result", {}), + } + + return app + + +def _register_octet_route(app: FastAPI, path_suffix: str) -> None: + """Wire up one of the DSSST1 octet-chunked routes. + + All three (fwd_bwd / fwd_no_grad / generate) share the same chunk-reassembly + logic; response shape is the only difference. + """ + + async def handler( + database: str, schema: str, endpoint: str, job_id: str, request: Request + ) -> dict: + raw = await request.body() + desc = wire.read_byte_chunk_metadata(raw) or {"total_chunks": 1, "chunk_idx": 0} + group_id = desc.get("chunk_group_id") or "single" + key = (job_id, path_suffix, group_id) + buf = STATE.chunk_buffers.setdefault(key, []) + buf.append(raw) + total = int(desc.get("total_chunks", 1)) + if len(buf) < total: + # Intermediate chunk: return empty body (no request_id yet), per + # `_post_octet_request_chunks` in the client. + return {} + + # Final chunk. Reassemble and inspect the payload shape. + try: + frame = wire.decode_byte_chunks(buf) if total > 1 else buf[0] + decoded = wire.loads(frame) + finally: + STATE.chunk_buffers.pop(key, None) + + result = _shape_octet_response(path_suffix, decoded) + req_id = _register_result(result, wire_encoded=(path_suffix != "forward-backward")) + return {"request_id": req_id} + + # FastAPI can't share a handler between routes with the same path template; + # register under the concrete suffix. + app.add_api_route( + _PREFIX + f"/{{job_id}}/{path_suffix}", + handler, + methods=["POST"], + name=f"octet_{path_suffix.replace('-', '_')}", + ) + + +# ─── Fake response shaping ──────────────────────────────────────────────── + + +def _register_result(result: dict, *, wire_encoded: bool = False) -> str: + """Register a canned response for a synthesized request_id. + + - `wire_encoded=False`: return `result` inline in the request-status JSON. + - `wire_encoded=True`: encode via DSSST1 + base64 and put it in the + request-status `result` field, per `_decode_result_payload` in the client. + """ + req_id = f"fake-req-{uuid.uuid4().hex[:8]}" + if wire_encoded: + frame = wire.dumps(result) + STATE.requests[req_id] = { + "result": { + "wire_format": wire.WIRE_FORMAT_VERSION, + "encoding": "base64", + "payload_b64": base64.b64encode(frame).decode("ascii"), + } + } + else: + STATE.requests[req_id] = {"result": result} + return req_id + + +def _shape_octet_response(path_suffix: str, decoded: Any) -> dict: + """Produce a response for one of the octet-chunked routes based on the + decoded request payload's shape. Correctness is not the point — we just + need a shape the client's response shim + verl/SkyRL adapters can parse. + """ + body = decoded if isinstance(decoded, dict) else {} + kwargs = body.get("kwargs") or {} + input_ids = kwargs.get("input_ids") + if torch.is_tensor(input_ids): + n_rows, seq_len = int(input_ids.shape[0]), int(input_ids.shape[1]) + else: + n_rows, seq_len = 1, 8 + + if path_suffix == "forward-backward": + # verl/SkyRL read `.get("loss")` / `.get("grad_norm")`; keep it JSON. + return _fake_step_result() + if path_suffix == "forward-no-grad": + return _fake_log_probs_result(n_rows=n_rows, seq_len=seq_len) + if path_suffix == "generate": + # `generate` body ships `{"prompts": [...]}` (no input_ids), so + # size the fake results off the prompts list. + prompts = body.get("prompts") or [] + return _fake_generate_result(n_prompts=max(1, len(prompts))) + return {} + + +def _fake_step_result(*, learning_rate: float | None = None) -> dict: + loss = round(random.uniform(0.1, 1.5), 4) + return { + "loss": loss, + "avg_loss": loss, + "metrics": { + "loss": loss, + "grad_norm": round(random.uniform(0.1, 2.0), 4), + "ppo_kl": round(random.uniform(0.0, 0.05), 4), + "pg_loss": loss, + "pg_clipfrac_lower": 0.0, + "kl_loss": 0.001, + "kl_coef": 0.001, + "last_lr": learning_rate or 1e-5, + }, + } + + +def _fake_log_probs_result(*, n_rows: int, seq_len: int) -> dict: + """Return a DSSST1-encodable dict with `model_outputs.logprobs` of the + right shape. `CortexTransport._shape_train_response` aliases + `model_outputs` -> `batch`, and verl's `_send_compute_log_prob` then + reads `response["batch"]["log_probs"]` after renaming `logprobs` -> + `log_probs`. + """ + logprobs = -torch.rand(n_rows, seq_len).abs() # sign-correct: log-probs ≤ 0 + entropy = torch.rand(n_rows, seq_len).abs() + return { + "model_outputs": { + "logprobs": logprobs, + "entropy": entropy, + }, + "metrics": {}, + } + + +def _fake_generate_result(*, n_prompts: int) -> dict: + # `_generate` in the client returns `{"results": [...]}`; each entry + # ships prompt / response ids + text. Keep it minimal. + results = [] + for _ in range(max(1, n_prompts)): + n_toks = random.randint(4, 12) + results.append( + { + "response_ids": [random.randint(1000, 50000) for _ in range(n_toks)], + "response_text": "fake", + "logprobs": [-abs(random.random()) for _ in range(n_toks)], + "finish_reason": "stop", + } + ) + return {"results": results} + + +def _completion_len_hint(completions: list) -> int | None: + if not completions: + return None + first = completions[0] + if isinstance(first, list): + return len(first) + if isinstance(first, str): + return max(1, len(first.split())) + return None + + +# ─── Fixture-friendly runner ────────────────────────────────────────────── + + +def serve_in_background(host: str = "127.0.0.1", port: int = 8080) -> threading.Thread: + """Start the fake GS in a daemon thread. Returns the thread so tests can + join or drop it. Uvicorn's built-in threading integration is used so we + don't need pytest-anyio. + """ + app = _make_app() + config = uvicorn.Config(app, host=host, port=port, log_level="warning", access_log=False) + server = uvicorn.Server(config) + + def _run() -> None: + server.run() + + thread = threading.Thread(target=_run, daemon=True, name=f"fake_cortex_gs:{port}") + thread.start() + # Wait for uvicorn to bind before returning. + import time + + for _ in range(200): + if server.started: + return thread + time.sleep(0.02) + raise RuntimeError("fake_cortex_gs failed to start within 4s") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8080) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s") + log.info("fake_cortex_gs listening on http://%s:%d", args.host, args.port) + uvicorn.run(_make_app(), host=args.host, port=args.port, log_level="info") + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/run_skyrl_against_mock.sh b/tests/e2e/run_skyrl_against_mock.sh new file mode 100755 index 0000000..32ca85e --- /dev/null +++ b/tests/e2e/run_skyrl_against_mock.sh @@ -0,0 +1,178 @@ +#!/usr/bin/env bash +# End-to-end validation: SkyRL gsm8k GRPO recipe driven against `fake_cortex_gs`. +# +# What this proves (or breaks): the whole client chain — SkyRL launcher → +# `integrations/arctic_rl/config.py::build_rl_config` (hardcoded `backend=local`) +# → `create_arctic_rl_client` → `ARCTIC_RL_BACKEND=cortex` env-override rewrite +# → `_CortexClientShim` → `arctic_platform.client.ArcticRLClient` → +# `CortexTransport` → HTTP → local fake Cortex GS → back — actually runs a +# real training step. No mocks anywhere in the client / transport / adapter +# code paths; only the *server* is fake. +# +# Convergence is NOT validated — the fake GS returns random losses. This is +# a plumbing gate: prove the driver → wire round-trip works and both +# integrations reach the Cortex path with zero adapter change. +# +# Usage: +# export SKYRL_ROOT=/path/to/SkyRL +# bash tests/e2e/run_skyrl_against_mock.sh +# +# Optional env: +# CACHE_ROOT — where uv/HF/tmp caches live (default: $HOME/.cache/cortex-e2e). +# Point at a scratch disk if $HOME is size-constrained; uv +# will pull ~15 GiB of wheels on first resolution. +# STEPS — training steps to run (default 3). +# GPUS — GPUs to split policy/generator across (default 4). + +set -euo pipefail + +STEPS=${STEPS:-3} +GPUS=${GPUS:-4} + +# --------------------------------------------------------------------------- +# Env caches — uv will otherwise fill $HOME with vllm + torch wheels (~15 GiB). +# Override CACHE_ROOT to a scratch disk (e.g. /data-fast/$USER/cortex-e2e) if +# $HOME is size-constrained. +# --------------------------------------------------------------------------- +CACHE_ROOT=${CACHE_ROOT:-${HOME}/.cache/cortex-e2e} +export UV_CACHE_DIR=${CACHE_ROOT}/uv-cache +export UV_PYTHON_INSTALL_DIR=${CACHE_ROOT}/uv-python +export XDG_CACHE_HOME=${CACHE_ROOT}/cache +export HF_HOME=${HF_HOME:-${CACHE_ROOT}/hf-cache} +export PIP_CACHE_DIR=${CACHE_ROOT}/pip-cache +export TMPDIR=${TMPDIR:-${CACHE_ROOT}/tmp} +mkdir -p "$UV_CACHE_DIR" "$UV_PYTHON_INSTALL_DIR" "$XDG_CACHE_HOME" "$HF_HOME" \ + "$PIP_CACHE_DIR" "$TMPDIR" + +# --------------------------------------------------------------------------- +# Paths. +# --------------------------------------------------------------------------- +E2E_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ARCTIC_PLATFORM_ROOT="$(cd "$E2E_DIR/../.." && pwd)" +DATA_DIR=${DATA_DIR:-${CACHE_ROOT}/gsm8k} + +# SkyRL checkout — required, no default. Point this at your local clone of +# https://github.com/NovaSky-AI/SkyRL (any branch that has +# integrations/arctic_rl/ works). +if [[ -z "${SKYRL_ROOT:-}" ]]; then + cat <&2 +ERR: SKYRL_ROOT is not set. Export it to your local SkyRL checkout, e.g.: + export SKYRL_ROOT=/path/to/your/SkyRL +EOF + exit 2 +fi +if [[ ! -d "$SKYRL_ROOT" ]]; then + echo "ERR: SkyRL checkout not found at SKYRL_ROOT=$SKYRL_ROOT" >&2 + exit 2 +fi + +# --------------------------------------------------------------------------- +# Prep GSM8K parquets under $DATA_DIR if missing. +# --------------------------------------------------------------------------- +if [[ ! -f "${DATA_DIR}/train.parquet" || ! -f "${DATA_DIR}/validation.parquet" ]]; then + echo "[e2e] GSM8K parquets missing under $DATA_DIR; using SkyRL prep script..." + cd "$SKYRL_ROOT" + uv run --isolated examples/train/gsm8k/gsm8k_dataset.py --output_dir "$DATA_DIR" +fi + +# --------------------------------------------------------------------------- +# Start the fake Cortex GS on a random port in the background. It runs from +# the local Arctic-Platform checkout — SkyRL's uv env doesn't need +# arctic-platform installed for the mock's sake (we're driving verl/SkyRL, +# not the mock, through the uv env). +# --------------------------------------------------------------------------- +MOCK_LOG="${CACHE_ROOT}/fake_cortex_gs.log" + +MOCK_PORT=$(python -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1",0)); print(s.getsockname()[1]); s.close()') +echo "[e2e] starting fake_cortex_gs on port $MOCK_PORT (log: $MOCK_LOG)" + +# Use the arctic_platform checkout's own env (dev conda) to run the mock — +# it only needs fastapi + uvicorn + safetensors + torch which are already +# there. Don't share the SkyRL uv env with this. +( + cd "$ARCTIC_PLATFORM_ROOT" + # `tests.e2e` is intentionally not a package; drive the file directly. + exec python "$E2E_DIR/fake_cortex_gs.py" --port "$MOCK_PORT" --host 127.0.0.1 +) >"$MOCK_LOG" 2>&1 & +MOCK_PID=$! +cleanup() { + if kill -0 "$MOCK_PID" 2>/dev/null; then + echo "[e2e] stopping fake_cortex_gs (pid=$MOCK_PID)" + kill "$MOCK_PID" 2>/dev/null || true + fi +} +trap cleanup EXIT + +# Wait for it to bind (openapi is served by FastAPI as soon as uvicorn is up). +for _ in $(seq 1 60); do + if curl -sf "http://127.0.0.1:${MOCK_PORT}/openapi.json" >/dev/null; then + echo "[e2e] fake_cortex_gs is up" + break + fi + sleep 0.2 +done +if ! curl -sf "http://127.0.0.1:${MOCK_PORT}/openapi.json" >/dev/null; then + echo "ERR: fake_cortex_gs failed to start within 12s" >&2 + tail -50 "$MOCK_LOG" >&2 || true + exit 3 +fi + +# --------------------------------------------------------------------------- +# Flip the launcher to Cortex via the env-var override. +# --------------------------------------------------------------------------- +export ARCTIC_RL_BACKEND=cortex +export CORTEX_BASE_URL="http://127.0.0.1:${MOCK_PORT}" +export CORTEX_DATABASE=e2e_db +export CORTEX_SCHEMA=e2e_sch +export CORTEX_ENDPOINT=cortex-training +export CORTEX_MAX_SEQ_LEN=1536 + +# Point uv at a local wheel of *this* Arctic-Platform checkout so `--with +# arctic-platform` in the SkyRL recipe resolves to our Cortex changes (0.1.3.dev0 +# beats PyPI's 0.1.2). We do this via env vars — no edits to the recipe. +WHEEL_DIR=${WHEEL_DIR:-${CACHE_ROOT}/wheels} +if ! ls "${WHEEL_DIR}"/arctic_platform-*.whl >/dev/null 2>&1; then + echo "[e2e] building local arctic-platform wheel into $WHEEL_DIR" + ( cd "$ARCTIC_PLATFORM_ROOT" && uv build --wheel -o "$WHEEL_DIR" ) >/dev/null +fi +export UV_FIND_LINKS="file://${WHEEL_DIR}" +export UV_PRERELEASE=allow +export UV_INDEX_STRATEGY=unsafe-best-match + +# --------------------------------------------------------------------------- +# Launch SkyRL. Its bundled recipe already handles isolated uv resolution +# for arctic-inference[vllm] + flash-attn. +# --------------------------------------------------------------------------- +cd "$SKYRL_ROOT" + +RECIPE=integrations/arctic_rl/examples/run_gsm8k_grpo_4gpu.sh +if [[ ! -f "$RECIPE" ]]; then + echo "ERR: expected recipe not found at $SKYRL_ROOT/$RECIPE" >&2 + exit 4 +fi + +RUN_LOG="${CACHE_ROOT}/skyrl_run.log" +echo "[e2e] launching SkyRL recipe (log: $RUN_LOG)" +echo "[e2e] steps=${STEPS} gpus=${GPUS} mock=http://127.0.0.1:${MOCK_PORT}" + +set +e +DATA_DIR="$DATA_DIR" \ + bash "$RECIPE" \ + trainer.epochs=1 \ + trainer.total_training_steps="$STEPS" \ + trainer.placement.policy_num_gpus_per_node=$((GPUS / 2)) \ + generator.inference_engine.num_engines=$((GPUS / 2)) \ + trainer.policy_mini_batch_size=4 \ + trainer.train_batch_size=64 \ + generator.n_samples_per_prompt=4 \ + trainer.eval_before_train=false \ + trainer.eval_interval=999999 \ + 2>&1 | tee "$RUN_LOG" +STATUS=$? +set -e + +echo "[e2e] SkyRL exit=$STATUS" +echo "[e2e] mock log tail:" +tail -30 "$MOCK_LOG" || true + +exit "$STATUS" diff --git a/tests/e2e/test_cortex_transport_smoke.py b/tests/e2e/test_cortex_transport_smoke.py new file mode 100644 index 0000000..736fc0b --- /dev/null +++ b/tests/e2e/test_cortex_transport_smoke.py @@ -0,0 +1,173 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end plumbing smoke: `CortexTransport` against `fake_cortex_gs`. + +Validates that: + +- `CortexTransport.initialize()` calls CreateJob + polls GetJob and produces + `JobHandles` covering training/sampling/log_prob sub-jobs. +- Every `call(Request)` op the transport dispatches (fwd_bwd, fwd_no_grad, + step, save, generate, sync_weights, reset_prefix_cache, log_probs, all the + colocation lifecycle no-ops) round-trips through the fake GS without + raising. +- Response shapes returned by `_shape_train_response` match what verl and + SkyRL then read from them: `.get("grad_norm")`, `.get("avg_loss")`, + `response["batch"]["log_probs"]` after the `model_outputs -> batch` alias. + +Nothing here validates *training correctness* — losses are random. It's a +plumbing gate: prove the client's REST + wire path works end-to-end without +mocks at any layer. +""" + +from __future__ import annotations + +import socket +import sys +from pathlib import Path + +import pytest +import torch + +# tests/ isn't a package in this repo, so relative-style imports don't work +# out of the box; put the current directory on the path so `import +# fake_cortex_gs` resolves the sibling module. +sys.path.insert(0, str(Path(__file__).parent)) + +from arctic_platform.client.config import ArcticRLClientConfig as UnifiedClientConfig +from arctic_platform.client.transport import Request +from arctic_platform.client.transports.cortex import CortexTransport + + +def _free_port() -> int: + """Grab an unused TCP port so parallel runs don't collide.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +@pytest.fixture(scope="module") +def fake_gs_url() -> str: + from fake_cortex_gs import serve_in_background + + port = _free_port() + serve_in_background(host="127.0.0.1", port=port) + return f"http://127.0.0.1:{port}" + + +@pytest.fixture +def transport(fake_gs_url: str) -> CortexTransport: + cfg = UnifiedClientConfig( + backend="cortex", + model_name="Qwen/Qwen3-0.6B", + training_gpus=1, + sampling_gpus=1, + log_prob_gpus=0, + cortex_base_url=fake_gs_url, + cortex_database="fake_db", + cortex_schema="fake_sch", + cortex_endpoint="cortex-training", + max_seq_len=512, + ) + tp = CortexTransport(cfg) + tp.initialize() + return tp + + +class TestInitialize: + def test_creates_job_and_captures_sub_jobs(self, transport: CortexTransport): + assert transport.job_id is not None + # Training + sampling sub-jobs should be recognized. Log-prob was + # 0 GPUs so it's the fallback synthesized handle. + assert "training" in transport.sub_jobs + assert "sampling" in transport.sub_jobs + + +class TestCoreOps: + def test_fwd_bwd_returns_loss_and_grad_norm(self, transport: CortexTransport): + # Build a verl-GRPO-shaped body: {batch: {...}, meta: {...}, processing: {...}} + batch = { + "input_ids": torch.zeros(2, 8, dtype=torch.long), + "attention_mask": torch.ones(2, 8, dtype=torch.long), + } + body = {"batch": batch, "meta": {"rollout_n": 1}, "processing": {"loss_fn": "verl_grpo"}} + resp = transport.call(Request(op="fwd-bwd", body=body)) + # verl reads `.get("avg_loss")`; SkyRL reads `.get("grad_norm")`. + assert "avg_loss" in resp + assert resp["metrics"].get("grad_norm") is not None + + def test_fwd_no_grad_returns_batch_log_probs(self, transport: CortexTransport): + batch = {"input_ids": torch.zeros(2, 16, dtype=torch.long)} + body = {"batch": batch, "meta": {}, "reference_model": True} + resp = transport.call(Request(op="fwd-no-grad", body=body)) + # _shape_train_response aliases model_outputs -> batch; verl then + # renames `logprobs` -> `log_probs`. Assert the alias took. + assert "batch" in resp + assert "logprobs" in resp["batch"] + lp = resp["batch"]["logprobs"] + assert torch.is_tensor(lp) + assert lp.shape == (2, 16) + + def test_step_returns_metrics(self, transport: CortexTransport): + resp = transport.call(Request(op="step", body={"learning_rate": 1e-5})) + assert "metrics" in resp + assert resp["metrics"].get("grad_norm") is not None + + def test_save_checkpoint(self, transport: CortexTransport): + resp = transport.call(Request(op="save-checkpoint", body={})) + assert isinstance(resp, dict) + assert "path" in resp + + def test_generate_returns_results_list(self, transport: CortexTransport): + body = { + "prompts": ["hello world", "second prompt"], + "sampling_params": {"temperature": 0.7, "max_tokens": 8}, + "routing_key": None, + } + resp = transport.call(Request(op="generate", body=body)) + assert isinstance(resp, dict) + assert "results" in resp + assert len(resp["results"]) >= 1 + + def test_log_probs(self, transport: CortexTransport): + body = { + "prompts": ["hello world"], + "completions": [[1, 2, 3, 4, 5]], + "top_k": None, + } + resp = transport.call(Request(op="log-probs", body=body)) + assert isinstance(resp, dict) + assert "model_outputs" in resp + assert "logprobs" in resp["model_outputs"] + + def test_sync_weights(self, transport: CortexTransport): + resp = transport.call(Request(op="sync-weights", body={"cuda_ipc": False, "low_memory": False})) + assert isinstance(resp, dict) + + def test_reset_prefix_cache(self, transport: CortexTransport): + resp = transport.call(Request(op="reset-prefix-cache", body={"drain": False, "timeout_s": 5.0})) + assert isinstance(resp, dict) + + +class TestColocationNoops: + """Every colocation lifecycle op is a no-op on Cortex (returns {}). + SkyRL calls these unconditionally when `colocate=True`; verl calls + `wake_inference` / `sleep_inference`. Regression here means SkyRL's + colocated recipe would raise on the very first weight sync.""" + + @pytest.mark.parametrize( + "op", + [ + "wake-training", + "sleep-training", + "wake-inference", + "sleep-inference", + "wake-log-prob", + "sleep-log-prob", + "empty-training-cache", + "weight-norm", + "save-weights", + ], + ) + def test_op_is_noop(self, transport: CortexTransport, op: str): + resp = transport.call(Request(op=op, body={})) + assert resp == {}