diff --git a/examples/dynamo/README.md b/examples/dynamo/README.md new file mode 100644 index 0000000000..067945e157 --- /dev/null +++ b/examples/dynamo/README.md @@ -0,0 +1,80 @@ +# Dynamo native-gRPC deployment requirements + +The recipes in this directory use Dynamo for inference and Prime only for the +trainer/orchestrator. Start from Dynamo's `sidecar_agg.yaml` or +`sidecar_disagg.yaml`, then apply the following RL overlay. The stock manifests +are serving examples and do not enable Prime worker discovery or vLLM's admin +control routes by themselves. + +## Frontend + +Set these variables on the Dynamo frontend and expose both container ports: + +```yaml +env: + - name: DYN_ENABLE_RL + value: "true" + - name: DYN_RL_PORT + value: "8001" +ports: + - name: http + containerPort: 8000 + - name: rl-discovery + containerPort: 8001 +``` + +The frontend Kubernetes Service must also map ports 8000 and 8001. Prime's +`base_url` targets 8000; `dynamo_discovery_url` targets 8001. + +## Every vLLM engine and sidecar pair + +The engine HTTP address published by discovery must be reachable from the +trainer, so bind vLLM to the pod network rather than loopback: + +```text +vllm-rs serve --host 0.0.0.0 --port 8000 --grpc-port 50051 -- \ + --worker-extension-cls prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker \ + +``` + +Install the matching Prime source in the engine image so Python can import the +worker extension. Set this environment variable on the `vllm-engine` +container to expose `/pause`, `/resume`, and `/collective_rpc`: + +```yaml +env: + - name: VLLM_SERVER_DEV_MODE + value: "1" +``` + +Set the following variables on each `dynamo-vllm-sidecar` container. `POD_IP` +must appear before `VLLM_HTTP_ENDPOINT` so Kubernetes expands it: + +```yaml +env: + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + - name: VLLM_HTTP_ENDPOINT + value: "http://$(POD_IP):8000" + - name: DYN_ENABLE_RL + value: "true" +``` + +Keep the existing `--grpc-endpoint 127.0.0.1:50051`: gRPC stays pod-local, +while discovery publishes the pod-reachable HTTP admin address. No per-worker +Kubernetes Service is required when trainer-to-pod networking is routable. + +After startup, `/v1/rl/workers` must return every expected worker with a +non-null `admin_base_url`, positive `world_size`, and no `error` before Prime is +launched. + +## Recipes + +- [`qwen3_06b_math`](qwen3_06b_math): single-GPU trainer and aggregate Dynamo + inference smoke test. +- [`qwen3_30b_Thinking`](qwen3_30b_Thinking): Qwen3-30B Thinking math with an + external prefill/decode deployment. +- [`glm52_fp8_r2e`](glm52_fp8_r2e): multi-node GLM-5.2 FP8 R2E training with a + separately managed DGD. diff --git a/examples/dynamo/glm52_fp8_r2e/README.md b/examples/dynamo/glm52_fp8_r2e/README.md new file mode 100644 index 0000000000..41932321ef --- /dev/null +++ b/examples/dynamo/glm52_fp8_r2e/README.md @@ -0,0 +1,128 @@ +# GLM-5.2 FP8 R2E with external Dynamo inference + +This three-step smoke recipe runs a distributed Prime trainer and orchestrator +against a separately managed Dynamo deployment serving +`zai-org/GLM-5.2-FP8`. Dynamo owns the frontend, vLLM engines, and one native +gRPC sidecar per engine group; Prime discovers the mutable engine control +surface through `/v1/rl/workers`. + +The recipe is split into trainer and orchestrator files because the external +Dynamo DGD and the multi-node trainer have independent lifecycles. It does not +add a Prime inference configuration or require Prime's launcher to manage the +DGD. + +## Reference topology + +| Component | Shape | GPUs | +|---|---|---:| +| Dynamo prefill | 2 nodes, DP4 x TP2 x PP1 x EP8 | 8 | +| Dynamo decode | 2 nodes, DP4 x TP2 x PP1 x EP8 | 8 | +| Prime trainer | 4 nodes, FSDP16 x CP4 x EP8 | 16 | + +The checked-in configuration targets this topology; it is not a claim that +every cluster can use these parallelism dimensions unchanged. The two +discovery records must report `prefill` and `backend` components with a +combined `world_size` of 16. If the DGD topology changes, update +`weight_broadcast.inference_world_size` in both TOML files to the atomic sum +returned by the same `/v1/rl/workers` response. + +The inference engines must load +`prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker`, expose vLLM's +admin routes, and run the version-matched `vllm-rs` and +`dynamo-vllm-sidecar` binaries described in [`../README.md`](../README.md). +For mutable GLM weight reloads, launch every vLLM rank with `--enforce-eager`. +When serving a snapshot path, set `--served-model-name +zai-org/GLM-5.2-FP8`. The tested GLM entrypoint also uses the `glm47` tool +parser, `glm45` reasoning parser, the model chat template, and complementary +NIXL `kv_producer`/`kv_consumer` roles for prefill/decode. +The trainer and every inference engine must also use a compatible NCCL +transport. Apply cluster-specific settings such as `NCCL_IB_DISABLE`, +`NCCL_SOCKET_IFNAME`, and the NCCL network plugin consistently on both sides; +do not force Socket on the trainer while allowing inference to select IB. +The filesystem rollout transport requires the orchestrator and all trainer +nodes to mount the same read-write shared output root. + +## Configure + +Initialize the R2E environment submodule and install its workspace package: + +```bash +git submodule update --init -- deps/research-environments +uv sync --package prime-rl --package r2e-gym-v1 +``` + +The taskset intentionally uses the current `r2e-gym-v1` default, +`PrimeIntellect/R2E-Gym-Subset-Verified`, instead of pinning an older dataset +override in this recipe. + +Replace the checked-in service names when the DGD and trainer use different +DNS names: + +- `model.client.base_url`: Dynamo OpenAI frontend on port 8000; +- `model.client.dynamo_discovery_url`: Dynamo RL discovery on port 8001; +- `weight_broadcast.host`: trainer rank zero, reachable from every inference + engine. + +The R2E harness uses Prime sandboxes. Configure the normal Prime credentials, +or replace the runtime with the sandbox backend used by your cluster. Model +and dataset caches should be shared across the trainer and inference nodes. + +Multi-turn affinity is not enabled merely by sending a session header. Start +the Dynamo frontend with `--router-session-affinity-ttl-secs ` (or +`DYN_ROUTER_SESSION_AFFINITY_TTL_SECS`) and choose an idle TTL longer than the +longest expected R2E turn gap. Prime maps each trajectory ID to the canonical +`X-Dynamo-Session-ID` header in `orchestrator.toml`. + +Verify both the model and the complete atomic worker snapshot before starting +Prime. A successful HTTP status alone is insufficient: + +```bash +MODEL=zai-org/GLM-5.2-FP8 +curl -fsS http://dynamo-frontend:8000/v1/models | + jq -e --arg model "$MODEL" '.data | any(.id == $model)' +curl -fsS http://dynamo-frontend:8001/v1/rl/workers | + jq -e --arg model "$MODEL" ' + .protocol_version == 1 and + (.workers | length == 2) and + (all(.workers[]; .model == $model and + ((.error // "") == "") and + (.instance_id != null) and + ((.admin_base_url // "") != ""))) and + ([.workers[].instance_id] | unique | length == 2) and + ([.workers[].admin_base_url] | unique | length == 2) and + ([.workers[] | select(.model == $model) | .component] | sort == ["backend", "prefill"]) and + ([.workers[] | select(.model == $model) | .world_size] | add == 16) + ' +``` + +## Run + +Launch the trainer on four 4-GPU nodes with the cluster's distributed runner. +For example, rank zero's rendezvous address can be passed to `torchrun` while +all ranks consume the same trainer file: + +```bash +uv run torchrun \ + --nnodes=4 --nproc-per-node=4 \ + --rdzv-backend=c10d --rdzv-endpoint="$TRAINER_RANK_ZERO:29501" \ + --node-rank="$NODE_RANK" \ + -m prime_rl.trainer.rl.train \ + @ examples/dynamo/glm52_fp8_r2e/trainer.toml \ + --output-dir /shared/glm52-dynamo-r2e/train +``` + +After trainer rank zero opens port 29500, launch the orchestrator once: + +```bash +uv run orchestrator \ + @ examples/dynamo/glm52_fp8_r2e/orchestrator.toml \ + --output-dir /shared/glm52-dynamo-r2e/train/run_0 +``` + +The gate succeeds when the first optimizer step completes, policy version 1 +settles on all 16 inference ranks through NCCL, and a later multi-turn rollout +completes without changing its Dynamo session assignment. Three steps are +required because finite NCCL runs skip broadcasts once +`step >= max_steps - 1`; this leaves step 1 as the first non-final broadcast +slot. The disabled post-batch zero-advantage filter keeps this small smoke run +from stalling on a homogeneous batch; enable it for a production training run. diff --git a/examples/dynamo/glm52_fp8_r2e/orchestrator.toml b/examples/dynamo/glm52_fp8_r2e/orchestrator.toml new file mode 100644 index 0000000000..74a0c0cf60 --- /dev/null +++ b/examples/dynamo/glm52_fp8_r2e/orchestrator.toml @@ -0,0 +1,60 @@ +max_steps = 3 +batch_size = 2 +group_size = 2 +seq_len = 32768 +max_inflight_episodes = 2 +max_off_policy_steps = 1 +tasks_per_minute = 1 + +[model] +name = "zai-org/GLM-5.2-FP8" + +[model.client] +base_url = ["http://dynamo-frontend:8000/v1"] +dynamo_discovery_url = "http://dynamo-frontend:8001" +wait_for_ready_timeout = 7200 + +[model.client.extra_headers_from_state] +X-Dynamo-Session-ID = "trajectory_id" + +[tokenizer] +name = "zai-org/GLM-5.2-FP8" + +[renderer] +name = "glm-5.1" +clear_thinking = false + +[train.sampling] +temperature = 1.0 +max_completion_tokens = 2048 +extra_body = { chat_template_kwargs = { clear_thinking = false } } + +[[train.source]] +name = "r2e" +group_size = 2 +serve.pool = { type = "static", num_workers = 2 } +env.taskset = { id = "r2e-gym-v1" } +env.agent.max_turns = 64 +env.agent.max_input_tokens = 30720 +env.agent.max_output_tokens = 16384 +env.agent.max_total_tokens = 32768 +env.agent.timeout = { setup = 600, rollout = 1800, finalize = 300, scoring = 600 } +env.agent.harness = { id = "bash", edit = true } +env.agent.runtime = { type = "prime", labels = ["glm52-dynamo"], cpu = 4, creates_per_min = 2 } + +[[post_batch_filters]] +type = "zero_advantage" +enforce = false + +[weight_broadcast] +type = "nccl" +host = "trainer-0.trainer-headless" +port = 29500 +timeout = 12000 +inference_world_size = 16 + +[rollout_transport] +type = "filesystem" + +[log] +level = "debug" diff --git a/examples/dynamo/glm52_fp8_r2e/trainer.toml b/examples/dynamo/glm52_fp8_r2e/trainer.toml new file mode 100644 index 0000000000..d7476a734d --- /dev/null +++ b/examples/dynamo/glm52_fp8_r2e/trainer.toml @@ -0,0 +1,46 @@ +max_steps = 3 +dist_timeout_seconds = 12000 + +[model] +name = "zai-org/GLM-5.2-FP8" +seq_len = 32768 +impl = "custom" +attn = "flash_attention_2" +dp_replicate = 1 +cp = 4 +ep = 8 +optimization_dtype = "bfloat16" +reduce_dtype = "bfloat16" +moe_router_dtype = "float32" +optim_cpu_offload = true +fused_lm_head_token_chunk_size = 1024 + +[model.ac] +freq = 1 + +[model.ac_offloading] +max_inflight_activations = 1 + +[tokenizer] +name = "zai-org/GLM-5.2-FP8" + +[optim] +type = "sign_sgd" +lr = 1e-6 +weight_decay = 0.0 + +[scheduler] +type = "constant" + +[weight_broadcast] +type = "nccl" +host = "0.0.0.0" +port = 29500 +timeout = 12000 +inference_world_size = 16 + +[rollout_transport] +type = "filesystem" + +[log] +level = "debug" diff --git a/examples/dynamo/qwen3_06b_math/README.md b/examples/dynamo/qwen3_06b_math/README.md new file mode 100644 index 0000000000..dcfae8c576 --- /dev/null +++ b/examples/dynamo/qwen3_06b_math/README.md @@ -0,0 +1,45 @@ +# Qwen3 0.6B math with external Dynamo inference + +This four-step smoke recipe runs the Prime trainer and orchestrator locally while generation is served by an already-running Dynamo frontend, vLLM sidecar, and vLLM engine. Prime does not launch local inference, so the configuration intentionally has no `[inference]` block. + +## Prerequisites + +Use this version-matched native-gRPC source set: + +- Dynamo `feat/dyn-pi-sidecar-v2-review-001` at `836fe81012` +- vLLM `feat/dyn-pi-sidecar-v2-review-001` at `e56ee21b2c` +- Prime `feat/dyn-pi-sidecar-v2-review-001` + +Build `vllm-rs` and `dynamo-vllm-sidecar` from those revisions into the same +runtime image. For Kubernetes, start from Dynamo's +`examples/backends/vllm/deploy/sidecar_agg.yaml`; its adjacent `README.md` +documents the paired-binary image. Then apply the required Prime RL discovery +and admin overlay in [`../README.md`](../README.md). This contract requires both +native gRPC and the Dynamo `/v1/rl/workers` endpoint; a standard Python-only +vLLM worker is not compatible. + +Install the math environment: + +```bash +prime env install primeintellect/math-env +``` + +Start an aggregated DP1 Dynamo deployment for `Qwen/Qwen3-0.6B`. The orchestrator waits for both model publication and worker discovery. These requests are useful diagnostics: + +```bash +curl http://127.0.0.1:8000/v1/models +curl http://127.0.0.1:8001/v1/rl/workers +``` + +The checked-in URLs assume Dynamo is reachable from the trainer through localhost, as in a shared dev pod. For a remote DGD, replace both URLs with its frontend services. Also replace `weight_broadcast.host` with a trainer hostname or IP reachable from every sidecar; localhost is not valid across pods or nodes. + +`inference_world_size` must equal the sum of `world_size` in one `/v1/rl/workers` response. This recipe assumes one aggregated DP1 engine and therefore uses `1`. + +## Run + +```bash +uv run rl @ examples/dynamo/qwen3_06b_math/rl.toml \ + --output-dir outputs/dynamo-qwen3-06b-math +``` + +The run is successful when four optimizer steps complete, the verifier reports math rewards, weight versions advance after each update, and the Dynamo workers remain healthy. diff --git a/examples/dynamo/qwen3_06b_math/rl.toml b/examples/dynamo/qwen3_06b_math/rl.toml new file mode 100644 index 0000000000..d26bf43ba2 --- /dev/null +++ b/examples/dynamo/qwen3_06b_math/rl.toml @@ -0,0 +1,43 @@ +max_steps = 4 +seq_len = 2048 + +[model] +name = "Qwen/Qwen3-0.6B" + +[deployment] +type = "single_node" +num_train_gpus = 1 +num_infer_gpus = 0 + +[weight_broadcast] +type = "nccl" +host = "127.0.0.1" +inference_world_size = 1 + +[trainer] + +[orchestrator] +batch_size = 32 +group_size = 4 + +[orchestrator.model.client] +base_url = ["http://127.0.0.1:8000/v1"] +dynamo_discovery_url = "http://127.0.0.1:8001" +wait_for_ready_timeout = 1800 + +[orchestrator.model.client.extra_headers_from_state] +X-Session-ID = "trajectory_id" +X-Dynamo-Session-ID = "trajectory_id" + +[orchestrator.renderer] +name = "auto" +thinking_retention = "all" + +[orchestrator.train.sampling] +max_completion_tokens = 2048 + +[[orchestrator.train.source]] +name = "math" +env.taskset = { id = "math-env-v1", dataset_name = "openai/gsm8k", dataset_subset = "main" } +env.agent.harness = { id = "null" } +env.agent.runtime = { type = "subprocess" } diff --git a/examples/dynamo/qwen3_30b_Thinking/README.md b/examples/dynamo/qwen3_30b_Thinking/README.md new file mode 100644 index 0000000000..d01b182882 --- /dev/null +++ b/examples/dynamo/qwen3_30b_Thinking/README.md @@ -0,0 +1,53 @@ +# Qwen3 30B Thinking math with external Dynamo inference + +This four-step scale recipe runs a two-GPU Prime trainer against an external 1-prefill/1-decode Dynamo deployment serving `Qwen/Qwen3-30B-A3B-Thinking-2507`. Prime launches no inference process; Dynamo owns the frontend, sidecars, and vLLM engines. + +The two-GPU trainer uses BF16 optimization and reduction plus CPU optimizer +offload. Leaving Prime's FP32 optimization default in place exceeds two GB200s +when Adam allocates its first-step state; larger production recipes use the +existing eight-GPU training shape instead. + +The same model is used by the existing public examples `qwen30b_math`, `qwen30b_swe`, `multinode/rl.toml`, and `multinode/sft.toml`. Those examples remain the source of truth for larger training and non-Dynamo deployment settings; this recipe adds only the external-Dynamo client shape. + +## Prerequisites + +Use this version-matched native-gRPC source set: + +- Dynamo `feat/dyn-pi-sidecar-v2-review-001` at `836fe81012` +- vLLM `feat/dyn-pi-sidecar-v2-review-001` at `e56ee21b2c` +- Prime `feat/dyn-pi-sidecar-v2-review-001` + +Build `vllm-rs` and `dynamo-vllm-sidecar` from those revisions into the same +runtime image. For Kubernetes, start from Dynamo's +`examples/backends/vllm/deploy/sidecar_disagg.yaml` and change the model plus +parallelism/resources for the 30B topology; its adjacent `README.md` documents +the paired-binary image. Then apply the required Prime RL discovery and admin +overlay in [`../README.md`](../README.md). This contract requires both native +gRPC and the Dynamo `/v1/rl/workers` endpoint; a standard Python-only vLLM +worker is not compatible. + +Install the math environment: + +```bash +prime env install primeintellect/math-env +``` + +Start a Dynamo 1P/1D deployment and verify its public endpoints: + +```bash +curl http://127.0.0.1:8000/v1/models +curl http://127.0.0.1:8001/v1/rl/workers +``` + +The checked-in localhost URLs are for a colocated dev-pod run. For a DGD, replace them with the frontend services. Replace `weight_broadcast.host` with a trainer address reachable from every sidecar and allow the configured NCCL port through the network policy. + +`inference_world_size` must equal the sum of `world_size` in one `/v1/rl/workers` response. This recipe assumes 1P/1D with one rank per engine, for a total of `2`; TP, PP, or managed-DP topologies require the corresponding larger sum. + +## Run + +```bash +uv run rl @ examples/dynamo/qwen3_30b_Thinking/rl.toml \ + --output-dir outputs/dynamo-qwen3-30b-thinking-math +``` + +The run is successful when four optimizer steps complete, math rewards are emitted, all worker weight versions advance, and both prefill and decode workers remain healthy. diff --git a/examples/dynamo/qwen3_30b_Thinking/rl.toml b/examples/dynamo/qwen3_30b_Thinking/rl.toml new file mode 100644 index 0000000000..12c6995bf6 --- /dev/null +++ b/examples/dynamo/qwen3_30b_Thinking/rl.toml @@ -0,0 +1,56 @@ +max_steps = 4 +seq_len = 2048 + +[model] +name = "Qwen/Qwen3-30B-A3B-Thinking-2507" + +[deployment] +type = "single_node" +num_train_gpus = 2 +num_infer_gpus = 0 + +[weight_broadcast] +type = "nccl" +host = "127.0.0.1" +timeout = 1800 +inference_world_size = 2 + +[trainer.model] +impl = "custom" +attn = "flash_attention_3" +ep = 2 +optim_cpu_offload = true +optimization_dtype = "bfloat16" +reduce_dtype = "bfloat16" + +[trainer.model.ac] +freq = 1 + +[orchestrator] +batch_size = 2 +group_size = 2 +max_inflight_episodes = 2 +max_off_policy_steps = 0 + +[orchestrator.model.client] +base_url = ["http://127.0.0.1:8000/v1"] +dynamo_discovery_url = "http://127.0.0.1:8001" +wait_for_ready_timeout = 3600 + +[orchestrator.model.client.extra_headers_from_state] +X-Session-ID = "trajectory_id" +X-Dynamo-Session-ID = "trajectory_id" + +[orchestrator.renderer] +name = "qwen3" +enable_thinking = true + +[orchestrator.train.sampling] +temperature = 1.0 +max_completion_tokens = 2048 + +[[orchestrator.train.source]] +name = "math" +env.taskset = { id = "math-env-v1", dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", task = { math_verify_timeout = 60 } } +env.agent.harness = { id = "null" } +env.agent.runtime = { type = "subprocess" } diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index e9fdbf2a6e..a4958724e4 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -403,6 +403,9 @@ class ZeroAdvantageFilterConfig(BaseConfig): class FileSystemWeightBroadcastConfig(BaseConfig): type: Literal["filesystem"] = "filesystem" + inference_world_size: int | None = Field(None, ge=1) + """Expected inference ranks for Dynamo discovery completeness; unused by filesystem transfer itself.""" + class InMemoryWeightBroadcastConfig(BaseConfig): host: str = "localhost" @@ -572,6 +575,17 @@ def auto_setup_session_headers(self): self.model.client.extra_headers_from_state.setdefault("X-Session-ID", "trajectory_id") return self + @model_validator(mode="after") + def validate_dynamo_world_size(self): + if not self.model.client.is_dynamo: + return self + if ( + self.weight_broadcast.inference_world_size is None + or "inference_world_size" not in self.weight_broadcast.model_fields_set + ): + raise ValueError("Dynamo inference requires an explicit weight_broadcast.inference_world_size") + return self + @model_validator(mode="after") def auto_setup_prime_monitor_run_name(self): """Default ``prime_monitor.run_name`` to the W&B run name when monitoring diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index 8a4a02ae21..06b408f457 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -138,6 +138,9 @@ class SharedNCCLWeightBroadcastConfig(SharedInMemoryWeightBroadcastConfig): quantize_in_weight_transfer: bool = False """Use kernel-format FP8 quantized NCCL transfer for weight updates. When disabled, uses default HF checkpoint-format transfer.""" + inference_world_size: int | None = Field(None, ge=1) + """Expected inference ranks when inference is managed externally.""" + class SharedNIXLWeightBroadcastConfig(SharedInMemoryWeightBroadcastConfig): type: Literal["nixl"] = "nixl" @@ -148,10 +151,16 @@ class SharedNIXLWeightBroadcastConfig(SharedInMemoryWeightBroadcastConfig): session_id: str = "default" """ModelExpress session ID.""" + inference_world_size: int | None = Field(None, ge=1) + """Expected inference ranks when inference is managed externally.""" + class SharedFileSystemWeightBroadcastConfig(BaseConfig): type: Literal["filesystem"] = "filesystem" + inference_world_size: int | None = Field(None, ge=1) + """Expected inference ranks when inference is managed externally (e.g. Dynamo LoRA over filesystem).""" + SharedWeightBroadcastConfig: TypeAlias = Annotated[ SharedFileSystemWeightBroadcastConfig | SharedNCCLWeightBroadcastConfig | SharedNIXLWeightBroadcastConfig, @@ -323,16 +332,6 @@ def validate_deployment(self): ) return self - @model_validator(mode="after") - def validate_enough_devices_for_nccl(self): - if self.deployment.type == "single_node": - if self.trainer.weight_broadcast.type == "nccl": - if self.deployment.num_train_gpus + self.deployment.num_infer_gpus < 2: - raise ValueError( - "NCCL weight broadcast requires at least 2 GPUs to build the broadcast process group." - ) - return self - @model_validator(mode="after") def validate_quantize_in_weight_transfer(self): if not isinstance(self.weight_broadcast, SharedNCCLWeightBroadcastConfig): @@ -393,13 +392,18 @@ def auto_setup_weight_broadcast(self): "Set weight_broadcast.type = 'filesystem'." ) if self.weight_broadcast.type in ("nccl", "nixl"): - inference_world_size = self.inference.parallel.dp * self.inference.parallel.tp if self.inference else 1 + inference_world_size = ( + self.inference.parallel.dp * self.inference.parallel.tp + if self.inference + else self.weight_broadcast.inference_world_size + ) common_config = dict( host=self.weight_broadcast.host, port=self.weight_broadcast.port, timeout=self.weight_broadcast.timeout, - inference_world_size=inference_world_size, ) + if inference_world_size is not None: + common_config["inference_world_size"] = inference_world_size if self.weight_broadcast.type == "nccl": transport_config = dict( quantize_in_weight_transfer=self.weight_broadcast.quantize_in_weight_transfer, @@ -414,7 +418,9 @@ def auto_setup_weight_broadcast(self): self.orchestrator.weight_broadcast = orchestrator_config_type(**common_config, **transport_config) elif self.weight_broadcast.type == "filesystem": self.trainer.weight_broadcast = TrainerFileSystemWeightBroadcastConfig() - self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig() + self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig( + inference_world_size=self.weight_broadcast.inference_world_size + ) if self.inference is not None: self.inference.weight_broadcast = InferenceWeightBroadcastConfig(type=self.weight_broadcast.type) @@ -444,6 +450,19 @@ def auto_setup_rollout_transport(self): self.rollout_transport = self.trainer.rollout_transport return self + @model_validator(mode="after") + def validate_enough_devices_for_nccl(self): + if self.deployment.type != "single_node" or self.trainer.weight_broadcast.type != "nccl": + return self + if self.inference is None and self.weight_broadcast.inference_world_size is not None: + return self + local_inference_gpus = self.deployment.num_infer_gpus if self.inference is not None else 0 + if self.deployment.num_train_gpus + local_inference_gpus < 2: + raise ValueError( + "NCCL weight broadcast requires at least 2 local GPUs or an explicit external inference_world_size." + ) + return self + @model_validator(mode="after") def validate_eplb_requires_quantized_weight_transfer(self): if self.inference is None or not self.inference.enable_eplb: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 74f2a859b8..4243fa1cc5 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -1,6 +1,6 @@ import os from pathlib import Path -from typing import Annotated, Literal, TypeAlias +from typing import Annotated, Literal, Self, TypeAlias from pydantic import AfterValidator, Field, model_validator @@ -144,17 +144,33 @@ class ClientConfig(BaseConfig): admin_base_url: list[str] | None = None """Separate base URLs for admin operations (weight updates, health checks). When set, admin clients bypass routers and hit each server directly — used in disaggregated P/D deployments where the router must not handle admin traffic.""" + dynamo_discovery_url: str | None = None + """Dynamo discovery URL. When set, Prime discovers vLLM admin endpoints and per-engine world sizes from ``/v1/rl/workers`` instead of requiring ``admin_base_url`` entries.""" + elastic: ElasticConfig | None = None """Elastic inference pool config for DNS-based service discovery. When set, ``base_url`` is ignored and inference servers are discovered dynamically via DNS.""" router_url: str | None = None """vllm-router URL for load-aware inference routing. With elastic mode, inference requests go through the router while admin ops still hit discovered pods directly.""" + @model_validator(mode="after") + def validate_pool_mode(self) -> Self: + if self.dynamo_discovery_url is not None and self.admin_base_url is not None: + raise ValueError("dynamo_discovery_url cannot be combined with admin_base_url") + if self.dynamo_discovery_url is not None and self.elastic is not None: + raise ValueError("dynamo_discovery_url cannot be combined with elastic discovery") + return self + @property def is_elastic(self) -> bool: """Check if elastic mode is enabled.""" return self.elastic is not None + @property + def is_dynamo(self) -> bool: + """Check if Dynamo worker discovery is enabled.""" + return self.dynamo_discovery_url is not None + class LogConfig(BaseConfig): level: str = Field(default_factory=lambda: os.environ.get("PRIME_LOG_LEVEL", "info")) diff --git a/packages/prime-rl-configs/src/prime_rl/utils/validation.py b/packages/prime-rl-configs/src/prime_rl/utils/validation.py index 6dac87c12e..5d507424c8 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/validation.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/validation.py @@ -129,6 +129,9 @@ def propagate(shared_path: str, *targets: str) -> None: # [rollout_transport] → both sub-configs (host is launcher-injected for zmq multi-node). propagate("rollout_transport", "trainer.rollout_transport", "orchestrator.rollout_transport") + # The orchestrator validates external inference topology during construction. + propagate("weight_broadcast", "orchestrator.weight_broadcast") + # Top-level scalars. propagate("max_steps", "trainer.max_steps", "orchestrator.max_steps") propagate("seq_len", "trainer.model.seq_len", "orchestrator.seq_len") diff --git a/src/prime_rl/inference/vllm/ranks.py b/src/prime_rl/inference/vllm/ranks.py new file mode 100644 index 0000000000..8de43f4ace --- /dev/null +++ b/src/prime_rl/inference/vllm/ranks.py @@ -0,0 +1,39 @@ +def global_inference_rank( + *, + rank_offset: int, + data_parallel_index: int, + data_parallel_size: int, + worker_rank: int, + tensor_parallel_size: int, + pipeline_parallel_size: int, + inference_world_size: int, + prefill_context_parallel_size: int = 1, + engine_world_size: int | None = None, +) -> int: + """Map one vLLM worker to its rank in Prime's inference NCCL group.""" + model_parallel_size = tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size + logical_data_parallel_size = data_parallel_size + if engine_world_size is not None: + if engine_world_size <= 0 or engine_world_size % model_parallel_size: + raise ValueError( + f"engine world size {engine_world_size} is not divisible by model parallel size {model_parallel_size}" + ) + logical_data_parallel_size = engine_world_size // model_parallel_size + # Dense vLLM EngineCore processes retain their global DP index but rewrite + # data_parallel_size to one. MoE EngineCore processes preserve the logical + # size, so keep validating that value rather than masking bad discovery. + if data_parallel_size != 1 and data_parallel_size != logical_data_parallel_size: + raise ValueError( + f"data parallel size {data_parallel_size} does not match engine-derived size " + f"{logical_data_parallel_size}" + ) + if not 0 <= data_parallel_index < logical_data_parallel_size: + raise ValueError( + f"data parallel index {data_parallel_index} is outside logical data parallel size " + f"{logical_data_parallel_size}" + ) + + rank = rank_offset + data_parallel_index * model_parallel_size + worker_rank % model_parallel_size + if not 0 <= rank < inference_world_size: + raise ValueError(f"calculated inference rank {rank} is outside inference world size {inference_world_size}") + return rank diff --git a/src/prime_rl/inference/vllm/routed_experts.py b/src/prime_rl/inference/vllm/routed_experts.py index 5f73512fe5..a32d93ded5 100644 --- a/src/prime_rl/inference/vllm/routed_experts.py +++ b/src/prime_rl/inference/vllm/routed_experts.py @@ -1,11 +1,10 @@ from __future__ import annotations -from collections.abc import AsyncIterator +import io from typing import Any import numpy as np import pybase64 -from vllm.outputs import RequestOutput def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, Any] | None: @@ -33,16 +32,9 @@ def serialize_routed_experts(routed_experts: Any, start: int = 0) -> dict[str, A } -class RoutedExpertsCapture: - def __init__(self, generator: AsyncIterator[RequestOutput], start: int = 0): - self._generator = generator - self._start = start - self.routed_experts: dict[int, dict[str, Any]] = {} - - async def __aiter__(self): - async for request_output in self._generator: - for output in request_output.outputs: - encoded = serialize_routed_experts(getattr(output, "routed_experts", None), start=self._start) - if encoded is not None: - self.routed_experts[output.index] = encoded - yield request_output +def compact_vllm_routed_experts(encoded: str | None, start: int = 0) -> dict[str, Any] | None: + """Convert vLLM's base64 ``.npy`` payload to Prime's compact payload.""" + if encoded is None: + return None + array = np.load(io.BytesIO(pybase64.b64decode(encoded)), allow_pickle=False) + return serialize_routed_experts(array, start=start) diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index 28e2f572ff..e69b651b3f 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -143,11 +143,21 @@ async def init_broadcaster(request: Request): timeout = data.get("timeout") rank_offset = data.get("rank_offset") inference_world_size = data.get("inference_world_size") + engine_world_size = data.get("engine_world_size") quantize_in_weight_transfer = data.get("quantize_in_weight_transfer", False) session_id = data.get("session_id", "default") await engine_client(request).collective_rpc( "init_broadcaster", - args=(host, port, rank_offset, inference_world_size, timeout, quantize_in_weight_transfer, session_id), + args=( + host, + port, + rank_offset, + inference_world_size, + timeout, + quantize_in_weight_transfer, + session_id, + engine_world_size, + ), ) return {"status": "ok"} @@ -158,28 +168,25 @@ async def custom_init_app_state( args: Namespace, supported_tasks: tuple, ): - """ - Modifies init_app_state: - 1. Call the original init_app_state to set up standard state, including - vLLM 0.20's ``serving_tokens`` for ``/inference/v1/generate``. - 2. Replace ``serving_tokens`` with ``PrimeRlServingTokens`` so DP-rank - routing and ``routed_experts`` export survive the migration off the - legacy ``/v1/generate`` endpoint. - """ + """Initialize vLLM app state and install Prime's token response adapter.""" await init_app_state(engine_client, state, args, supported_tasks) state.liveness_timeout_seconds = args.liveness_timeout_seconds - # Swap in our ServingTokens subclass for /inference/v1/generate so the - # X-data-parallel-rank header and routed_experts response field — both - # used by prime-RL's renderer / router-replay paths — keep working. if "generate" in supported_tasks and state.serving_tokens is not None: from prime_rl.inference.vllm.serving_tokens import PrimeRlServingTokens upstream = state.serving_tokens - prime_serving = object.__new__(PrimeRlServingTokens) - prime_serving.__dict__.update(upstream.__dict__) - state.serving_tokens = prime_serving + state.serving_tokens = PrimeRlServingTokens( + upstream.engine_client, + upstream.models, + upstream.online_renderer, + request_logger=upstream.request_logger, + return_tokens_as_token_ids=upstream.return_tokens_as_token_ids, + force_no_detokenize=upstream.force_no_detokenize, + enable_prompt_tokens_details=True, + enable_log_outputs=upstream.enable_log_outputs, + ) import vllm.entrypoints.openai.api_server diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index e14a5ac83e..77caea00ce 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -1,53 +1,21 @@ -"""Prime-RL extensions to vLLM's `/inference/v1/generate` handler. - -vLLM ships a generic tokens-in / tokens-out handler at -``vllm.entrypoints.scale_out.token_in_token_out.serving.ServingTokens`` that covers -prefix-cache salting, lora dispatch, multimodal features, prompt logprobs, -priority, ``data_parallel_rank`` header routing and server-side ``max_tokens`` -defaulting. We subclass it for the bits still missing from the upstream handler: - -1. ``data_parallel_rank`` routing — read from the ``X-data-parallel-rank`` - header and forwarded to ``engine_client.generate``. Upstream ``ServingTokens`` - now does this too; we keep the equivalent path for the DP-replicated - inference servers prime-RL runs. - -2. Compact ``routed_experts`` export — when the engine emits routing - decisions, surface them as base64 raw-byte payloads without requiring a vLLM - source fork. - -3. Server-side ``max_tokens`` defaulting — upstream ``ServingTokens`` now applies - this itself (via ``GenerateRequest.is_sampling_param_provided`` + - ``get_max_tokens``); we keep an equivalent guard so callers that omit - ``max_tokens`` don't truncate at vLLM's 16-token ``SamplingParams`` default. - -Everything else (request/response schema, sampling params, error handling) -delegates to upstream so we track future vLLM changes for free. -""" +"""Small Prime extensions to vLLM's canonical token-in/token-out handler.""" from __future__ import annotations -from collections.abc import AsyncGenerator, AsyncIterable -from functools import cached_property +from collections.abc import AsyncGenerator from typing import Any from fastapi import Request -from vllm.entrypoints.openai.engine.protocol import ( - ErrorResponse, - PromptTokenUsageInfo, - RequestResponseMetadata, - UsageInfo, -) +from vllm.entrypoints.openai.engine.protocol import ErrorResponse, RequestResponseMetadata from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateRequest, GenerateResponse, GenerateResponseChoice, ) from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens -from vllm.entrypoints.serve.utils.api_utils import get_max_tokens from vllm.outputs import RequestOutput -from vllm.sampling_params import RequestOutputKind, SamplingParams -from prime_rl.inference.vllm.routed_experts import RoutedExpertsCapture +from prime_rl.inference.vllm.routed_experts import compact_vllm_routed_experts class PrimeRlGenerateResponseChoice(GenerateResponseChoice): @@ -56,255 +24,26 @@ class PrimeRlGenerateResponseChoice(GenerateResponseChoice): class PrimeRlGenerateResponse(GenerateResponse): choices: list[PrimeRlGenerateResponseChoice] - # Upstream ``GenerateResponse`` doesn't declare a ``usage`` field, so the - # parent ``ServingTokens.serve_tokens_full_generator`` constructs it and - # Pydantic silently drops it on serialization. Declare it here so the - # router can extract per-run token counts (and cached-prefix tokens) for - # platform billing — see https://github.com/PrimeIntellect-ai/router/pull/43. - usage: UsageInfo | None = None - - -class _GenerateRoutedExpertsCapture(RoutedExpertsCapture): - def post_process(self, response: GenerateResponse) -> PrimeRlGenerateResponse: - choices = [ - PrimeRlGenerateResponseChoice( - **choice.model_dump(exclude={"routed_experts"}), - routed_experts=self.routed_experts.get(choice.index), - ) - for choice in response.choices - ] - return PrimeRlGenerateResponse( - request_id=response.request_id, - choices=choices, - prompt_logprobs=response.prompt_logprobs, - kv_transfer_params=response.kv_transfer_params, - ) - - -class _FinalOutputCapture: - """Wraps a ``RequestOutput`` async generator to record the last yielded item. - - Needed so the response builder can construct a ``usage`` block from - ``final_res.prompt_token_ids`` / ``output.token_ids`` / ``num_cached_tokens`` - after delegating iteration to upstream. - """ - - def __init__(self, source: AsyncIterable[RequestOutput]) -> None: - # ``source`` may be any async-iterable — including - # ``_GenerateRoutedExpertsCapture``, which exposes the protocol via - # ``async def __aiter__`` (an async generator function) and has no - # ``__anext__`` method. Drive it through ``async for`` rather than - # poking ``__anext__`` directly so both shapes work. - self._source = source - self.final_res: RequestOutput | None = None - - async def __aiter__(self) -> AsyncGenerator[RequestOutput, None]: - async for item in self._source: - self.final_res = item - yield item - - -def _build_usage(final_res: RequestOutput) -> UsageInfo: - assert final_res.prompt_token_ids is not None - num_prompt_tokens = len(final_res.prompt_token_ids) - if final_res.encoder_prompt_token_ids is not None: - num_prompt_tokens += len(final_res.encoder_prompt_token_ids) - num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs) - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=num_generated_tokens, - total_tokens=num_prompt_tokens + num_generated_tokens, - ) - # Always emit cached tokens when vLLM reports any. Upstream gates this on - # ``enable_prompt_tokens_details`` (default False) for OpenAI-API compat, - # but ``/inference/v1/generate`` is prime-rl internal — the cache-discount - # billing pipeline always wants the cached subset surfaced. - if final_res.num_cached_tokens: - usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens) - return usage - - -async def _client_set_max_tokens(raw_request: Request | None) -> bool: - """Whether the inbound JSON body carried ``sampling_params.max_tokens``. - - ``GenerateRequest.sampling_params`` is parsed into a ``SamplingParams`` - instance, which means an unset ``max_tokens`` is indistinguishable from - an explicit ``max_tokens=16`` once the request reaches the handler — - both surface as ``sampling_params.max_tokens == 16``. We re-read the - cached body to recover that distinction. When we can't (no raw_request, - non-JSON body, or read error), pessimistically assume the client did - set it so we never clobber an explicit value. - """ - if raw_request is None: - return True - try: - body = await raw_request.json() - except Exception: - return True - if not isinstance(body, dict): - return True - sp = body.get("sampling_params") - return isinstance(sp, dict) and "max_tokens" in sp class PrimeRlServingTokens(ServingTokens): - """ServingTokens + DP-rank routing + compact routed experts + max_tokens defaulting.""" - - @cached_property - def _max_tokens_defaults(self) -> tuple[dict, int | None]: - """Server-side ``max_tokens`` defaulting inputs, mirroring upstream ``ServingTokens``. - - Computed lazily because ``custom_init_app_state`` swaps in this - subclass via ``object.__new__`` + ``__dict__.update`` (so our - ``__init__`` never runs). - """ - diff = self.model_config.get_diff_sampling_param() - mc = self.model_config - override = ( - diff.get("max_tokens") - if mc.generation_config not in ("auto", "vllm") - # Upstream uses ``getattr(..., {})`` directly. Defensive ``or {}`` - # in case a downstream caller ever sets the attribute to ``None`` - # (``getattr``'s default only fires when the attribute is missing, - # not when it exists with a ``None`` value). - else (getattr(mc, "override_generation_config", None) or {}).get("max_new_tokens") - ) - return diff, override + """Add KV handoff and Prime's compact routed-expert response encoding.""" async def serve_tokens( self, request: GenerateRequest, raw_request: Request | None = None, - ) -> PrimeRlGenerateResponse | ErrorResponse | AsyncGenerator[str, None]: - # Mirrors upstream ``ServingTokens.serve_tokens``. Diffs: - # (a) inject ``data_parallel_rank`` from the inbound header into - # ``engine_client.generate``; (b) default ``sampling_params.max_tokens`` - # to ``max_model_len - prompt_len`` when the caller didn't set it; and - # (c) dispatch to our overridden response builder so ``routed_experts`` - # makes it into the JSON. - error_check_ret = await self._check_model(request) - if error_check_ret is not None: - return error_check_ret - - if self.engine_client.errored: - raise self.engine_client.dead_error - - lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True) - model_name = self.models.model_name(lora_request) + ) -> GenerateResponse | ErrorResponse | AsyncGenerator[str, None]: + if request.kv_transfer_params is None: + return await super().serve_tokens(request, raw_request) - request_id = f"generate-tokens-{self._base_request_id(raw_request, request.request_id)}" - request_metadata = RequestResponseMetadata(request_id=request_id) - if raw_request: - raw_request.state.request_metadata = request_metadata + forwarded = request.model_copy(deep=True) + extra_args = dict(forwarded.sampling_params.extra_args or {}) + extra_args["kv_transfer_params"] = forwarded.kv_transfer_params + forwarded.sampling_params.extra_args = extra_args + return await super().serve_tokens(forwarded, raw_request) - # Build the engine input — features-aware (MM) or text-only fallback. - # Identical to upstream so we keep tracking it. - if features := request.features: - from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item - from vllm.inputs import mm_input - from vllm.multimodal.inputs import ( - MultiModalKwargsItem, - MultiModalKwargsItems, - PlaceholderRange, - ) - - mm_placeholders = { - modality: [PlaceholderRange(offset=p.offset, length=p.length) for p in ranges] - for modality, ranges in features.mm_placeholders.items() - } - mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {} - if features.kwargs_data is not None: - for modality, items in features.kwargs_data.items(): - mm_kwargs[modality] = [decode_mm_kwargs_item(item) if item is not None else None for item in items] - else: - for modality, hashes in features.mm_hashes.items(): - mm_kwargs[modality] = [None] * len(hashes) - engine_input = mm_input( - prompt_token_ids=request.token_ids, - mm_kwargs=MultiModalKwargsItems(mm_kwargs), - mm_hashes=features.mm_hashes, - mm_placeholders=mm_placeholders, - cache_salt=request.cache_salt, - ) - else: - (engine_input,) = await self.online_renderer.preprocess_completion( - request, - prompt_input=request.token_ids, - prompt_embeds=None, - skip_mm_cache=True, - ) - - sampling_params: SamplingParams = request.sampling_params - - # Upstream ``ServingTokens.serve_tokens`` parses ``request.kv_transfer_params`` - # but never threads it into the engine, so PD disagg never fires on - # ``/inference/v1/generate`` — decode receives an empty NIXL handshake - # and ends up re-prefilling the prompt locally (~100× slower under - # concurrency). Bridge it through ``sampling_params.extra_args`` so the - # engine's KV connector picks the params up. - # - # Upstream fix: https://github.com/vllm-project/vllm/pull/42644 — drop - # this block once we pin a vLLM version that includes it. - if request.kv_transfer_params is not None: - extra = sampling_params.extra_args or {} - extra["kv_transfer_params"] = request.kv_transfer_params - sampling_params.extra_args = extra - - # Server-side ``max_tokens`` defaulting — see module docstring. Upstream - # ``ServingTokens`` now does this too; kept here so callers that omit - # ``max_tokens`` don't get capped at vLLM's 16-token ``SamplingParams`` - # default. - if not await _client_set_max_tokens(raw_request): - diff_sp, override = self._max_tokens_defaults - sampling_params.max_tokens = get_max_tokens( - max_model_len=self.model_config.max_model_len, - max_tokens=None, - input_length=len(request.token_ids), - default_sampling_params=diff_sp, - override_max_tokens=override, - ) - - if self.force_no_detokenize: - sampling_params.detokenize = False - if request.stream: - sampling_params.output_kind = RequestOutputKind.DELTA - - self._log_inputs( - request_id, - engine_input, - params=sampling_params, - lora_request=lora_request, - ) - - trace_headers = None if raw_request is None else await self._get_trace_headers(raw_request.headers) - - result_generator = self.engine_client.generate( - engine_input, - sampling_params, - request_id, - lora_request=lora_request, - trace_headers=trace_headers, - priority=request.priority, - data_parallel_rank=self._get_data_parallel_rank(raw_request), - ) - - if request.stream: - # Streaming path: defer to upstream — prime-RL's renderer client - # only consumes the full response, so adding routed_experts to the - # streaming choice schema is unnecessary churn. - return self.serve_tokens_stream_generator( - request, - result_generator, - request_id, - model_name, - request_metadata, - ) - - return await self.serve_tokens_full_generator( - request, result_generator, request_id, model_name, request_metadata - ) - - async def serve_tokens_full_generator( # type: ignore[override] + async def serve_tokens_full_generator( self, request: GenerateRequest, result_generator: AsyncGenerator[RequestOutput, None], @@ -312,45 +51,25 @@ async def serve_tokens_full_generator( # type: ignore[override] model_name: str, request_metadata: RequestResponseMetadata, ) -> ErrorResponse | GenerateResponse: - # Capture routed_experts as vLLM streams request outputs, then post-process - # the final response into our GenerateResponse subclass so the encoded - # experts surface in the JSON. - capture: _GenerateRoutedExpertsCapture | None = None - if self.model_config.enable_return_routed_experts: - start = request.sampling_params.routed_experts_prompt_start - capture = _GenerateRoutedExpertsCapture( - result_generator, - start=start, - ) - result_generator = capture - - # Always capture the final ``RequestOutput`` so we can attach a - # ``usage`` block to the response. The router parses ``usage`` for - # per-run billing metrics; without it the cache-discount counter - # (``vllm_router_run_cached_prompt_tokens_total``) stays at zero. - final_capture = _FinalOutputCapture(result_generator) - result_generator = final_capture - response = await super().serve_tokens_full_generator( - request, result_generator, request_id, model_name, request_metadata + request, + result_generator, + request_id, + model_name, + request_metadata, ) - - if not isinstance(response, GenerateResponse): + if not isinstance(response, GenerateResponse) or not any( + choice.routed_experts is not None for choice in response.choices + ): return response - - if capture is not None: - response = capture.post_process(response) - elif not isinstance(response, PrimeRlGenerateResponse): - # Upgrade to the prime-rl subclass so the declared ``usage`` field - # actually surfaces in JSON (the parent class would drop it). - response = PrimeRlGenerateResponse( - request_id=response.request_id, - choices=[PrimeRlGenerateResponseChoice(**choice.model_dump()) for choice in response.choices], - prompt_logprobs=response.prompt_logprobs, - kv_transfer_params=response.kv_transfer_params, - ) - - if final_capture.final_res is not None: - response.usage = _build_usage(final_capture.final_res) - - return response + start = request.sampling_params.routed_experts_prompt_start or 0 + return PrimeRlGenerateResponse( + **response.model_dump(exclude={"choices"}), + choices=[ + PrimeRlGenerateResponseChoice( + **choice.model_dump(exclude={"routed_experts"}), + routed_experts=compact_vllm_routed_experts(choice.routed_experts, start=start), + ) + for choice in response.choices + ], + ) diff --git a/src/prime_rl/inference/vllm/worker/nccl.py b/src/prime_rl/inference/vllm/worker/nccl.py index 686bb856cf..468be61280 100644 --- a/src/prime_rl/inference/vllm/worker/nccl.py +++ b/src/prime_rl/inference/vllm/worker/nccl.py @@ -7,6 +7,7 @@ from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger +from prime_rl.inference.vllm.ranks import global_inference_rank from prime_rl.inference.vllm.worker.weight_transfer import ( load_weights_checkpoint_layerwise, load_weights_kernel, @@ -78,15 +79,22 @@ def __init__( self.communicator = PyNcclCommunicator(pg, device=device) @torch.no_grad() - def receive_state_dict(self): - """Receives the state dict of a model from the trainer master rank using NCCL communicator.""" + def receive_state_dicts( + self, + ) -> Generator[Generator[tuple[str, torch.Tensor], None, None], None, None]: + """Receive each trainer-broadcast state dict as a separate stream.""" logger.info("Receiving weights from trainer") num_state_dict_to_receive = receive_integer(self.communicator) logger.info(f"Receiving {num_state_dict_to_receive} layer state dicts") for layer_id in range(num_state_dict_to_receive): logger.info(f"Receiving state dict {layer_id + 1}/{num_state_dict_to_receive}") - for key, value in receive_state_dict(self.communicator): - yield key, value + yield receive_state_dict(self.communicator) + + @torch.no_grad() + def receive_state_dict(self): + """Receive trainer weights as one flat stream for kernel-format loading.""" + for state_dict in self.receive_state_dicts(): + yield from state_dict class NCCLWeightUpdateWorker(Worker): @@ -101,27 +109,37 @@ def init_broadcaster( timeout: int, quantize_in_weight_transfer: bool = False, session_id: str = "default", + engine_world_size: int | None = None, ) -> None: """Initialize the NCCL broadcast receiver. Args: rank_offset: Starting GPU offset for this server in the global inference group. inference_world_size: Total number of inference GPUs across all servers. + engine_world_size: Number of inference GPUs assigned to this server. """ del session_id self.quantize_in_weight_transfer = quantize_in_weight_transfer - # Use the worker's device index directly as the local rank. - # The previous dp_group-based computation broke in vLLM v1 multiprocess - # DP mode where each worker is a separate process with a singleton - # DP group (rank_in_group is always 0). - local_rank = self.device.index - global_rank_inference = rank_offset + local_rank + if engine_world_size is None: + global_rank_inference = rank_offset + self.device.index + else: + parallel_config = self.parallel_config + global_rank_inference = global_inference_rank( + rank_offset=rank_offset, + data_parallel_index=parallel_config.data_parallel_index, + data_parallel_size=parallel_config.data_parallel_size, + worker_rank=self.rank, + tensor_parallel_size=parallel_config.tensor_parallel_size, + pipeline_parallel_size=parallel_config.pipeline_parallel_size, + prefill_context_parallel_size=parallel_config.prefill_context_parallel_size, + inference_world_size=inference_world_size, + engine_world_size=engine_world_size, + ) logger.info( - f"Worker [local_rank={local_rank} rank_offset={rank_offset}] " + f"Worker [worker_rank={self.rank} rank_offset={rank_offset}] " f"-> [global_rank={global_rank_inference} inference_world_size={inference_world_size}]" ) - self.nccl_broadcast_receiver = NCCLWeightBroadcastReceiver( host=host, port=port, @@ -144,15 +162,14 @@ def update_weights_from_path(self, weight_dir: str) -> None: model = model_runner.model assert isinstance(model, Module) - state_iter = self.nccl_broadcast_receiver.receive_state_dict() if self.quantize_in_weight_transfer: - load_weights_kernel(model, state_iter) + load_weights_kernel(model, self.nccl_broadcast_receiver.receive_state_dict()) update_mla_absorbed_weights(model) - return - - load_weights_checkpoint_layerwise( - model, - state_iter, - self.model_runner.model_config, - self.vllm_config, - ) + else: + for state_iter in self.nccl_broadcast_receiver.receive_state_dicts(): + load_weights_checkpoint_layerwise( + model, + state_iter, + self.model_runner.model_config, + self.vllm_config, + ) diff --git a/src/prime_rl/inference/vllm/worker/nixl.py b/src/prime_rl/inference/vllm/worker/nixl.py index d6990211a1..5959b3a6bc 100644 --- a/src/prime_rl/inference/vllm/worker/nixl.py +++ b/src/prime_rl/inference/vllm/worker/nixl.py @@ -18,6 +18,7 @@ from vllm.config import set_current_vllm_config from vllm.logger import init_logger +from prime_rl.inference.vllm.ranks import global_inference_rank from prime_rl.inference.vllm.worker.weight_transfer import update_mla_absorbed_weights from prime_rl.trainer.rl.broadcast.nixl.agent import MemDesc, NixlAgent, make_agent_name, set_ucx_env_defaults from prime_rl.trainer.rl.broadcast.nixl.cuda_malloc_memory import ( @@ -96,9 +97,24 @@ def init_broadcaster( timeout: int, quantize_in_weight_transfer: bool = False, session_id: str = "default", + engine_world_size: int | None = None, ) -> None: - del inference_world_size, quantize_in_weight_transfer - global_rank = rank_offset + self.device.index + del quantize_in_weight_transfer + if engine_world_size is None: + global_rank = rank_offset + self.device.index + else: + parallel_config = self.parallel_config + global_rank = global_inference_rank( + rank_offset=rank_offset, + data_parallel_index=parallel_config.data_parallel_index, + data_parallel_size=parallel_config.data_parallel_size, + worker_rank=self.rank, + tensor_parallel_size=parallel_config.tensor_parallel_size, + pipeline_parallel_size=parallel_config.pipeline_parallel_size, + prefill_context_parallel_size=parallel_config.prefill_context_parallel_size, + inference_world_size=inference_world_size, + engine_world_size=engine_world_size, + ) server_url = f"{host}:{port}" set_ucx_env_defaults() self.nixl_agent = NixlAgent(make_agent_name("inference", global_rank)) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index d6b8d07e19..17e04aaf6c 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -76,7 +76,6 @@ from prime_rl.trainer.rl.broadcast.nixl.model_express import ModelExpressSession from prime_rl.transport import TrainingBatch, setup_training_batch_sender from prime_rl.utils.async_utils import EventLoopLagMonitor, EventLoopLagStats, safe_cancel -from prime_rl.utils.client import init_nccl_broadcast, init_nixl_broadcast from prime_rl.utils.config import to_toml_dict from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.logger import format_time, get_logger, setup_logger @@ -311,22 +310,20 @@ async def setup(self) -> None: get_logger().info(f"Initializing weight broadcast ({config.weight_broadcast})") if config.weight_broadcast.type == "nccl": - await init_nccl_broadcast( - self.policy_inference.admin_clients, - config.weight_broadcast.host, - config.weight_broadcast.port, - config.weight_broadcast.timeout, + await self.policy_inference.init_nccl_broadcast( + host=config.weight_broadcast.host, + port=config.weight_broadcast.port, + timeout=config.weight_broadcast.timeout, inference_world_size=config.weight_broadcast.inference_world_size, quantize_in_weight_transfer=config.weight_broadcast.quantize_in_weight_transfer, ) elif config.weight_broadcast.type == "nixl": - await init_nixl_broadcast( - self.policy_inference.admin_clients, - config.weight_broadcast.host, - config.weight_broadcast.port, - config.weight_broadcast.timeout, - config.weight_broadcast.inference_world_size, - config.weight_broadcast.session_id, + await self.policy_inference.init_nixl_broadcast( + host=config.weight_broadcast.host, + port=config.weight_broadcast.port, + timeout=config.weight_broadcast.timeout, + inference_world_size=config.weight_broadcast.inference_world_size, + session_id=config.weight_broadcast.session_id, ) self.model_express = ModelExpressSession( client=MxClient(server_url=f"{config.weight_broadcast.host}:{config.weight_broadcast.port}"), diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 374a1070f0..4929207c8a 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -45,6 +45,7 @@ async def setup_policy_inference_pool(*, config: OrchestratorConfig, tokenizer): train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=config.renderer, + expected_inference_world_size=config.weight_broadcast.inference_world_size, ) return renderer, inference_pool diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 2d35f122ef..d673cb0115 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -5,7 +5,7 @@ from collections.abc import Mapping from itertools import cycle from pathlib import Path -from typing import Protocol, runtime_checkable +from typing import Protocol, cast, runtime_checkable import httpx import verifiers.v1 as vf @@ -66,6 +66,30 @@ async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> N """Wait for inference pool to be ready.""" ... + async def init_nccl_broadcast( + self, + *, + host: str, + port: int, + timeout: int, + inference_world_size: int | None, + quantize_in_weight_transfer: bool, + ) -> None: + """Initialize the inference workers' NCCL receivers.""" + ... + + async def init_nixl_broadcast( + self, + *, + host: str, + port: int, + timeout: int, + inference_world_size: int, + session_id: str, + ) -> None: + """Initialize the inference workers' NIXL receivers.""" + ... + async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: """Update weights on all inference servers.""" ... @@ -122,6 +146,8 @@ def __init__( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, + *, + admin_clients: list[AsyncClient] | None = None, ): renderer_model_name = model_name if train_client_type == "renderer" else None self._train_clients = setup_clients( @@ -131,7 +157,7 @@ def __init__( renderer_model_name=renderer_model_name, ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) - self._admin_clients = setup_admin_clients(client_config) + self._admin_clients = setup_admin_clients(client_config) if admin_clients is None else admin_clients # When admin URLs bypass a router, also health-check the client-facing # (router) endpoint - it only starts serving once its workers are healthy. self._router_clients = ( @@ -175,9 +201,32 @@ async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> N ) await maybe_check_has_model(self._admin_clients, model_name, skip_model_check=self._skip_model_check) + async def init_nccl_broadcast( + self, + *, + host: str, + port: int, + timeout: int, + inference_world_size: int | None, + quantize_in_weight_transfer: bool, + ) -> None: + await init_nccl_broadcast( + self._admin_clients, + host=host, + port=port, + timeout=timeout, + inference_world_size=inference_world_size, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) + async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step) + async def init_nixl_broadcast( + self, *, host: str, port: int, timeout: int, inference_world_size: int, session_id: str + ) -> None: + await init_nixl_broadcast(self._admin_clients, host, port, timeout, inference_world_size, session_id) + async def score(self, token_ids: list[int]) -> list[float]: """Prefill-score ``token_ids`` under this pool's model (one logprob per token, 0.0 for the leading token). Delegates to the shared scorer.""" @@ -193,6 +242,7 @@ async def setup_inference_pool( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, + expected_inference_world_size: int | None = None, ) -> InferencePool: """Create an inference pool from config (static or elastic).""" if client_config.is_elastic: @@ -206,6 +256,18 @@ async def setup_inference_pool( renderer_config=renderer_config, ) + if client_config.is_dynamo: + from prime_rl.utils.dynamo import DynamoInferencePool + + return await DynamoInferencePool.from_config( + client_config, + model_name=model_name, + train_client_type=train_client_type, + eval_client_type=eval_client_type, + renderer_config=renderer_config, + expected_inference_world_size=cast(int, expected_inference_world_size), + ) + return StaticInferencePool( client_config, model_name=model_name, @@ -396,6 +458,8 @@ async def update_weights( weight_dir: Path | None, lora_name: str | None = None, step: int = 0, + *, + use_native_collective_rpc: bool = False, ) -> None: """Update weights on static inference servers. @@ -425,12 +489,18 @@ async def update_weights( nccl_ready_file.touch() logger.debug(f"Created NCCL_READY marker at {nccl_ready_file}") + update_path = "/collective_rpc" if use_native_collective_rpc else "/update_weights" + payload = ( + {"method": "update_weights_from_path", "args": [weight_dir_posix]} + if use_native_collective_rpc + else {"weight_dir": weight_dir_posix} + ) await asyncio.gather( *[ _admin_post( admin_client, - "/update_weights", - json={"weight_dir": weight_dir_posix}, + update_path, + json=payload, timeout_s=UPDATE_WEIGHTS_TIMEOUT_S, ) for admin_client in admin_clients @@ -515,6 +585,9 @@ async def init_nccl_broadcast( timeout: int, inference_world_size: int | None = None, quantize_in_weight_transfer: bool = False, + *, + engine_world_sizes: list[int] | None = None, + use_native_collective_rpc: bool = False, ) -> None: """Initialize NCCL broadcast on all inference servers. @@ -524,42 +597,65 @@ async def init_nccl_broadcast( """ logger = get_logger() + has_explicit_engine_world_sizes = engine_world_sizes is not None if inference_world_size is None: - inference_world_size = len(admin_clients) + if engine_world_sizes is not None: + inference_world_size = sum(engine_world_sizes) + else: + inference_world_size = len(admin_clients) logger.warning( f"inference_world_size not provided, defaulting to {inference_world_size} (one GPU per admin client)" ) - gpus_per_server = inference_world_size // len(admin_clients) + if engine_world_sizes is None: + if inference_world_size % len(admin_clients) != 0: + raise ValueError("inference_world_size must be divisible by the number of admin clients") + engine_world_sizes = [inference_world_size // len(admin_clients)] * len(admin_clients) + if len(engine_world_sizes) != len(admin_clients): + raise ValueError("one engine world size is required for each admin client") + rank_offsets = _rank_offsets(engine_world_sizes, inference_world_size) logger.info( f"Initializing NCCL broadcast: {len(admin_clients)} servers, " - f"inference_world_size={inference_world_size}, gpus_per_server={gpus_per_server}" + f"inference_world_size={inference_world_size}, engine_world_sizes={engine_world_sizes}" ) - async def _init_nccl_broadcast(admin_client: AsyncClient, rank_offset: int) -> None: - try: + async def _init_nccl_broadcast( + admin_client: AsyncClient, + rank_offset: int, + engine_world_size: int, + ) -> None: + payload = { + "host": host, + "port": port, + "rank_offset": rank_offset, + "inference_world_size": inference_world_size, + "timeout": timeout, + "quantize_in_weight_transfer": quantize_in_weight_transfer, + } + if has_explicit_engine_world_sizes: + payload["engine_world_size"] = engine_world_size + if use_native_collective_rpc: response = await admin_client.post( - "/init_broadcaster", - json={ - "host": host, - "port": port, - "rank_offset": rank_offset, - "inference_world_size": inference_world_size, - "timeout": timeout, - "quantize_in_weight_transfer": quantize_in_weight_transfer, - }, + "/collective_rpc", + json={"method": "init_broadcaster", "kwargs": payload}, ) + else: + response = await admin_client.post("/init_broadcaster", json=payload) + try: response.raise_for_status() - except httpx.HTTPStatusError as e: - if e.response.status_code == 404: + except httpx.HTTPStatusError as error: + if not use_native_collective_rpc and error.response.status_code == 404: logger.warning("The route /init_broadcaster does not exist. Skipping NCCL broadcast initialization.") return + raise await asyncio.gather( *[ - _init_nccl_broadcast(admin_client, client_num * gpus_per_server) - for client_num, admin_client in enumerate(admin_clients) + _init_nccl_broadcast(admin_client, rank_offset, engine_world_size) + for admin_client, rank_offset, engine_world_size in zip( + admin_clients, rank_offsets, engine_world_sizes, strict=True + ) ] ) @@ -571,31 +667,69 @@ async def init_nixl_broadcast( timeout: int, inference_world_size: int, session_id: str, + *, + engine_world_sizes: list[int] | None = None, + use_native_collective_rpc: bool = False, ) -> None: """Configure every vLLM worker for NIXL + ModelExpress pulls.""" - workers_per_server = inference_world_size // len(admin_clients) - - async def initialize(admin_client: AsyncClient, rank_offset: int) -> None: - await _admin_post( - admin_client, - "/init_broadcaster", - timeout_s=max(ADMIN_TIMEOUT_S, timeout), - json={ - "host": host, - "port": port, - "rank_offset": rank_offset, - "inference_world_size": inference_world_size, - "timeout": timeout, - "quantize_in_weight_transfer": False, - "session_id": session_id, - }, - ) + has_explicit_engine_world_sizes = engine_world_sizes is not None + if engine_world_sizes is None: + if inference_world_size % len(admin_clients) != 0: + raise ValueError("inference_world_size must be divisible by the number of admin clients") + engine_world_sizes = [inference_world_size // len(admin_clients)] * len(admin_clients) + if len(engine_world_sizes) != len(admin_clients): + raise ValueError("one engine world size is required for each admin client") + rank_offsets = _rank_offsets(engine_world_sizes, inference_world_size) + + async def initialize(admin_client: AsyncClient, rank_offset: int, engine_world_size: int) -> None: + payload = { + "host": host, + "port": port, + "rank_offset": rank_offset, + "inference_world_size": inference_world_size, + "timeout": timeout, + "quantize_in_weight_transfer": False, + "session_id": session_id, + } + if has_explicit_engine_world_sizes: + payload["engine_world_size"] = engine_world_size + if use_native_collective_rpc: + response = await admin_client.post( + "/collective_rpc", + json={"method": "init_broadcaster", "kwargs": payload}, + ) + response.raise_for_status() + else: + await _admin_post( + admin_client, + "/init_broadcaster", + timeout_s=max(ADMIN_TIMEOUT_S, timeout), + json=payload, + ) await asyncio.gather( - *[initialize(admin_client, index * workers_per_server) for index, admin_client in enumerate(admin_clients)] + *[ + initialize(admin_client, rank_offset, engine_world_size) + for admin_client, rank_offset, engine_world_size in zip( + admin_clients, rank_offsets, engine_world_sizes, strict=True + ) + ] ) +def _rank_offsets(engine_world_sizes: list[int], inference_world_size: int) -> list[int]: + if not engine_world_sizes or any(isinstance(size, bool) or size <= 0 for size in engine_world_sizes): + raise ValueError("engine world sizes must be positive integers") + if sum(engine_world_sizes) != inference_world_size: + raise ValueError("engine world sizes do not match inference_world_size") + offsets: list[int] = [] + offset = 0 + for world_size in engine_world_sizes: + offsets.append(offset) + offset += world_size + return offsets + + async def prefill_logprobs(openai: AsyncOpenAI, model: str, token_ids: list[int]) -> list[float]: """Prefill-score ``token_ids`` under ``model`` via ``/inference/v1/generate`` + ``prompt_logprobs`` (the prime-rl server-side extension in diff --git a/src/prime_rl/utils/dynamo.py b/src/prime_rl/utils/dynamo.py new file mode 100644 index 0000000000..b8ac801c77 --- /dev/null +++ b/src/prime_rl/utils/dynamo.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any, cast + +import httpx +from httpx import AsyncClient +from pydantic import BaseModel, ConfigDict, Field +from tenacity import AsyncRetrying, retry, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential + +from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.client import ( + LORA_LOAD_READ_TIMEOUT_S, + LORA_LOAD_TOTAL_TIMEOUT_S, + StaticInferencePool, + _is_retryable_lora_error, + _pause_engines, + _resume_engines, + init_nccl_broadcast, + init_nixl_broadcast, + setup_admin_clients, + update_weights, +) + +DYNAMO_RL_DISCOVERY_PROTOCOL_VERSION = 1 +DYNAMO_READINESS_REQUEST_TIMEOUT_S = 30.0 + + +class DiscoveredDynamoWorker(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + component: str = Field(min_length=1) + instance_id: int = Field(ge=0, strict=True) + model: str + admin_base_url: str = Field(min_length=1) + world_size: int = Field(gt=0, strict=True) + system_url: str | None = Field(None, min_length=1) + system_routes: tuple[str, ...] = () + + +class DynamoDiscoverySnapshot(BaseModel): + model_config = ConfigDict(extra="ignore") + + protocol_version: int = Field( + strict=True, + ge=DYNAMO_RL_DISCOVERY_PROTOCOL_VERSION, + le=DYNAMO_RL_DISCOVERY_PROTOCOL_VERSION, + ) + workers: list[dict[str, Any]] + + +class DynamoDiscoveryPending(ValueError): + """A well-formed discovery snapshot that is not ready yet.""" + + +def _is_retryable_dynamo_error(exception: BaseException) -> bool: + if isinstance(exception, httpx.HTTPStatusError): + return exception.response.status_code == 429 or exception.response.status_code >= 500 + return isinstance(exception, (DynamoDiscoveryPending, httpx.TransportError)) + + +def _parse_dynamo_workers(payload: object, model_name: str) -> tuple[DiscoveredDynamoWorker, ...]: + snapshot = DynamoDiscoverySnapshot.model_validate(payload) + workers = [] + for raw_worker in snapshot.workers: + if raw_worker.get("model") not in (None, model_name): + continue + if error := raw_worker.get("error"): + raise DynamoDiscoveryPending(f"Dynamo RL worker probe is not ready: {error}") + workers.append(DiscoveredDynamoWorker.model_validate(raw_worker)) + if not workers: + raise DynamoDiscoveryPending("Dynamo RL discovery returned no workers yet") + + identities = [(worker.component, worker.instance_id) for worker in workers] + admin_urls = [worker.admin_base_url for worker in workers] + if len(set(identities)) != len(identities): + raise ValueError("Dynamo RL discovery returned duplicate worker identities") + if len(set(admin_urls)) != len(admin_urls): + raise ValueError("Dynamo RL discovery returned duplicate admin endpoints") + lora_workers = [worker for worker in workers if "update/load_lora" in worker.system_routes] + if lora_workers and len(lora_workers) != len(workers): + raise ValueError("Dynamo RL discovery returned a partial update/load_lora capability snapshot") + if any(worker.system_url is None for worker in lora_workers): + raise ValueError("Dynamo RL discovery returned update/load_lora without a system_url") + return tuple(sorted(workers, key=lambda worker: (worker.component, worker.instance_id))) + + +def _setup_control_clients(urls: list[str]) -> list[AsyncClient]: + return [ + AsyncClient( + base_url=url.rstrip("/"), + limits=httpx.Limits(max_connections=4, max_keepalive_connections=1), + timeout=httpx.Timeout(None), + ) + for url in urls + ] + + +async def _load_lora_adapter(update_clients: list[AsyncClient], lora_name: str, lora_path: Path) -> None: + timeout = httpx.Timeout(connect=10.0, read=LORA_LOAD_READ_TIMEOUT_S, write=60.0, pool=10.0) + payload = { + "lora_name": lora_name, + "source": {"uri": lora_path.resolve().as_uri()}, + "load_inplace": True, + } + + @retry( + retry=retry_if_exception(_is_retryable_lora_error), + stop=stop_after_delay(LORA_LOAD_TOTAL_TIMEOUT_S) | stop_after_attempt(10), + wait=wait_exponential(multiplier=1, min=1, max=10), + reraise=True, + ) + async def load(update_client: AsyncClient) -> None: + response = await update_client.post("/v1/loras", json=payload, timeout=timeout) + response.raise_for_status() + result = response.json() + if isinstance(result, dict) and result.get("status") == "error": + raise RuntimeError(result.get("message") or "Dynamo LoRA update failed") + + await asyncio.gather(*(load(update_client) for update_client in update_clients)) + + +async def _wait_for_model(clients: list[AsyncClient], model_name: str, timeout: float) -> None: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + async with asyncio.timeout(timeout): + async for attempt in AsyncRetrying( + stop=stop_after_delay(timeout), + wait=wait_exponential(multiplier=0.1, min=0.1, max=1), + retry=retry_if_exception(_is_retryable_dynamo_error), + reraise=True, + ): + with attempt: + remaining = deadline - loop.time() + if remaining <= 0: + raise TimeoutError + request_timeout = httpx.Timeout(min(DYNAMO_READINESS_REQUEST_TIMEOUT_S, remaining)) + responses = await asyncio.gather( + *(client.get("/v1/models", timeout=request_timeout) for client in clients) + ) + for response in responses: + response.raise_for_status() + models = response.json().get("data", []) + if not any(model.get("id") == model_name for model in models): + raise DynamoDiscoveryPending(f"Dynamo frontend has not published model {model_name!r}") + + +class DynamoInferencePool(StaticInferencePool): + """Static request pool whose direct admin clients come from Dynamo discovery.""" + + def __init__(self, client_config: ClientConfig, workers: tuple[DiscoveredDynamoWorker, ...], **kwargs): + admin_clients = _setup_control_clients([worker.admin_base_url for worker in workers]) + super().__init__(client_config, admin_clients=admin_clients, **kwargs) + self._admin_world_sizes = [worker.world_size for worker in workers] + self._lora_update_clients: list[AsyncClient] = [] + if all("update/load_lora" in worker.system_routes for worker in workers): + system_urls = [worker.system_url for worker in workers if worker.system_url is not None] + self._lora_update_clients = _setup_control_clients(system_urls) + self._frontend_model_clients = setup_admin_clients(client_config) + self._readiness_deadline: float | None = None + + async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: + effective_timeout = self._wait_for_ready_timeout if timeout is None else timeout + loop = asyncio.get_running_loop() + deadline = ( + self._readiness_deadline + if timeout is None and self._readiness_deadline is not None + else loop.time() + effective_timeout + ) + remaining = max(0.0, deadline - loop.time()) + try: + async with asyncio.timeout(remaining): + await super().wait_for_ready(model_name, timeout=remaining) + if not self._skip_model_check: + await _wait_for_model( + self._frontend_model_clients, + model_name, + timeout=max(0.0, deadline - loop.time()), + ) + finally: + self._readiness_deadline = None + + async def init_nccl_broadcast( + self, + *, + host: str, + port: int, + timeout: int, + inference_world_size: int | None, + quantize_in_weight_transfer: bool, + ) -> None: + await init_nccl_broadcast( + self._admin_clients, + host=host, + port=port, + timeout=timeout, + inference_world_size=inference_world_size, + engine_world_sizes=self._admin_world_sizes, + quantize_in_weight_transfer=quantize_in_weight_transfer, + use_native_collective_rpc=True, + ) + + async def init_nixl_broadcast( + self, *, host: str, port: int, timeout: int, inference_world_size: int, session_id: str + ) -> None: + await init_nixl_broadcast( + self._admin_clients, + host, + port, + timeout, + inference_world_size, + session_id, + engine_world_sizes=self._admin_world_sizes, + use_native_collective_rpc=True, + ) + + async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: + if lora_name is None or weight_dir is None: + await update_weights( + self._admin_clients, + weight_dir, + lora_name=lora_name, + step=step, + use_native_collective_rpc=True, + ) + return + if not self._lora_update_clients: + raise RuntimeError("Dynamo LoRA update requires every worker to advertise system_url and update/load_lora") + try: + await _pause_engines(self._admin_clients, step=step) + await _load_lora_adapter(self._lora_update_clients, lora_name, weight_dir) + await _wait_for_model( + self._frontend_model_clients, + lora_name, + timeout=self._wait_for_ready_timeout, + ) + finally: + await _resume_engines(self._admin_clients) + + async def stop(self) -> None: + await super().stop() + await asyncio.gather( + *( + client.aclose() + for client in [*self._admin_clients, *self._lora_update_clients, *self._frontend_model_clients] + ) + ) + + @classmethod + async def from_config( + cls, + client_config: ClientConfig, + model_name: str, + expected_inference_world_size: int, + **kwargs, + ) -> DynamoInferencePool: + discovery_url = cast(str, client_config.dynamo_discovery_url).rstrip("/").removesuffix("/v1") + loop = asyncio.get_running_loop() + deadline = loop.time() + client_config.wait_for_ready_timeout + async with asyncio.timeout(client_config.wait_for_ready_timeout): + async with AsyncClient(timeout=httpx.Timeout(None)) as client: + workers: tuple[DiscoveredDynamoWorker, ...] = () + async for attempt in AsyncRetrying( + stop=stop_after_delay(client_config.wait_for_ready_timeout), + wait=wait_exponential(multiplier=0.1, min=0.1, max=1), + retry=retry_if_exception(_is_retryable_dynamo_error), + reraise=True, + ): + with attempt: + remaining = deadline - loop.time() + if remaining <= 0: + raise TimeoutError + response = await client.get( + f"{discovery_url}/v1/rl/workers", + timeout=httpx.Timeout(min(DYNAMO_READINESS_REQUEST_TIMEOUT_S, remaining)), + ) + response.raise_for_status() + workers = _parse_dynamo_workers(response.json(), model_name) + discovered_world_size = sum(worker.world_size for worker in workers) + if discovered_world_size != expected_inference_world_size: + raise DynamoDiscoveryPending( + "Dynamo RL discovery returned " + f"inference_world_size={discovered_world_size}; " + f"waiting for expected inference_world_size={expected_inference_world_size}" + ) + pool = cls(client_config, workers, model_name=model_name, **kwargs) + pool._readiness_deadline = deadline + return pool diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index 5d64c48794..13484094f1 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -26,6 +26,8 @@ ClientIdentity, PrefillScorer, client_identity, + init_nccl_broadcast, + init_nixl_broadcast, load_lora_adapter, setup_admin_clients, setup_clients, @@ -169,6 +171,31 @@ async def from_config( def update_model_name(self, model_name: str) -> None: self.model_name = model_name + async def init_nccl_broadcast( + self, + *, + host: str, + port: int, + timeout: int, + inference_world_size: int | None, + quantize_in_weight_transfer: bool, + ) -> None: + await init_nccl_broadcast( + list(self._admin_clients.values()), + host=host, + port=port, + timeout=timeout, + inference_world_size=inference_world_size, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) + + async def init_nixl_broadcast( + self, *, host: str, port: int, timeout: int, inference_world_size: int, session_id: str + ) -> None: + await init_nixl_broadcast( + list(self._admin_clients.values()), host, port, timeout, inference_world_size, session_id + ) + def _build_url(self, ip: str) -> str: return f"http://{ip}:{self.port}" diff --git a/tests/unit/inference/test_nccl_rank.py b/tests/unit/inference/test_nccl_rank.py new file mode 100644 index 0000000000..dc2955814d --- /dev/null +++ b/tests/unit/inference/test_nccl_rank.py @@ -0,0 +1,90 @@ +import pytest + +from prime_rl.inference.vllm.ranks import global_inference_rank + + +def test_dense_aggregate_uses_engine_size_when_vllm_rewrites_dp_size(): + actual = { + global_inference_rank( + rank_offset=0, + data_parallel_index=dp_index, + data_parallel_size=1, + worker_rank=tp_rank, + tensor_parallel_size=2, + pipeline_parallel_size=1, + inference_world_size=8, + engine_world_size=8, + ) + for dp_index in range(4) + for tp_rank in range(2) + } + + assert actual == set(range(8)) + + +def test_already_global_moe_worker_ranks_are_not_double_counted(): + actual = { + global_inference_rank( + rank_offset=0, + data_parallel_index=dp_index, + data_parallel_size=4, + worker_rank=dp_index * 2 + tp_rank, + tensor_parallel_size=2, + pipeline_parallel_size=1, + inference_world_size=8, + ) + for dp_index in range(4) + for tp_rank in range(2) + } + + assert actual == set(range(8)) + + +def test_rank_offset_composes_with_pipeline_parallel_rank(): + actual = { + global_inference_rank( + rank_offset=8, + data_parallel_index=dp_index, + data_parallel_size=2, + worker_rank=model_parallel_rank, + tensor_parallel_size=2, + pipeline_parallel_size=2, + inference_world_size=16, + ) + for dp_index in range(2) + for model_parallel_rank in range(4) + } + + assert actual == set(range(8, 16)) + + +def test_global_inference_rank_rejects_out_of_bounds_rank(): + with pytest.raises(ValueError, match="outside inference world size"): + global_inference_rank( + rank_offset=2, + data_parallel_index=3, + data_parallel_size=4, + worker_rank=0, + tensor_parallel_size=2, + pipeline_parallel_size=1, + inference_world_size=8, + ) + + +def test_prefill_context_parallel_ranks_are_unique(): + actual = { + global_inference_rank( + rank_offset=0, + data_parallel_index=0, + data_parallel_size=1, + worker_rank=worker_rank, + tensor_parallel_size=1, + pipeline_parallel_size=1, + prefill_context_parallel_size=2, + inference_world_size=2, + engine_world_size=2, + ) + for worker_rank in range(2) + } + + assert actual == {0, 1} diff --git a/tests/unit/inference/test_nccl_weight_update.py b/tests/unit/inference/test_nccl_weight_update.py new file mode 100644 index 0000000000..155fd98cf0 --- /dev/null +++ b/tests/unit/inference/test_nccl_weight_update.py @@ -0,0 +1,15 @@ +from prime_rl.inference.vllm.worker import nccl + + +def test_receiver_preserves_state_dict_boundaries(monkeypatch): + receiver = object.__new__(nccl.NCCLWeightBroadcastReceiver) + receiver.communicator = object() + streams = [iter([("layer.0", 0)]), iter([("layer.1", 1)])] + + monkeypatch.setattr(nccl, "receive_integer", lambda _communicator: len(streams)) + monkeypatch.setattr(nccl, "receive_state_dict", lambda _communicator: streams.pop(0)) + + assert [list(stream) for stream in receiver.receive_state_dicts()] == [ + [("layer.0", 0)], + [("layer.1", 1)], + ] diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index 951ef5c9d8..71c90e0377 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -1,66 +1,40 @@ -"""Sanity tests for the prime-RL ``ServingTokens`` subclass. - -The full happy-path is owned upstream by vLLM's -``vllm/entrypoints/serve/disagg`` test suite. We only cover the prime-RL -deltas here: - * ``serialize_routed_experts`` round-trips a compact raw-byte payload. - * The subclass attaches its overrides without monkey-patching the parent. - * ``_client_set_max_tokens`` distinguishes raw-body shapes correctly. -""" - from __future__ import annotations import asyncio +import io import numpy as np import pybase64 -from vllm.entrypoints.openai.engine.protocol import UsageInfo -from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateResponse, GenerateResponseChoice +from vllm.entrypoints.openai.engine.protocol import RequestResponseMetadata, UsageInfo +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + GenerateRequest, + GenerateResponse, + GenerateResponseChoice, +) +from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens +from vllm.sampling_params import SamplingParams -from prime_rl.inference.vllm.routed_experts import serialize_routed_experts -from prime_rl.inference.vllm.serving_tokens import ( - PrimeRlGenerateResponse, - PrimeRlGenerateResponseChoice, - PrimeRlServingTokens, - _build_usage, - _client_set_max_tokens, - _FinalOutputCapture, - _GenerateRoutedExpertsCapture, +from prime_rl.inference.vllm.routed_experts import ( + compact_vllm_routed_experts, + serialize_routed_experts, ) +from prime_rl.inference.vllm.serving_tokens import PrimeRlServingTokens -def _decode_routed_experts(encoded: dict) -> np.ndarray: +def _decode_compact(encoded: dict) -> np.ndarray: return np.frombuffer( pybase64.b64decode_as_bytearray(encoded["data"]), dtype=np.uint8, ).reshape(encoded["shape"]) -class _FakeRawRequest: - def __init__(self, body): - self._body = body - self._raise = isinstance(body, Exception) - - async def json(self): - if self._raise: - raise self._body - return self._body +def _encode_vllm(array: np.ndarray) -> str: + buffer = io.BytesIO() + np.save(buffer, array) + return pybase64.b64encode(buffer.getvalue()).decode("ascii") -async def _empty_request_outputs(): - if False: - yield - - -def test_subclass_only_overrides_serve_tokens(): - assert PrimeRlServingTokens.serve_tokens is not PrimeRlServingTokens.__mro__[1].serve_tokens - assert ( - PrimeRlServingTokens.serve_tokens_full_generator - is not PrimeRlServingTokens.__mro__[1].serve_tokens_full_generator - ) - - -def test_serialize_routed_experts_uses_compact_raw_payload(): +def test_routed_experts_round_trip_both_wire_formats(): routed_experts = np.array( [ [[1, 2], [3, 4]], @@ -69,201 +43,96 @@ def test_serialize_routed_experts_uses_compact_raw_payload(): dtype=np.int64, ) - encoded = serialize_routed_experts(routed_experts) - assert encoded is not None + compact = serialize_routed_experts(routed_experts, start=2) + converted = compact_vllm_routed_experts(_encode_vllm(routed_experts), start=2) + + assert compact is not None + assert converted is not None + assert converted["start"] == 2 + np.testing.assert_array_equal(_decode_compact(compact), routed_experts) + np.testing.assert_array_equal(_decode_compact(converted), routed_experts) + + +def test_serve_tokens_forwards_kv_transfer_params_without_mutating_request(monkeypatch): + expected = object() + observed_request = None + + async def upstream(_self, request, _raw_request=None): + nonlocal observed_request + observed_request = request + assert request.sampling_params.extra_args == { + "existing": True, + "kv_transfer_params": {"remote": "metadata"}, + } + return expected + + monkeypatch.setattr(ServingTokens, "serve_tokens", upstream) + server = object.__new__(PrimeRlServingTokens) + request = GenerateRequest( + token_ids=[1], + sampling_params=SamplingParams(max_tokens=1, extra_args={"existing": True}), + kv_transfer_params={"remote": "metadata"}, + ) - decoded = _decode_routed_experts(encoded) - assert decoded.dtype == np.uint8 - np.testing.assert_array_equal(decoded, routed_experts) + assert asyncio.run(server.serve_tokens(request)) is expected + assert observed_request is not request + assert request.sampling_params.extra_args == {"existing": True} -def test_generate_response_post_process_replaces_upstream_routed_experts(): - compact_routed_experts = {"data": "AQID", "shape": [1, 1, 3], "start": 0} - capture = _GenerateRoutedExpertsCapture(_empty_request_outputs()) - capture.routed_experts[0] = compact_routed_experts - response = GenerateResponse( - request_id="request-id", +def test_full_generator_preserves_all_upstream_response_fields(monkeypatch): + routed_experts = np.array([[[1, 2, 3]]], dtype=np.uint8) + usage = UsageInfo( + prompt_tokens=3, + completion_tokens=2, + total_tokens=5, + prompt_tokens_details={"cached_tokens": 2}, + ) + upstream_response = GenerateResponse( + request_id="canonical-request-id", + model="served-model", + created=123456789, + usage=usage, choices=[ GenerateResponseChoice( index=0, - token_ids=[1, 2, 3], - routed_experts="upstream-npy-payload", + token_ids=[10, 11], + routed_experts=_encode_vllm(routed_experts), ) ], ) - processed = capture.post_process(response) - - assert processed.choices[0].routed_experts == compact_routed_experts - - -def test_client_set_max_tokens_recognizes_explicit_value(): - body = {"token_ids": [1, 2, 3], "sampling_params": {"max_tokens": 256}} - assert asyncio.run(_client_set_max_tokens(_FakeRawRequest(body))) is True - - -def test_client_set_max_tokens_detects_unset(): - body = {"token_ids": [1, 2, 3], "sampling_params": {}} - assert asyncio.run(_client_set_max_tokens(_FakeRawRequest(body))) is False - - body_without_sp = {"token_ids": [1, 2, 3]} - assert asyncio.run(_client_set_max_tokens(_FakeRawRequest(body_without_sp))) is False - - -class _FakeOutput: - def __init__(self, token_ids): - self.token_ids = token_ids - - -class _FakeRequestOutput: - """Minimal stand-in for ``vllm.outputs.RequestOutput``. - - ``_build_usage`` only touches four attributes; constructing a real - ``RequestOutput`` would require a full ``CompletionOutput`` graph and - isn't worth it for a serialization-shape test. - """ - - def __init__(self, prompt_token_ids, output_token_ids_list, num_cached_tokens=0, encoder_prompt_token_ids=None): - self.prompt_token_ids = prompt_token_ids - self.encoder_prompt_token_ids = encoder_prompt_token_ids - self.outputs = [_FakeOutput(t) for t in output_token_ids_list] - self.num_cached_tokens = num_cached_tokens - - -def test_prime_rl_generate_response_serializes_usage_block(): - # Regression for prime-rl PR #2408: parent ``GenerateResponse`` doesn't - # declare ``usage``, so the field must be declared on the subclass for - # Pydantic to emit it in JSON. Without this the router can't extract - # per-run token / cache counts for billing. - response = PrimeRlGenerateResponse( - request_id="req-1", - choices=[PrimeRlGenerateResponseChoice(index=0, token_ids=[1, 2, 3])], - usage=UsageInfo(prompt_tokens=4, completion_tokens=3, total_tokens=7), - ) - payload = response.model_dump(mode="json") - assert payload["usage"] == { - "prompt_tokens": 4, - "completion_tokens": 3, - "total_tokens": 7, - "prompt_tokens_details": None, - } - - -def test_build_usage_sums_prompt_and_completion_tokens(): - final_res = _FakeRequestOutput( - prompt_token_ids=[1, 2, 3, 4, 5], - output_token_ids_list=[[10, 11], [20, 21, 22]], - ) - usage = _build_usage(final_res) - assert usage.prompt_tokens == 5 - assert usage.completion_tokens == 5 # 2 + 3 - assert usage.total_tokens == 10 - assert usage.prompt_tokens_details is None - - -def test_build_usage_includes_encoder_prompt_tokens(): - final_res = _FakeRequestOutput( - prompt_token_ids=[1, 2, 3], - output_token_ids_list=[[10]], - encoder_prompt_token_ids=[100, 101], - ) - usage = _build_usage(final_res) - assert usage.prompt_tokens == 5 # 3 + 2 - assert usage.total_tokens == 6 - - -def test_build_usage_reports_cached_tokens_unconditionally(): - # Unlike upstream's ``enable_prompt_tokens_details`` gate, prime-rl always - # surfaces cached tokens — the cache-discount billing pipeline needs them. - final_res = _FakeRequestOutput( - prompt_token_ids=[1, 2, 3, 4], - output_token_ids_list=[[10, 11]], - num_cached_tokens=3, + async def upstream(_self, _request, result_generator, *_args): + async for _ in result_generator: + pass + return upstream_response + + monkeypatch.setattr(ServingTokens, "serve_tokens_full_generator", upstream) + server = object.__new__(PrimeRlServingTokens) + server.enable_prompt_tokens_details = True + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=2, routed_experts_prompt_start=1), ) - usage = _build_usage(final_res) - assert usage.prompt_tokens_details is not None - assert usage.prompt_tokens_details.cached_tokens == 3 - -def test_build_usage_skips_cached_tokens_when_zero(): - # Don't emit a details block with cached=0, which would be misleading - # to the router's billing extractor. - final_res = _FakeRequestOutput( - prompt_token_ids=[1, 2, 3, 4], - output_token_ids_list=[[10, 11]], - num_cached_tokens=0, + async def outputs(): + if False: + yield + + response = asyncio.run( + server.serve_tokens_full_generator( + request, + outputs(), + "input-request-id", + "input-model", + RequestResponseMetadata(request_id="input-request-id"), + ) ) - usage = _build_usage(final_res) - assert usage.prompt_tokens_details is None - - -def test_final_output_capture_records_last_item(): - async def _gen(): - for r in [ - _FakeRequestOutput(prompt_token_ids=[1], output_token_ids_list=[[1]]), - _FakeRequestOutput(prompt_token_ids=[1, 2], output_token_ids_list=[[1, 2]]), - _FakeRequestOutput(prompt_token_ids=[1, 2, 3], output_token_ids_list=[[1, 2, 3]]), - ]: - yield r - - async def _drain(capture): - async for _ in capture: - pass - - capture = _FinalOutputCapture(_gen()) - asyncio.run(_drain(capture)) - assert capture.final_res is not None - assert capture.final_res.prompt_token_ids == [1, 2, 3] - - -def test_final_output_capture_works_over_async_def_aiter_source(): - # ``_GenerateRoutedExpertsCapture`` exposes the async-iterator protocol - # via ``async def __aiter__`` (an async generator function) and has no - # ``__anext__``. The wrapper must drive it through ``async for`` rather - # than poking ``__anext__`` directly, or routed-experts runs raise - # AttributeError before the response is built. - - class _AsyncGenAiterSource: - def __init__(self, items): - self._items = items - - async def __aiter__(self): - for item in self._items: - yield item - - items = [ - _FakeRequestOutput(prompt_token_ids=[1], output_token_ids_list=[[1]]), - _FakeRequestOutput(prompt_token_ids=[1, 2], output_token_ids_list=[[1, 2]]), - ] - capture = _FinalOutputCapture(_AsyncGenAiterSource(items)) - - async def _drain(): - async for _ in capture: - pass - - asyncio.run(_drain()) - assert capture.final_res is not None - assert capture.final_res.prompt_token_ids == [1, 2] - - -def test_final_output_capture_handles_empty_stream(): - capture = _FinalOutputCapture(_empty_request_outputs()) - - async def _drain(): - async for _ in capture: - pass - - asyncio.run(_drain()) - assert capture.final_res is None - - -def test_client_set_max_tokens_assumes_set_when_body_unreadable(): - # No raw_request → can't tell, don't override. - assert asyncio.run(_client_set_max_tokens(None)) is True - - # body read raises → can't tell, don't override. - err = ValueError("bad json") - assert asyncio.run(_client_set_max_tokens(_FakeRawRequest(err))) is True - # non-dict body → can't tell, don't override. - assert asyncio.run(_client_set_max_tokens(_FakeRawRequest([1, 2, 3]))) is True + assert response.request_id == "canonical-request-id" + assert response.model == "served-model" + assert response.created == 123456789 + assert response.usage == usage + encoded = response.choices[0].routed_experts + assert isinstance(encoded, dict) + assert encoded["start"] == 1 + np.testing.assert_array_equal(_decode_compact(encoded), routed_experts) diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index 60f13e01aa..48254c2fc7 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -18,6 +18,7 @@ async def run() -> None: ), renderer=renderer_settings, any_policy_sourced=True, + weight_broadcast=SimpleNamespace(type="filesystem", inference_world_size=None), ) renderer = object() inference_pool = object() @@ -43,6 +44,7 @@ async def run() -> None: train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=renderer_settings, + expected_inference_world_size=None, ) asyncio.run(run()) @@ -64,6 +66,7 @@ async def run() -> None: ), renderer=renderer_settings, any_policy_sourced=False, + weight_broadcast=SimpleNamespace(inference_world_size=8), ) renderer = object() inference_pool = object() @@ -89,6 +92,7 @@ async def run() -> None: train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=renderer_settings, + expected_inference_world_size=8, ) asyncio.run(run()) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 53492118fb..0bca37733a 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -215,6 +215,116 @@ def test_trainer_enable_token_export_cli_flag(): assert cli(TrainerConfig, args=["--enable-token-export"]).enable_token_export +def test_external_dynamo_world_size_survives_rl_config_resolution(): + config = RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": { + "model": { + "client": { + "base_url": ["http://frontend:8000/v1"], + "dynamo_discovery_url": "http://frontend:8001", + } + } + }, + "inference": None, + "weight_broadcast": { + "type": "nccl", + "host": "trainer.service", + "inference_world_size": 8, + }, + } + ) + + assert config.trainer.weight_broadcast.inference_world_size == 8 + assert config.trainer.weight_broadcast.host == "trainer.service" + assert config.orchestrator.weight_broadcast.inference_world_size == 8 + assert config.orchestrator.weight_broadcast.host == "trainer.service" + + +def test_external_dynamo_nccl_does_not_require_a_local_inference_gpu(): + config = RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": { + "model": { + "client": { + "base_url": ["http://frontend:8000/v1"], + "dynamo_discovery_url": "http://frontend:8001", + } + } + }, + "inference": None, + "deployment": { + "type": "single_node", + "num_train_gpus": 1, + "num_infer_gpus": 0, + }, + "weight_broadcast": { + "type": "nccl", + "host": "trainer.service", + "inference_world_size": 1, + }, + } + ) + + assert config.deployment.num_train_gpus == 1 + assert config.deployment.num_infer_gpus == 0 + assert config.trainer.weight_broadcast.inference_world_size == 1 + + +def test_default_nccl_world_size_does_not_bypass_local_gpu_guard(): + with pytest.raises(ValueError, match="NCCL weight broadcast requires at least 2"): + RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": {}, + "inference": None, + "deployment": { + "type": "single_node", + "num_train_gpus": 1, + "num_infer_gpus": 0, + }, + "weight_broadcast": {"type": "nccl"}, + } + ) + + +def test_external_dynamo_lora_world_size_survives_filesystem_config_resolution(): + config = RLConfig.model_validate( + { + "trainer": {"model": {"lora": {}}}, + "orchestrator": { + "model": { + "client": { + "base_url": ["http://frontend:8000/v1"], + "dynamo_discovery_url": "http://frontend:8001", + } + } + }, + "inference": None, + "weight_broadcast": {"type": "filesystem", "inference_world_size": 8}, + } + ) + + assert config.orchestrator.weight_broadcast.inference_world_size == 8 + + +def test_dynamo_orchestrator_requires_explicit_inference_world_size(): + with pytest.raises(ValueError, match="inference_world_size"): + OrchestratorConfig.model_validate( + { + "model": { + "client": { + "base_url": ["http://frontend:8000/v1"], + "dynamo_discovery_url": "http://frontend:8001", + } + }, + "weight_broadcast": {"type": "filesystem"}, + } + ) + + def test_single_node_auto_inference_ports_follow_server_port(): config = RLConfig.model_validate( { diff --git a/tests/unit/utils/test_dynamo.py b/tests/unit/utils/test_dynamo.py new file mode 100644 index 0000000000..d6e0c844eb --- /dev/null +++ b/tests/unit/utils/test_dynamo.py @@ -0,0 +1,119 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from prime_rl.configs.shared import ClientConfig, ElasticConfig +from prime_rl.utils.dynamo import DynamoInferencePool, _parse_dynamo_workers + +MODEL = "Qwen/Qwen3-0.6B" + + +def worker(**updates): + value = { + "component": "backend", + "instance_id": 10, + "model": MODEL, + "admin_base_url": "http://decode:8120", + "world_size": 2, + } + return {**value, **updates} + + +def payload(*workers): + return {"protocol_version": 1, "workers": list(workers)} + + +def response(body): + result = MagicMock() + result.raise_for_status = MagicMock() + result.json.return_value = body + return result + + +def test_parse_workers_orders_identity_and_preserves_topology(): + workers = _parse_dynamo_workers( + payload( + worker(component="prefill", instance_id=20, admin_base_url="http://prefill:8121"), + worker(), + ), + MODEL, + ) + + assert [(item.component, item.instance_id) for item in workers] == [ + ("backend", 10), + ("prefill", 20), + ] + assert [item.world_size for item in workers] == [2, 2] + + +@pytest.mark.parametrize( + "workers", + [ + [], + [worker(error="probe timed out")], + [worker(admin_base_url=None)], + [worker(world_size=0)], + [worker(model="other/model")], + [worker(), worker(instance_id=11)], + [worker(), worker(component="prefill", instance_id=20, admin_base_url="http://decode:8120")], + ], +) +def test_parse_workers_rejects_incomplete_or_duplicate_snapshots(workers): + with pytest.raises(ValueError): + _parse_dynamo_workers(payload(*workers), MODEL) + + +@pytest.mark.parametrize( + "conflict", + [ + {"admin_base_url": ["http://worker:8120"]}, + {"elastic": ElasticConfig(hostname="workers")}, + ], +) +def test_discovery_config_rejects_other_pool_modes(conflict): + with pytest.raises(ValueError, match="dynamo_discovery_url"): + ClientConfig(dynamo_discovery_url="http://frontend:8001", **conflict) + + +def test_discovery_retries_until_expected_world_size_is_complete(): + transient = MagicMock() + transient.raise_for_status.side_effect = httpx.HTTPStatusError( + "Service unavailable", + request=httpx.Request("GET", "http://frontend:8001/v1/rl/workers"), + response=httpx.Response(503), + ) + discovery_client = AsyncMock() + discovery_client.get.side_effect = [ + transient, + response(payload(worker())), + response( + payload( + worker(), + worker(component="prefill", instance_id=20, admin_base_url="http://prefill:8121"), + ) + ), + ] + context = AsyncMock() + context.__aenter__.return_value = discovery_client + + class DiscoveryOnlyPool(DynamoInferencePool): + def __init__(self, _config, workers, **_kwargs): + self.workers = workers + + with patch("prime_rl.utils.dynamo.AsyncClient", return_value=context): + pool = asyncio.run( + DiscoveryOnlyPool.from_config( + ClientConfig( + base_url=["http://frontend:8000/v1"], + dynamo_discovery_url="http://frontend:8001", + wait_for_ready_timeout=1, + ), + model_name=MODEL, + expected_inference_world_size=4, + ) + ) + + assert discovery_client.get.await_count == 3 + assert [item.component for item in pool.workers] == ["backend", "prefill"] diff --git a/tests/unit/utils/test_dynamo_inmemory.py b/tests/unit/utils/test_dynamo_inmemory.py new file mode 100644 index 0000000000..f296665054 --- /dev/null +++ b/tests/unit/utils/test_dynamo_inmemory.py @@ -0,0 +1,81 @@ +import asyncio +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +from prime_rl.utils.client import init_nccl_broadcast, update_weights +from prime_rl.utils.dynamo import DynamoInferencePool + + +def test_native_nccl_initialization_uses_collective_rpc(): + clients = [AsyncMock(), AsyncMock()] + for client in clients: + response = MagicMock() + response.raise_for_status = MagicMock() + client.post.return_value = response + + asyncio.run( + init_nccl_broadcast( + clients, + host="127.0.0.1", + port=29519, + timeout=1200, + inference_world_size=4, + engine_world_sizes=[2, 2], + use_native_collective_rpc=True, + ) + ) + + assert [client.post.await_args.args[0] for client in clients] == ["/collective_rpc", "/collective_rpc"] + assert [client.post.await_args.kwargs["json"]["kwargs"]["rank_offset"] for client in clients] == [0, 2] + + +def test_dynamo_pool_passes_discovered_topology_to_nccl(): + pool = DynamoInferencePool.__new__(DynamoInferencePool) + pool._admin_clients = [AsyncMock(), AsyncMock()] + pool._admin_world_sizes = [1, 3] + + with patch("prime_rl.utils.dynamo.init_nccl_broadcast", new=AsyncMock()) as initialize: + asyncio.run( + pool.init_nccl_broadcast( + host="trainer", + port=29501, + timeout=1200, + inference_world_size=4, + quantize_in_weight_transfer=False, + ) + ) + + assert initialize.await_args.kwargs["engine_world_sizes"] == [1, 3] + assert initialize.await_args.kwargs["use_native_collective_rpc"] is True + + +def test_dynamo_pool_uses_native_full_weight_update(): + pool = DynamoInferencePool.__new__(DynamoInferencePool) + pool._admin_clients = [AsyncMock()] + + with patch("prime_rl.utils.dynamo.update_weights", new=AsyncMock()) as update: + asyncio.run(pool.update_weights(Path("/weights"), step=2)) + + assert update.await_args.kwargs["use_native_collective_rpc"] is True + + +def test_native_full_weight_update_uses_positional_path(tmp_path): + client = AsyncMock() + response = MagicMock() + response.raise_for_status = MagicMock() + client.post.return_value = response + + asyncio.run( + update_weights( + [client], + tmp_path, + step=2, + use_native_collective_rpc=True, + ) + ) + + collective_calls = [call for call in client.post.await_args_list if call.args[0] == "/collective_rpc"] + assert collective_calls[0].kwargs["json"] == { + "method": "update_weights_from_path", + "args": [tmp_path.as_posix()], + } diff --git a/tests/unit/utils/test_dynamo_lora.py b/tests/unit/utils/test_dynamo_lora.py new file mode 100644 index 0000000000..d286bf1706 --- /dev/null +++ b/tests/unit/utils/test_dynamo_lora.py @@ -0,0 +1,75 @@ +import asyncio +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest + +from prime_rl.utils.dynamo import DynamoInferencePool, _parse_dynamo_workers + +MODEL = "Qwen/Qwen3-0.6B" + + +def worker(**updates): + value = { + "component": "backend", + "instance_id": 1, + "model": MODEL, + "admin_base_url": "http://worker:8120", + "world_size": 1, + "system_url": "http://worker:8181", + "system_routes": ["update/load_lora"], + } + return {**value, **updates} + + +def pool(): + value = DynamoInferencePool.__new__(DynamoInferencePool) + value._admin_clients = [AsyncMock()] + value._lora_update_clients = [AsyncMock()] + value._frontend_model_clients = [AsyncMock()] + value._wait_for_ready_timeout = 1 + return value + + +def test_discovery_rejects_partial_lora_capability(): + payload = { + "protocol_version": 1, + "workers": [ + worker(), + worker(component="prefill", instance_id=2, admin_base_url="http://prefill:8120", system_routes=[]), + ], + } + + with pytest.raises(ValueError, match="partial update/load_lora"): + _parse_dynamo_workers(payload, MODEL) + + +def test_lora_update_resumes_after_publication(): + inference_pool = pool() + + with ( + patch("prime_rl.utils.dynamo._pause_engines", new=AsyncMock()) as pause, + patch("prime_rl.utils.dynamo._load_lora_adapter", new=AsyncMock()) as load, + patch("prime_rl.utils.dynamo._wait_for_model", new=AsyncMock()) as wait, + patch("prime_rl.utils.dynamo._resume_engines", new=AsyncMock()) as resume, + ): + asyncio.run(inference_pool.update_weights(Path("/weights/adapter"), lora_name="policy", step=3)) + + pause.assert_awaited_once() + load.assert_awaited_once() + wait.assert_awaited_once() + resume.assert_awaited_once() + + +def test_lora_update_resumes_after_failure(): + inference_pool = pool() + + with ( + patch("prime_rl.utils.dynamo._pause_engines", new=AsyncMock()), + patch("prime_rl.utils.dynamo._load_lora_adapter", new=AsyncMock(side_effect=RuntimeError("failed"))), + patch("prime_rl.utils.dynamo._resume_engines", new=AsyncMock()) as resume, + pytest.raises(RuntimeError, match="failed"), + ): + asyncio.run(inference_pool.update_weights(Path("/weights/adapter"), lora_name="policy", step=3)) + + resume.assert_awaited_once() diff --git a/tests/unit/utils/test_external_engine_topology.py b/tests/unit/utils/test_external_engine_topology.py new file mode 100644 index 0000000000..591bd8c78a --- /dev/null +++ b/tests/unit/utils/test_external_engine_topology.py @@ -0,0 +1,44 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock + +from prime_rl.utils.client import _rank_offsets, init_nccl_broadcast + + +def test_rank_offsets_support_heterogeneous_engines(): + assert _rank_offsets([1, 3, 2], inference_world_size=6) == [0, 1, 4] + + +def test_nccl_broadcast_forwards_explicit_engine_sizes(): + clients = [AsyncMock(), AsyncMock()] + for client in clients: + response = MagicMock() + response.raise_for_status = MagicMock() + client.post.return_value = response + + asyncio.run( + init_nccl_broadcast( + clients, + host="127.0.0.1", + port=29519, + timeout=1200, + inference_world_size=4, + engine_world_sizes=[1, 3], + ) + ) + + assert [call.kwargs["json"]["rank_offset"] for client in clients for call in client.post.await_args_list] == [0, 1] + assert [call.kwargs["json"]["engine_world_size"] for client in clients for call in client.post.await_args_list] == [ + 1, + 3, + ] + + +def test_nccl_broadcast_preserves_legacy_payload(): + client = AsyncMock() + response = MagicMock() + response.raise_for_status = MagicMock() + client.post.return_value = response + + asyncio.run(init_nccl_broadcast([client], "127.0.0.1", 29519, 1200, 1)) + + assert "engine_world_size" not in client.post.await_args.kwargs["json"]