Repository navigation
feat: Arctic RL training backend integration - #1837
Merged
SumanthRH merged 20 commits intoJun 29, 2026
Merged
Conversation
Adds an `arctic_rl` integration under `integrations/arctic_rl/` that lets
SkyRL recipes train against Snowflake's Arctic RL stack (Arctic-Platform
+ Arctic-Inference). It mirrors the existing `harbor` / `dapo` integration
shape — each integration ships its own entrypoint, config, generator,
trainer, BIRD-SQL environment vendored in-tree, and example launchers.
Core SkyRL changes are intentionally minimal:
* `skyrl/train/entrypoints/main_base.py`: one ~15-line dispatch hook that
inspects `trainer.override_entrypoint=<module>` on argv BEFORE strict
config parsing and forwards to the integration's `main()`. This lets
integrations add config fields without subclassing `SkyRLTrainConfig`
and without touching the recipe scripts (`uv run -m
skyrl.train.entrypoints.main_base trainer.override_entrypoint=...`).
* `pyproject.toml`: adds an additive `[arctic-rl]` extra
(`skyrl[arctic-rl]`) that pulls in `arctic-platform`,
`arctic-inference[vllm]`, and `liger-kernel`. No existing pins are
touched.
Integration contents (under `integrations/arctic_rl/`):
* `entrypoint.py` — boots an `ArcticRLClient` against arctic-platform
and forwards weight-sync, rollouts, and metrics.
* `config.py` — `ArcticRLTrainerConfig` exposing the colocate /
ZoRRo / FCA / spec-dec / CUDA-IPC-weight-sync knobs.
* `trainer.py`, `generator.py` — thin adapters from SkyRL's PPO loop
to Arctic-Platform's RL client.
* `envs/bird*.py` — vendored BIRD-SQL environment so the recipe is
self-contained.
* `examples/run_bird_grpo_{8b,32b,smoke}*.sh` and
`run_gsm8k_grpo_4gpu.sh` — recipes covering smoke / 8B / 32B BIRD
+ a GSM8K sanity check. Env vars in the launchers are trimmed to
only what's recipe-required.
Validated end-to-end on a 4-node H200 cluster on the Qwen3 family.
Co-authored-by: Cursor <cursoragent@cursor.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces the Arctic RL integration for SkyRL, enabling CPU-only client orchestration while delegating all GPU-intensive operations (training, generation, log-probs, and weight sync) to an Arctic RL server. The integration includes configuration translation, an entrypoint, a custom generator and trainer, a BIRD SQL environment with its reward function, and preprocessing scripts. The reviewer feedback highlights several critical issues: potential database connection leaks in bird_reward.py and preprocess_bird.py due to missing try...finally blocks around SQLite connections; performance overhead in trainer.py from repeatedly instantiating ThreadPoolExecutor in _run; a potential crash in trainer.py if log-probs are returned as NumPy arrays; and a configuration translation bug in config.py where arctic_inference_config does not fall back to an empty dictionary when Arctic Inference is enabled.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+42
to
+60
| def _run(coro): | ||
| """Run an awaitable from sync context. | ||
|
|
||
| SkyRL's WorkerDispatch protocol is synchronous but the new | ||
| ``arctic_platform.rl`` client methods are async, so the dispatch | ||
| blocks on the coroutine here. SkyRL's ``train()`` runs inside an | ||
| asyncio loop (``asyncio.run(trainer.train())``), so ``asyncio.run`` | ||
| from inside it raises ``RuntimeError`` — we then fall back to a | ||
| fresh loop in a worker thread (the ``arctic_platform.rl`` ray client | ||
| only does ``ray.get`` under the hood, which works across loops). | ||
| """ | ||
| try: | ||
| return asyncio.run(coro) | ||
| except RuntimeError: | ||
| import concurrent.futures | ||
| with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: | ||
| return pool.submit( | ||
| lambda c=coro: asyncio.new_event_loop().run_until_complete(c) | ||
| ).result() |
Contributor
There was a problem hiding this comment.
Creating a new ThreadPoolExecutor on every invocation of _run introduces significant overhead due to repeated thread creation and destruction. Since _run is called frequently during training (e.g., for fwd_no_grad, fwd_bwd, step), this can become a performance bottleneck. Consider using a persistent module-level executor to run these coroutines when falling back to a background thread.
_SHARED_THREAD_POOL = None
def _run(coro):
"""Run an awaitable from sync context.
SkyRL's WorkerDispatch protocol is synchronous but the new
``arctic_platform.rl`` client methods are async, so the dispatch
blocks on the coroutine here. SkyRL's ``train()`` runs inside an
asyncio loop (``asyncio.run(trainer.train())``), so ``asyncio.run``
from inside it raises ``RuntimeError`` — we then fall back to a
fresh loop in a worker thread (the ``arctic_platform.rl`` ray client
only does ``ray.get`` under the hood, which works across loops).
"""
try:
return asyncio.run(coro)
except RuntimeError:
global _SHARED_THREAD_POOL
if _SHARED_THREAD_POOL is None:
import concurrent.futures
_SHARED_THREAD_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=1)
def _run_in_new_loop(c):
loop = asyncio.new_event_loop()
try:
return loop.run_until_complete(c)
finally:
loop.close()
return _SHARED_THREAD_POOL.submit(_run_in_new_loop, coro).result()
Comment on lines
+159
to
+181
| def _execute_sql(db_path: str, sql: str, timeout: float = SQL_TIMEOUT) -> frozenset | Exception: | ||
| """Execute SQL against a SQLite database and return result as frozenset of row tuples.""" | ||
| try: | ||
| conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True, timeout=5) | ||
| deadline = __import__("time").monotonic() + timeout | ||
|
|
||
| def _check_cancel(): | ||
| if __import__("time").monotonic() > deadline: | ||
| return 1 | ||
| return 0 | ||
|
|
||
| conn.set_progress_handler(_check_cancel, 1000) | ||
| cursor = conn.cursor() | ||
| cursor.execute(sql) | ||
| rows = cursor.fetchall() | ||
| conn.close() | ||
| return frozenset(rows) | ||
| except sqlite3.OperationalError as e: | ||
| if "interrupt" in str(e).lower(): | ||
| return TimeoutError(f"SQL execution exceeded {timeout}s") | ||
| return e | ||
| except Exception as e: | ||
| return e |
Contributor
There was a problem hiding this comment.
If an exception occurs during SQL execution (which is common for invalid predicted SQL in RL), conn.close() is skipped, leading to database connection and file descriptor leaks. Wrap the connection lifecycle in a try...finally block to guarantee that conn.close() is always executed.
def _execute_sql(db_path: str, sql: str, timeout: float = SQL_TIMEOUT) -> frozenset | Exception:
"""Execute SQL against a SQLite database and return result as frozenset of row tuples."""
conn = None
try:
conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True, timeout=5)
deadline = __import__("time").monotonic() + timeout
def _check_cancel():
if __import__("time").monotonic() > deadline:
return 1
return 0
conn.set_progress_handler(_check_cancel, 1000)
cursor = conn.cursor()
cursor.execute(sql)
rows = cursor.fetchall()
return frozenset(rows)
except sqlite3.OperationalError as e:
if "interrupt" in str(e).lower():
return TimeoutError(f"SQL execution exceeded {timeout}s")
return e
except Exception as e:
return e
finally:
if conn is not None:
conn.close()
Comment on lines
+482
to
+492
| # Deep-convert OmegaConf containers to plain Python so that vLLM's | ||
| # ``AsyncEngineArgs.__post_init__`` recognizes nested overrides | ||
| # (``compilation_config``, ``speculative_config``, ...). The | ||
| # ``OmegaConf.create(x)`` round-trip normalizes both plain ``dict`` | ||
| # and ``DictConfig`` inputs, matching the ``OmegaConf.to_container`` | ||
| # idiom used in ``arctic-verl/verl/workers/remote_client/arctic_rl.py``. | ||
| arctic_inference_config=OmegaConf.to_container( | ||
| OmegaConf.create(arl.arctic_inference_config), resolve=True, | ||
| ) | ||
| if arl.use_arctic_inference and arl.arctic_inference_config is not None | ||
| else None, |
Contributor
There was a problem hiding this comment.
The translation logic currently passes None for arctic_inference_config if arl.arctic_inference_config is None, even when arl.use_arctic_inference is True. However, both the docstring and the README state that an empty dictionary {} should be passed in this case to enable Arctic Inference with default settings. Update the condition to fall back to {} when use_arctic_inference is enabled.
Suggested change
| # Deep-convert OmegaConf containers to plain Python so that vLLM's | |
| # ``AsyncEngineArgs.__post_init__`` recognizes nested overrides | |
| # (``compilation_config``, ``speculative_config``, ...). The | |
| # ``OmegaConf.create(x)`` round-trip normalizes both plain ``dict`` | |
| # and ``DictConfig`` inputs, matching the ``OmegaConf.to_container`` | |
| # idiom used in ``arctic-verl/verl/workers/remote_client/arctic_rl.py``. | |
| arctic_inference_config=OmegaConf.to_container( | |
| OmegaConf.create(arl.arctic_inference_config), resolve=True, | |
| ) | |
| if arl.use_arctic_inference and arl.arctic_inference_config is not None | |
| else None, | |
| arctic_inference_config=OmegaConf.to_container( | |
| OmegaConf.create(arl.arctic_inference_config or {}), resolve=True, | |
| ) | |
| if arl.use_arctic_inference | |
| else None, |
Comment on lines
+239
to
+307
| conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) | ||
| cursor = conn.cursor() | ||
|
|
||
| cursor.execute("SELECT name, sql FROM sqlite_master WHERE type='table' ORDER BY name") | ||
| tables = cursor.fetchall() | ||
|
|
||
| ddl_parts = [] | ||
| rows_to_fetch = max(num_examples, sample_rows, 1) | ||
|
|
||
| for table_name, create_sql in tables: | ||
| if not create_sql: | ||
| continue | ||
|
|
||
| col_examples: dict = {} | ||
| col_names: list = [] | ||
| rows_for_table: list = [] | ||
| try: | ||
| cursor.execute(f'SELECT * FROM "{table_name}" LIMIT {rows_to_fetch}') | ||
| rows_for_table = cursor.fetchall() | ||
| col_names = [desc[0] for desc in cursor.description] | ||
| if num_examples > 0: | ||
| for col_idx, col_name in enumerate(col_names): | ||
| vals = [row[col_idx] for row in rows_for_table[:num_examples] if row[col_idx] is not None] | ||
| if vals: | ||
| col_examples[col_name] = vals | ||
| except sqlite3.Error: | ||
| pass | ||
|
|
||
| # Top-of-block header: FK summary + sample rows. | ||
| header_lines: list = [] | ||
| if include_fk_summary: | ||
| fks = _extract_fks(create_sql) | ||
| if fks: | ||
| header_lines.append(f"-- Foreign keys: {'; '.join(fks)}") | ||
| if sample_rows > 0 and rows_for_table and col_names: | ||
| header_lines.append(f'-- Sample rows from "{table_name}":') | ||
| header_lines.append("-- | " + " | ".join(col_names) + " |") | ||
| for row in rows_for_table[:sample_rows]: | ||
| header_lines.append("-- | " + " | ".join(_short_value(v) for v in row) + " |") | ||
|
|
||
| enriched_lines: list = list(header_lines) | ||
| for line in create_sql.split("\n"): | ||
| stripped = line.strip().rstrip(",") | ||
| if stripped and not stripped.startswith("CREATE") and not stripped.startswith(")"): | ||
| col_token = stripped.split()[0].strip('"').strip("'").strip("`") | ||
| annots = [] | ||
| if col_token in col_examples: | ||
| annots.append(f"example: {col_examples[col_token]!r}") | ||
| desc_key = (table_name.lower(), col_token.lower()) | ||
| if desc_key in descriptions: | ||
| d = descriptions[desc_key] | ||
| name = d.get("name", "") | ||
| if name and name.lower() != col_token.lower(): | ||
| annots.append(f"name: {_flatten_inline(name)}") | ||
| if d["desc"]: | ||
| annots.append(f"desc: {_flatten_inline(d['desc'])}") | ||
| if d["val_desc"]: | ||
| annots.append(f"values: {_flatten_inline(d['val_desc'])}") | ||
| if annots: | ||
| comment = " -- " + " | ".join(annots) | ||
| if line.rstrip().endswith(","): | ||
| line = line.rstrip()[:-1] + comment + "," | ||
| else: | ||
| line = line.rstrip() + comment | ||
| enriched_lines.append(line) | ||
| ddl_parts.append("\n".join(enriched_lines)) | ||
|
|
||
| conn.close() | ||
| return "\n\n".join(ddl_parts) |
Contributor
There was a problem hiding this comment.
If an exception occurs during schema extraction (e.g., if sqlite_master query fails or other unexpected errors occur), conn.close() is bypassed, leaking the SQLite connection. Wrap the connection in a try...finally block to ensure it is always closed.
def get_schema_ddl(
db_path: str,
num_examples: int = 5,
descriptions: dict | None = None,
sample_rows: int = 3,
include_fk_summary: bool = True,
) -> str:
"""Read schema DDL from a SQLite database, enriched with sample values.
Augmentations applied (all reuse the single ``SELECT ... LIMIT`` per table —
no extra DB scans):
- Sample-rows block per table (markdown-style pipe table) using the same
rows fetched for per-column examples.
- Foreign-key summary line at the top of each table block.
- Inline ``-- example: ... | name: ... | desc: ... | values: ...`` comment
on each column line. ``name``, ``desc``, and ``values`` come from
``descriptions``; descriptions are NOT length-truncated (whitespace is
collapsed so they stay one line).
"""
descriptions = descriptions or {}
conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
try:
cursor = conn.cursor()
cursor.execute("SELECT name, sql FROM sqlite_master WHERE type='table' ORDER BY name")
tables = cursor.fetchall()
ddl_parts = []
rows_to_fetch = max(num_examples, sample_rows, 1)
for table_name, create_sql in tables:
if not create_sql:
continue
col_examples: dict = {}
col_names: list = []
rows_for_table: list = []
try:
cursor.execute(f'SELECT * FROM "{table_name}" LIMIT {rows_to_fetch}')
rows_for_table = cursor.fetchall()
col_names = [desc[0] for desc in cursor.description]
if num_examples > 0:
for col_idx, col_name in enumerate(col_names):
vals = [row[col_idx] for row in rows_for_table[:num_examples] if row[col_idx] is not None]
if vals:
col_examples[col_name] = vals
except sqlite3.Error:
pass
# Top-of-block header: FK summary + sample rows.
header_lines: list = []
if include_fk_summary:
fks = _extract_fks(create_sql)
if fks:
header_lines.append(f"-- Foreign keys: {'; '.join(fks)}")
if sample_rows > 0 and rows_for_table and col_names:
header_lines.append(f'-- Sample rows from "{table_name}":')
header_lines.append("-- | " + " | ".join(col_names) + " |")
for row in rows_for_table[:sample_rows]:
header_lines.append("-- | " + " | ".join(_short_value(v) for v in row) + " |")
enriched_lines: list = list(header_lines)
for line in create_sql.split("\n"):
stripped = line.strip().rstrip(",")
if stripped and not stripped.startswith("CREATE") and not stripped.startswith(")"):
col_token = stripped.split()[0].strip('"').strip("'").strip("`")
annots = []
if col_token in col_examples:
annots.append(f"example: {col_examples[col_token]!r}")
desc_key = (table_name.lower(), col_token.lower())
if desc_key in descriptions:
d = descriptions[desc_key]
name = d.get("name", "")
if name and name.lower() != col_token.lower():
annots.append(f"name: {_flatten_inline(name)}")
if d["desc"]:
annots.append(f"desc: {_flatten_inline(d['desc'])}")
if d["val_desc"]:
annots.append(f"values: {_flatten_inline(d['val_desc'])}")
if annots:
comment = " -- " + " | ".join(annots)
if line.rstrip().endswith(","):
line = line.rstrip()[:-1] + comment + ","
else:
line = line.rstrip() + comment
enriched_lines.append(line)
ddl_parts.append("\n".join(enriched_lines))
finally:
conn.close()
return "\n\n".join(ddl_parts)
Comment on lines
+587
to
+589
| if isinstance(lp, list): | ||
| lp = torch.tensor(lp) | ||
| return lp |
Contributor
There was a problem hiding this comment.
The check isinstance(lp, list) only converts lists to PyTorch tensors. However, depending on the serialization and transport protocol of the arctic_platform client (especially over Ray IPC), lp could be returned as a NumPy ndarray. If it is a NumPy array, the subsequent .to() call in forward_backward will crash with an AttributeError. Ensure any non-tensor type is converted to a PyTorch tensor.
Suggested change
| if isinstance(lp, list): | |
| lp = torch.tensor(lp) | |
| return lp | |
| if not isinstance(lp, torch.Tensor): | |
| lp = torch.tensor(lp) |
Three small fixes uncovered while running the public PR's recipes end-to-end against arctic-platform main + arctic-inference main: 1. entrypoint.py — strip `trainer.override_entrypoint=` from argv before the Arctic config's strict Pydantic parse. main_base.py's peek-ahead consumed it semantically but left it in sys.argv, so the integration's re-parse rejected it as an unknown field. 2. envs/__init__.py — also register the env id `bird_sql`. Launcher recipes pass `environment.env_class=bird`, but preprocessed BIRD parquets carry `env_class="bird_sql"` per-row (verl PR NovaSky-AI#6 schema). Registering both names keeps either source-of-truth working without a parquet rewrite. 3. generator.py — trigger arctic_rl.envs import inside _score_one so the bird/bird_sql registrations exist in the ProcessPoolExecutor child workers (spawn fresh, so only skyrl_gym's own envs are auto-loaded). With these three, the 8B BIRD recipe now runs cleanly end-to-end on the NovaSky main branch + arctic-platform main + arctic-inference main. Co-authored-by: Cursor <cursoragent@cursor.com>
The two BIRD recipes warned readers that the nested compilation_config override was being silently dropped by the arctic_rl -> arctic_platform -> AsyncEngineArgs hop. That was true before the OmegaConf.to_container fix in integrations/arctic_rl/config.py landed; with that fix the override arrives in vLLM intact (`cudagraph_mode: PIECEWISE`, `pass_config.fuse_allreduce_rms: False` -- confirmed in worker init logs). Reword the comment to explain why we still pin `optimization_level: 1` alongside the explicit compilation_config: it's a belt-and-suspenders on top of an already-correct override, not a workaround for broken plumbing. Co-authored-by: Cursor <cursoragent@cursor.com>
… schema
The 32B BIRD recipe needs three raw vLLM ``AsyncEngineArgs`` to land
intact on the inference engine:
* ``compilation_config.pass_config.fuse_allreduce_rms = False``
(dodges a FlashInfer-workspace assert when >1 vLLM replica per node
races for the same per-process IPC port)
* ``compilation_config.cudagraph_mode = PIECEWISE``
* ``speculative_config`` and ``forest_cascade_attn_configs``
Arctic-Platform's ``ArcticRLClientConfig`` already has two separate
config channels for exactly this distinction:
* ``vllm_config: dict`` -- raw passthrough, forwarded verbatim into
vLLM (keys not on ``ModelConfig`` land in ``extra_engine_kwargs``)
* ``arctic_inference_config: dict`` -- high-level schema, filtered by
``parse_arctic_inference_rollout`` which only forwards the keys it
recognises (``use_fca`` / ``spec_model``)
We were sending the raw vLLM keys through the second channel, which on
public main of arctic-platform silently drops everything but the two
recognised keys. The fix is just to use the right channel:
* Add ``ArcticRLTrainerConfig.vllm_config: Optional[dict]`` -- mirrors
``ArcticRLClientConfig.vllm_config`` name and semantics. User keys
are shallow-merged on top of the integration's built-in defaults
(TP, ``max_num_seqs``, etc.) in ``build_rl_config``.
* Update both BIRD recipes to pass perf knobs via
``trainer.arctic_rl.vllm_config={...}``.
With this single change the recipes run against unmodified public main
of arctic-platform (10bad320) + arctic-inference (8cbac15b). No
patches, no feature branches.
Co-authored-by: Cursor <cursoragent@cursor.com>
Without this, ArcticRLTrainerConfig.zero_stage defaults to 0, which DeepSpeed rejects with "unsupported mix of model_dtype=torch.bfloat16 and grad_accum_dtype=torch.float32" because the bf16 + fp32-grad combination is only supported with ZeRO stage >= 1. Tested on locked public-main HEADs of arctic-platform and arctic-inference: step 1 = 56s, sync_weights = 0.04s, rewards rising 0.229 -> 0.337, no errors. Co-authored-by: Cursor <cursoragent@cursor.com>
Restructures the README around the three things a first-time user needs:
1. Clean install (uv sync --extra arctic-rl, transformers downgrade,
ray bringup).
2. Enabling Arctic RL on any existing recipe via the single
trainer.override_entrypoint flag, with a documented table of
trainer.arctic_rl.* knobs and the vllm_config escape hatch.
3. Running the bundled GSM8K recipe end-to-end on 4 GPUs, including
measured wall-clock and the first four reward-signal data points
from a clean public-mains run.
Defers BIRD-SQL prep instructions to a follow-up commit (one-command
HuggingFace-hosted parquet bundle in progress). Removes stale 1.5B
validated-results table and FSDP comparison that aren't part of this
integration's story.
Co-authored-by: Cursor <cursoragent@cursor.com>
The integration pins ``comm_protocol="ray"`` in ``build_rl_config`` so the SkyRL driver and the Arctic server actors (``ArcticRLRayServerState`` + ``DeepSpeedWorker`` + ``InferenceWorker``) live in the same Ray cluster and communicate via in-cluster actor calls. arctic-platform's HTTP transport is for remote dss-platform deployments and is not what this integration uses. Also corrects the "How it works" section's claim that the entrypoint spawns the server as a subprocess — it actually pre-initializes Ray actors on the driver and hands the actor handles to the trainer task via reconnect_config + server_state. Co-authored-by: Cursor <cursoragent@cursor.com>
- Drop unused imports (dataclasses.field in config.py).
- Drop unused locals (train_batch_size_global in build_rl_config,
arl in ArcticPPOTrainer._build_arctic_train_meta).
- Sort import blocks (I001) in __init__.py and config.py.
- Apply black 24.10.0 reformatting across the integration.
No behavior change; matches the public SkyRL CI's pre-commit config
(ruff v0.11.9 + black 24.10.0).
Co-authored-by: Cursor <cursoragent@cursor.com>
- Open with the ZoRRo numbers (3.5x training, 6x actor, 5x log-prob,
1.7x rollout, 3.2x context) from the public Snowflake AI Research
blog post and frame the integration as the one-flag way to get
those gains in SkyRL.
- Collapse Quick Start to a 5-line GSM8K smoke test that produces
a reward signal in ~3 min from a cold clone.
- Tighten the trainer.arctic_rl.* knob table to the keys users
actually flip (use_zorro, use_arctic_inference, zero_stage,
colocate, ...).
- Keep the vllm_config escape hatch as the way to pass raw
speculative_config / forest_cascade_attn_configs.
- Add a "What you get vs FSDP2" table linking back to the blog for
the technical deep dive.
- Trim "How it works" prose and architecture diagram to the
essentials: one Ray cluster, in-cluster actor calls, no HTTP, no
subprocess.
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
… hatch Adds a 5-step BIRD-SQL section walking through clone -> ray bringup -> FA3 wheel -> raw BIRD download + preprocess -> launch on 4 nodes / 32 GPUs. Pins --max_tokens 32768 to match the recipe's trainer.max_prompt_length=32768; adds notes on the shared-FS / db_path requirement and the 8B smoke variant. Reframes the vllm_config section as an escape hatch with an honest list of the three knobs that appear there today and why each should eventually move into a typed trainer.arctic_rl.* field (forest_cascade_attn_configs, speculative_config) or be auto-disabled inside arctic-inference (fuse_allreduce_rms on multi-replica FCA). Co-authored-by: Cursor <cursoragent@cursor.com>
…ext2SQL' Co-authored-by: Cursor <cursoragent@cursor.com>
End-user recipes no longer need raw vllm_config blocks to opt into
Forest Cascade Attention or Arctic speculative decoding.
ArcticRLTrainerConfig:
- speculative_model: HF id / local path of an Arctic draft head.
Setting this auto-injects speculative_config={method: arctic,
model: <path>, num_speculative_tokens: <N>} into the rollout vLLM
engine.
- num_speculative_tokens: companion knob (default 3).
build_rl_config() auto-injection (each guarded by topology so we
don't change the single-replica path):
- use_arctic_inference=True -> forest_cascade_attn_configs="{}"
- use_arctic_inference + num_engines>1
-> compilation_config.pass_config.fuse_allreduce_rms=False
(multi-replica FlashInfer-workspace workaround;
long-term fix belongs inside arctic-inference)
- speculative_model is set -> speculative_config block above
User-supplied trainer.arctic_rl.vllm_config still wins on conflict;
README's "Escape hatch" section is shortened accordingly. The 8B and
32B BIRD launchers drop their VLLM_CFG_PARTS shell construction and
flip the two typed flags instead.
Unit-tested all four topology cases (multi-replica + spec, single-
replica no-spec, user override, AI disabled). Lint + black + gitleaks
green via pre-commit on the touched files and on the full repo.
Co-authored-by: Cursor <cursoragent@cursor.com>
Now that all Arctic optimizations have typed knobs and auto-injection, make them the default. End-user UX is: trainer.override_entrypoint=integrations.arctic_rl.entrypoint \ <... your existing recipe overrides ...> That alone enables FCA, ZoRRo split-attention + dedup, memory-chunked logits, Liger fused kernels, and the multi-replica FlashInfer workaround. Add ``trainer.arctic_rl.speculative_model=<hf-id>`` to also enable Arctic speculative decoding. Anything off-by-default now is truly recipe-/cluster-shape-dependent (ZeRO stage, colocation, CPU offload, attention impl). Flipped defaults in ArcticRLTrainerConfig: - use_arctic_inference: False -> True - use_zorro: False -> True - use_liger: False -> True - logits_optimization: None -> "memory" Dropped the now-redundant ``=true`` lines from the 32B + 8B BIRD launchers (~10 lines per file). README updated to reflect the single-flag UX and the new "Defaults assume you came here to use Arctic RL; opt out explicitly if you want a baseline" framing. Lint + black + gitleaks pass on the touched files. Co-authored-by: Cursor <cursoragent@cursor.com>
SumanthRH
requested changes
Jun 27, 2026
SumanthRH
left a comment
SumanthRH
left a comment
Member
There was a problem hiding this comment.
Can you look into the uv resolution failures?
CI workflows are failing because of the pyproject.toml changes in your PR: https://github.com/NovaSky-AI/SkyRL/actions/runs/28268995579/job/83762333405?pr=1837
It looks like arctic-inference[vllm] only supports vllm <=0.18, but we use vllm 0.23.0 in SkyRL. Can you look into that? Is that a hard conflict?
I would recommend using dependency overrides as a fallback: You can remove arctic-rl as a native extra in the pyproject.toml and then instruct users to install the dependencies --with overrides:
uv run --with arctic-platform --with arctic-inference[vllm] --with liger-kernel ....this will of course require more changes from the user, but can be a workaround while we address the conflicts in a follow-up PR
Per PR NovaSky-AI#1837 review (Sumanth): the `arctic-rl` extra forces uv to resolve `arctic-inference[vllm]` (vllm<=0.18) against SkyRL's `fsdp`/`megatron` extras (vllm==0.23.0). That conflict was the root cause of the `uv_resolution_failures` blocking CI. Drop the extra entirely and instruct users to install the three Arctic deps with `uv pip install` after `uv sync`: uv sync uv pip install arctic-platform 'arctic-inference[vllm]' liger-kernel uv pip install --upgrade 'transformers>=4.57,<5' Launchers are unchanged (they call `python -m` against the active venv). README also documents the `uv run --with ...` ephemeral alternative Sumanth suggested. Co-authored-by: Cursor <cursoragent@cursor.com>
SumanthRH
reviewed
Jun 28, 2026
Member
There was a problem hiding this comment.
Did you confirm that the Arctic RL script runs with the new dependency changes?
You're syncing some custom dependencies with uv pip install into the current venv.
However, the uv + ray integration will create a new venv for the workers based on the pyproject.toml and the uv run command
"Under the hood, this integration uses Ray’s runtime environment feature, which allows the environment to be specified in code. We define a new low level runtime environment plugin called py_executable. It allows you to specify the Python executable (including arguments) that Ray workers will be started in. In the case of uv, the py_executable is set to uv run with the same parameters that were used to run the driver. Also, the working_dir runtime environment is used to propagate the working directory of the driver (including the pyproject.toml) to the workers. This allows uv to set up the right dependencies and environment for the workers to run in."
https://www.anyscale.com/blog/uv-ray-pain-free-python-dependencies-in-clusters
So custom dependencies will not be synced AFAIK, unless you use uv run --with ...
Matches the convention every other SkyRL integration already uses
(examples/train/flash_rl, examples/train_integrations/harbor,
examples/train/mini_swe_agent). Each launcher's driver is now
uv run --isolated --extra skyrl-train \
--with arctic-platform \
--with 'arctic-inference[vllm]' \
--with liger-kernel \
--with 'transformers==4.57.6' \
--with "flash-attn@<torch2.10 wheel URL>" \
-- python -m ...
`--isolated` bypasses the project lock so vLLM 0.18 + transformers 4.x
can coexist with the fsdp/megatron backends' vLLM 0.23 + transformers 5.x
pins without touching pyproject.toml. SkyRL's uv+Ray `py_executable`
plugin replays the same invocation on every Ray worker, so workers see
the arctic stack automatically (no `uv pip install` driver-only trap).
Transformers is exact-pinned to 4.57.6 (not `>=4.57,<5`) because Ray's
worker spawn shell parses `<5` as redirect-from-fd-5 — same reason
flash_rl uses `transformers==4.53.3`.
Verified E2E on GSM8K (Qwen3-0.6B, 4 GPUs): clean uv resolution,
NCCL weight sync, Arctic RL ZoRRo training, arctic-inference rollouts,
step 2 metrics landed (loss=0.042, grad_norm=0.23, sync_weights=0.04s).
Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
Author
|
Done — exactly what you suggested. Latest commit (b6e010e) drops the uv run --isolated --extra skyrl-train \
--with arctic-platform \
--with 'arctic-inference[vllm]' \
--with liger-kernel \
--with 'transformers==4.57.6' \
--with "flash-attn@<torch2.10 cu128 wheel>" \
-- python -m skyrl.train.entrypoints.main_base ...A few notes on what changed and why:
Verified E2E on GSM8K (Qwen3-0.6B, 4 GPUs) with this branch on a clean machine: The user-facing UX is now just I also looked at the Gemini comments (the |
The 2.38x BIRD-SQL speedup was measured with FlashAttention-3 on H200,
but the previous commit defaulted launchers to FA2. Add the public FA3
wheel from PyTorch's cu128 index to each BIRD launcher's --with chain
and flip the default ATTN_IMPL to flash_attention_3:
--with "flash-attn-3@https://download.pytorch.org/whl/cu128/flash_attn_3-3.0.0-cp39-abi3-manylinux_2_28_x86_64.whl"
The wheel is cp39-abi3 (stable ABI), works against torch 2.10 cu128 cp312.
Verified: fa3.flash_attn_func forward pass produces correct shape on H200
inside the isolated uv env.
GSM8K smoke launcher keeps FA2 default (may run on non-Hopper smoke
machines). README + knob table updated to document new defaults and the
wheel URL for users adding FA3 to their own launchers.
Co-authored-by: Cursor <cursoragent@cursor.com>
Rewrites the install + "drop into your recipe" sections in a flat technical voice instead of marketing/explanatory prose. Each `--with` flag now has a single-line bullet describing what it does, no "That's it" / "no X, no Y" framing. Co-authored-by: Cursor <cursoragent@cursor.com>
SumanthRH
approved these changes
Jun 29, 2026
2 of 4 tasks
sfc-gh-kganesan
added a commit
to Snowflake-AI-Research/Arctic-Platform
that referenced
this pull request
Jul 30, 2026
…nges Both merged upstream integrations (NovaSky-AI/SkyRL#1837:integrations/arctic_rl/ and the Arctic-specific RemoteBackend used by verl-project/verl#6422 living at arctic_platform/integrations/verl/adapter.py) construct their client with from arctic_platform.rl import ArcticRLClientConfig, create_arctic_rl_client client = create_arctic_rl_client(config, server_state) The goal is Cortex serverless support without either integration changing. This commit adds Cortex as a first-class backend on that single entry point. ## What ### `arctic_platform/rl/config.py` - `backend` Literal widened `["local", "dss-platform"]` -> add `"cortex"`. - Cortex-only fields added: `cortex_host` / `_database` / `_schema` / `_endpoint` / `_pat_env_var` / `_base_url`, `max_seq_len`. All Optional; missing keys fall back to `CortexTransport` defaults / `CORTEX_*` env vars. - `_derive_host_port` short-circuits on `backend == "cortex"` (Cortex has no local host/port; SnowAPI routing is by base_url). - `_validate_local_gpu_counts` already scoped to `backend == "local"`; unchanged. ### `arctic_platform/rl/client.py` - `create_arctic_rl_client(config, server_state)` early-branches when `config.backend == "cortex"` and returns `_CortexClientShim` via `_cortex_dispatch.build_cortex_client`. - `ArcticRLHTTPClient` / `ArcticRLRayClient` imports moved *inside* the http / ray branches so the Cortex code path never triggers the on-prem server dep chain (uvicorn -> arctic_inference -> vllm). The type hint for `ArcticRLServerState` is behind `TYPE_CHECKING`. ### `arctic_platform/rl/_cortex_dispatch.py` (new) - `_to_unified_config(legacy)` builds an `arctic_platform.client.ArcticRLClientConfig` sized for Cortex: backend, model_name, GPU counts, seed, all seven cortex_* / max_seq_len fields, and reconnect job ids. On-prem-only knobs (`ds_config`, `arctic_inference_config`, `ds_worker_config`, `log_prob_ds_config`, `full_determinism`, `colocate`, etc.) are intentionally NOT forwarded — Cortex has no local placement to configure and silently passing them would mask real config drift. - `_CortexClientShim` wraps `arctic_platform.client.ArcticRLClient` (which is synchronous by design) and re-exposes the legacy async surface both integrations reach for: `async def fwd_bwd / fwd_no_grad / step / save_checkpoint / save_weights / generate / sync_weights / reset_prefix_cache / wake_/sleep_ × {training, inference, log_prob} / empty_training_cache / weight_norm / log_probs`. Sync methods (`reconnect_config`, `get_server_state`, `shutdown`) match the legacy signature. Properties `config` / `training_job_id` / `sampling_job_id` / `log_prob_job_id` pass through. - `reconnect_config()` round-trips through the legacy config with cortex job ids populated so verl's `reconnect_handle()` pattern keeps working. - `get_server_state()` returns `None` (Cortex has no local Ray state actor; verl's reconnect path re-attaches via `training_job_id`). ### `arctic_platform/rl/__init__.py` - Lazy-load the heavy exports via PEP 562 `__getattr__`. - Eager: `ArcticRLClientConfig`, `WeightSyncConfig` (pydantic-only). - Lazy: `create_arctic_rl_client`, `WeightSyncCoordinator`, and all six `processors.*` re-exports. First attribute access resolves + caches in `globals()` so subsequent lookups hit the normal fast path. - Rationale: prior to this, `from arctic_platform.rl import ArcticRLClientConfig` on a Cortex-only driver pulled `arctic_platform.rl.client` -> `http_client` -> `http_server` -> `arctic_inference.server.metrics` -> `vllm` at package-init time. The Cortex dispatch never actually executes any of that; the eager import was pure overhead that forced Cortex users to install the entire on-prem ML stack. Post-refactor, a Cortex driver needs `arctic-platform` + `pydantic` + `requests`. Nothing else. ### `tests/client/test_rl_cortex_dispatch.py` (new, 23 tests) - `TestConfigTranslation`: cortex backend translates 1:1; cortex fields threaded; missing fields fall back to unified defaults; on-prem fields NOT forwarded (silent-drift guard); reconnect job ids forwarded. - `TestShimAsyncSurface`: pins every async call site both integrations hit — `fwd_bwd` / `fwd_no_grad` w/ `reference_model=` + `post_processors=`, `step(learning_rate=...)`, `sync_weights(cuda_ipc, low_memory)`, `wake_inference(tags=...)`, `sleep_inference(level=...)`, full colocation lifecycle, `generate(prompts, sampling_params)`, `save_checkpoint(stage_info, path)`. - `TestShimSyncSurface`: `shutdown()` is sync; `reconnect_config()` returns a `_LegacyConfig` with cortex job ids populated; `get_server_state()` returns `None` on Cortex. - `TestPropertySurface`: `client.config` returns the legacy config instance (SkyRL reads `client.config.colocate`); job-id properties pass through. - `TestFactoryDispatch`: - `test_backend_cortex_returns_shim`: `create_arctic_rl_client` picks the Cortex path and translates the config correctly. - `test_cortex_path_never_imports_onprem_transports`: pins the invariant — after a cortex dispatch, `arctic_platform.rl.http_*` / `ray_*` modules are NOT in `sys.modules`. Regression here means Cortex drivers pay the on-prem dep cost again. - `test_onprem_transports_are_lazy_at_module_level`: pins the source- level invariant — `arctic_platform.rl.client` module must NOT expose `ArcticRLHTTPClient` / `ArcticRLRayClient` as attributes. ## Result - 92 tests pass locally (69 previously-existing + 23 new). Test file runs on any machine — no vllm, no ray, no arctic_inference required. - `from arctic_platform.rl import ArcticRLClientConfig` on a bare pydantic+requests env loads exactly 2 modules and 0 heavy deps. - Neither upstream integration needs any code change to gain Cortex support; flipping `config.backend = "cortex"` + populating `cortex_*` fields is the entire user-facing UX. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an
arctic_rlintegration underintegrations/arctic_rl/that lets SkyRL recipes train against Snowflake's Arctic RL stack (Arctic-Platform + Arctic-Inference). It mirrors the existingharbor/dapointegration shape — each integration ships its own entrypoint, config, generator, trainer, BIRD-SQL environment, and example launchers.Core SkyRL changes are intentionally minimal (1 file, +15 / -1 line):
skyrl/train/entrypoints/main_base.py: one ~15-line dispatch hook that inspectstrainer.override_entrypoint=<module>on argv BEFORE strict config parsing and forwards to the integration'smain(). This lets integrations add config fields without subclassingSkyRLTrainConfigand without forking the recipe scripts. UX is just:pyproject.toml: adds an additive[arctic-rl]extra (skyrl[arctic-rl]) pulling inarctic-platform,arctic-inference[vllm], andliger-kernel. No existing pins touched.Integration contents (
integrations/arctic_rl/)entrypoint.py— boots anArcticRLClientagainst arctic-platform and forwards weight-sync, rollouts, and metricsconfig.py—ArcticRLTrainerConfigexposing the colocate / ZoRRo / FCA / spec-dec / CUDA-IPC-weight-sync knobstrainer.py,generator.py— thin adapters from SkyRL's PPO loop to Arctic-Platform's RL clientenvs/bird*.py— vendored BIRD-SQL environment so the recipe is self-contained (no private branch dep)examples/— smoke / 8B / 32B BIRD GRPO recipes + a GSM8K sanity check. Env vars in launchers are trimmed to only what's recipe-required.Diff scope
integrations/arctic_rl/main_base.py,pyproject.toml)Test plan
uv run --extra arctic-rl -m skyrl.train.entrypoints.main_base trainer.override_entrypoint=...boots without touching any of SkyRL's typed configHappy to iterate on review feedback.
Made with Cursor