feat(dynamo): support LoRA lifecycle

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

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:

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"))

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")

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
 

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,