diff --git a/examples/train_integrations/harbor/README.md b/examples/train_integrations/harbor/README.md index 61d1450db7..9421d3a183 100644 --- a/examples/train_integrations/harbor/README.md +++ b/examples/train_integrations/harbor/README.md @@ -14,7 +14,8 @@ examples/train_integrations/harbor/ entrypoints/ main_harbor.py # Full training entrypoint main_harbor_generate.py # Generation-only debug entrypoint - run_codecontest.sh # Code contest training (Qwen3-8B) + run_codecontest.sh # Code contest training (Qwen3-8B, FSDP baseline) + run_codecontest_arctic.sh # Same recipe, routed through the Arctic RL backend run_harbor_gen.sh # Debug generation-only ``` @@ -39,3 +40,23 @@ uv run examples/train_integrations/harbor/prepare_harbor_dataset.py \ # 3. Launch training bash examples/train_integrations/harbor/run_codecontest.sh ``` + +### Arctic RL backend + +Any Harbor recipe can be routed through the Arctic RL server (ZoRRo / FCA / Arctic speculative decoding) by adding one CLI flag to `main_harbor`: + +``` +trainer.override_entrypoint=integrations.arctic_rl.harbor_entrypoint +``` + +`run_codecontest_arctic.sh` next to `run_codecontest.sh` is that flag applied to the CodeContests recipe: + +```bash +export WANDB_API_KEY=... DAYTONA_API_KEY=... +bash examples/train_integrations/harbor/run_codecontest_arctic.sh +# smaller smoke: +NUM_POLICY_GPUS=4 MODEL=Qwen/Qwen3-0.6B MAX_MODEL_LEN=8192 \ + bash examples/train_integrations/harbor/run_codecontest_arctic.sh +``` + +Setup, `trainer.arctic_rl.*` knobs, and design notes: [`integrations/arctic_rl/README.md`](../../../integrations/arctic_rl/README.md) and [`integrations/arctic_rl/docs/HARBOR_DESIGN.md`](../../../integrations/arctic_rl/docs/HARBOR_DESIGN.md). diff --git a/examples/train_integrations/harbor/entrypoints/main_harbor.py b/examples/train_integrations/harbor/entrypoints/main_harbor.py index d7aad69455..4ee74f30af 100644 --- a/examples/train_integrations/harbor/entrypoints/main_harbor.py +++ b/examples/train_integrations/harbor/entrypoints/main_harbor.py @@ -97,6 +97,19 @@ def skyrl_entrypoint(cfg): def main() -> None: + # Same override pattern as ``skyrl.train.entrypoints.main_base``: peek at + # ``trainer.override_entrypoint=`` before parsing so an integration (e.g. + # ``integrations.arctic_rl.harbor_entrypoint``) can add its own config + # fields and take over dispatch. Users toggle a Harbor recipe to Arctic + # RL by appending exactly one CLI flag to any Harbor launcher; no other + # changes required. + for arg in sys.argv[1:]: + if arg.startswith("trainer.override_entrypoint="): + override = arg.split("=", 1)[1] + from importlib import import_module + + return import_module(override).main() + cfg = HarborSkyRLConfig.from_cli_overrides(sys.argv[1:]) # Load harbor defaults and merge CLI overrides on top diff --git a/examples/train_integrations/harbor/run_codecontest_arctic.sh b/examples/train_integrations/harbor/run_codecontest_arctic.sh new file mode 100755 index 0000000000..f8b0b76e65 --- /dev/null +++ b/examples/train_integrations/harbor/run_codecontest_arctic.sh @@ -0,0 +1,195 @@ +set -ex + +# run_codecontest.sh with the Arctic RL backend spliced in. Diff versus the +# FSDP baseline: swap uv extras, add `trainer.override_entrypoint=...`, add +# `trainer.arctic_rl.*`. See integrations/arctic_rl/docs/HARBOR_DESIGN.md. + +# export WANDB_API_KEY=... +# export DAYTONA_API_KEY=... # or MODAL_TOKEN_ID/_SECRET, E2B_API_KEY +: "${WANDB_API_KEY:?WANDB_API_KEY not set}" +: "${DAYTONA_API_KEY:?DAYTONA_API_KEY not set (or pick another sandbox provider)}" + +# Some envs firewall SSH; force plain HTTPS for uv git clones. +export GIT_CONFIG_GLOBAL="${GIT_CONFIG_GLOBAL:-/dev/null}" + +# Bind + advertise addresses for the Arctic OpenAI shim (openai_bridge.py). +export ARCTIC_HARBOR_SHIM_HOST="${ARCTIC_HARBOR_SHIM_HOST:-0.0.0.0}" +export ARCTIC_HARBOR_SHIM_PORT="${ARCTIC_HARBOR_SHIM_PORT:-8000}" + +# Repo root on PYTHONPATH so the driver Ray task can import integrations.arctic_rl.*. +_REPO_ROOT="$(cd "$(dirname "$0")"/../../.. && pwd)" +export PYTHONPATH="${_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +#----------------------- +# Dataset setup +#----------------------- +# Prepare datasets first (downloads from HuggingFace and extracts tasks): +# uv run examples/train_integrations/harbor/prepare_harbor_dataset.py --dataset open-thoughts/CodeContests +# uv run examples/train_integrations/harbor/prepare_harbor_dataset.py --dataset open-thoughts/OpenThoughts-TB-dev +DATA_DIR="$HOME/data/harbor" +TRAIN_DATA="['$DATA_DIR/CodeContests']" +EVAL_DATA="['$DATA_DIR/OpenThoughts-TB-dev']" + +#----------------------- +# Directory setup +#----------------------- +RUN_NAME="codecontest_arctic" +STORAGE_ROOT="/mnt/local_storage/$RUN_NAME" +TRIALS_DIR="$STORAGE_ROOT/trials_run" +CKPTS_DIR="$STORAGE_ROOT/ckpts" +EXPORTS_DIR="$STORAGE_ROOT/exports" +LOG_DIR="$STORAGE_ROOT/logs" + +#----------------------- +# Training setup +#----------------------- +# Env-overridable knobs. Smoke on a smaller model with e.g. +# NUM_POLICY_GPUS=4 MODEL=Qwen/Qwen3-0.6B MAX_MODEL_LEN=8192 bash $0 +: "${MODEL:=Qwen/Qwen3-8B}" +: "${SERVED_MODEL_NAME:=$(basename "${MODEL}")}" +: "${N_SAMPLES_PER_PROMPT:=8}" +: "${MINI_BATCH_SIZE:=32}" +: "${MAX_MODEL_LEN:=32768}" + +# Algorithmic parameters +LOSS_REDUCTION="token_mean" # with step-wise training, we have to use token_mean to be prefix-merge-invariant +GRPO_NORM_BY_STD=false +USE_KL_LOSS=false +# Keep overlong trajectories in the batch under Arctic. Upstream +# arctic_platform.rl.utils.batch.split_dict raises IndexError if the surviving +# batch shrinks below num_workers, so we can't afford to drop samples. Overlong +# trajectories still contribute a zero-reward learning signal. +APPLY_OVERLONG_FILTERING=false + +# Essentially achieves interleaved thinking (does not strip thinking tokens). Allows our step-wise +# training to be able to merge more step-wise outputs and hence speed up training. +# If you change the model you train, please change it accordingly, and decide if you need to make +# modifications. +CHAT_TEMPLATE_PATH="$(dirname "$0")/../../../skyrl/train/utils/templates/qwen3_acc_thinking.jinja2" + +# TIS corrections +TIS_TYPE=token +TIS_IMP_RATIO_CAP=2.0 + +#---------------- +# Infrastructure setup +#---------------- +: "${NUM_POLICY_GPUS:=8}" +: "${NUM_INFERENCE_ENGINES:=${NUM_POLICY_GPUS}}" # arctic runs one vLLM replica per GPU +: "${TP_SIZE:=1}" +: "${ENABLE_RATE_LIMITING:=true}" +: "${TRAJECTORIES_PER_SECOND:=5}" +: "${MAX_CONCURRENCY:=512}" + +#---------------- +# Arctic RL knobs — see integrations/arctic_rl/README.md for the full table. +#---------------- +COLOCATE=true +ZERO_STAGE=3 +USE_LIGER=true +USE_ARCTIC_INFERENCE=true +USE_ZORRO=false # off for Harbor: prompts don't share the (problem x N samples) prefix ZoRRo dedupes on +CUDA_IPC_WEIGHT_SYNC=true +LOW_MEM_WEIGHT_SYNC=false +VLLM_ENFORCE_EAGER=false +VLLM_GPU_MEM_UTIL=0.7 +ATTN_IMPL=flash_attention_2 + +# torch-2.10 flash-attn wheel matching arctic-inference's vLLM 0.18 pin +# (SkyRL's default lock ships a torch-2.11 wheel). +FLASH_ATTN_WHL="https://github.com/lesj0610/flash-attention/releases/download/v2.8.3-cu12-torch2.10-cp312/flash_attn-2.8.3%2Bcu12torch2.10cxx11abiTRUE-cp312-cp312-linux_x86_64.whl" + +# Preflight: surface leaked ports / busy GPUs from a previous crashed run so we +# don't fail 5 minutes into engine init with a cryptic EADDRINUSE. +if command -v ss >/dev/null 2>&1 && ss -Hltn "sport = :${ARCTIC_HARBOR_SHIM_PORT}" 2>/dev/null | grep -q .; then + echo "!! Port ${ARCTIC_HARBOR_SHIM_PORT} (Arctic OpenAI shim) is already in use." >&2 + echo " Rerun with ARCTIC_HARBOR_SHIM_PORT=, or clean stale workers:" >&2 + echo " pkill -9 -f 'InferenceWorker|DeepSpeedWorker|VLLM::|ArcticRLRayServerState|skyrl_entrypoint' && sleep 5" >&2 + exit 1 +fi +if command -v nvidia-smi >/dev/null 2>&1; then + busy=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | awk '$1 > 1024' | wc -l) + if [[ "${busy}" -gt 0 ]]; then + echo "!! WARNING: ${busy} GPU(s) already have >1 GiB memory used." >&2 + echo " If a previous Arctic run crashed, free them with:" >&2 + echo " pkill -9 -f 'InferenceWorker|DeepSpeedWorker|VLLM::|ArcticRLRayServerState|skyrl_entrypoint' && sleep 5" >&2 + fi +fi + +uv run --isolated --extra skyrl-train --extra harbor \ + --with arctic-platform \ + --with 'arctic-inference[vllm]' \ + --with liger-kernel \ + --with 'transformers==4.57.6' \ + --with "flash-attn@${FLASH_ATTN_WHL}" \ + -m examples.train_integrations.harbor.entrypoints.main_harbor \ + trainer.override_entrypoint=integrations.arctic_rl.harbor_entrypoint \ + data.train_data=$TRAIN_DATA \ + data.val_data=$EVAL_DATA \ + trainer.policy.model.path=$MODEL \ + generator.inference_engine.served_model_name=$SERVED_MODEL_NAME \ + harbor_trial_config.trials_dir=$TRIALS_DIR \ + trainer.export_path=$EXPORTS_DIR \ + trainer.ckpt_path=$CKPTS_DIR \ + trainer.log_path=$LOG_DIR \ + trainer.algorithm.advantage_estimator=grpo \ + trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ + trainer.algorithm.grpo_norm_by_std=$GRPO_NORM_BY_STD \ + trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ + trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ + trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ + trainer.strategy=fsdp2 \ + trainer.placement.colocate_all=false \ + trainer.placement.policy_num_nodes=1 \ + trainer.placement.ref_num_nodes=1 \ + trainer.placement.policy_num_gpus_per_node=$NUM_POLICY_GPUS \ + trainer.placement.ref_num_gpus_per_node=$NUM_POLICY_GPUS \ + trainer.arctic_rl.colocate=$COLOCATE \ + trainer.arctic_rl.zero_stage=$ZERO_STAGE \ + trainer.arctic_rl.use_liger=$USE_LIGER \ + trainer.arctic_rl.use_arctic_inference=$USE_ARCTIC_INFERENCE \ + trainer.arctic_rl.use_zorro=$USE_ZORRO \ + trainer.arctic_rl.cuda_ipc_weight_sync=$CUDA_IPC_WEIGHT_SYNC \ + trainer.arctic_rl.low_memory_weight_sync=$LOW_MEM_WEIGHT_SYNC \ + trainer.arctic_rl.attn_implementation=$ATTN_IMPL \ + trainer.arctic_rl.vllm_enforce_eager=$VLLM_ENFORCE_EAGER \ + trainer.arctic_rl.vllm_max_model_len=$MAX_MODEL_LEN \ + trainer.arctic_rl.vllm_enable_prefix_caching=true \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$TP_SIZE \ + generator.inference_engine.gpu_memory_utilization=$VLLM_GPU_MEM_UTIL \ + generator.inference_engine.engine_init_kwargs.chat_template=$CHAT_TEMPLATE_PATH \ + generator.inference_engine.engine_init_kwargs.max_model_len=$MAX_MODEL_LEN \ + generator.inference_engine.engine_init_kwargs.enable_log_requests=false \ + trainer.epochs=3 \ + trainer.eval_batch_size=128 \ + trainer.eval_before_train=false \ + trainer.eval_interval=100 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=$MINI_BATCH_SIZE \ + trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.ckpt_interval=5 \ + trainer.max_ckpts_to_keep=5 \ + trainer.hf_save_interval=5 \ + trainer.algorithm.max_seq_len=$MAX_MODEL_LEN \ + trainer.policy.optimizer_config.lr=1.0e-6 \ + generator.step_wise_trajectories=true \ + generator.merge_stepwise_output=true \ + generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ + generator.eval_n_samples_per_prompt=2 \ + generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ + trainer.logger=wandb \ + trainer.project_name=harbor \ + trainer.run_name=$RUN_NAME \ + trainer.resume_mode=latest \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.batched=false \ + generator.inference_engine.enforce_eager=$VLLM_ENFORCE_EAGER \ + generator.rate_limit.enabled=$ENABLE_RATE_LIMITING \ + generator.rate_limit.trajectories_per_second=$TRAJECTORIES_PER_SECOND \ + generator.rate_limit.max_concurrency=$MAX_CONCURRENCY \ + "$@" diff --git a/integrations/arctic_rl/README.md b/integrations/arctic_rl/README.md index a5ac2e1023..06dfbfde6e 100644 --- a/integrations/arctic_rl/README.md +++ b/integrations/arctic_rl/README.md @@ -83,6 +83,68 @@ uv run --isolated --extra skyrl-train \ Add `trainer.arctic_rl.speculative_model=` to enable Arctic speculative decoding. +## Harbor recipes + +The same override flag routes any Harbor recipe through Arctic RL. Target Harbor's own entrypoint and point `trainer.override_entrypoint` at `integrations.arctic_rl.harbor_entrypoint`: + +``` +uv run --isolated --extra harbor \ + --with arctic-platform --with 'arctic-inference[vllm]' --with liger-kernel \ + --with 'transformers==4.57.6' \ + --with "flash-attn@${FLASH_ATTN_WHL}" \ + -- python -m examples.train_integrations.harbor.entrypoints.main_harbor \ + trainer.override_entrypoint=integrations.arctic_rl.harbor_entrypoint \ + +``` + +`harbor_entrypoint` spins up a small OpenAI-compatible FastAPI shim +(`openai_bridge.ArcticInferenceEngineAdapter`) in front of +`arctic_client.generate()` before constructing `HarborGenerator`, so +Harbor's `terminus-2` agent inside the sandbox reaches the exact same +model weights the trainer is updating — the shim is served in-process +on the driver and reachable at `http://127.0.0.1:8000/v1` by default. +Weight sync (NCCL or CUDA-IPC in colocated mode) still flows through +the Arctic RL server, not through the shim. + +The shim's bind/advertise addresses are env-only (no new SkyRL config +fields): + +| env var | default | notes | +|--------------------------------------|------------|------------------------------------------------------------| +| `ARCTIC_HARBOR_SHIM_HOST` | `0.0.0.0` | bind address (needed on `0.0.0.0` for off-driver sandboxes)| +| `ARCTIC_HARBOR_SHIM_PORT` | `8000` | starting port; auto-bumps if busy | +| `ARCTIC_HARBOR_SHIM_ADVERTISED_HOST` | `127.0.0.1`| host clients (sandbox) should call; set to the driver node IP if the sandbox runs off-node | + +### Harbor CodeContests recipe + +Two launchers, same underlying entrypoint: + +- [`examples/train_integrations/harbor/run_codecontest_arctic.sh`](../../examples/train_integrations/harbor/run_codecontest_arctic.sh) — Qwen3-8B / 8-GPU, mirrors Harbor's FSDP `run_codecontest.sh`. +- [`examples/run_codecontest_arctic_harbor.sh`](examples/run_codecontest_arctic_harbor.sh) — env-driven, defaults to a Qwen3-0.6B / 4-GPU smoke (~5 min from a cold uv cache). + +```bash +export WANDB_API_KEY=... # WANDB_BASE_URL auto-selects the Snowflake host +export DAYTONA_API_KEY=... # or MODAL_TOKEN_ID/_SECRET, E2B_API_KEY + +uv run --extra harbor examples/train_integrations/harbor/prepare_harbor_dataset.py \ + --dataset open-thoughts/CodeContests \ + --output_dir /data/skyrl-runs/arctic_harbor/data/CodeContests + +bash integrations/arctic_rl/examples/run_codecontest_arctic_harbor.sh +``` + +Scale up via env vars (defaults already match the FSDP baseline recipe — Qwen3-8B, 32K context, colocated CUDA-IPC): + +```bash +NUM_GPUS=8 MODEL=Qwen/Qwen3-8B \ + TRAIN_BATCH_SIZE=4 N_SAMPLES_PER_PROMPT=8 MAX_CONCURRENCY=32 \ + bash integrations/arctic_rl/examples/run_codecontest_arctic_harbor.sh +``` + +Design + change list: [`docs/HARBOR_DESIGN.md`](docs/HARBOR_DESIGN.md) (see **Launcher invariants** for the two settings — `APPLY_OVERLONG_FILTERING=false` and `max_input_tokens`/`max_output_tokens` sourcing — that must not be casually flipped). + +Sandbox concurrency: Harbor spins one sandbox per trial. On Daytona's free tier (10 CPUs total), keep `MAX_CONCURRENCY ≤ 8` and `train_batch_size × n_samples_per_prompt ≤ 16`. If a launcher crashes mid-run, the preflight tells you exactly which port / GPUs are stuck and prints the `pkill` line to free them. + ### `trainer.arctic_rl.*` knobs | Knob | Default | Effect | @@ -157,17 +219,29 @@ GPU layout: ``` integrations/arctic_rl/ ├── README.md -├── trainer.py ArcticPPOTrainer — routes training to Arctic server actors -├── generator.py ArcticGenerator — routes generation to vLLM, scores via skyrl-gym -├── config.py ArcticRLTrainerConfig + build_rl_config -├── entrypoint.py Dispatched from main_base via trainer.override_entrypoint +├── trainer.py ArcticPPOTrainer — routes training to Arctic server actors +├── generator.py ArcticGenerator — routes generation to vLLM, scores via skyrl-gym +├── config.py ArcticRLTrainerConfig + build_rl_config +├── entrypoint.py Dispatched from main_base via trainer.override_entrypoint +├── harbor_entrypoint.py Same, dispatched from main_harbor — Harbor rollouts + Arctic training +├── openai_bridge.py OpenAI-compatible HTTP shim over arctic_client.generate() +├── docs/ +│ └── HARBOR_DESIGN.md ├── envs/ │ ├── bird.py │ ├── bird_reward.py │ └── preprocess_bird.py └── examples/ ├── run_gsm8k_grpo_4gpu.sh - └── run_bird_grpo_*.sh (BIRD recipes — manual data prep today) + ├── run_bird_grpo_*.sh (BIRD recipes — manual data prep today) + └── run_codecontest_arctic_harbor.sh (Harbor recipe — env-driven, 0.6B smoke -> 8B) +``` + +Harbor's own recipe folder also carries a fixed-config companion: + +``` +examples/train_integrations/harbor/ + └── run_codecontest_arctic.sh (Qwen3-8B / 8-GPU parity with run_codecontest.sh) ``` ## BIRD-SQL: 32B Qwen3 on 32 × H200 diff --git a/integrations/arctic_rl/docs/HARBOR_DESIGN.md b/integrations/arctic_rl/docs/HARBOR_DESIGN.md new file mode 100644 index 0000000000..bc66744043 --- /dev/null +++ b/integrations/arctic_rl/docs/HARBOR_DESIGN.md @@ -0,0 +1,298 @@ +# Harbor recipes on the Arctic RL backend — design + +Companion doc to PR **NovaSky-AI/SkyRL#1879**. + +## Goal + +Let any Harbor recipe opt into the Arctic RL backend the same way stock +SkyRL recipes do — one CLI flag, no changes to Arctic Platform or +Harbor: + +``` +trainer.override_entrypoint=integrations.arctic_rl.harbor_entrypoint +``` + +Prior state: + +- `integrations.arctic_rl.entrypoint` (already in `main`) dispatches + from `main_base.py` and gets `ArcticGenerator`, which calls + `arctic_client.generate()` directly (no HTTP). +- Harbor recipes use `main_harbor.py` and `HarborGenerator`, which + drives `terminus-2` inside a sandbox VM that talks to the model over + `POST /v1/chat/completions` via LiteLLM. + +The two paths never met: Harbor's agent lives inside a Daytona/Modal +sandbox and can only reach the model over HTTP; Arctic's sampling side +is a `ReplicaPool` reached only via async +`arctic_client.generate(prompts, sampling_params)`. This PR is the +adapter that closes that gap without touching either side. + +Upstream SkyRL recently removed its OpenAI HTTP shim +(`inference_engine_client_http_endpoint.serve`) — vLLM's `VLLMServerActor` +now serves OpenAI natively. This PR replaces the shim we used to lean on +with a self-contained FastAPI app in `openai_bridge.py`, so nothing here +depends on an internal SkyRL module that can move again. + +## Change surface + +``` ++ integrations/arctic_rl/harbor_entrypoint.py NEW entrypoint + config ++ integrations/arctic_rl/openai_bridge.py NEW InferenceEngineInterface + + self-contained FastAPI shim ++ integrations/arctic_rl/examples/ + run_codecontest_arctic_harbor.sh NEW env-driven launcher + (0.6B smoke -> 8B) ++ integrations/arctic_rl/docs/HARBOR_DESIGN.md NEW THIS FILE ++ examples/train_integrations/harbor/ + run_codecontest_arctic.sh NEW Qwen3-8B / 8-GPU parity + with run_codecontest.sh +~ examples/train_integrations/harbor/entrypoints/ + main_harbor.py +13 override_entrypoint peek +~ examples/train_integrations/harbor/README.md +18 Arctic RL backend + + companion launcher +~ integrations/arctic_rl/README.md +20 Harbor recipes section +``` + +No changes to `pyproject.toml`. No changes to any core `skyrl/` module. +No changes to `arctic-platform`, `arctic-inference`, or `harbor`. All +integration code lives under `integrations/arctic_rl/`. + +## Architecture + +``` + Single Ray cluster (Arctic in-process; HTTP only inside sandbox) ++---------------------------------------------------------------------------------------+ +| | +| SkyRL driver (num_gpus=0) | +| | +| ArcticHarborExp --> HarborGenerator --> Harbor Trial | +| | (Daytona / Modal sandbox VM) | +| | | | +| | | POST /v1/chat/completions | +| | v (via LiteLLM hosted_vllm/...) | +| | +---------------------------------+ | +| +--------->| ArcticInferenceEngineAdapter | | +| get_endpoint_url | - FastAPI on a daemon thread | | +| finish_session | - Implements | | +| | InferenceEngineInterface | | +| +---------------------------------+ | +| | | +| | await arctic_client.generate | +| v | +| Arctic RL Ray actors: | +| - DeepSpeedWorker (x N training GPUs) --NCCL/CUDA-IPC--> InferenceWorker | +| - InferenceWorker (x M sampling GPUs, ArcticInference vLLM) | +| | ++---------------------------------------------------------------------------------------+ +``` + +## `openai_bridge.py` — `ArcticInferenceEngineAdapter` + +Implements +`skyrl.backends.skyrl_train.inference_servers.base.InferenceEngineInterface` +so `HarborGenerator` can call the two methods it actually needs +(`get_endpoint_url()`, `finish_session()`) on this object directly, with +no wrapper. + +Responsibilities: + +1. Build a FastAPI app (`/v1/chat/completions`, `/v1/completions`, + `/v1/models`, `/health`). +2. Serve it via `uvicorn.Server` in a daemon thread on a port reserved + at startup (`ARCTIC_HARBOR_SHIM_PORT`, default 8000, auto-bumps if + busy). +3. Translate OpenAI requests into `arctic_client.generate(...)` and + the results back into OpenAI-shaped JSON. +4. Forward or no-op the rest of the interface. Weight sync, sleep/wake, + and pause/resume all go through `arctic_client` directly, not this + adapter. + +### Response shape + +Harbor's `terminus-2` uses `collect_rollout_details=True`, which makes +LiteLLM request `logprobs=True` + `extra_body.return_token_ids=True` +and then read three vLLM-specific response fields — matching vLLM's +`ChatCompletionResponseChoice` (`vllm.entrypoints.openai.protocol`): + +- `choices[i].token_ids` (direct choice field, NOT nested inside + `provider_specific_fields`) — LiteLLM auto-folds any non-`Choices` + field into `provider_specific_fields`, so terminus-2 gets it at + `choice.provider_specific_fields["token_ids"]`. +- `response.prompt_token_ids` (top-level) — read via `getattr` on + LiteLLM's `ModelResponse`; that model inherits OpenAI's + `BaseModel(extra="allow")` so unknown top-level keys survive. +- `choices[i].logprobs.content[j].logprob` — standard OpenAI shape, + filled in from Arctic's per-token logprobs when the request enables + them. + +Without any of those, `terminus-2._collect_subagent_rollout_details` +short-circuits, `chat.rollout_details` stays empty, and +`HarborGenerator._harbor_agent_loop` fails every trajectory with +"empty/missing rollout_details". The bridge emits all three; a +standalone LiteLLM round-trip test in the run log verifies the shape. + +### What the shim deliberately does *not* touch + +`HarborGenerator` reads `completion_token_ids` / `prompt_token_ids` / +`logprobs` straight out of `rollout_details` — it never re-tokenizes +the response text. That means anything the shim does to `content` +(strip whitespace, inject `reasoning_content: None`, etc.) is invisible +to the trainer: gradients depend only on the token IDs vLLM's sampler +emitted, which the shim forwards verbatim from +`arctic_client.generate()`. The FSDP baseline runs the same recipe +with vLLM's native OpenAI endpoint and the same LiteLLM +``-stripping behavior; it trains fine without any of those +mutations, and so does the Arctic path. Keeping the shim thin means +every future SkyRL integration on top of Arctic inherits a bridge that +just shape-shifts, not one that carries workarounds for one caller's +quirks. + +## `harbor_entrypoint.py` — `ArcticHarborExp` + `main()` + +- **`ArcticHarborSkyRLConfig(SkyRLTrainConfig)`** — adds Harbor's + `harbor_trial_config: dict` and points `generator` at + `HarborGeneratorConfig` (defined by Harbor). `trainer` is bumped to + `ArcticTrainerConfig` so `trainer.arctic_rl.*` is available. **No + `from __future__ import annotations`** — SkyRL's + `build_nested_dataclass` introspects `dataclasses.fields(cls)[i].type` + and expects a concrete class, not a string. +- **`ArcticHarborExp(ArcticRLExp)`** — reuses `ArcticRLExp.__init__` + (which pre-inits the Arctic client) verbatim. Only overrides: + - `get_train_dataset` / `get_eval_dataset` → `HarborTaskDataset`. + - `get_generator` → `HarborGenerator`, passing the adapter as + `inference_engine_client` (so `HarborGenerator` picks up the shim + URL via `.get_endpoint_url()` and bakes it into LiteLLM's + `api_base`). + - `_setup_trainer` → spins up the OpenAI HTTP shim **before** + constructing `HarborGenerator` (order matters — `HarborGenerator` + reads the URL at `__init__`), then hands the same + `_ArcticInferenceEngineStub` upstream uses to `ArcticPPOTrainer` so + colocated sleep / wake still route to the Arctic server. +- **`main()`** — mirrors `integrations.arctic_rl.entrypoint.main()`: + pre-init the Arctic client on the driver, forward sandbox-provider + credentials + `PYTHONPATH` onto Ray workers, dispatch the remote + entrypoint task. Environment-variable prefixes forwarded: + `ARCTIC_`, `WANDB_`, `DAYTONA_`, `MODAL_`, `E2B_`, `RUNLOOP_`, + `GKE_`, `OPENAI_`, `HF_`. + +## `main_harbor.py` edit + ++13 lines: the same `trainer.override_entrypoint=` peek +`skyrl.train.entrypoints.main_base` already uses. Any existing Harbor +launcher becomes an Arctic launcher by appending the flag to its CLI. + +```python +def main() -> None: + for arg in sys.argv[1:]: + if arg.startswith("trainer.override_entrypoint="): + override = arg.split("=", 1)[1] + from importlib import import_module + return import_module(override).main() + ... +``` + +## Reference launchers + +Both invoke Harbor's own `main_harbor` with the override flag: + +- `examples/train_integrations/harbor/run_codecontest_arctic.sh` — + mirrors the FSDP baseline `run_codecontest.sh` (Qwen3-8B / 8 GPUs), + discoverable next to it. +- `integrations/arctic_rl/examples/run_codecontest_arctic_harbor.sh` — + env-driven, defaults to a Qwen3-0.6B / 4-GPU smoke, scales up via + `NUM_GPUS=... MODEL=...`. + +Both use `uv run --isolated --with arctic-platform --with 'arctic-inference[vllm]' --with liger-kernel --with 'transformers==4.57.6' --with ''`. +Harbor's step-wise training requirements +(`generator.step_wise_trajectories=true`, +`generator.merge_stepwise_output=true`) are baked into both so +prefix-merge invariance holds. + +Shim host/port + `trainer.arctic_rl.*` knob tables live in the arctic +[`README.md`](../README.md) — no duplication here. + +### Launcher invariants + +Two settings are locked to non-default values in both launchers; changing +either without understanding the interaction below will crash Harbor +training under Arctic on any multi-turn dataset. + +- **`APPLY_OVERLONG_FILTERING=false`.** Upstream + `arctic_platform.rl.utils.batch.split_dict` raises + `IndexError: list index out of range` when the surviving batch shrinks + below `num_workers` (== `NUM_GPUS` under colocation). Overlong-filtering + strips trajectories that exceed the context budget, and CodeContests + routinely produces enough of them to trigger the split. Keeping overlong + trajectories in the batch (as zero-reward samples) is the right learning + signal *and* sidesteps the upstream bug — no arctic-platform patch + required. +- **`harbor_trial_config.agent.kwargs.model_info.max_input_tokens` and + `max_output_tokens` point at `MAX_PROMPT_LENGTH` / `MAX_GENERATE_LENGTH`, + not `MAX_MODEL_LEN`.** Otherwise Harbor packs a full-context prompt and + then requests another full-context response, so any 2+ turn trial + eventually blows past the vLLM engine's `max_model_len` and gets rejected + with `VLLMValidationError: input N > max_model_len`. Enough rejections in + one batch → same `split_dict` crash. The Harbor-side companion launcher + matches the FSDP baseline recipe and simply omits these overrides so + Harbor's own defaults apply. + +### Preflight + +Both launchers do a non-destructive preflight before spinning up the uv +env: (a) fail with an actionable message if the Arctic OpenAI shim port +is already bound, (b) warn (don't fail) if any GPU has >1 GiB in use, and +(c) print the exact `pkill` incantation to free stale +`InferenceWorker` / `DeepSpeedWorker` / `VLLM::` / `ArcticRLRayServerState` +processes from a previous crashed run. Turns the failure mode from a +cryptic `EADDRINUSE` five minutes into engine init into an actionable +one-line error. + +## Smoke result — Qwen3-8B, 8×H100, CodeContests + +The end-to-end smoke on `upstream/main + this PR` (`arctic_smoke_*.log`): + +``` +Started: 'step' +Started: 'generate' +Finished: 'generate', time cost: 131.71s +Finished: 'sync_weights', time cost: 5.87s (CUDA IPC, 8 replicas, 399 params) +Finished: 'step', time cost: 142.27s +{'actor/pg_loss': -0.0955..., ...} +``` + +Zero "empty/missing rollout_details" warnings, non-zero policy gradient +loss, weight sync using colocated CUDA IPC. Convergence run (~40 steps) +is queued after this smoke completes. + +## Risks and known limits + +- **`use_zorro=false` default.** ZoRRo's prompt-group dedup path + assumes the "same problem × N samples" prefix pattern; Harbor's + agentic prompts don't share prefixes. Off by default; flip when the + prompt distribution allows. +- **Sandbox concurrency.** Daytona's free tier caps concurrent + sandboxes at 10 CPUs. On that tier keep + `train_batch_size × n_samples_per_prompt ≤ 16` and + `MAX_CONCURRENCY ≤ 8`; the launcher runs unmodified but back-pressures + on sandbox creation. +- **Chat template path.** The adapter re-applies the chat template + server-side so training-time re-tokenization sees the same prompt + bytes as sampling. Both Harbor and Arctic read the template from + `generator.inference_engine.engine_init_kwargs.chat_template` — must + point at the same file. + +## Alternatives considered + +- **In-process transport, no HTTP.** Harbor's agent runs inside a + sandbox VM, so it needs a real network endpoint. +- **Reuse `inference_engine_client_http_endpoint`.** Upstream removed + it along with the `enable_http_endpoint` config surface; reimplementing + the FastAPI app in-integration insulates this PR from further churn. +- **Fork `HarborGenerator` to call `arctic_client` directly.** Would + fork Harbor's `terminus-2` LiteLLM path — hard to keep in sync with + upstream Harbor. The adapter preserves the LiteLLM path verbatim. +- **New extra in `pyproject.toml`.** `uv` still evaluates the full + `[project.optional-dependencies]` graph under `--isolated`, so an + `arctic-inference[vllm]==0.18` pin would block launchers that also + need harbor / megatron / fsdp (vLLM 0.19+). `uv run --isolated + --with ...` (used by every arctic launcher) sidesteps this. diff --git a/integrations/arctic_rl/examples/run_codecontest_arctic_harbor.sh b/integrations/arctic_rl/examples/run_codecontest_arctic_harbor.sh new file mode 100755 index 0000000000..f72ec4062f --- /dev/null +++ b/integrations/arctic_rl/examples/run_codecontest_arctic_harbor.sh @@ -0,0 +1,267 @@ +#!/usr/bin/env bash +# Arctic RL + Harbor CodeContests GRPO smoke. +# +# Dispatch mechanism: this script invokes Harbor's existing entrypoint +# (``examples.train_integrations.harbor.entrypoints.main_harbor``) with one +# extra CLI flag — +# trainer.override_entrypoint=integrations.arctic_rl.harbor_entrypoint +# — mirroring the way any SkyRL recipe toggles to Arctic RL. All the +# Harbor knobs below are the same knobs the FSDP baseline +# (``examples/train_integrations/harbor/run_codecontest.sh``) accepts; +# switching backend is a one-line change in your own launcher. + +set -euo pipefail +export PYTHONUNBUFFERED=1 + +# ---------- credentials ---------- +: "${WANDB_API_KEY:?WANDB_API_KEY not set; export it or source env.sh}" +: "${DAYTONA_API_KEY:?DAYTONA_API_KEY not set (or pick another sandbox provider)}" +: "${WANDB_BASE_URL:=https://snowflake.wandb.io}" +: "${WANDB_PROJECT:=arctic_rl_harbor}" +export WANDB_API_KEY WANDB_BASE_URL WANDB_PROJECT DAYTONA_API_KEY +# Bypass user's global gitconfig url.insteadOf that rewrites github HTTPS -> SSH +# (SSH port is firewalled in the target env). Required for the uv resolver to +# clone arctic-platform / arctic-inference / harbor / megatron-core. +export GIT_CONFIG_GLOBAL="${GIT_CONFIG_GLOBAL:-/dev/null}" +# Reuse a persistent HF cache across runs so we don't re-download Qwen weights +# on every smoke. +export HF_HOME="${HF_HOME:-/data/hf_cache}" +export HF_HUB_ENABLE_HF_TRANSFER="${HF_HUB_ENABLE_HF_TRANSFER:-1}" + +# SkyRL's driver task deserializes ``integrations.arctic_rl.*`` + +# ``examples.train_integrations.harbor.*`` inside a fresh Ray worker; add the +# repo root to PYTHONPATH so those imports resolve without relying on the +# entrypoint's own env-forwarding fallback (which only fires on the driver +# task, not on the transient uv resolver process that spawns it). +_REPO_ROOT="$(cd "$(dirname "$0")"/../../.. && pwd)" +export PYTHONPATH="${_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +# ---------- topology ---------- +: "${NUM_GPUS:=4}" +: "${MODEL:=Qwen/Qwen3-0.6B}" +: "${SERVED_MODEL_NAME:=$(basename "${MODEL}")}" +# Match the Harbor FSDP baseline recipe (examples/train_integrations/harbor/ +# run_codecontest.sh): 32K vLLM context, with 24K reserved for the packed +# multi-turn prompt and 8K for the response so Harbor's agent never +# overshoots the vLLM engine's max_model_len mid-trial. +: "${MAX_MODEL_LEN:=32768}" +: "${MAX_PROMPT_LENGTH:=24576}" +: "${MAX_GENERATE_LENGTH:=8192}" + +# ---------- run + data layout ---------- +RUN_NAME="${RUN_NAME:-arctic_harbor_codecontest_$(basename "${MODEL}")_$(date -u +%Y%m%dT%H%M%SZ)}" +RUN_ROOT="${RUN_ROOT:-/data/skyrl-runs/arctic_harbor}" +# Persistent, deduplicated Harbor task cache (prepare_harbor_dataset.py extracts +# each dataset once into a subdir). Falls back to $HOME/data/harbor for parity +# with the upstream FSDP launcher. +: "${DATA_DIR:=${RUN_ROOT}/data}" +: "${TRAIN_DATA:=[\"${DATA_DIR}/CodeContests\"]}" +: "${EVAL_DATA:=[\"${DATA_DIR}/CodeContests\"]}" # reuse train for smoke; override for real eval + +TRIALS_DIR="${TRIALS_DIR:-${RUN_ROOT}/trials/${RUN_NAME}}" +CKPTS_DIR="${CKPTS_DIR:-${RUN_ROOT}/ckpts/${RUN_NAME}}" +EXPORTS_DIR="${EXPORTS_DIR:-${RUN_ROOT}/exports/${RUN_NAME}}" +LOG_DIR="${LOG_DIR:-${RUN_ROOT}/logs/${RUN_NAME}}" +mkdir -p "${TRIALS_DIR}" "${CKPTS_DIR}" "${EXPORTS_DIR}" "${LOG_DIR}" + +# ---------- batch math ---------- +# Small-by-default to keep the smoke under 10 min on Qwen3-0.6B. Scale up +# via `trainer.train_batch_size=... trainer.policy_mini_batch_size=...` on +# the CLI when moving off smoke. +: "${TRAIN_BATCH_SIZE:=8}" +: "${MINI_BATCH_SIZE:=8}" +: "${N_SAMPLES_PER_PROMPT:=4}" +: "${LR:=1.0e-6}" +: "${EPOCHS:=1}" +: "${EVAL_INTERVAL:=0}" # 0 = never; smoke doesn't run Harbor eval +: "${EVAL_BEFORE_TRAIN:=false}" + +# ---------- rate limiting (Harbor sandbox back-pressure) ---------- +: "${ENABLE_RATE_LIMITING:=true}" +: "${TRAJECTORIES_PER_SECOND:=5}" +: "${MAX_CONCURRENCY:=64}" + +# ---------- Arctic RL knobs ---------- +# Defaults match the arctic BIRD-8B recipe (run_bird_grpo_8b_32gpu.sh) with +# Harbor-specific safety flips (ZoRRO off — untested through the OpenAI shim). +: "${COLOCATE:=true}" +: "${ZERO_STAGE:=3}" # 8B under colocation needs ZeRO-3 to fit vLLM alongside +: "${USE_LIGER:=true}" +: "${USE_ARCTIC_INFERENCE:=true}" # FCA + Arctic spec-dec on server-side vLLM +: "${USE_ZORRO:=false}" # prompt-group dedup + chunked logits; enable once shape assertions verified +: "${CUDA_IPC_WEIGHT_SYNC:=true}" # near-zero-copy weight transfer in colocated mode +: "${LOW_MEM_WEIGHT_SYNC:=false}" +: "${VLLM_ENFORCE_EAGER:=false}" # graph-compile pays off after step 1 +: "${VLLM_GPU_MEM_UTIL:=0.7}" +: "${ATTN_IMPL:=flash_attention_2}" # FA3 requires the FA3 wheel; FA2 is safe on Hopper + +# ---------- Dr. GRPO / loss knobs (parity with the FSDP baseline recipe) ---------- +# Note: with step-wise Harbor training, upstream's run_codecontest.sh uses +# ``token_mean`` for prefix-merge invariance. If you flip this on your own +# launcher, keep it consistent with generator.merge_stepwise_output below. +: "${LOSS_REDUCTION:=token_mean}" +: "${GRPO_NORM_BY_STD:=false}" +: "${USE_KL_LOSS:=false}" +# Keep overlong trajectories in the batch under Arctic: upstream +# arctic_platform.rl.utils.batch.split_dict raises IndexError if the surviving +# batch shrinks below num_workers, so we can't afford to drop samples. Trajectories +# that overflow still get a zero reward, which is the right learning signal. +: "${APPLY_OVERLONG_FILTERING:=false}" +: "${MICRO_FWD_BATCH_PER_GPU:=1}" +: "${MICRO_TRAIN_BATCH_PER_GPU:=1}" + +# ---------- Harbor HTTP endpoint (Arctic OpenAI shim serves here) ---------- +# The shim binds on HTTP_ENDPOINT_HOST and starts port-searching from +# HTTP_ENDPOINT_PORT (upstream removed the corresponding SkyRL config keys; +# these now flow through env vars so we don't add new config surface). +: "${HTTP_ENDPOINT_HOST:=0.0.0.0}" +: "${HTTP_ENDPOINT_PORT:=8000}" +export ARCTIC_HARBOR_SHIM_HOST="${HTTP_ENDPOINT_HOST}" +export ARCTIC_HARBOR_SHIM_PORT="${HTTP_ENDPOINT_PORT}" + +# Chat template — use SkyRL's bundled Qwen3 thinking template for parity with +# the FSDP Harbor baseline (see examples/train_integrations/harbor/run_codecontest.sh). +CHAT_TEMPLATE_PATH="${CHAT_TEMPLATE_PATH:-$(cd "$(dirname "$0")"/../../.. && pwd)/skyrl/train/utils/templates/qwen3_acc_thinking.jinja2}" + +if [[ ! -d "${DATA_DIR}/CodeContests" ]]; then + echo "!! CodeContests tasks not found at ${DATA_DIR}/CodeContests." >&2 + echo "!! Run: uv run --extra harbor examples/train_integrations/harbor/prepare_harbor_dataset.py --dataset open-thoughts/CodeContests --output_dir ${DATA_DIR}/CodeContests" >&2 + exit 1 +fi + +# ---------- preflight ---------- +# vLLM / Ray / OpenAI shim all bind ports on startup; if a previous run crashed +# they leak. Surface that up-front (with the exact cleanup command) instead of +# failing 5 minutes into engine init with a cryptic EADDRINUSE. +if command -v ss >/dev/null 2>&1 && ss -Hltn "sport = :${HTTP_ENDPOINT_PORT}" 2>/dev/null | grep -q .; then + echo "!! Port ${HTTP_ENDPOINT_PORT} (Arctic OpenAI shim) is already in use." >&2 + echo " Rerun with HTTP_ENDPOINT_PORT=, or clean stale workers:" >&2 + echo " pkill -9 -f 'InferenceWorker|DeepSpeedWorker|VLLM::|ArcticRLRayServerState|skyrl_entrypoint' && sleep 5" >&2 + exit 1 +fi +if command -v nvidia-smi >/dev/null 2>&1; then + busy=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | awk '$1 > 1024' | wc -l) + if [[ "${busy}" -gt 0 ]]; then + echo "!! WARNING: ${busy}/8 GPU(s) already have >1 GiB memory used." >&2 + echo " If a previous Arctic run crashed, free them with:" >&2 + echo " pkill -9 -f 'InferenceWorker|DeepSpeedWorker|VLLM::|ArcticRLRayServerState|skyrl_entrypoint' && sleep 5" >&2 + echo " Continuing anyway; vLLM will fail if it can't fit the ${MODEL} weights." >&2 + fi +fi + +echo "==> Launching Arctic RL + Harbor CodeContests" +echo " RUN_NAME=${RUN_NAME}" +echo " MODEL=${MODEL} (served as '${SERVED_MODEL_NAME}')" +echo " NUM_GPUS=${NUM_GPUS} colocate=${COLOCATE} zero_stage=${ZERO_STAGE}" +echo " train_batch=${TRAIN_BATCH_SIZE} x ${N_SAMPLES_PER_PROMPT} samples/prompt" +echo " context: max_model_len=${MAX_MODEL_LEN} max_prompt=${MAX_PROMPT_LENGTH} max_generate=${MAX_GENERATE_LENGTH}" +echo " knobs: apply_overlong_filtering=${APPLY_OVERLONG_FILTERING} max_concurrency=${MAX_CONCURRENCY} traj/sec=${TRAJECTORIES_PER_SECOND}" +echo " HTTP endpoint: ${HTTP_ENDPOINT_HOST}:${HTTP_ENDPOINT_PORT}" +echo " Trials dir: ${TRIALS_DIR}" +echo " Log dir: ${LOG_DIR}" +echo + +# Driver: same shape as the stock arctic_rl launchers +# (integrations/arctic_rl/examples/run_gsm8k_grpo_4gpu.sh) so no changes to +# pyproject.toml are needed. +# - ``--isolated`` bypasses the project lock: arctic-inference[vllm] needs +# vLLM 0.18 + torch 2.10, which conflict with the fsdp/megatron pins in +# ``[tool.uv.sources]`` (vLLM 0.23 / torch 2.11). +# - ``--extra skyrl-train`` pulls the base training deps (ray, deepspeed, +# fastapi, wandb, ...); ``--extra harbor`` layers harbor[daytona,modal] +# on top. Harbor doesn't pin vLLM/torch so it composes cleanly. +# - ``--with arctic-platform arctic-inference[vllm] liger-kernel`` provides +# the arctic stack. +# - ``--with 'transformers==4.57.6'`` is exact-pinned (not ``<5``): raylet +# re-spawns workers via ``bash -c``, and the shell would parse the +# unquoted ``<5`` as a redirect from fd 5 and kill the worker. +# - ``--with "flash-attn@"`` overrides SkyRL's default +# torch-2.11 flash-attn wheel to match arctic-inference's torch 2.10 pin. +# SkyRL's uv+ray plugin replays this exact invocation on every Ray worker via +# ``py_executable``, so workers get the same env automatically — no +# ``uv pip install`` bootstrap needed. +FLASH_ATTN_WHL="https://github.com/lesj0610/flash-attention/releases/download/v2.8.3-cu12-torch2.10-cp312/flash_attn-2.8.3%2Bcu12torch2.10cxx11abiTRUE-cp312-cp312-linux_x86_64.whl" + +exec uv run --isolated --extra skyrl-train --extra harbor \ + --with arctic-platform \ + --with 'arctic-inference[vllm]' \ + --with liger-kernel \ + --with 'transformers==4.57.6' \ + --with "flash-attn@${FLASH_ATTN_WHL}" \ + -m examples.train_integrations.harbor.entrypoints.main_harbor \ + trainer.override_entrypoint=integrations.arctic_rl.harbor_entrypoint \ + \ + data.train_data="${TRAIN_DATA}" \ + data.val_data="${EVAL_DATA}" \ + trainer.policy.model.path="${MODEL}" \ + generator.inference_engine.served_model_name="${SERVED_MODEL_NAME}" \ + trainer.max_prompt_length="${MAX_PROMPT_LENGTH}" \ + generator.sampling_params.max_generate_length="${MAX_GENERATE_LENGTH}" \ + trainer.algorithm.max_seq_len="${MAX_MODEL_LEN}" \ + \ + trainer.export_path="${EXPORTS_DIR}" \ + trainer.ckpt_path="${CKPTS_DIR}" \ + trainer.log_path="${LOG_DIR}" \ + trainer.resume_mode=null \ + \ + trainer.algorithm.advantage_estimator=grpo \ + trainer.algorithm.loss_reduction="${LOSS_REDUCTION}" \ + trainer.algorithm.grpo_norm_by_std="${GRPO_NORM_BY_STD}" \ + trainer.algorithm.use_kl_loss="${USE_KL_LOSS}" \ + trainer.algorithm.use_kl_in_reward=false \ + trainer.update_epochs_per_batch=1 \ + trainer.epochs="${EPOCHS}" \ + trainer.eval_batch_size=8 \ + trainer.eval_before_train="${EVAL_BEFORE_TRAIN}" \ + trainer.eval_interval="${EVAL_INTERVAL}" \ + trainer.train_batch_size="${TRAIN_BATCH_SIZE}" \ + trainer.policy_mini_batch_size="${MINI_BATCH_SIZE}" \ + trainer.micro_forward_batch_size_per_gpu="${MICRO_FWD_BATCH_PER_GPU}" \ + trainer.micro_train_batch_size_per_gpu="${MICRO_TRAIN_BATCH_PER_GPU}" \ + trainer.policy.optimizer_config.lr="${LR}" \ + \ + trainer.strategy=fsdp2 \ + trainer.placement.colocate_all=false \ + trainer.placement.policy_num_gpus_per_node="${NUM_GPUS}" \ + trainer.placement.ref_num_gpus_per_node="${NUM_GPUS}" \ + generator.inference_engine.num_engines="${NUM_GPUS}" \ + generator.inference_engine.tensor_parallel_size=1 \ + generator.inference_engine.gpu_memory_utilization="${VLLM_GPU_MEM_UTIL}" \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.enforce_eager="${VLLM_ENFORCE_EAGER}" \ + generator.inference_engine.engine_init_kwargs.chat_template="${CHAT_TEMPLATE_PATH}" \ + generator.inference_engine.engine_init_kwargs.max_model_len="${MAX_MODEL_LEN}" \ + generator.inference_engine.engine_init_kwargs.enable_log_requests=false \ + generator.batched=false \ + generator.step_wise_trajectories=true \ + generator.merge_stepwise_output=true \ + generator.n_samples_per_prompt="${N_SAMPLES_PER_PROMPT}" \ + generator.eval_n_samples_per_prompt=1 \ + generator.apply_overlong_filtering="${APPLY_OVERLONG_FILTERING}" \ + generator.rate_limit.enabled="${ENABLE_RATE_LIMITING}" \ + generator.rate_limit.trajectories_per_second="${TRAJECTORIES_PER_SECOND}" \ + generator.rate_limit.max_concurrency="${MAX_CONCURRENCY}" \ + \ + trainer.arctic_rl.colocate="${COLOCATE}" \ + trainer.arctic_rl.zero_stage="${ZERO_STAGE}" \ + trainer.arctic_rl.use_liger="${USE_LIGER}" \ + trainer.arctic_rl.use_arctic_inference="${USE_ARCTIC_INFERENCE}" \ + trainer.arctic_rl.use_zorro="${USE_ZORRO}" \ + trainer.arctic_rl.cuda_ipc_weight_sync="${CUDA_IPC_WEIGHT_SYNC}" \ + trainer.arctic_rl.low_memory_weight_sync="${LOW_MEM_WEIGHT_SYNC}" \ + trainer.arctic_rl.attn_implementation="${ATTN_IMPL}" \ + trainer.arctic_rl.vllm_enforce_eager="${VLLM_ENFORCE_EAGER}" \ + trainer.arctic_rl.vllm_max_model_len="${MAX_MODEL_LEN}" \ + trainer.arctic_rl.vllm_enable_prefix_caching=true \ + \ + harbor_trial_config.trials_dir="${TRIALS_DIR}" \ + harbor_trial_config.agent.kwargs.model_info.max_input_tokens="${MAX_PROMPT_LENGTH}" \ + harbor_trial_config.agent.kwargs.model_info.max_output_tokens="${MAX_GENERATE_LENGTH}" \ + \ + trainer.logger=wandb \ + trainer.project_name="${WANDB_PROJECT}" \ + trainer.run_name="${RUN_NAME}" \ + \ + "$@" diff --git a/integrations/arctic_rl/harbor_entrypoint.py b/integrations/arctic_rl/harbor_entrypoint.py new file mode 100644 index 0000000000..c2fde62171 --- /dev/null +++ b/integrations/arctic_rl/harbor_entrypoint.py @@ -0,0 +1,284 @@ +"""``main_harbor`` override entrypoint: routes Harbor recipes through Arctic RL. + +User flow (same for any existing Harbor launcher — no other changes needed): + + python -m examples.train_integrations.harbor.entrypoints.main_harbor \\ + trainer.override_entrypoint=integrations.arctic_rl.harbor_entrypoint \\ + + +Everything reused verbatim from upstream: + * ``ArcticRLExp`` / ``ArcticPPOTrainer`` — training routed to Arctic server. + * ``HarborGenerator`` / ``HarborTaskDataset`` — Harbor Trials + reward. + +Only new piece: ``ArcticInferenceEngineAdapter`` (openai_bridge.py) exposes an +OpenAI-compatible HTTP shim in front of the same Arctic ``ReplicaPool`` that +serves training rollouts, so Harbor's LiteLLM agent (inside the sandbox) sees +the up-to-date policy after every weight sync. + +Do NOT ``from __future__ import annotations``: SkyRL's +``build_nested_dataclass`` introspects ``dataclasses.fields(cls)[i].type`` and +expects a concrete class, not a string. With future-annotations, nested +configs (e.g. ``generator: HarborGeneratorConfig``) get parsed as raw dicts +and ``SkyRLTrainConfig.__post_init__`` then fails with +``AttributeError: 'dict' object has no attribute 'max_input_length'``. +""" + +import os +import sys +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + +import ray +import yaml +from arctic_platform.rl import ArcticRLClientConfig, create_arctic_rl_client +from loguru import logger + +from examples.train_integrations.harbor.dataset import HarborTaskDataset # noqa: E402 +from examples.train_integrations.harbor.entrypoints.main_harbor import ( # noqa: E402 + HARBOR_DEFAULT_CONFIG, + HarborGeneratorConfig, + _deep_merge, +) +from examples.train_integrations.harbor.harbor_generator import ( + HarborGenerator, # noqa: E402 +) +from skyrl.train.config import SkyRLTrainConfig +from skyrl.train.utils import validate_cfg +from skyrl.train.utils.utils import prepare_runtime_environment + +from . import ArcticPPOTrainer +from .config import ArcticRLTrainerConfig, ArcticTrainerConfig, build_rl_config +from .entrypoint import ArcticRLExp +from .openai_bridge import ArcticInferenceEngineAdapter +from .trainer import _ArcticInferenceEngineStub + +# Shim bind address env vars. Kept out of the config schema so this recipe +# adds no new SkyRL config fields. Full table lives in the arctic README. +_SHIM_HOST_ENV = "ARCTIC_HARBOR_SHIM_HOST" +_SHIM_PORT_ENV = "ARCTIC_HARBOR_SHIM_PORT" + + +@dataclass +class ArcticHarborSkyRLConfig(SkyRLTrainConfig): + """SkyRL config with Harbor's ``harbor_trial_config`` and Arctic RL's + ``trainer.arctic_rl`` sub-config both slotted in. + """ + + harbor_trial_config: Dict[str, Any] = field(default_factory=dict) + generator: HarborGeneratorConfig = field(default_factory=HarborGeneratorConfig) + trainer: ArcticTrainerConfig = field(default_factory=ArcticTrainerConfig) + + +class ArcticHarborExp(ArcticRLExp): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._openai_adapter: Optional[ArcticInferenceEngineAdapter] = None + + # -- data --------------------------------------------------------------- + def get_train_dataset(self): + ds = HarborTaskDataset(data_files=self.cfg.data.train_data) + assert len(ds) >= self.cfg.trainer.train_batch_size, ( + f"HarborTaskDataset size ({len(ds)}) < train_batch_size " + f"({self.cfg.trainer.train_batch_size})" + ) + return ds + + def get_eval_dataset(self): + if self.cfg.trainer.eval_interval > 0 and self.cfg.data.val_data: + return HarborTaskDataset(data_files=self.cfg.data.val_data) + return None + + # -- generator --------------------------------------------------------- + def get_generator(self, cfg, tokenizer, inference_engine_client): + # HarborGenerator looks up ``inference_engine_client.get_endpoint_url()`` + # to bake the LiteLLM ``api_base`` into every Trial's config; we pass + # our adapter (which exposes the OpenAI shim we just started). + return HarborGenerator( + generator_cfg=cfg.generator, + harbor_cfg=cfg.harbor_trial_config, + inference_engine_client=inference_engine_client, + tokenizer=tokenizer, + max_seq_len=cfg.trainer.algorithm.max_seq_len, + ) + + def get_trainer( + self, + cfg, + tracker, + tokenizer, + train_dataset, + eval_dataset, + inference_engine_client, + generator, + colocate_pg, + ): + return ArcticPPOTrainer( + cfg=cfg, + tracker=tracker, + tokenizer=tokenizer, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + inference_engine_client=inference_engine_client, + generator=generator, + colocate_pg=colocate_pg, + arctic_client=self.arctic_client, + ) + + # -- assembly ---------------------------------------------------------- + def _setup_trainer(self): + logger.info( + "Setting up ArcticHarbor trainer (GPU work delegated to Arctic RL server, " + "rollouts served over an OpenAI shim for Harbor's LiteLLM agent)" + ) + os.makedirs(self.cfg.trainer.export_path, exist_ok=True) + os.makedirs(self.cfg.trainer.ckpt_path, exist_ok=True) + + arl_cfg = getattr(self.cfg.trainer, "arctic_rl", None) + cfg_colocate = bool(getattr(arl_cfg, "colocate", False)) if arl_cfg else False + + ie_cfg = self.cfg.generator.inference_engine + assert ie_cfg.served_model_name is not None, ( + "generator.inference_engine.served_model_name must be set — Harbor " + "sends model=hosted_vllm/." + ) + default_max_tokens = int( + getattr(self.cfg.generator.sampling_params, "max_generate_length", 4096) + or 4096 + ) + + shim_host = os.environ.get(_SHIM_HOST_ENV, "0.0.0.0") + shim_port = int(os.environ.get(_SHIM_PORT_ENV, "8000")) + chat_template_path = None + engine_init_kwargs = getattr(ie_cfg, "engine_init_kwargs", None) or {} + if hasattr(engine_init_kwargs, "get"): + chat_template_path = engine_init_kwargs.get("chat_template") + chat_template = None + if chat_template_path: + with open(str(chat_template_path)) as f: + chat_template = f.read() + + # HarborGenerator bakes the shim URL into LiteLLM's ``api_base`` at + # __init__ time — start the shim first so the URL is real. + self._openai_adapter = ArcticInferenceEngineAdapter( + arctic_client=self.arctic_client, + tokenizer=self.tokenizer, + model_name=ie_cfg.served_model_name, + default_max_tokens=default_max_tokens, + host=shim_host, + port=shim_port, + chat_template=chat_template, + ) + self._openai_adapter.spin_up_http_endpoint() + + # Trainer-side inference-engine handle: same stub upstream uses, so + # colocated sleep/wake still routes to the Arctic RL server. + ie_stub = _ArcticInferenceEngineStub( + client=self.arctic_client, colocate=cfg_colocate + ) + tracker = self.get_tracker() + generator = self.get_generator( + self.cfg, self.tokenizer, self._openai_adapter + ) + trainer = self.get_trainer( + cfg=self.cfg, + tracker=tracker, + tokenizer=self.tokenizer, + train_dataset=self.train_dataset, + eval_dataset=self.eval_dataset, + inference_engine_client=ie_stub, + generator=generator, + colocate_pg=self.colocate_pg, + ) + trainer.build_models() + if cfg_colocate: + trainer.colocate_all = True + return trainer + + +@ray.remote(num_cpus=1) +def skyrl_entrypoint( + cfg: ArcticHarborSkyRLConfig, + reconnect_config: Optional[ArcticRLClientConfig] = None, + server_state: Optional[Any] = None, +): + exp = ArcticHarborExp( + cfg, reconnect_config=reconnect_config, server_state=server_state + ) + exp.run() + + +def main() -> None: + # ``main_harbor.py`` already stripped its ``override_entrypoint=`` peek + # arg before dispatching here, but SkyRL's strict config parser would + # still reject any leftover ``trainer.override_entrypoint=…`` in argv, so + # defensively drop it here too. + argv = [ + a for a in sys.argv[1:] if not a.startswith("trainer.override_entrypoint=") + ] + cfg = ArcticHarborSkyRLConfig.from_cli_overrides(argv) + with open(HARBOR_DEFAULT_CONFIG) as f: + defaults = yaml.safe_load(f) + cfg.harbor_trial_config = _deep_merge(defaults, cfg.harbor_trial_config) + + if cfg.trainer.arctic_rl is None: + cfg.trainer.arctic_rl = ArcticRLTrainerConfig() + validate_cfg(cfg) + if cfg.trainer.algorithm.max_seq_len is None: + raise ValueError( + "trainer.algorithm.max_seq_len must be set — Harbor uses it to " + "truncate responses." + ) + + # ``create_arctic_rl_client`` starts a Ray cluster on first call, so + # pre-init BEFORE ``ray.init``; ``ignore_reinit_error=True`` reuses that + # cluster on the follow-up ``ray.init``. + rl_config = build_rl_config(cfg) + logger.info("Pre-initializing ArcticRL jobs…") + pre_client = create_arctic_rl_client(rl_config) + reconnect_cfg = pre_client.reconnect_config() + server_state = ( + pre_client.get_server_state() if rl_config.comm_protocol == "ray" else None + ) + logger.info( + "ArcticRL ready — training=%s, sample=%s, log_prob=%s", + pre_client.training_job_id, + pre_client.sampling_job_id, + pre_client.log_prob_job_id, + ) + + # Ray workers reconstruct the driver task, so they need our imports on + # PYTHONPATH and the sandbox-provider creds forwarded. + env_vars = prepare_runtime_environment(cfg) + for prefix in ( + "ARCTIC_", + "WANDB_", + "DAYTONA_", + "MODAL_", + "E2B_", + "RUNLOOP_", + "GKE_", + "OPENAI_", + "HF_", + ): + env_vars.update( + {k: v for k, v in os.environ.items() if k.startswith(prefix)} + ) + _repo_root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + _existing_pp = env_vars.get("PYTHONPATH") or os.environ.get("PYTHONPATH", "") + env_vars["PYTHONPATH"] = _repo_root + ( + os.pathsep + _existing_pp if _existing_pp else "" + ) + runtime_env = {"env_vars": env_vars} + + ray.init(num_gpus=0, runtime_env=runtime_env, ignore_reinit_error=True) + ray.get( + skyrl_entrypoint.options(runtime_env=runtime_env).remote( + cfg, reconnect_config=reconnect_cfg, server_state=server_state + ) + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/arctic_rl/openai_bridge.py b/integrations/arctic_rl/openai_bridge.py new file mode 100644 index 0000000000..9550720862 --- /dev/null +++ b/integrations/arctic_rl/openai_bridge.py @@ -0,0 +1,627 @@ +"""OpenAI HTTP shim in front of ``arctic_client.generate()``. + +Harbor's terminus-2 agent lives inside a sandbox VM and only reaches the model +over HTTP (LiteLLM's ``hosted_vllm/`` → ``POST /v1/chat/completions``). +Arctic's sampling side is a ``ReplicaPool`` reached in-process via +``arctic_client.generate(prompts, sampling_params)``. This module bridges the +two by: + + 1. Building a small FastAPI app that exposes ``/v1/chat/completions`` and + ``/v1/completions``. + 2. Serving it in a daemon thread via ``uvicorn.Server`` on a port we own. + 3. Implementing SkyRL's ``InferenceEngineInterface`` so ``HarborGenerator`` + can call ``get_endpoint_url()`` / ``finish_session()`` on the same object + with no wrapper. + +Only the OpenAI paths call into Arctic; every other interface method is a +no-op or forwards to the underlying ``arctic_client`` (weight-sync goes +through the Arctic RL server directly, not through this shim). +""" + +import asyncio +import logging +import socket +import threading +import time +import uuid +from typing import Any, Dict, List, Optional, Tuple + +from transformers import PreTrainedTokenizerBase + +from skyrl.backends.skyrl_train.inference_servers.base import ( + InferenceEngineInput, + InferenceEngineInterface, + InferenceEngineOutput, +) + +logger = logging.getLogger(__name__) + + +# Fields the arctic sampling worker forwards to ``vllm.SamplingParams`` +# (see ``arctic_inference.server.worker.InferenceWorker.generate``). Anything +# not in this set is dropped from the payload so we don't hit vLLM's strict +# kwargs check when LiteLLM adds a new knob. +_VLLM_SAMPLING_KEYS = frozenset( + { + "n", + "best_of", + "temperature", + "top_p", + "top_k", + "min_p", + "max_tokens", + "min_tokens", + "seed", + "stop", + "stop_token_ids", + "presence_penalty", + "frequency_penalty", + "repetition_penalty", + "logprobs", + "prompt_logprobs", + "logit_bias", + "include_stop_str_in_output", + "skip_special_tokens", + "spaces_between_special_tokens", + "ignore_eos", + } +) + + +def _map_finish_reason(reason: Optional[str]) -> str: + # vLLM emits stop|length|abort; OpenAI accepts stop|length|tool_calls|... + # Fold abort -> length (token-budget cutoff), unknown -> stop. + if reason == "length" or reason == "abort": + return "length" + return "stop" + + +def _build_sampling_params_from_openai( + body: Dict[str, Any], + *, + default_max_tokens: int, +) -> Dict[str, Any]: + out: Dict[str, Any] = {k: v for k, v in body.items() if k in _VLLM_SAMPLING_KEYS} + + # ``max_completion_tokens`` is OpenAI's newer name for ``max_tokens``. + if "max_completion_tokens" in body and "max_tokens" not in out: + out["max_tokens"] = body["max_completion_tokens"] + out.setdefault("max_tokens", int(default_max_tokens)) + + # LiteLLM's /chat/completions ``logprobs`` is a bool; vLLM wants an int + # (top-k). Harbor doesn't consume logprobs, so drop the bool form. + if isinstance(out.get("logprobs"), bool): + out["logprobs"] = 1 if out["logprobs"] else None + + return {k: v for k, v in out.items() if v is not None} + + +def _openai_logprobs_from_arctic( + arctic_result: Dict[str, Any], + token_ids: List[int], + tokenizer: PreTrainedTokenizerBase, +) -> Optional[Dict[str, Any]]: + """Translate ``arctic_client.generate()``'s logprob payload to OpenAI + ``chat.logprobs``. Assumes vLLM's canonical + ``list[dict[token_id, {"logprob": float, ...}]]`` shape from + ``SamplingOutput.logprobs``; anything else falls through to zeros. + """ + if not token_ids: + return None + + raw = arctic_result.get("logprobs") or [] + content: List[Dict[str, Any]] = [] + for i, tid in enumerate(token_ids): + tok = tokenizer.decode([tid], skip_special_tokens=False) + lp = 0.0 + if i < len(raw) and isinstance(raw[i], dict) and raw[i]: + entry = raw[i].get(tid) or next(iter(raw[i].values())) + if isinstance(entry, dict): + lp = float(entry.get("logprob", 0.0)) + content.append( + { + "token": tok, + "logprob": lp, + "bytes": list(tok.encode("utf-8", errors="ignore")), + "top_logprobs": [], + } + ) + return {"content": content} + + +def _bad_request(message: str, code: int = 400) -> Dict[str, Any]: + return {"error": {"message": message, "type": "BadRequest", "code": code}} + + +def _server_error(message: str) -> Dict[str, Any]: + return {"error": {"message": message, "type": "InternalServerError", "code": 500}} + + +def _pick_port(start_port: int, host: str = "0.0.0.0") -> Tuple[int, socket.socket]: + """Reserve a free TCP port at or above ``start_port`` (bind-and-listen).""" + end = start_port + 128 + port = start_port + while port < end: + try: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind((host, port)) + sock.listen(1) + return port, sock + except OSError: + port += 1 + raise RuntimeError(f"No available port in [{start_port}, {end}) on {host}") + + +class ArcticInferenceEngineAdapter(InferenceEngineInterface): + """Adapts ``arctic_client`` to SkyRL's ``InferenceEngineInterface``. + + Serves an OpenAI-compatible HTTP endpoint (FastAPI + uvicorn) that + Harbor's LiteLLM agent (inside the Daytona/Modal sandbox) can reach. + Weight-sync and other server-side operations bypass this class — + ``arctic_client`` owns those paths and the vLLM engine the shim forwards to. + """ + + def __init__( + self, + *, + arctic_client: Any, + tokenizer: PreTrainedTokenizerBase, + model_name: str, + default_max_tokens: int = 4096, + host: str = "0.0.0.0", + port: int = 8000, + chat_template: Optional[str] = None, + ) -> None: + self._arctic_client = arctic_client + self.tokenizer = tokenizer + self._model_name = model_name + self._default_max_tokens = int(default_max_tokens) + self._host = host + self._preferred_port = int(port) + self._custom_chat_template = chat_template + + self._server: Optional[Any] = None # uvicorn.Server + self._server_thread: Optional[threading.Thread] = None + self._reserved_sock: Optional[socket.socket] = None + self._bound_port: Optional[int] = None + + # ------------------------------------------------------------------ # + # InferenceEngineInterface — required properties # + # ------------------------------------------------------------------ # + + @property + def model_name(self) -> str: + return self._model_name + + def get_endpoint_url(self) -> str: + if self._bound_port is None: + raise RuntimeError( + "OpenAI shim has not been started; call spin_up_http_endpoint() first." + ) + # Harbor's sandbox reaches this via the host's routable IP; bind on + # 0.0.0.0 (host arg) but advertise 127.0.0.1 by default so single-host + # runs work without egress reconfig. Override with `.set_advertised_host()` + # if the sandbox is on a different machine. + advertised_host = self._advertised_host() + return f"http://{advertised_host}:{self._bound_port}" + + def _advertised_host(self) -> str: + if self._host in ("0.0.0.0", "::"): + # Prefer explicit override; else localhost (sandbox → host loopback + # via port forwarding, or localhost if driver + sandbox share a node). + import os + + return os.environ.get("ARCTIC_HARBOR_SHIM_ADVERTISED_HOST", "127.0.0.1") + return self._host + + # ------------------------------------------------------------------ # + # InferenceEngineInterface — OpenAI paths # + # ------------------------------------------------------------------ # + + def _messages_to_prompt(self, messages: List[Dict[str, Any]]) -> str: + return self.tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=False, + chat_template=self._custom_chat_template, + ) + + def _count_tokens(self, text: str) -> int: + try: + return len(self.tokenizer.encode(text, add_special_tokens=False)) + except Exception: + return 0 + + async def chat_completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]: + body = request_payload["json"] if "json" in request_payload else request_payload + messages = body.get("messages", []) + if not messages: + return _bad_request("empty messages in /v1/chat/completions") + + prompt_text = self._messages_to_prompt(messages) + prompt_ids = self.tokenizer.encode(prompt_text, add_special_tokens=False) + + sampling_params = _build_sampling_params_from_openai( + body, default_max_tokens=self._default_max_tokens + ) + # Arctic's InferenceWorker.generate returns choice.outputs[0] only, so + # n>1 is expressed by duplicating the prompt and stitching results + # back into ``choices``. + n = int(sampling_params.pop("n", 1) or 1) + prompts = [prompt_text] * n + + want_logprobs = bool(body.get("logprobs")) or ("logprobs" in sampling_params) + if want_logprobs: + # Arctic's ``generate`` accepts vLLM sampling kwargs; asking for the + # top-1 logprob populates per-token ``logprobs`` in the result dict. + sampling_params.setdefault("logprobs", 1) + + try: + results = await self._arctic_client.generate( + prompts=prompts, sampling_params=sampling_params + ) + except Exception as exc: + logger.error("chat_completion arctic_client.generate failed: %s", exc) + return _server_error(f"arctic_client.generate failed: {exc}") + + if not results: + return _server_error("arctic_client.generate returned no results") + + prompt_tokens = int(results[0].get("prompt_len") or len(prompt_ids)) + completion_tokens = int( + sum(r.get("generation_len", 0) or 0 for r in results) + ) + choices = [] + for i, r in enumerate(results): + text = r.get("text", "") + token_ids = list(r.get("token_ids") or []) + if not token_ids and text: + token_ids = self.tokenizer.encode(text, add_special_tokens=False) + + # ``token_ids`` sits directly on the choice — matches vLLM's + # ``ChatCompletionResponseChoice``. LiteLLM's ``convert_dict_to_response`` + # folds non-``_CHOICES_FIELDS`` keys into ``provider_specific_fields``, + # which is where Harbor's ``_extract_token_ids`` reads them. + choice: Dict[str, Any] = { + "index": i, + "message": {"role": "assistant", "content": text}, + "finish_reason": _map_finish_reason(r.get("finish_reason")), + "logprobs": _openai_logprobs_from_arctic( + r, token_ids, self.tokenizer + ) + if want_logprobs + else None, + "token_ids": token_ids, + } + choices.append(choice) + + return { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": int(time.time()), + "model": self._model_name, + "choices": choices, + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + # vLLM emits the full prompt token IDs at the top level; Harbor's + # LiteLLM reads them via ``response.prompt_token_ids`` (attribute + # lookup on the pydantic ModelResponse — unknown JSON keys become + # attributes via ``__pydantic_extra__``). + "prompt_token_ids": prompt_ids, + } + + async def completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]: + body = request_payload["json"] if "json" in request_payload else request_payload + raw_prompt = body.get("prompt") + + prompts: List[str] + if isinstance(raw_prompt, str): + prompts = [raw_prompt] + elif isinstance(raw_prompt, list): + if not raw_prompt: + return _bad_request("empty prompt in /v1/completions") + if isinstance(raw_prompt[0], str): + prompts = list(raw_prompt) + elif isinstance(raw_prompt[0], int): + prompts = [self.tokenizer.decode(raw_prompt, skip_special_tokens=False)] + elif isinstance(raw_prompt[0], list): + prompts = [ + self.tokenizer.decode(p, skip_special_tokens=False) + for p in raw_prompt + ] + else: + return _bad_request( + f"unsupported prompt type {type(raw_prompt[0]).__name__}" + ) + else: + return _bad_request("missing `prompt` in /v1/completions") + + sampling_params = _build_sampling_params_from_openai( + body, default_max_tokens=self._default_max_tokens + ) + sampling_params.pop("n", None) + + try: + results = await self._arctic_client.generate( + prompts=prompts, sampling_params=sampling_params + ) + except Exception as exc: + logger.error("completion arctic_client.generate failed: %s", exc) + return _server_error(f"arctic_client.generate failed: {exc}") + + texts = [r.get("text", "") for r in results] + finish_reasons = [_map_finish_reason(r.get("finish_reason")) for r in results] + prompt_tokens = int(sum(r.get("prompt_len", 0) or 0 for r in results)) or sum( + self._count_tokens(p) for p in prompts + ) + completion_tokens = int( + sum(r.get("generation_len", 0) or 0 for r in results) + ) + return { + "id": "cmpl-" + uuid.uuid4().hex, + "object": "text_completion", + "created": int(time.time()), + "model": self._model_name, + "choices": [ + {"index": i, "text": t, "finish_reason": f, "logprobs": None} + for i, (t, f) in enumerate(zip(texts, finish_reasons)) + ], + "usage": { + "prompt_tokens": int(prompt_tokens), + "completion_tokens": int(completion_tokens), + "total_tokens": int(prompt_tokens + completion_tokens), + }, + } + + async def render_chat_completion( + self, request_payload: Dict[str, Any] + ) -> Dict[str, Any]: + """Apply the chat template + tokenize without generating. + + Only used on HarborGenerator's multimodal path; mirrors vLLM's + renderer output shape. + """ + body = request_payload["json"] if "json" in request_payload else request_payload + messages = body.get("messages", []) + prompt_text = self._messages_to_prompt(messages) + token_ids = self.tokenizer.encode(prompt_text, add_special_tokens=False) + return {"prompt": prompt_text, "prompt_token_ids": token_ids} + + # ------------------------------------------------------------------ # + # InferenceEngineInterface — SkyRL trainer paths (unused for Harbor) # + # ------------------------------------------------------------------ # + + async def generate( + self, + input_batch: InferenceEngineInput, + model: Optional[str] = None, + ) -> InferenceEngineOutput: + """Interface method — unused on the Harbor path (rollouts go over + HTTP), implemented so eval / dry-run callers still work. + """ + prompts = input_batch.get("prompts") or [] + prompt_texts: List[str] = [] + for prompt in prompts: + if isinstance(prompt, list): + prompt_texts.append( + self.tokenizer.apply_chat_template( + prompt, + add_generation_prompt=True, + tokenize=False, + chat_template=self._custom_chat_template, + ) + ) + else: + prompt_texts.append(str(prompt)) + + sampling_params = dict(input_batch.get("sampling_params") or {}) + sampling_params.setdefault("max_tokens", self._default_max_tokens) + + results = await self._arctic_client.generate( + prompts=prompt_texts, sampling_params=sampling_params + ) + + responses: List[str] = [] + response_ids: List[List[int]] = [] + stop_reasons: List[str] = [] + for r in results: + text = r.get("text", "") + token_ids = r.get("token_ids") or self.tokenizer.encode( + text, add_special_tokens=False + ) + responses.append(text) + response_ids.append(list(token_ids)) + stop_reasons.append( + "completed" + if _map_finish_reason(r.get("finish_reason")) == "stop" + else "length" + ) + + return InferenceEngineOutput( # type: ignore[typeddict-item] + responses=responses, + response_ids=response_ids, + stop_reasons=stop_reasons, + response_logprobs=None, + prompt_logprobs=None, + rollout_expert_indices=None, + ) + + # Weight sync flows through the Arctic RL server (arctic_client owns the + # DeepSpeed <-> vLLM handshake); the SkyRL trainer never calls these on + # this adapter — we return no-ops for interface completeness. + async def init_weight_update_communicator(self, init_info): # type: ignore[override] + return None + + async def update_named_weights(self, request): # type: ignore[override] + return None + + async def reset_prefix_cache(self, reset_running_requests: bool = False): + # Best-effort: arctic_client may expose this; if not, silently skip. + fn = getattr(self._arctic_client, "reset_prefix_cache", None) + if fn is None: + return None + try: + return await fn(reset_running_requests=reset_running_requests) + except TypeError: + return await fn() + + async def sleep(self, *args: Any, **kwargs: Any): + fn = getattr(self._arctic_client, "sleep_inference", None) + if fn is None: + return None + level = kwargs.get("level", 2) + return await fn(level=level) + + async def wake_up(self, *args: Any, **kwargs: Any): + fn = getattr(self._arctic_client, "wake_inference", None) + if fn is None: + return None + tags = kwargs.get("tags") + try: + return await fn(tags=tags) + except TypeError: + return await fn() + + async def pause_generation(self) -> None: + return None + + async def resume_generation(self) -> None: + return None + + async def finish_session(self, session_id: str) -> None: + # Arctic has no session concept; best-effort no-op. + return None + + async def get_world_size(self) -> Tuple[int, int]: + # Not consumed on the Harbor path (used by weight-sync setup only). + # Return (1, 1) so any accidental consumer gets a sane placeholder. + return (1, 1) + + async def teardown(self) -> None: + srv = self._server + if srv is not None: + try: + srv.should_exit = True + except Exception: + pass + thread = self._server_thread + if thread is not None and thread.is_alive(): + thread.join(timeout=5) + + # ------------------------------------------------------------------ # + # FastAPI + uvicorn lifecycle # + # ------------------------------------------------------------------ # + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + # uvicorn / socket / thread aren't picklable; strip before Ray serdes. + state["_server"] = None + state["_server_thread"] = None + state["_reserved_sock"] = None + return state + + def _build_app(self): + from fastapi import FastAPI, Request + from fastapi.responses import JSONResponse + + app = FastAPI(title="arctic-harbor-openai-shim") + + @app.get("/v1/models") + async def _models() -> Dict[str, Any]: + return { + "object": "list", + "data": [ + { + "id": self._model_name, + "object": "model", + "created": int(time.time()), + "owned_by": "arctic-rl", + } + ], + } + + @app.post("/v1/chat/completions") + async def _chat(request: Request) -> Any: + body = await request.json() + result = await self.chat_completion({"json": body}) + status = 400 if isinstance(result, dict) and "error" in result else 200 + return JSONResponse(result, status_code=status) + + @app.post("/v1/completions") + async def _completions(request: Request) -> Any: + body = await request.json() + result = await self.completion({"json": body}) + status = 400 if isinstance(result, dict) and "error" in result else 200 + return JSONResponse(result, status_code=status) + + @app.get("/health") + async def _health() -> Dict[str, str]: + return {"status": "ok"} + + return app + + def spin_up_http_endpoint(self) -> None: + if self._server_thread is not None and self._server_thread.is_alive(): + return + + # Reserve a port up front so we know the bound port before starting + # uvicorn (needed for get_endpoint_url() ordering). + port, sock = _pick_port(self._preferred_port, host=self._host) + # uvicorn creates its own socket; free ours before it binds. + sock.close() + self._bound_port = port + + import uvicorn + + app = self._build_app() + config = uvicorn.Config( + app, + host=self._host, + port=port, + log_level="warning", + access_log=False, + loop="asyncio", + ) + self._server = uvicorn.Server(config) + + def _run() -> None: + asyncio.run(self._server.serve()) + + self._server_thread = threading.Thread( + target=_run, name="arctic-harbor-openai-shim", daemon=True + ) + self._server_thread.start() + self._wait_ready(port) + logger.info( + "arctic-harbor OpenAI shim ready at http://%s:%s/v1 (model=%s)", + self._advertised_host(), + port, + self._model_name, + ) + + def _wait_ready(self, port: int, timeout_s: float = 60.0) -> None: + import urllib.error + import urllib.request + + url = f"http://127.0.0.1:{port}/health" + deadline = time.monotonic() + timeout_s + last_err: Optional[Exception] = None + while time.monotonic() < deadline: + try: + with urllib.request.urlopen(url, timeout=1) as resp: + if resp.status == 200: + return + except (urllib.error.URLError, OSError, ConnectionError) as exc: + last_err = exc + time.sleep(0.2) + raise RuntimeError( + f"arctic-harbor OpenAI shim on :{port} failed to become ready in {timeout_s}s " + f"(last error: {last_err!r})" + ) + + +__all__ = ["ArcticInferenceEngineAdapter"]