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/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..e652f3cc2c 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 @@ -122,6 +122,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 +133,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 = ( @@ -193,6 +195,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 +209,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, diff --git a/src/prime_rl/utils/dynamo.py b/src/prime_rl/utils/dynamo.py new file mode 100644 index 0000000000..8cb950a0cc --- /dev/null +++ b/src/prime_rl/utils/dynamo.py @@ -0,0 +1,246 @@ +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, + setup_admin_clients, +) + +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 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 super().update_weights(weight_dir, lora_name=lora_name, step=step) + 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/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_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()