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/server.py b/src/prime_rl/inference/vllm/server.py index 28e2f572ff..5077f3003b 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"} 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..93e0b25018 --- /dev/null +++ b/src/prime_rl/utils/dynamo.py @@ -0,0 +1,230 @@ +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_if_exception, stop_after_delay, wait_exponential + +from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.client import ( + StaticInferencePool, + 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) + + +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") + 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 _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._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: + await update_weights( + self._admin_clients, + weight_dir, + lora_name=lora_name, + step=step, + use_native_collective_rpc=True, + ) + + async def stop(self) -> None: + await super().stop() + await asyncio.gather(*(client.aclose() for client in [*self._admin_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/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_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"]