From 3690dc064da218afdfcc8298e2e5b4b3725c8678 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 3 Jul 2026 12:24:02 -0700 Subject: [PATCH 01/18] feat(inference): add Dynamo backend configuration Signed-off-by: Biswa Panda --- .../src/prime_rl/configs/inference.py | 44 ++++++++++- .../src/prime_rl/configs/rl.py | 26 ++++++- .../src/prime_rl/configs/shared.py | 6 ++ tests/unit/test_configs.py | 74 +++++++++++++++++++ 4 files changed, 145 insertions(+), 5 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index 20ac441af0..e5e84c21b7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -304,10 +304,10 @@ class DisaggregatedInferenceDeploymentConfig(BaseInferenceDeploymentConfig): """Extra environment variables exported only on decode nodes.""" prefill_vllm_overrides: dict[str, Any] = {} - """Extra vLLM config options merged into --vllm-extra only for prefill ranks (SLURM only).""" + """Extra vLLM config options merged into the resolved config only for prefill workers.""" decode_vllm_overrides: dict[str, Any] = {} - """Extra vLLM config options merged into --vllm-extra only for decode ranks (SLURM only).""" + """Extra vLLM config options merged into the resolved config only for decode workers.""" @property def num_prefill_nodes(self) -> int: @@ -328,7 +328,24 @@ def num_nodes(self) -> int: ] +class VllmInferenceBackendConfig(BaseConfig): + type: Literal["vllm"] = "vllm" + + +class DynamoInferenceBackendConfig(BaseConfig): + type: Literal["dynamo"] = "dynamo" + + +InferenceBackendConfig: TypeAlias = Annotated[ + VllmInferenceBackendConfig | DynamoInferenceBackendConfig, + Field(discriminator="type"), +] + + class InferenceConfig(BaseConfig): + backend: InferenceBackendConfig = VllmInferenceBackendConfig() + """Serving backend. Existing configs default to Prime's native vLLM launcher.""" + server: ServerConfig = ServerConfig() model: ModelConfig = Field(default_factory=ModelConfig) @@ -428,10 +445,31 @@ class InferenceConfig(BaseConfig): @model_validator(mode="after") def validate_multi_node_requires_slurm(self): - if self.deployment.type in ("multi_node", "disaggregated") and self.slurm is None: + if self.deployment.type == "multi_node" and self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + if self.deployment.type == "disaggregated" and self.slurm is None and self.backend.type != "dynamo": raise ValueError("Must use SLURM for multi-node / disaggregated deployment.") return self + @model_validator(mode="after") + def validate_dynamo_backend(self): + if self.backend.type != "dynamo": + return self + if self.slurm is not None: + raise ValueError( + "Dynamo is launched locally or through a DynamoGraphDeployment, not Prime's SLURM template." + ) + if self.deployment.type == "multi_node": + raise ValueError("Dynamo multi-node inference must use a DynamoGraphDeployment.") + router = getattr(self.deployment, "router", None) + if router is not None and router.type == "llm-d": + raise ValueError("The Dynamo backend owns request routing and cannot use the llm-d router.") + if self.deployment.type == "disaggregated": + if self.enable_prefix_caching is False: + raise ValueError("Dynamo disaggregated inference requires prefix caching for exact KV-aware routing.") + self.enable_prefix_caching = True + return self + @model_validator(mode="after") def validate_llmd_no_routed_experts(self): """Reject routed-expert return with the llm-d router (breaks P/D, unverified for multi-node).""" 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 c92e53889a..b7e9d92768 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -534,7 +534,23 @@ def auto_setup_deployment(self): # fill up inference capacity with dp ranks if self.inference is not None: num_infer_gpus = self.deployment.num_infer_gpus - if num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp: + is_dynamo_disaggregated = ( + self.inference.backend.type == "dynamo" and self.inference.deployment.type == "disaggregated" + ) + if is_dynamo_disaggregated: + infer_deploy = self.inference.deployment + expected_infer_gpus = infer_deploy.num_nodes * infer_deploy.gpus_per_node + if num_infer_gpus != expected_infer_gpus: + raise ValueError( + "deployment.num_infer_gpus must equal the Dynamo prefill/decode topology GPU count " + f"({expected_infer_gpus}), got {num_infer_gpus}." + ) + if self.weight_broadcast is not None and self.weight_broadcast.type == "nccl": + assert self.trainer.weight_broadcast.type == "nccl" + self.trainer.weight_broadcast.inference_world_size = expected_infer_gpus + assert self.orchestrator.weight_broadcast.type == "nccl" + self.orchestrator.weight_broadcast.inference_world_size = expected_infer_gpus + elif num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp: assert num_infer_gpus % self.inference.parallel.tp == 0, ( "Number of inference GPUs must be divisible by the tensor parallel size" ) @@ -675,8 +691,14 @@ def auto_setup_inference_client(self): if self.inference is None: return self client = self.orchestrator.model.client + if "admin_api" in client.model_fields_set and client.admin_api != self.inference.backend.type: + raise ValueError( + "orchestrator.model.client.admin_api conflicts with inference.backend.type; " + "configure the backend only under inference." + ) + client.admin_api = self.inference.backend.type if "dp_rank_count" not in client.model_fields_set: - if self.deployment.type == "multi_node": + if self.inference.backend.type == "dynamo" or self.deployment.type == "multi_node": client.dp_rank_count = 1 else: client.dp_rank_count = self.inference.data_parallel_size_local or self.inference.parallel.dp 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 c63bcc9c53..3485d05f0d 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -147,6 +147,12 @@ 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.""" + admin_api: Literal["vllm", "dynamo"] = "vllm" + """Admin protocol used for health and weight updates. Auto-derived from the local inference backend.""" + + rl_base_url: list[str] | None = None + """Dynamo RL worker-discovery URLs. When omitted, they are derived from the frontend URL or ``DYN_RL_DISCOVERY_URL``.""" + 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.""" diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 49243a9f98..8624eca192 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -211,6 +211,80 @@ def test_single_node_auto_inference_client_dp_rank_count_matches_local_dp(): assert config.orchestrator.model.client.dp_rank_count == 2 +def test_inference_backend_defaults_to_vllm(): + assert InferenceConfig().backend.type == "vllm" + + +def test_dynamo_disaggregated_config_is_local_and_enables_prefix_caching(): + config = InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_nodes": 2, + "num_decode_nodes": 2, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + ) + + assert config.enable_prefix_caching is True + assert config.use_pd_kv_transfer is True + + +def test_dynamo_disaggregated_config_rejects_disabled_prefix_caching(): + with pytest.raises(ValidationError, match="requires prefix caching"): + InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "enable_prefix_caching": False, + "deployment": {"type": "disaggregated"}, + } + ) + + +def test_native_disaggregated_config_still_requires_slurm(): + with pytest.raises(ValidationError, match="Must use SLURM"): + InferenceConfig.model_validate({"deployment": {"type": "disaggregated"}}) + + +def test_single_node_dynamo_disaggregated_keeps_per_worker_dp_and_sets_nccl_world_size(): + config = RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": {}, + "inference": { + "backend": {"type": "dynamo"}, + "parallel": {"tp": 1}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_nodes": 2, + "num_decode_nodes": 2, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + }, + "weight_broadcast": {"type": "nccl"}, + "deployment": { + "type": "single_node", + "gpus_per_node": 8, + "num_train_gpus": 4, + "num_infer_gpus": 4, + }, + } + ) + + assert config.inference is not None + assert config.inference.parallel.dp == 1 + assert config.orchestrator.model.client.admin_api == "dynamo" + assert config.orchestrator.model.client.dp_rank_count == 1 + assert config.trainer.weight_broadcast.inference_world_size == 4 + assert config.orchestrator.weight_broadcast.inference_world_size == 4 + + def test_multi_node_auto_inference_client_dp_rank_count_uses_router_url(): config = RLConfig.model_validate( { From c58de96d4be80ea35435c680ba89eb4fc254ea93 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 3 Jul 2026 12:27:31 -0700 Subject: [PATCH 02/18] feat(inference): launch Dynamo disaggregated workers Signed-off-by: Biswa Panda --- src/prime_rl/entrypoints/inference.py | 13 ++ src/prime_rl/inference/dynamo.py | 319 ++++++++++++++++++++++++++ tests/unit/inference/test_dynamo.py | 146 ++++++++++++ 3 files changed, 478 insertions(+) create mode 100644 src/prime_rl/inference/dynamo.py create mode 100644 tests/unit/inference/test_dynamo.py diff --git a/src/prime_rl/entrypoints/inference.py b/src/prime_rl/entrypoints/inference.py index efc120f96e..474cff9a76 100644 --- a/src/prime_rl/entrypoints/inference.py +++ b/src/prime_rl/entrypoints/inference.py @@ -148,6 +148,13 @@ def inference_local(config: InferenceConfig): logger = setup_logger(config.log.level, json_logging=config.log.json_logging) if config.dry_run: + if config.backend.type == "dynamo": + from prime_rl.inference.dynamo import build_frontend_command, build_local_worker_specs + + specs = build_local_worker_specs(config) + logger.info(f"Dynamo frontend: {' '.join(build_frontend_command(config))}") + for spec in specs: + logger.info(f"Dynamo {spec.name}: {' '.join(spec.command())}") logger.success("Dry run complete. To start inference locally, remove --dry-run from your command.") return @@ -162,6 +169,12 @@ def inference_local(config: InferenceConfig): setup_vllm_env(config) + if config.backend.type == "dynamo": + from prime_rl.inference.dynamo import run_dynamo_local + + run_dynamo_local(config) + return + from prime_rl.inference.vllm.server import server # pyright: ignore server(config, vllm_extra=config.vllm_extra) diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py new file mode 100644 index 0000000000..f0b01d4cb9 --- /dev/null +++ b/src/prime_rl/inference/dynamo.py @@ -0,0 +1,319 @@ +"""Translate Prime inference config into Dynamo worker processes.""" + +from __future__ import annotations + +import json +import os +import signal +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any, Literal + +from prime_rl.configs.inference import DisaggregatedInferenceDeploymentConfig, InferenceConfig +from prime_rl.utils.pathing import get_config_dir + +Role = Literal["agg", "prefill", "decode"] + +ENGINE_CONFIG_DIR = "dynamo" +PREFILL_ENGINE_CONFIG = "prefill-engine.json" +DECODE_ENGINE_CONFIG = "decode-engine.json" +AGG_ENGINE_CONFIG = "agg-engine.json" + +_ENGINE_CONFIG_EXCLUDED = frozenset( + { + "api_server_count", + "chat_template", + "enable_auto_tool_choice", + "host", + "liveness_timeout_seconds", + "port", + "reasoning_parser", + "tool_call_parser", + } +) +_RESERVED_ENGINE_KEYS = frozenset( + { + "disaggregation_mode", + "enable_rl", + "kv_events_config", + "kv_transfer_config", + "worker_extension_cls", + } +) +_WORKER_EXTENSION_CLS = { + "nccl": "prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker", + "filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker", +} + + +@dataclass(frozen=True) +class DynamoWorkerSpec: + name: str + role: Role + gpu_ids: tuple[str, ...] + system_port: int + nixl_port: int + kv_events_port: int | None + engine_config: Path + + def command(self) -> list[str]: + return [ + sys.executable, + "-m", + "dynamo.vllm", + "--engine-config-json", + str(self.engine_config), + "--disaggregation-mode", + self.role, + "--enable-rl", + ] + + +def _json_default(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, Enum): + return value.value + raise TypeError(f"Cannot serialize {type(value).__name__}") + + +def _role_overrides(config: InferenceConfig, role: Role) -> dict[str, Any]: + if config.deployment.type != "disaggregated": + return {} + if role == "prefill": + return config.deployment.prefill_vllm_overrides + if role == "decode": + return config.deployment.decode_vllm_overrides + return {} + + +def _validate_overrides(source: str, values: dict[str, Any]) -> None: + conflicts = sorted(_RESERVED_ENGINE_KEYS & values.keys()) + if conflicts: + raise ValueError(f"{source} cannot override Dynamo-managed engine keys: {conflicts}") + + +def build_engine_config( + config: InferenceConfig, + role: Role, + *, + kv_events_port: int | None = None, +) -> dict[str, Any]: + """Build one deterministic vLLM ``AsyncEngineArgs`` object.""" + _validate_overrides("vllm_extra", config.vllm_extra) + overrides = _role_overrides(config, role) + _validate_overrides(f"{role}_vllm_overrides", overrides) + + values = vars(config.to_vllm()).copy() + for key in _ENGINE_CONFIG_EXCLUDED: + values.pop(key, None) + values.update(config.vllm_extra) + values.update(overrides) + + if config.deployment.type == "disaggregated": + # Each generated worker is an independent vLLM server. Preserve local + # DP within a worker, but never turn the P/D worker count into vLLM DP. + local_dp = config.deployment.gpus_per_node // config.parallel.tp + values["data_parallel_size"] = local_dp + if local_dp == 1: + values.pop("data_parallel_size_local", None) + values.pop("data_parallel_rpc_port", None) + else: + values["data_parallel_size_local"] = local_dp + + if role in ("prefill", "agg") and kv_events_port is not None: + values["kv_events_config"] = { + "publisher": "zmq", + "topic": "kv-events", + "endpoint": f"tcp://*:{kv_events_port}", + "enable_kv_cache_events": True, + } + else: + values.pop("kv_events_config", None) + + values["worker_extension_cls"] = _WORKER_EXTENSION_CLS[config.weight_broadcast.type] + return {key: value for key, value in values.items() if value is not None} + + +def _write_json(path: Path, value: dict[str, Any]) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, default=_json_default, indent=2, sort_keys=True) + "\n") + return path + + +def write_role_engine_configs(config: InferenceConfig, output_dir: Path | None = None) -> dict[Role, Path]: + """Write canonical role configs used by DGD and dry-run inspection.""" + config_dir = output_dir or (get_config_dir(config.output_dir) / ENGINE_CONFIG_DIR) + if config.deployment.type == "disaggregated": + return { + "prefill": _write_json( + config_dir / PREFILL_ENGINE_CONFIG, + build_engine_config(config, "prefill", kv_events_port=20080), + ), + "decode": _write_json(config_dir / DECODE_ENGINE_CONFIG, build_engine_config(config, "decode")), + } + return { + "agg": _write_json( + config_dir / AGG_ENGINE_CONFIG, + build_engine_config(config, "agg", kv_events_port=20080), + ) + } + + +def _visible_gpu_ids() -> list[str]: + configured = os.environ.get("CUDA_VISIBLE_DEVICES") + if configured: + return [gpu.strip() for gpu in configured.split(",") if gpu.strip()] + try: + output = subprocess.check_output( + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError) as exc: + raise RuntimeError("Cannot discover GPUs; set CUDA_VISIBLE_DEVICES explicitly") from exc + return [line.strip() for line in output.splitlines() if line.strip()] + + +def build_local_worker_specs( + config: InferenceConfig, + output_dir: Path | None = None, + gpu_ids: list[str] | None = None, +) -> list[DynamoWorkerSpec]: + """Allocate local workers and write instance-specific engine configs.""" + config_dir = output_dir or (get_config_dir(config.output_dir) / ENGINE_CONFIG_DIR) + available = gpu_ids or _visible_gpu_ids() + + if config.deployment.type == "disaggregated": + deployment: DisaggregatedInferenceDeploymentConfig = config.deployment + if deployment.num_prefill_nodes != deployment.num_prefill_replicas: + raise ValueError("Local Dynamo requires one prefill node per prefill replica") + if deployment.num_decode_nodes != deployment.num_decode_replicas: + raise ValueError("Local Dynamo requires one decode node per decode replica") + roles: list[Role] = ["decode"] * deployment.num_decode_replicas + ["prefill"] * deployment.num_prefill_replicas + gpus_per_worker = deployment.gpus_per_node + else: + roles = ["agg"] + gpus_per_worker = config.parallel.tp * config.parallel.dp + + required = len(roles) * gpus_per_worker + if len(available) < required: + raise ValueError(f"Dynamo topology requires {required} GPUs, but only {len(available)} are visible") + + specs: list[DynamoWorkerSpec] = [] + role_indexes: dict[Role, int] = {"agg": 0, "prefill": 0, "decode": 0} + for worker_index, role in enumerate(roles): + role_index = role_indexes[role] + role_indexes[role] += 1 + start = worker_index * gpus_per_worker + worker_gpus = tuple(available[start : start + gpus_per_worker]) + kv_events_port = 20080 + role_index if role in ("prefill", "agg") else None + name = f"{role}-{role_index}" + engine_path = _write_json( + config_dir / f"{name}-engine.json", + build_engine_config(config, role, kv_events_port=kv_events_port), + ) + specs.append( + DynamoWorkerSpec( + name=name, + role=role, + gpu_ids=worker_gpus, + system_port=8081 + worker_index, + nixl_port=20100 + worker_index, + kv_events_port=kv_events_port, + engine_config=engine_path, + ) + ) + return specs + + +def build_frontend_command(config: InferenceConfig) -> list[str]: + return [ + sys.executable, + "-m", + "dynamo.frontend", + "--http-host", + config.server.host or "0.0.0.0", + "--http-port", + str(config.server.port), + "--router-mode", + "kv", + "--router-reset-states", + "--enable-engine-apis", + ] + + +def build_worker_environment( + config: InferenceConfig, + spec: DynamoWorkerSpec, + base_environment: dict[str, str], +) -> dict[str, str]: + environment = base_environment | { + "CUDA_VISIBLE_DEVICES": ",".join(spec.gpu_ids), + "DYN_SYSTEM_PORT": str(spec.system_port), + "VLLM_NIXL_SIDE_CHANNEL_HOST": "127.0.0.1", + "VLLM_NIXL_SIDE_CHANNEL_PORT": str(spec.nixl_port), + "VLLM_PLUGINS": "prime_rl", + } + if config.deployment.type == "disaggregated": + role_environment = ( + config.deployment.prefill_env_vars if spec.role == "prefill" else config.deployment.decode_env_vars + ) + environment.update(role_environment) + return environment + + +def _terminate(process: subprocess.Popen) -> None: + if process.poll() is not None: + return + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + return + try: + process.wait(timeout=15) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + + +def run_dynamo_local(config: InferenceConfig) -> None: + """Run a Dynamo frontend and all configured workers until one exits.""" + specs = build_local_worker_specs(config) + environment = os.environ.copy() + environment.setdefault("DYN_DISCOVERY_BACKEND", "file") + environment.setdefault("DYN_EVENT_PLANE", "zmq") + environment.setdefault("DYN_FILE_KV_TTL_SECS", "1800") + environment.setdefault("DYN_NAMESPACE", f"prime-rl-{os.getpid()}") + environment.setdefault("DYN_ENABLE_RL", "1") + environment.setdefault("DYN_RL_PORT", "8001") + environment.setdefault("PYTHONHASHSEED", "0") + + def request_stop(_signum, _frame): + raise KeyboardInterrupt + + signal.signal(signal.SIGTERM, request_stop) + processes: list[subprocess.Popen] = [] + with tempfile.TemporaryDirectory(prefix="prime-dynamo-") as temporary_dir: + environment.setdefault("DYN_FILE_KV", str(Path(temporary_dir) / "discovery")) + frontend_env = environment | {"CUDA_VISIBLE_DEVICES": ""} + frontend_env.pop("DYN_SYSTEM_PORT", None) + + try: + processes.append(subprocess.Popen(build_frontend_command(config), env=frontend_env, start_new_session=True)) + for spec in specs: + worker_env = build_worker_environment(config, spec, environment) + processes.append(subprocess.Popen(spec.command(), env=worker_env, start_new_session=True)) + + while all(process.poll() is None for process in processes): + time.sleep(0.2) + raise SystemExit(next((process.returncode for process in processes if process.returncode), 1)) + except KeyboardInterrupt: + return + finally: + for process in reversed(processes): + _terminate(process) diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py new file mode 100644 index 0000000000..d6d77df4cb --- /dev/null +++ b/tests/unit/inference/test_dynamo.py @@ -0,0 +1,146 @@ +import json +from pathlib import Path + +import pytest + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference import dynamo +from prime_rl.inference.dynamo import ( + build_engine_config, + build_local_worker_specs, + build_worker_environment, + write_role_engine_configs, +) + + +def disaggregated_config(**overrides) -> InferenceConfig: + data = { + "backend": {"type": "dynamo"}, + "weight_broadcast": {"type": "nccl"}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_nodes": 2, + "num_decode_nodes": 2, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + data.update(overrides) + return InferenceConfig.model_validate(data) + + +def test_role_engine_configs_share_nixl_and_only_prefill_publishes_events(tmp_path: Path): + paths = write_role_engine_configs(disaggregated_config(), tmp_path) + prefill = json.loads(paths["prefill"].read_text()) + decode = json.loads(paths["decode"].read_text()) + + assert prefill["kv_transfer_config"] == decode["kv_transfer_config"] + assert prefill["kv_transfer_config"]["kv_connector"] == "NixlConnector" + assert prefill["kv_events_config"]["enable_kv_cache_events"] is True + assert "kv_events_config" not in decode + assert prefill["worker_extension_cls"].endswith("NCCLWeightUpdateWorker") + assert decode["worker_extension_cls"] == prefill["worker_extension_cls"] + + +def test_role_overrides_are_isolated(): + config = disaggregated_config( + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_nodes": 2, + "num_decode_nodes": 2, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + "prefill_vllm_overrides": {"max_num_batched_tokens": 8192}, + "decode_vllm_overrides": {"max_num_seqs": 64}, + } + ) + + prefill = build_engine_config(config, "prefill", kv_events_port=20080) + decode = build_engine_config(config, "decode") + + assert prefill["max_num_batched_tokens"] == 8192 + assert "max_num_batched_tokens" not in decode + assert decode["max_num_seqs"] == 64 + assert "max_num_seqs" not in prefill + + +@pytest.mark.parametrize("key", ["kv_transfer_config", "kv_events_config", "worker_extension_cls"]) +def test_reserved_engine_override_is_rejected(key: str): + config = disaggregated_config(vllm_extra={key: {}}) + with pytest.raises(ValueError, match="Dynamo-managed"): + build_engine_config(config, "prefill", kv_events_port=20080) + + +def test_local_specs_allocate_four_workers_and_unique_ports(tmp_path: Path): + specs = build_local_worker_specs(disaggregated_config(), tmp_path, gpu_ids=["4", "5", "6", "7"]) + + assert [spec.role for spec in specs] == ["decode", "decode", "prefill", "prefill"] + assert [spec.gpu_ids for spec in specs] == [("4",), ("5",), ("6",), ("7",)] + assert len({spec.system_port for spec in specs}) == 4 + assert len({spec.nixl_port for spec in specs}) == 4 + assert [spec.kv_events_port for spec in specs if spec.role == "prefill"] == [20080, 20081] + assert all("--enable-rl" in spec.command() for spec in specs) + + +def test_wrapper_options_are_not_written_to_engine_json(): + engine = build_engine_config(disaggregated_config(), "prefill", kv_events_port=20080) + assert "disaggregation_mode" not in engine + assert "enable_rl" not in engine + + +def test_worker_environment_applies_only_matching_role_overrides(tmp_path: Path): + config = disaggregated_config( + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_nodes": 1, + "num_decode_nodes": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + "prefill_env_vars": {"ROLE_SETTING": "prefill"}, + "decode_env_vars": {"ROLE_SETTING": "decode"}, + } + ) + decode, prefill = build_local_worker_specs(config, tmp_path, gpu_ids=["3", "7"]) + + decode_env = build_worker_environment(config, decode, {"COMMON": "value"}) + prefill_env = build_worker_environment(config, prefill, {"COMMON": "value"}) + + assert decode_env["ROLE_SETTING"] == "decode" + assert prefill_env["ROLE_SETTING"] == "prefill" + assert decode_env["CUDA_VISIBLE_DEVICES"] == "3" + assert prefill_env["CUDA_VISIBLE_DEVICES"] == "7" + assert decode_env["VLLM_PLUGINS"] == "prime_rl" + + +def test_child_failure_tears_down_complete_process_group(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + config = disaggregated_config(output_dir=tmp_path) + processes = [] + terminated = [] + + class FakeProcess: + def __init__(self, returncode): + self.pid = 1000 + len(processes) + self.returncode = returncode + + def poll(self): + return self.returncode + + def popen(*_args, **_kwargs): + process = FakeProcess(7 if not processes else None) + processes.append(process) + return process + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3") + monkeypatch.setattr(dynamo.subprocess, "Popen", popen) + monkeypatch.setattr(dynamo.signal, "signal", lambda *_args: None) + monkeypatch.setattr(dynamo, "_terminate", terminated.append) + + with pytest.raises(SystemExit) as exc: + dynamo.run_dynamo_local(config) + + assert exc.value.code == 7 + assert len(processes) == 5 + assert terminated == list(reversed(processes)) From 7c1a61d1b35e9fd553497e7efa1baea14c181aa2 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 3 Jul 2026 12:30:32 -0700 Subject: [PATCH 03/18] feat(orchestrator): administer Dynamo inference workers Signed-off-by: Biswa Panda --- src/prime_rl/inference/dynamo_admin.py | 220 ++++++++++++++++++++++ src/prime_rl/orchestrator/orchestrator.py | 4 +- src/prime_rl/utils/client.py | 97 +++++++++- src/prime_rl/utils/elastic.py | 18 ++ tests/unit/inference/test_dynamo_admin.py | 132 +++++++++++++ 5 files changed, 461 insertions(+), 10 deletions(-) create mode 100644 src/prime_rl/inference/dynamo_admin.py create mode 100644 tests/unit/inference/test_dynamo_admin.py diff --git a/src/prime_rl/inference/dynamo_admin.py b/src/prime_rl/inference/dynamo_admin.py new file mode 100644 index 0000000000..b63a690d05 --- /dev/null +++ b/src/prime_rl/inference/dynamo_admin.py @@ -0,0 +1,220 @@ +"""Dynamo worker discovery and engine administration.""" + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path +from urllib.parse import urlsplit, urlunsplit + +import httpx +from httpx import AsyncClient +from tenacity import AsyncRetrying, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential + +from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.logger import get_logger + +NCCL_READY_MARKER = "NCCL_READY" +ADMIN_TIMEOUT_S = 300.0 +UPDATE_WEIGHTS_TIMEOUT_S = 720.0 +_REQUIRED_ROUTES = frozenset( + { + "init_weights_update_group", + "pause_generation", + "resume_generation", + "update_weights_from_disk", + "update_weights_from_distributed", + } +) + + +def _root_url(url: str) -> str: + return url.rstrip("/").removesuffix("/v1") + + +def discovery_urls(config: ClientConfig) -> list[str]: + if config.rl_base_url: + return [_root_url(url) for url in config.rl_base_url] + + configured = os.getenv("DYN_RL_DISCOVERY_URL") + if configured: + return [_root_url(url.strip()) for url in configured.split(",") if url.strip()] + + port = int(os.getenv("DYN_RL_PORT", "8001")) + urls: list[str] = [] + for base_url in config.base_url: + parsed = urlsplit(_root_url(base_url)) + host = parsed.hostname or "localhost" + if ":" in host: + host = f"[{host}]" + netloc = f"{host}:{port}" + urls.append(urlunsplit((parsed.scheme or "http", netloc, "", "", ""))) + return urls + + +async def discover_worker_urls( + discovery_clients: list[AsyncClient], + timeout: int, + model_name: str | None = None, +) -> list[str]: + """Wait for a complete Dynamo worker set and return stable system URLs.""" + logger = get_logger() + last_error: Exception | None = None + deadline = asyncio.get_running_loop().time() + timeout + + while asyncio.get_running_loop().time() < deadline: + try: + results = await asyncio.gather(*(client.get("/v1/rl/workers") for client in discovery_clients)) + workers: list[dict] = [] + for response in results: + response.raise_for_status() + payload = response.json() + if not isinstance(payload, dict) or not isinstance(payload.get("workers"), list): + raise ValueError("Dynamo worker discovery returned an invalid response") + workers.extend(payload["workers"]) + + urls: list[str] = [] + for worker in workers: + if worker.get("error"): + raise RuntimeError( + f"Dynamo worker {worker.get('component', '')} is unhealthy: {worker['error']}" + ) + system_url = worker.get("system_url") + if not system_url: + raise ValueError("Dynamo worker discovery response is missing system_url") + worker_model = worker.get("model") + if model_name is not None and worker_model not in (None, model_name): + raise ValueError(f"Dynamo worker {system_url} serves {worker_model!r}, expected {model_name!r}") + missing = _REQUIRED_ROUTES - set(worker.get("routes", [])) + if missing: + raise ValueError(f"Dynamo worker {system_url} is missing RL routes: {sorted(missing)}") + urls.append(_root_url(system_url)) + + if urls: + if len(set(urls)) != len(urls): + raise ValueError("Dynamo worker discovery returned duplicate system URLs") + resolved = sorted(urls) + logger.info(f"Discovered {len(resolved)} Dynamo inference worker(s)") + return resolved + last_error = RuntimeError("Dynamo worker discovery returned no workers") + except Exception as exc: + last_error = exc + await asyncio.sleep(1) + + raise TimeoutError(f"Dynamo workers were not ready after {timeout} seconds: {last_error}") + + +def validate_worker_membership(expected: tuple[str, ...], discovered: list[str]) -> None: + if tuple(discovered) != expected: + raise RuntimeError( + f"Dynamo worker membership changed after initialization: expected {list(expected)}, discovered {discovered}" + ) + + +class DynamoAdminAPI: + """Typed adapter for Dynamo's per-worker ``/engine`` endpoints.""" + + def __init__(self) -> None: + self._distributed_updates = False + + @staticmethod + def _retryable(exception: BaseException) -> bool: + if isinstance(exception, httpx.HTTPStatusError): + return exception.response.status_code >= 500 + return isinstance(exception, (httpx.TimeoutException, httpx.TransportError)) + + async def _post( + self, + client: AsyncClient, + method: str, + body: dict | None = None, + *, + timeout_s: float = ADMIN_TIMEOUT_S, + ) -> dict: + async for attempt in AsyncRetrying( + retry=retry_if_exception(self._retryable), + stop=stop_after_delay(2 * timeout_s) | stop_after_attempt(10), + wait=wait_exponential(multiplier=1, min=1, max=10), + reraise=True, + ): + with attempt: + response = await client.post( + f"/engine/{method}", + json=body or {}, + timeout=httpx.Timeout(connect=10.0, read=timeout_s, write=60.0, pool=10.0), + ) + response.raise_for_status() + payload = response.json() + if not isinstance(payload, dict): + raise ValueError(f"Dynamo /engine/{method} returned a non-object response") + if payload.get("status") != "ok": + raise RuntimeError(payload.get("message", f"Dynamo /engine/{method} failed")) + return payload + raise AssertionError("unreachable") + + async def initialize_nccl( + self, + clients: list[AsyncClient], + *, + host: str, + port: int, + timeout: int, + inference_world_size: int | None, + quantize_in_weight_transfer: bool, + ) -> None: + world_size = inference_world_size or len(clients) + if not clients or world_size % len(clients) != 0: + raise ValueError(f"inference_world_size={world_size} must be divisible by {len(clients)} Dynamo workers") + gpus_per_worker = world_size // len(clients) + await asyncio.gather( + *( + self._post( + client, + "init_weights_update_group", + { + "host": host, + "port": port, + "rank_offset": index * gpus_per_worker, + "inference_world_size": world_size, + "timeout": timeout, + "quantize_in_weight_transfer": quantize_in_weight_transfer, + "engine_rpc": "init_broadcaster", + }, + ) + for index, client in enumerate(clients) + ) + ) + self._distributed_updates = True + + async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | None, step: int) -> None: + try: + await asyncio.gather( + *(self._post(client, "pause_generation", {"mode": "keep", "clear_cache": False}) for client in clients) + ) + if weight_dir is not None: + marker = weight_dir / NCCL_READY_MARKER + marker.parent.mkdir(parents=True, exist_ok=True) + marker.touch() + + if self._distributed_updates: + body = { + "weight_version": str(step), + "weight_dir": weight_dir.as_posix() if weight_dir is not None else None, + "engine_rpc": "update_weights_from_path", + } + method = "update_weights_from_distributed" + else: + if weight_dir is None: + raise ValueError("Dynamo filesystem weight updates require weight_dir") + body = { + "model_path": str(weight_dir.resolve()), + "weight_version": str(step), + "engine_rpc": "update_weights_from_path", + } + method = "update_weights_from_disk" + + await asyncio.gather( + *(self._post(client, method, body, timeout_s=UPDATE_WEIGHTS_TIMEOUT_S) for client in clients) + ) + finally: + await asyncio.gather(*(self._post(client, "resume_generation") for client in clients)) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index c17ed9e2c7..4cad1fcbe4 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -71,7 +71,6 @@ from prime_rl.trainer.model import setup_tokenizer 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 from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.logger import format_time, get_logger, setup_logger from prime_rl.utils.monitor import setup_monitor @@ -293,8 +292,7 @@ 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, + await self.policy_inference.init_nccl_broadcast( config.weight_broadcast.host, config.weight_broadcast.port, config.weight_broadcast.timeout, diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index ae8c1cde74..0de76c0e27 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -16,6 +16,12 @@ from verifiers.v1.clients.config import EvalClientConfig, TrainClientConfig from prime_rl.configs.shared import ClientConfig +from prime_rl.inference.dynamo_admin import ( + DynamoAdminAPI, + discover_worker_urls, + discovery_urls, + validate_worker_membership, +) from prime_rl.utils.logger import get_logger # Identity tuple used by ``select_train_client`` to key load counts. ``base_url`` @@ -73,6 +79,17 @@ async def update_weights(self, weight_dir: Path | None, lora_name: str | None = """Update weights on all inference servers.""" ... + async def init_nccl_broadcast( + self, + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, + ) -> None: + """Initialize weight broadcast on all inference servers.""" + ... + async def score(self, token_ids: list[int]) -> list[float]: """Prefill-score ``token_ids`` under the pool's model — one logprob per token.""" ... @@ -136,7 +153,14 @@ def __init__( pool_size=pool_size, ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) - self._admin_clients = setup_admin_clients(client_config) + self._client_config = client_config + self._frontend_admin_clients = setup_admin_clients(client_config, urls=client_config.base_url) + self._dynamo_admin = DynamoAdminAPI() if client_config.admin_api == "dynamo" else None + self._discovery_clients = ( + setup_admin_clients(client_config, urls=discovery_urls(client_config)) if self._dynamo_admin else [] + ) + self._dynamo_worker_urls: tuple[str, ...] = () + self._admin_clients = [] if self._dynamo_admin else setup_admin_clients(client_config) self._skip_model_check = client_config.skip_model_check self._wait_for_ready_timeout = client_config.wait_for_ready_timeout self._eval_cycle = cycle(self._eval_clients) @@ -167,14 +191,65 @@ async def select_train_client(self, load: Mapping[ClientIdentity, int]) -> vf.Cl return min(self.train_clients, key=lambda c: load[client_identity(c)]) async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: + ready_timeout = timeout if timeout is not None else self._wait_for_ready_timeout await check_health( - self._admin_clients, timeout=timeout if timeout is not None else self._wait_for_ready_timeout + self._frontend_admin_clients, + timeout=ready_timeout, + ) + await maybe_check_has_model( + self._frontend_admin_clients, + model_name, + skip_model_check=self._skip_model_check, ) - await maybe_check_has_model(self._admin_clients, model_name, skip_model_check=self._skip_model_check) + if self._dynamo_admin is not None: + worker_urls = await discover_worker_urls(self._discovery_clients, ready_timeout, model_name=model_name) + self._dynamo_worker_urls = tuple(worker_urls) + self._admin_clients = setup_admin_clients(self._client_config, urls=worker_urls) + await check_health(self._admin_clients, timeout=ready_timeout, strict=True) + + async def _validate_dynamo_membership(self) -> None: + if self._dynamo_admin is None: + return + discovered = await discover_worker_urls(self._discovery_clients, 30, model_name=self.model_name) + validate_worker_membership(self._dynamo_worker_urls, discovered) async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: + if self._dynamo_admin is not None: + if lora_name is not None: + raise ValueError("Dynamo backend does not yet support Prime LoRA weight updates") + await self._validate_dynamo_membership() + await self._dynamo_admin.update_weights(self._admin_clients, weight_dir, step) + return await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step) + async def init_nccl_broadcast( + self, + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, + ) -> None: + if self._dynamo_admin is not None: + await self._validate_dynamo_membership() + await self._dynamo_admin.initialize_nccl( + self._admin_clients, + host=host, + port=port, + timeout=timeout, + inference_world_size=inference_world_size, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) + return + await init_nccl_broadcast( + self._admin_clients, + host, + port, + timeout, + inference_world_size=inference_world_size, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) + 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.""" @@ -182,6 +257,8 @@ async def score(self, token_ids: list[int]) -> list[float]: async def stop(self) -> None: await self._scorer.aclose() + clients = {*self._frontend_admin_clients, *self._discovery_clients, *self._admin_clients} + await asyncio.gather(*(client.aclose() for client in clients)) async def setup_inference_pool( @@ -250,14 +327,14 @@ def setup_clients( return clients -def setup_admin_clients(client_config: ClientConfig) -> list[AsyncClient]: +def setup_admin_clients(client_config: ClientConfig, urls: list[str] | None = None) -> list[AsyncClient]: """Create dedicated admin clients for weight update operations. Uses a separate connection pool to avoid queueing behind streaming requests. When admin_base_url is set, uses those URLs instead of base_url, allowing weight updates to bypass routers in disaggregated P/D deployments. """ - urls = client_config.admin_base_url if client_config.admin_base_url else client_config.base_url + urls = urls or (client_config.admin_base_url if client_config.admin_base_url else client_config.base_url) def _setup_admin_client(base_url: str) -> httpx.AsyncClient: env_headers = { @@ -297,7 +374,11 @@ async def maybe_check_has_model( async def check_health( - admin_clients: list[AsyncClient], interval: int = 1, log_interval: int = 10, timeout: int = 1800 + admin_clients: list[AsyncClient], + interval: int = 1, + log_interval: int = 10, + timeout: int = 1800, + strict: bool = False, ) -> None: logger = get_logger() @@ -306,7 +387,9 @@ async def _check_health(admin_client: AsyncClient) -> None: logger.debug("Starting pinging /health to check health") while wait_time < timeout: try: - await admin_client.get("/health") + response = await admin_client.get("/health") + if strict: + response.raise_for_status() logger.debug(f"Inference pool is ready after {wait_time} seconds") return except NotFoundError: diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index ef47dee774..4266beffe4 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -26,6 +26,7 @@ ClientIdentity, PrefillScorer, client_identity, + init_nccl_broadcast, load_lora_adapter, setup_admin_clients, setup_clients, @@ -511,3 +512,20 @@ async def update_weights(self, weight_dir: Path | None, lora_name: str | None = if lora_name is None: raise ValueError("Elastic inference pool requires LoRA training (lora_name must be set)") await self.sync_weights(weight_dir, lora_name, step) + + async def init_nccl_broadcast( + self, + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, + ) -> None: + await init_nccl_broadcast( + self.admin_clients, + host, + port, + timeout, + inference_world_size=inference_world_size, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) diff --git a/tests/unit/inference/test_dynamo_admin.py b/tests/unit/inference/test_dynamo_admin.py new file mode 100644 index 0000000000..e1b240da3d --- /dev/null +++ b/tests/unit/inference/test_dynamo_admin.py @@ -0,0 +1,132 @@ +import json +from pathlib import Path + +import httpx +import pytest + +from prime_rl.configs.shared import ClientConfig +from prime_rl.inference.dynamo_admin import ( + DynamoAdminAPI, + discover_worker_urls, + discovery_urls, + validate_worker_membership, +) + + +def async_client(handler, base_url: str = "http://worker:8081") -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(handler), base_url=base_url) + + +def test_discovery_url_defaults_to_rl_listener_port(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DYN_RL_DISCOVERY_URL", raising=False) + monkeypatch.delenv("DYN_RL_PORT", raising=False) + config = ClientConfig(base_url=["http://frontend.example:8000/v1"]) + assert discovery_urls(config) == ["http://frontend.example:8001"] + + +@pytest.mark.asyncio +async def test_worker_discovery_validates_and_sorts_system_urls(): + routes = [ + "init_weights_update_group", + "pause_generation", + "resume_generation", + "update_weights_from_disk", + "update_weights_from_distributed", + ] + + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "namespace": "test", + "workers": [ + {"system_url": "http://worker-b:8081", "routes": routes}, + {"system_url": "http://worker-a:8081", "routes": routes}, + ], + }, + ) + + client = async_client(handler, "http://frontend:8001") + try: + assert await discover_worker_urls([client], timeout=1) == [ + "http://worker-a:8081", + "http://worker-b:8081", + ] + finally: + await client.aclose() + + +@pytest.mark.asyncio +async def test_worker_discovery_rejects_incomplete_admin_surface(): + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"workers": [{"system_url": "http://worker:8081", "routes": []}]}) + + client = async_client(handler, "http://frontend:8001") + try: + with pytest.raises(TimeoutError, match="missing RL routes"): + await discover_worker_urls([client], timeout=1) + finally: + await client.aclose() + + +def test_worker_membership_change_is_rejected(): + with pytest.raises(RuntimeError, match="membership changed"): + validate_worker_membership(("http://worker-a:8081",), ["http://worker-b:8081"]) + + +@pytest.mark.asyncio +async def test_nccl_initialization_and_update_use_engine_routes(tmp_path: Path): + requests: list[tuple[str, dict]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append((request.url.path, json.loads(request.content))) + return httpx.Response(200, json={"status": "ok"}) + + clients = [async_client(handler, f"http://worker-{index}:8081") for index in range(2)] + admin = DynamoAdminAPI() + try: + await admin.initialize_nccl( + clients, + host="localhost", + port=29511, + timeout=12000, + inference_world_size=4, + quantize_in_weight_transfer=False, + ) + await admin.update_weights(clients, tmp_path / "step_1", step=1) + finally: + for client in clients: + await client.aclose() + + init_bodies = [body for path, body in requests if path.endswith("/init_weights_update_group")] + assert [body["rank_offset"] for body in init_bodies] == [0, 2] + assert all(body["inference_world_size"] == 4 for body in init_bodies) + + updates = [body for path, body in requests if path.endswith("/update_weights_from_distributed")] + assert len(updates) == 2 + assert all(body["engine_rpc"] == "update_weights_from_path" for body in updates) + assert all(body["weight_version"] == "1" for body in updates) + assert (tmp_path / "step_1" / "NCCL_READY").exists() + + paths = [path for path, _body in requests] + assert paths.count("/engine/pause_generation") == 2 + assert paths.count("/engine/resume_generation") == 2 + + +@pytest.mark.asyncio +async def test_engine_status_error_is_not_accepted(): + paths = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + if request.url.path.endswith("resume_generation"): + return httpx.Response(200, json={"status": "ok"}) + return httpx.Response(200, json={"status": "error", "message": "not paused"}) + + client = async_client(handler) + try: + with pytest.raises(RuntimeError, match="not paused"): + await DynamoAdminAPI().update_weights([client], Path("weights"), step=1) + finally: + await client.aclose() + assert paths == ["/engine/pause_generation", "/engine/resume_generation"] From 908705df7e03603b749cc6213268695cb20785d9 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 3 Jul 2026 15:26:31 -0700 Subject: [PATCH 04/18] fix(config): stop injecting response-only token flag Signed-off-by: Biswa Panda --- .../prime-rl-configs/src/prime_rl/configs/orchestrator.py | 1 - tests/unit/test_configs.py | 6 ++++++ 2 files changed, 6 insertions(+), 1 deletion(-) 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 2c11a4ddb6..c50f5af42c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -728,7 +728,6 @@ def resolve_env_config(self): if env.algo.sampling.source == "policy": env.sampling.extra_body.setdefault("top_k", -1) env.sampling.extra_body.setdefault("min_p", 0.0) - env.sampling.extra_body.setdefault("return_token_ids", True) if env.is_legacy: # v0 env: cap per-turn response tokens to the training budget (the legacy # bridge applies extra_env_kwargs via env.set_kwargs). diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 8624eca192..f0e081198f 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -186,6 +186,12 @@ def test_env_algo_overrides_top_level(): assert reloaded.train.env[0].algo is not None and reloaded.train.env[0].algo.type == "grpo" +def test_policy_sampling_does_not_add_response_format_flags(): + config = OrchestratorConfig.model_validate({"train": {"env": [{"id": "math-env"}]}}) + + assert config.train.env[0].sampling.extra_body == {"top_k": -1, "min_p": 0.0} + + def test_trainer_enable_token_export_cli_flag(): assert not cli(TrainerConfig, args=[]).enable_token_export assert cli(TrainerConfig, args=["--enable-token-export"]).enable_token_export From 770b568c8dd0f9714a2dcecf94c507a9d2d3f711 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 3 Jul 2026 16:14:17 -0700 Subject: [PATCH 05/18] fix(client): preserve native admin routing Signed-off-by: Biswa Panda --- src/prime_rl/utils/client.py | 16 ++++++++++------ tests/unit/utils/test_client.py | 16 +++++++++++++++- 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 0de76c0e27..7b325a0a6d 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -154,13 +154,16 @@ def __init__( ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) self._client_config = client_config - self._frontend_admin_clients = setup_admin_clients(client_config, urls=client_config.base_url) self._dynamo_admin = DynamoAdminAPI() if client_config.admin_api == "dynamo" else None - self._discovery_clients = ( - setup_admin_clients(client_config, urls=discovery_urls(client_config)) if self._dynamo_admin else [] - ) self._dynamo_worker_urls: tuple[str, ...] = () - self._admin_clients = [] if self._dynamo_admin else setup_admin_clients(client_config) + if self._dynamo_admin is None: + self._admin_clients = setup_admin_clients(client_config) + self._frontend_admin_clients = self._admin_clients + self._discovery_clients = [] + else: + self._frontend_admin_clients = setup_admin_clients(client_config, urls=client_config.base_url) + self._discovery_clients = setup_admin_clients(client_config, urls=discovery_urls(client_config)) + self._admin_clients = [] self._skip_model_check = client_config.skip_model_check self._wait_for_ready_timeout = client_config.wait_for_ready_timeout self._eval_cycle = cycle(self._eval_clients) @@ -334,7 +337,8 @@ def setup_admin_clients(client_config: ClientConfig, urls: list[str] | None = No When admin_base_url is set, uses those URLs instead of base_url, allowing weight updates to bypass routers in disaggregated P/D deployments. """ - urls = urls or (client_config.admin_base_url if client_config.admin_base_url else client_config.base_url) + if urls is None: + urls = client_config.admin_base_url if client_config.admin_base_url else client_config.base_url def _setup_admin_client(base_url: str) -> httpx.AsyncClient: env_headers = { diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 40de4cfee6..8708d74d2c 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -6,7 +6,7 @@ from verifiers.v1.clients.config import EvalClientConfig from prime_rl.configs.shared import ClientConfig -from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter, setup_clients +from prime_rl.utils.client import StaticInferencePool, _is_retryable_lora_error, load_lora_adapter, setup_clients def test_is_retryable_lora_error_returns_true_for_404(): @@ -108,3 +108,17 @@ def test_setup_clients_preserves_chat_client_defaults(): headers={}, ) ] + + +def test_native_pool_preserves_admin_base_url_and_reuses_admin_clients(): + pool = StaticInferencePool( + ClientConfig( + base_url=["http://router:8000/v1"], + admin_base_url=["http://worker:8001/v1"], + ), + model_name="test-model", + ) + + assert pool._frontend_admin_clients is pool._admin_clients + assert [str(client.base_url).rstrip("/") for client in pool.admin_clients] == ["http://worker:8001"] + asyncio.run(pool.stop()) From 75100105040ec1c4e112be477315cd52cce21ee1 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 3 Jul 2026 16:14:17 -0700 Subject: [PATCH 06/18] fix(dynamo): preserve weight update failures Signed-off-by: Biswa Panda --- src/prime_rl/inference/dynamo_admin.py | 11 ++++++++++- tests/unit/inference/test_dynamo_admin.py | 19 +++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/inference/dynamo_admin.py b/src/prime_rl/inference/dynamo_admin.py index b63a690d05..30cfe93cea 100644 --- a/src/prime_rl/inference/dynamo_admin.py +++ b/src/prime_rl/inference/dynamo_admin.py @@ -187,6 +187,7 @@ async def initialize_nccl( self._distributed_updates = True async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | None, step: int) -> None: + primary_error: BaseException | None = None try: await asyncio.gather( *(self._post(client, "pause_generation", {"mode": "keep", "clear_cache": False}) for client in clients) @@ -216,5 +217,13 @@ async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | No await asyncio.gather( *(self._post(client, method, body, timeout_s=UPDATE_WEIGHTS_TIMEOUT_S) for client in clients) ) + except BaseException as exc: + primary_error = exc + raise finally: - await asyncio.gather(*(self._post(client, "resume_generation") for client in clients)) + try: + await asyncio.gather(*(self._post(client, "resume_generation") for client in clients)) + except BaseException as exc: + if primary_error is None: + raise + primary_error.add_note(f"Dynamo resume_generation cleanup also failed: {exc!r}") diff --git a/tests/unit/inference/test_dynamo_admin.py b/tests/unit/inference/test_dynamo_admin.py index e1b240da3d..3a4e9dd743 100644 --- a/tests/unit/inference/test_dynamo_admin.py +++ b/tests/unit/inference/test_dynamo_admin.py @@ -130,3 +130,22 @@ def handler(request: httpx.Request) -> httpx.Response: finally: await client.aclose() assert paths == ["/engine/pause_generation", "/engine/resume_generation"] + + +@pytest.mark.asyncio +async def test_resume_error_does_not_hide_primary_update_error(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + + async def post(_client, method, *_args, **_kwargs): + if method == "pause_generation": + raise RuntimeError("pause failed") + if method == "resume_generation": + raise RuntimeError("resume failed") + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + + with pytest.raises(RuntimeError, match="pause failed") as exc: + await admin.update_weights([object()], Path("weights"), step=1) + + assert exc.value.__notes__ == ["Dynamo resume_generation cleanup also failed: RuntimeError('resume failed')"] From beaec343a7346c2b7a65b47b4187dffed440edad Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sat, 4 Jul 2026 20:24:00 -0700 Subject: [PATCH 07/18] refactor(inference): isolate Dynamo runtime boundaries Signed-off-by: Biswa Panda --- src/prime_rl/inference/dynamo.py | 135 ++++++++++++++++++------- src/prime_rl/utils/client.py | 149 +++++++++++++++++++--------- tests/unit/inference/test_dynamo.py | 35 +++++++ tests/unit/utils/test_client.py | 24 ++++- 4 files changed, 262 insertions(+), 81 deletions(-) diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py index f0b01d4cb9..e17cd12e39 100644 --- a/src/prime_rl/inference/dynamo.py +++ b/src/prime_rl/inference/dynamo.py @@ -51,6 +51,19 @@ } +@dataclass(frozen=True) +class DynamoProcessSpec: + module: str + arguments: tuple[str, ...] + environment_items: tuple[tuple[str, str], ...] + + def command(self, executable: str = sys.executable) -> list[str]: + return [executable, "-m", self.module, *self.arguments] + + def environment(self, base: dict[str, str] | None = None) -> dict[str, str]: + return (base or {}) | dict(self.environment_items) + + @dataclass(frozen=True) class DynamoWorkerSpec: name: str @@ -60,18 +73,10 @@ class DynamoWorkerSpec: nixl_port: int kv_events_port: int | None engine_config: Path + process: DynamoProcessSpec def command(self) -> list[str]: - return [ - sys.executable, - "-m", - "dynamo.vllm", - "--engine-config-json", - str(self.engine_config), - "--disaggregation-mode", - self.role, - "--enable-rl", - ] + return self.process.command() def _json_default(value: Any) -> Any: @@ -98,6 +103,79 @@ def _validate_overrides(source: str, values: dict[str, Any]) -> None: raise ValueError(f"{source} cannot override Dynamo-managed engine keys: {conflicts}") +def _environment_items(values: dict[str, str]) -> tuple[tuple[str, str], ...]: + return tuple(sorted(values.items())) + + +def _role_environment(config: InferenceConfig, role: Role) -> dict[str, str]: + if config.deployment.type != "disaggregated": + return {} + if role == "prefill": + return config.deployment.prefill_env_vars + if role == "decode": + return config.deployment.decode_env_vars + return {} + + +def build_frontend_process( + config: InferenceConfig, + *, + host: str | None = None, + port: int | None = None, +) -> DynamoProcessSpec: + """Build the canonical Dynamo frontend process contract.""" + environment = { + **config.env_vars, + "DYN_ENABLE_RL": "1", + "DYN_RL_PORT": "8001", + } + return DynamoProcessSpec( + module="dynamo.frontend", + arguments=( + "--http-host", + host or config.server.host or "0.0.0.0", + "--http-port", + str(port or config.server.port), + "--router-mode", + "kv", + "--router-reset-states", + "--enable-engine-apis", + ), + environment_items=_environment_items(environment), + ) + + +def build_worker_process( + config: InferenceConfig, + role: Role, + engine_config: Path, + *, + nixl_host: str | None, + nixl_port: int, +) -> DynamoProcessSpec: + """Build the canonical Dynamo vLLM worker process contract.""" + environment = { + **config.env_vars, + "DYN_ENABLE_RL": "1", + "VLLM_NIXL_SIDE_CHANNEL_PORT": str(nixl_port), + "VLLM_PLUGINS": "prime_rl", + **_role_environment(config, role), + } + if nixl_host is not None: + environment["VLLM_NIXL_SIDE_CHANNEL_HOST"] = nixl_host + return DynamoProcessSpec( + module="dynamo.vllm", + arguments=( + "--engine-config-json", + str(engine_config), + "--disaggregation-mode", + role, + "--enable-rl", + ), + environment_items=_environment_items(environment), + ) + + def build_engine_config( config: InferenceConfig, role: Role, @@ -226,25 +304,20 @@ def build_local_worker_specs( nixl_port=20100 + worker_index, kv_events_port=kv_events_port, engine_config=engine_path, + process=build_worker_process( + config, + role, + engine_path, + nixl_host="127.0.0.1", + nixl_port=20100 + worker_index, + ), ) ) return specs def build_frontend_command(config: InferenceConfig) -> list[str]: - return [ - sys.executable, - "-m", - "dynamo.frontend", - "--http-host", - config.server.host or "0.0.0.0", - "--http-port", - str(config.server.port), - "--router-mode", - "kv", - "--router-reset-states", - "--enable-engine-apis", - ] + return build_frontend_process(config).command() def build_worker_environment( @@ -252,19 +325,10 @@ def build_worker_environment( spec: DynamoWorkerSpec, base_environment: dict[str, str], ) -> dict[str, str]: - environment = base_environment | { + return spec.process.environment(base_environment) | { "CUDA_VISIBLE_DEVICES": ",".join(spec.gpu_ids), "DYN_SYSTEM_PORT": str(spec.system_port), - "VLLM_NIXL_SIDE_CHANNEL_HOST": "127.0.0.1", - "VLLM_NIXL_SIDE_CHANNEL_PORT": str(spec.nixl_port), - "VLLM_PLUGINS": "prime_rl", } - if config.deployment.type == "disaggregated": - role_environment = ( - config.deployment.prefill_env_vars if spec.role == "prefill" else config.deployment.decode_env_vars - ) - environment.update(role_environment) - return environment def _terminate(process: subprocess.Popen) -> None: @@ -300,11 +364,12 @@ def request_stop(_signum, _frame): processes: list[subprocess.Popen] = [] with tempfile.TemporaryDirectory(prefix="prime-dynamo-") as temporary_dir: environment.setdefault("DYN_FILE_KV", str(Path(temporary_dir) / "discovery")) - frontend_env = environment | {"CUDA_VISIBLE_DEVICES": ""} + frontend = build_frontend_process(config) + frontend_env = frontend.environment(environment) | {"CUDA_VISIBLE_DEVICES": ""} frontend_env.pop("DYN_SYSTEM_PORT", None) try: - processes.append(subprocess.Popen(build_frontend_command(config), env=frontend_env, start_new_session=True)) + processes.append(subprocess.Popen(frontend.command(), env=frontend_env, start_new_session=True)) for spec in specs: worker_env = build_worker_environment(config, spec, environment) processes.append(subprocess.Popen(spec.command(), env=worker_env, start_new_session=True)) diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 7b325a0a6d..091620c52c 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -132,8 +132,8 @@ async def aclose(self) -> None: await asyncio.gather(*(c.close() for c in self._clients.values())) -class StaticInferencePool: - """Static inference pool with fixed client list.""" +class _StaticClientPool: + """Shared request-client behavior for fixed inference pools.""" def __init__( self, @@ -154,16 +154,6 @@ def __init__( ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) self._client_config = client_config - self._dynamo_admin = DynamoAdminAPI() if client_config.admin_api == "dynamo" else None - self._dynamo_worker_urls: tuple[str, ...] = () - if self._dynamo_admin is None: - self._admin_clients = setup_admin_clients(client_config) - self._frontend_admin_clients = self._admin_clients - self._discovery_clients = [] - else: - self._frontend_admin_clients = setup_admin_clients(client_config, urls=client_config.base_url) - self._discovery_clients = setup_admin_clients(client_config, urls=discovery_urls(client_config)) - self._admin_clients = [] self._skip_model_check = client_config.skip_model_check self._wait_for_ready_timeout = client_config.wait_for_ready_timeout self._eval_cycle = cycle(self._eval_clients) @@ -193,6 +183,34 @@ async def select_train_client(self, load: Mapping[ClientIdentity, int]) -> vf.Cl await asyncio.sleep(0.5) return min(self.train_clients, key=lambda c: load[client_identity(c)]) + async def score(self, token_ids: list[int]) -> list[float]: + """Prefill-score tokens under this pool's model.""" + return await self._scorer.score(self.train_clients, self.model_name, token_ids) + + +class StaticInferencePool(_StaticClientPool): + """Static native-vLLM inference pool with fixed clients.""" + + def __init__( + self, + client_config: ClientConfig, + model_name: str, + train_client_type: str = "openai_chat_completions", + eval_client_type: str = "openai_chat_completions", + renderer_config: RendererConfig | None = None, + pool_size: int | None = None, + ): + super().__init__( + client_config, + model_name, + train_client_type, + eval_client_type, + renderer_config, + pool_size, + ) + self._admin_clients = setup_admin_clients(client_config) + self._frontend_admin_clients = self._admin_clients + async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: ready_timeout = timeout if timeout is not None else self._wait_for_ready_timeout await check_health( @@ -204,25 +222,8 @@ async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> N model_name, skip_model_check=self._skip_model_check, ) - if self._dynamo_admin is not None: - worker_urls = await discover_worker_urls(self._discovery_clients, ready_timeout, model_name=model_name) - self._dynamo_worker_urls = tuple(worker_urls) - self._admin_clients = setup_admin_clients(self._client_config, urls=worker_urls) - await check_health(self._admin_clients, timeout=ready_timeout, strict=True) - - async def _validate_dynamo_membership(self) -> None: - if self._dynamo_admin is None: - return - discovered = await discover_worker_urls(self._discovery_clients, 30, model_name=self.model_name) - validate_worker_membership(self._dynamo_worker_urls, discovered) async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: - if self._dynamo_admin is not None: - if lora_name is not None: - raise ValueError("Dynamo backend does not yet support Prime LoRA weight updates") - await self._validate_dynamo_membership() - await self._dynamo_admin.update_weights(self._admin_clients, weight_dir, step) - return await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step) async def init_nccl_broadcast( @@ -233,17 +234,6 @@ async def init_nccl_broadcast( inference_world_size: int | None = None, quantize_in_weight_transfer: bool = False, ) -> None: - if self._dynamo_admin is not None: - await self._validate_dynamo_membership() - await self._dynamo_admin.initialize_nccl( - self._admin_clients, - host=host, - port=port, - timeout=timeout, - inference_world_size=inference_world_size, - quantize_in_weight_transfer=quantize_in_weight_transfer, - ) - return await init_nccl_broadcast( self._admin_clients, host, @@ -253,10 +243,78 @@ async def init_nccl_broadcast( quantize_in_weight_transfer=quantize_in_weight_transfer, ) - 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.""" - return await self._scorer.score(self.train_clients, self.model_name, token_ids) + async def stop(self) -> None: + await self._scorer.aclose() + clients = {*self._frontend_admin_clients, *self._admin_clients} + await asyncio.gather(*(client.aclose() for client in clients)) + + +class DynamoInferencePool(_StaticClientPool): + """Static Dynamo pool with worker discovery and Dynamo administration.""" + + def __init__( + self, + client_config: ClientConfig, + model_name: str, + train_client_type: str = "openai_chat_completions", + eval_client_type: str = "openai_chat_completions", + renderer_config: RendererConfig | None = None, + pool_size: int | None = None, + ): + super().__init__( + client_config, + model_name, + train_client_type, + eval_client_type, + renderer_config, + pool_size, + ) + self._frontend_admin_clients = setup_admin_clients(client_config, urls=client_config.base_url) + self._discovery_clients = setup_admin_clients(client_config, urls=discovery_urls(client_config)) + self._admin_clients: list[AsyncClient] = [] + self._worker_urls: tuple[str, ...] = () + self._admin = DynamoAdminAPI() + + async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: + ready_timeout = timeout if timeout is not None else self._wait_for_ready_timeout + await check_health(self._frontend_admin_clients, timeout=ready_timeout) + await maybe_check_has_model( + self._frontend_admin_clients, + model_name, + skip_model_check=self._skip_model_check, + ) + worker_urls = await discover_worker_urls(self._discovery_clients, ready_timeout, model_name=model_name) + self._worker_urls = tuple(worker_urls) + self._admin_clients = setup_admin_clients(self._client_config, urls=worker_urls) + await check_health(self._admin_clients, timeout=ready_timeout, strict=True) + + async def _validate_membership(self) -> None: + discovered = await discover_worker_urls(self._discovery_clients, 30, model_name=self.model_name) + validate_worker_membership(self._worker_urls, discovered) + + async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: + if lora_name is not None: + raise ValueError("Dynamo backend does not yet support Prime LoRA weight updates") + await self._validate_membership() + await self._admin.update_weights(self._admin_clients, weight_dir, step) + + async def init_nccl_broadcast( + self, + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, + ) -> None: + await self._validate_membership() + await self._admin.initialize_nccl( + 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 stop(self) -> None: await self._scorer.aclose() @@ -285,7 +343,8 @@ async def setup_inference_pool( pool_size=pool_size, ) - return StaticInferencePool( + pool_type = DynamoInferencePool if client_config.admin_api == "dynamo" else StaticInferencePool + return pool_type( client_config, model_name=model_name, train_client_type=train_client_type, diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index d6d77df4cb..39bab4392e 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -7,8 +7,10 @@ from prime_rl.inference import dynamo from prime_rl.inference.dynamo import ( build_engine_config, + build_frontend_process, build_local_worker_specs, build_worker_environment, + build_worker_process, write_role_engine_configs, ) @@ -90,6 +92,39 @@ def test_wrapper_options_are_not_written_to_engine_json(): assert "enable_rl" not in engine +def test_process_specs_own_canonical_commands_and_environment(tmp_path: Path): + config = disaggregated_config( + env_vars={"SHARED": "value"}, + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_nodes": 1, + "num_decode_nodes": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + "prefill_env_vars": {"ROLE": "prefill"}, + }, + ) + + frontend = build_frontend_process(config) + prefill = build_worker_process( + config, + "prefill", + tmp_path / "prefill.json", + nixl_host="127.0.0.1", + nixl_port=20100, + ) + + assert frontend.module == "dynamo.frontend" + assert frontend.arguments[-1] == "--enable-engine-apis" + assert frontend.environment()["DYN_ENABLE_RL"] == "1" + assert prefill.module == "dynamo.vllm" + assert prefill.arguments[-3:] == ("--disaggregation-mode", "prefill", "--enable-rl") + assert prefill.environment()["ROLE"] == "prefill" + assert prefill.environment()["VLLM_PLUGINS"] == "prime_rl" + assert prefill.environment()["VLLM_NIXL_SIDE_CHANNEL_PORT"] == "20100" + + def test_worker_environment_applies_only_matching_role_overrides(tmp_path: Path): config = disaggregated_config( deployment={ diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 8708d74d2c..d615715633 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -6,7 +6,14 @@ from verifiers.v1.clients.config import EvalClientConfig from prime_rl.configs.shared import ClientConfig -from prime_rl.utils.client import StaticInferencePool, _is_retryable_lora_error, load_lora_adapter, setup_clients +from prime_rl.utils.client import ( + DynamoInferencePool, + StaticInferencePool, + _is_retryable_lora_error, + load_lora_adapter, + setup_clients, + setup_inference_pool, +) def test_is_retryable_lora_error_returns_true_for_404(): @@ -122,3 +129,18 @@ def test_native_pool_preserves_admin_base_url_and_reuses_admin_clients(): assert pool._frontend_admin_clients is pool._admin_clients assert [str(client.base_url).rstrip("/") for client in pool.admin_clients] == ["http://worker:8001"] asyncio.run(pool.stop()) + + +def test_setup_inference_pool_selects_dynamo_pool_once(): + pool = asyncio.run( + setup_inference_pool( + ClientConfig( + base_url=["http://frontend:8000/v1"], + admin_api="dynamo", + ), + model_name="test-model", + ) + ) + + assert isinstance(pool, DynamoInferencePool) + asyncio.run(pool.stop()) From c6fab8f860461ccaa22407a783baaced61f4e5a5 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sat, 4 Jul 2026 20:40:47 -0700 Subject: [PATCH 08/18] refactor(inference): keep process environment canonical Signed-off-by: Biswa Panda --- src/prime_rl/inference/dynamo.py | 5 +---- tests/unit/inference/test_dynamo.py | 4 ++-- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py index e17cd12e39..e5c26cc3d2 100644 --- a/src/prime_rl/inference/dynamo.py +++ b/src/prime_rl/inference/dynamo.py @@ -321,7 +321,6 @@ def build_frontend_command(config: InferenceConfig) -> list[str]: def build_worker_environment( - config: InferenceConfig, spec: DynamoWorkerSpec, base_environment: dict[str, str], ) -> dict[str, str]: @@ -353,8 +352,6 @@ def run_dynamo_local(config: InferenceConfig) -> None: environment.setdefault("DYN_EVENT_PLANE", "zmq") environment.setdefault("DYN_FILE_KV_TTL_SECS", "1800") environment.setdefault("DYN_NAMESPACE", f"prime-rl-{os.getpid()}") - environment.setdefault("DYN_ENABLE_RL", "1") - environment.setdefault("DYN_RL_PORT", "8001") environment.setdefault("PYTHONHASHSEED", "0") def request_stop(_signum, _frame): @@ -371,7 +368,7 @@ def request_stop(_signum, _frame): try: processes.append(subprocess.Popen(frontend.command(), env=frontend_env, start_new_session=True)) for spec in specs: - worker_env = build_worker_environment(config, spec, environment) + worker_env = build_worker_environment(spec, environment) processes.append(subprocess.Popen(spec.command(), env=worker_env, start_new_session=True)) while all(process.poll() is None for process in processes): diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index 39bab4392e..9aef232176 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -140,8 +140,8 @@ def test_worker_environment_applies_only_matching_role_overrides(tmp_path: Path) ) decode, prefill = build_local_worker_specs(config, tmp_path, gpu_ids=["3", "7"]) - decode_env = build_worker_environment(config, decode, {"COMMON": "value"}) - prefill_env = build_worker_environment(config, prefill, {"COMMON": "value"}) + decode_env = build_worker_environment(decode, {"COMMON": "value"}) + prefill_env = build_worker_environment(prefill, {"COMMON": "value"}) assert decode_env["ROLE_SETTING"] == "decode" assert prefill_env["ROLE_SETTING"] == "prefill" From 5ce895940a3a3399c92491d6950747960f4573ae Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sat, 4 Jul 2026 20:43:12 -0700 Subject: [PATCH 09/18] refactor(inference): remove redundant process wrappers Signed-off-by: Biswa Panda --- src/prime_rl/entrypoints/inference.py | 6 +++--- src/prime_rl/inference/dynamo.py | 11 +---------- tests/unit/inference/test_dynamo.py | 2 +- 3 files changed, 5 insertions(+), 14 deletions(-) diff --git a/src/prime_rl/entrypoints/inference.py b/src/prime_rl/entrypoints/inference.py index 474cff9a76..8f5e11caf4 100644 --- a/src/prime_rl/entrypoints/inference.py +++ b/src/prime_rl/entrypoints/inference.py @@ -149,12 +149,12 @@ def inference_local(config: InferenceConfig): if config.dry_run: if config.backend.type == "dynamo": - from prime_rl.inference.dynamo import build_frontend_command, build_local_worker_specs + from prime_rl.inference.dynamo import build_frontend_process, build_local_worker_specs specs = build_local_worker_specs(config) - logger.info(f"Dynamo frontend: {' '.join(build_frontend_command(config))}") + logger.info(f"Dynamo frontend: {' '.join(build_frontend_process(config).command())}") for spec in specs: - logger.info(f"Dynamo {spec.name}: {' '.join(spec.command())}") + logger.info(f"Dynamo {spec.name}: {' '.join(spec.process.command())}") logger.success("Dry run complete. To start inference locally, remove --dry-run from your command.") return diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py index e5c26cc3d2..97907716f7 100644 --- a/src/prime_rl/inference/dynamo.py +++ b/src/prime_rl/inference/dynamo.py @@ -72,12 +72,8 @@ class DynamoWorkerSpec: system_port: int nixl_port: int kv_events_port: int | None - engine_config: Path process: DynamoProcessSpec - def command(self) -> list[str]: - return self.process.command() - def _json_default(value: Any) -> Any: if isinstance(value, Path): @@ -303,7 +299,6 @@ def build_local_worker_specs( system_port=8081 + worker_index, nixl_port=20100 + worker_index, kv_events_port=kv_events_port, - engine_config=engine_path, process=build_worker_process( config, role, @@ -316,10 +311,6 @@ def build_local_worker_specs( return specs -def build_frontend_command(config: InferenceConfig) -> list[str]: - return build_frontend_process(config).command() - - def build_worker_environment( spec: DynamoWorkerSpec, base_environment: dict[str, str], @@ -369,7 +360,7 @@ def request_stop(_signum, _frame): processes.append(subprocess.Popen(frontend.command(), env=frontend_env, start_new_session=True)) for spec in specs: worker_env = build_worker_environment(spec, environment) - processes.append(subprocess.Popen(spec.command(), env=worker_env, start_new_session=True)) + processes.append(subprocess.Popen(spec.process.command(), env=worker_env, start_new_session=True)) while all(process.poll() is None for process in processes): time.sleep(0.2) diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index 9aef232176..38c66303e1 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -83,7 +83,7 @@ def test_local_specs_allocate_four_workers_and_unique_ports(tmp_path: Path): assert len({spec.system_port for spec in specs}) == 4 assert len({spec.nixl_port for spec in specs}) == 4 assert [spec.kv_events_port for spec in specs if spec.role == "prefill"] == [20080, 20081] - assert all("--enable-rl" in spec.command() for spec in specs) + assert all("--enable-rl" in spec.process.command() for spec in specs) def test_wrapper_options_are_not_written_to_engine_json(): From 69986cce2a5214b4b9004de8cbba29cc5a1fe0c6 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sat, 4 Jul 2026 20:43:43 -0700 Subject: [PATCH 10/18] refactor(inference): derive worker port evidence Signed-off-by: Biswa Panda --- src/prime_rl/inference/dynamo.py | 4 ---- tests/unit/inference/test_dynamo.py | 10 ++++++++-- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py index 97907716f7..3c3592bfe6 100644 --- a/src/prime_rl/inference/dynamo.py +++ b/src/prime_rl/inference/dynamo.py @@ -70,8 +70,6 @@ class DynamoWorkerSpec: role: Role gpu_ids: tuple[str, ...] system_port: int - nixl_port: int - kv_events_port: int | None process: DynamoProcessSpec @@ -297,8 +295,6 @@ def build_local_worker_specs( role=role, gpu_ids=worker_gpus, system_port=8081 + worker_index, - nixl_port=20100 + worker_index, - kv_events_port=kv_events_port, process=build_worker_process( config, role, diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index 38c66303e1..69822db9e8 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -81,8 +81,14 @@ def test_local_specs_allocate_four_workers_and_unique_ports(tmp_path: Path): assert [spec.role for spec in specs] == ["decode", "decode", "prefill", "prefill"] assert [spec.gpu_ids for spec in specs] == [("4",), ("5",), ("6",), ("7",)] assert len({spec.system_port for spec in specs}) == 4 - assert len({spec.nixl_port for spec in specs}) == 4 - assert [spec.kv_events_port for spec in specs if spec.role == "prefill"] == [20080, 20081] + assert len({spec.process.environment()["VLLM_NIXL_SIDE_CHANNEL_PORT"] for spec in specs}) == 4 + prefill_configs = [ + json.loads(Path(spec.process.arguments[1]).read_text()) for spec in specs if spec.role == "prefill" + ] + assert [config["kv_events_config"]["endpoint"] for config in prefill_configs] == [ + "tcp://*:20080", + "tcp://*:20081", + ] assert all("--enable-rl" in spec.process.command() for spec in specs) From 1ac4298e6336982c322b9989df7f456edeb1639c Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 15:24:05 -0700 Subject: [PATCH 11/18] test: update Dynamo topology fixtures --- tests/unit/inference/test_dynamo.py | 16 ++++++++-------- tests/unit/test_configs.py | 8 ++++---- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index 69822db9e8..7bbc9ab287 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -22,8 +22,8 @@ def disaggregated_config(**overrides) -> InferenceConfig: "deployment": { "type": "disaggregated", "gpus_per_node": 1, - "num_prefill_nodes": 2, - "num_decode_nodes": 2, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, "num_prefill_replicas": 2, "num_decode_replicas": 2, }, @@ -50,8 +50,8 @@ def test_role_overrides_are_isolated(): deployment={ "type": "disaggregated", "gpus_per_node": 1, - "num_prefill_nodes": 2, - "num_decode_nodes": 2, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, "num_prefill_replicas": 2, "num_decode_replicas": 2, "prefill_vllm_overrides": {"max_num_batched_tokens": 8192}, @@ -104,8 +104,8 @@ def test_process_specs_own_canonical_commands_and_environment(tmp_path: Path): deployment={ "type": "disaggregated", "gpus_per_node": 1, - "num_prefill_nodes": 1, - "num_decode_nodes": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, "num_prefill_replicas": 1, "num_decode_replicas": 1, "prefill_env_vars": {"ROLE": "prefill"}, @@ -136,8 +136,8 @@ def test_worker_environment_applies_only_matching_role_overrides(tmp_path: Path) deployment={ "type": "disaggregated", "gpus_per_node": 1, - "num_prefill_nodes": 1, - "num_decode_nodes": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, "num_prefill_replicas": 1, "num_decode_replicas": 1, "prefill_env_vars": {"ROLE_SETTING": "prefill"}, diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index f0e081198f..4204fa3539 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -228,8 +228,8 @@ def test_dynamo_disaggregated_config_is_local_and_enables_prefix_caching(): "deployment": { "type": "disaggregated", "gpus_per_node": 1, - "num_prefill_nodes": 2, - "num_decode_nodes": 2, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, "num_prefill_replicas": 2, "num_decode_replicas": 2, }, @@ -267,8 +267,8 @@ def test_single_node_dynamo_disaggregated_keeps_per_worker_dp_and_sets_nccl_worl "deployment": { "type": "disaggregated", "gpus_per_node": 1, - "num_prefill_nodes": 2, - "num_decode_nodes": 2, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, "num_prefill_replicas": 2, "num_decode_replicas": 2, }, From 04eecebc41e01787c622c3431a846a8ec2d107b6 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 16:05:07 -0700 Subject: [PATCH 12/18] fix(dynamo): enforce runtime topology contracts --- .../src/prime_rl/configs/inference.py | 2 + .../src/prime_rl/configs/rl.py | 41 ++- .../src/prime_rl/configs/shared.py | 6 + src/prime_rl/entrypoints/inference.py | 4 +- src/prime_rl/inference/dynamo.py | 166 +++++++++-- src/prime_rl/inference/dynamo_admin.py | 275 ++++++++++++++---- src/prime_rl/utils/client.py | 124 ++++++-- tests/unit/inference/test_dynamo.py | 123 +++++++- tests/unit/inference/test_dynamo_admin.py | 252 ++++++++++++++-- tests/unit/test_configs.py | 68 +++++ tests/unit/utils/test_client.py | 61 ++++ 11 files changed, 985 insertions(+), 137 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index e5e84c21b7..0577fdd9b1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -461,6 +461,8 @@ def validate_dynamo_backend(self): ) if self.deployment.type == "multi_node": raise ValueError("Dynamo multi-node inference must use a DynamoGraphDeployment.") + if self.enable_lora: + raise ValueError("The Dynamo backend does not support LoRA weight updates.") router = getattr(self.deployment, "router", None) if router is not None and router.type == "llm-d": raise ValueError("The Dynamo backend owns request routing and cannot use the llm-d router.") 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 b7e9d92768..6abafc72ba 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -460,6 +460,12 @@ def auto_setup_lora(self): return self + @model_validator(mode="after") + def validate_auto_setup_does_not_enable_dynamo_lora(self): + if self.inference is not None and self.inference.backend.type == "dynamo" and self.inference.enable_lora: + raise ValueError("The Dynamo backend does not support LoRA weight updates.") + return self + @model_validator(mode="after") def auto_setup_router_replay(self): if self.trainer.enable_router_replay: @@ -691,21 +697,48 @@ def auto_setup_inference_client(self): if self.inference is None: return self client = self.orchestrator.model.client + client_updates: dict[str, Any] = {} if "admin_api" in client.model_fields_set and client.admin_api != self.inference.backend.type: raise ValueError( "orchestrator.model.client.admin_api conflicts with inference.backend.type; " "configure the backend only under inference." ) - client.admin_api = self.inference.backend.type + client_updates["admin_api"] = self.inference.backend.type if "dp_rank_count" not in client.model_fields_set: if self.inference.backend.type == "dynamo" or self.deployment.type == "multi_node": - client.dp_rank_count = 1 + client_updates["dp_rank_count"] = 1 + else: + client_updates["dp_rank_count"] = self.inference.data_parallel_size_local or self.inference.parallel.dp + if self.inference.backend.type == "dynamo": + if self.inference.deployment.type == "disaggregated": + deployment = self.inference.deployment + expected_roles = ("prefill",) * deployment.num_prefill_replicas + ( + "decode", + ) * deployment.num_decode_replicas + expected_gpus_per_worker = deployment.gpus_per_node else: - client.dp_rank_count = self.inference.data_parallel_size_local or self.inference.parallel.dp + expected_roles = ("agg",) + expected_gpus_per_worker = self.inference.parallel.tp * self.inference.parallel.dp + + expected_topology = { + "dynamo_worker_roles": expected_roles, + "dynamo_gpus_per_worker": expected_gpus_per_worker, + } + for field, expected in expected_topology.items(): + if field in client.model_fields_set and getattr(client, field) != expected: + raise ValueError( + f"orchestrator.model.client.{field} conflicts with the inference topology; " + "configure the topology only under inference." + ) + client_updates.update(expected_topology) if not self.orchestrator.any_policy_sourced and "base_url" not in client.model_fields_set: host = self.inference.server.host or "localhost" port = self.inference.server.port - client.base_url = [f"http://{host}:{port}/v1"] + client_updates["base_url"] = [f"http://{host}:{port}/v1"] + + updated_client = client.model_copy(update=client_updates) + updated_model = self.orchestrator.model.model_copy(update={"client": updated_client}) + self.orchestrator = self.orchestrator.model_copy(update={"model": updated_model}) return self @model_validator(mode="after") 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 3485d05f0d..fbc3e2f4c2 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -153,6 +153,12 @@ class ClientConfig(BaseConfig): rl_base_url: list[str] | None = None """Dynamo RL worker-discovery URLs. When omitted, they are derived from the frontend URL or ``DYN_RL_DISCOVERY_URL``.""" + dynamo_worker_roles: tuple[Literal["agg", "prefill", "decode"], ...] | None = None + """Exact Dynamo worker roles expected during readiness. Auto-derived from the local inference topology.""" + + dynamo_gpus_per_worker: int | None = Field(None, ge=1) + """GPUs owned by each discovered Dynamo worker. Auto-derived from the local inference topology.""" + 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.""" diff --git a/src/prime_rl/entrypoints/inference.py b/src/prime_rl/entrypoints/inference.py index 8f5e11caf4..3a6a62d637 100644 --- a/src/prime_rl/entrypoints/inference.py +++ b/src/prime_rl/entrypoints/inference.py @@ -149,9 +149,9 @@ def inference_local(config: InferenceConfig): if config.dry_run: if config.backend.type == "dynamo": - from prime_rl.inference.dynamo import build_frontend_process, build_local_worker_specs + from prime_rl.inference.dynamo import build_dry_run_worker_specs, build_frontend_process - specs = build_local_worker_specs(config) + specs = build_dry_run_worker_specs(config) logger.info(f"Dynamo frontend: {' '.join(build_frontend_process(config).command())}") for spec in specs: logger.info(f"Dynamo {spec.name}: {' '.join(spec.process.command())}") diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py index 3c3592bfe6..764e324e39 100644 --- a/src/prime_rl/inference/dynamo.py +++ b/src/prime_rl/inference/dynamo.py @@ -23,6 +23,7 @@ PREFILL_ENGINE_CONFIG = "prefill-engine.json" DECODE_ENGINE_CONFIG = "decode-engine.json" AGG_ENGINE_CONFIG = "agg-engine.json" +CHAT_TEMPLATE_ASSET = "chat-template.jinja" _ENGINE_CONFIG_EXCLUDED = frozenset( { @@ -49,6 +50,11 @@ "nccl": "prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker", "filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker", } +_WORKER_COMPONENT = { + "agg": "backend", + "prefill": "prefill", + "decode": "backend", +} @dataclass(frozen=True) @@ -111,11 +117,61 @@ def _role_environment(config: InferenceConfig, role: Role) -> dict[str, str]: return {} +def resolve_chat_template_content(config: InferenceConfig) -> str | None: + """Resolve a configured inline or file-backed chat template to immutable content.""" + template = config.model.chat_template + if template is None: + return None + template_source = Path(os.path.expanduser(template)) + return template_source.read_text(encoding="utf-8") if template_source.is_file() else template + + +def _materialize_chat_template( + config: InferenceConfig, + template_content: str, + output_dir: Path | None, +) -> Path: + config_dir = output_dir or (get_config_dir(config.output_dir) / ENGINE_CONFIG_DIR) + template_path = config_dir / CHAT_TEMPLATE_ASSET + template_path.parent.mkdir(parents=True, exist_ok=True) + template_path.write_text(template_content, encoding="utf-8") + return template_path + + +def _frontend_model_arguments( + config: InferenceConfig, + output_dir: Path | None, + runtime_chat_template_path: Path | None, +) -> tuple[str, ...]: + template_content = resolve_chat_template_content(config) + if template_content is None: + return () + template_path = runtime_chat_template_path or _materialize_chat_template(config, template_content, output_dir) + tool_arguments = ( + ("--enable-auto-tool-choice", "--tool-call-parser", config.model.tool_call_parser) + if config.model.tool_call_parser is not None + else () + ) + reasoning_arguments = ( + ("--reasoning-parser", config.model.reasoning_parser) if config.model.reasoning_parser is not None else () + ) + return ( + *tool_arguments, + *reasoning_arguments, + "--dyn-chat-processor", + "vllm", + "--chat-template", + str(template_path), + ) + + def build_frontend_process( config: InferenceConfig, *, host: str | None = None, port: int | None = None, + output_dir: Path | None = None, + runtime_chat_template_path: Path | None = None, ) -> DynamoProcessSpec: """Build the canonical Dynamo frontend process contract.""" environment = { @@ -123,22 +179,46 @@ def build_frontend_process( "DYN_ENABLE_RL": "1", "DYN_RL_PORT": "8001", } + arguments = ( + "--http-host", + host or config.server.host or "0.0.0.0", + "--http-port", + str(port or config.server.port), + "--router-mode", + "kv", + "--router-reset-states", + "--enable-engine-apis", + *_frontend_model_arguments(config, output_dir, runtime_chat_template_path), + ) return DynamoProcessSpec( module="dynamo.frontend", - arguments=( - "--http-host", - host or config.server.host or "0.0.0.0", - "--http-port", - str(port or config.server.port), - "--router-mode", - "kv", - "--router-reset-states", - "--enable-engine-apis", - ), + arguments=arguments, environment_items=_environment_items(environment), ) +def _worker_parser_arguments(config: InferenceConfig, role: Role) -> tuple[str, ...]: + if role == "prefill": + return () + tool_arguments = ( + ("--dyn-tool-call-parser", config.model.tool_call_parser) if config.model.tool_call_parser is not None else () + ) + reasoning_arguments = ( + ("--dyn-reasoning-parser", config.model.reasoning_parser) if config.model.reasoning_parser is not None else () + ) + return (*tool_arguments, *reasoning_arguments) + + +def _worker_endpoint_contract(namespace: str | None, component: str) -> tuple[dict[str, str], tuple[str, ...]]: + if namespace is None: + return {}, () + endpoint = f"dyn://{namespace}.{component}.generate" + return ( + {"DYN_NAMESPACE": namespace, "DYN_ENDPOINT": endpoint}, + ("--endpoint", endpoint), + ) + + def build_worker_process( config: InferenceConfig, role: Role, @@ -146,26 +226,34 @@ def build_worker_process( *, nixl_host: str | None, nixl_port: int, + namespace: str | None = None, ) -> DynamoProcessSpec: """Build the canonical Dynamo vLLM worker process contract.""" + resolved_namespace = namespace or config.env_vars.get("DYN_NAMESPACE") + component = _WORKER_COMPONENT[role] + endpoint_environment, endpoint_arguments = _worker_endpoint_contract(resolved_namespace, component) environment = { **config.env_vars, + **_role_environment(config, role), "DYN_ENABLE_RL": "1", + "DYN_COMPONENT": component, + **endpoint_environment, "VLLM_NIXL_SIDE_CHANNEL_PORT": str(nixl_port), + **({"VLLM_NIXL_SIDE_CHANNEL_HOST": nixl_host} if nixl_host is not None else {}), "VLLM_PLUGINS": "prime_rl", - **_role_environment(config, role), } - if nixl_host is not None: - environment["VLLM_NIXL_SIDE_CHANNEL_HOST"] = nixl_host + arguments = ( + "--engine-config-json", + str(engine_config), + *endpoint_arguments, + "--disaggregation-mode", + role, + "--enable-rl", + *_worker_parser_arguments(config, role), + ) return DynamoProcessSpec( module="dynamo.vllm", - arguments=( - "--engine-config-json", - str(engine_config), - "--disaggregation-mode", - role, - "--enable-rl", - ), + arguments=arguments, environment_items=_environment_items(environment), ) @@ -255,10 +343,12 @@ def build_local_worker_specs( config: InferenceConfig, output_dir: Path | None = None, gpu_ids: list[str] | None = None, + namespace: str | None = None, ) -> list[DynamoWorkerSpec]: """Allocate local workers and write instance-specific engine configs.""" config_dir = output_dir or (get_config_dir(config.output_dir) / ENGINE_CONFIG_DIR) - available = gpu_ids or _visible_gpu_ids() + available = gpu_ids if gpu_ids is not None else _visible_gpu_ids() + resolved_namespace = namespace or config.env_vars.get("DYN_NAMESPACE") or "dynamo" if config.deployment.type == "disaggregated": deployment: DisaggregatedInferenceDeploymentConfig = config.deployment @@ -301,12 +391,31 @@ def build_local_worker_specs( engine_path, nixl_host="127.0.0.1", nixl_port=20100 + worker_index, + namespace=resolved_namespace, ), ) ) return specs +def build_dry_run_worker_specs( + config: InferenceConfig, + output_dir: Path | None = None, +) -> list[DynamoWorkerSpec]: + """Build local specs without consulting host GPU hardware.""" + if config.deployment.type == "disaggregated": + worker_count = config.deployment.num_prefill_replicas + config.deployment.num_decode_replicas + gpu_count = worker_count * config.deployment.gpus_per_node + else: + gpu_count = config.parallel.tp * config.parallel.dp + return build_local_worker_specs( + config, + output_dir=output_dir, + gpu_ids=[f"" for index in range(gpu_count)], + namespace=config.env_vars.get("DYN_NAMESPACE") or "dynamo", + ) + + def build_worker_environment( spec: DynamoWorkerSpec, base_environment: dict[str, str], @@ -333,13 +442,14 @@ def _terminate(process: subprocess.Popen) -> None: def run_dynamo_local(config: InferenceConfig) -> None: """Run a Dynamo frontend and all configured workers until one exits.""" - specs = build_local_worker_specs(config) environment = os.environ.copy() environment.setdefault("DYN_DISCOVERY_BACKEND", "file") environment.setdefault("DYN_EVENT_PLANE", "zmq") environment.setdefault("DYN_FILE_KV_TTL_SECS", "1800") - environment.setdefault("DYN_NAMESPACE", f"prime-rl-{os.getpid()}") + namespace = config.env_vars.get("DYN_NAMESPACE") or environment.get("DYN_NAMESPACE") or f"prime-rl-{os.getpid()}" + environment["DYN_NAMESPACE"] = namespace environment.setdefault("PYTHONHASHSEED", "0") + specs = build_local_worker_specs(config, namespace=namespace) def request_stop(_signum, _frame): raise KeyboardInterrupt @@ -358,9 +468,15 @@ def request_stop(_signum, _frame): worker_env = build_worker_environment(spec, environment) processes.append(subprocess.Popen(spec.process.command(), env=worker_env, start_new_session=True)) - while all(process.poll() is None for process in processes): + exited_process = next((process for process in processes if process.poll() is not None), None) + while exited_process is None: time.sleep(0.2) - raise SystemExit(next((process.returncode for process in processes if process.returncode), 1)) + exited_process = next((process for process in processes if process.poll() is not None), None) + returncode = exited_process.returncode + if returncode is None: + raise RuntimeError("Dynamo child exit was observed without a return code") + # A clean child exit is still a service failure while its siblings are supervised. + raise SystemExit(returncode if returncode != 0 else 1) except KeyboardInterrupt: return finally: diff --git a/src/prime_rl/inference/dynamo_admin.py b/src/prime_rl/inference/dynamo_admin.py index 30cfe93cea..14fc9211a3 100644 --- a/src/prime_rl/inference/dynamo_admin.py +++ b/src/prime_rl/inference/dynamo_admin.py @@ -4,7 +4,11 @@ import asyncio import os +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass from pathlib import Path +from typing import Literal, TypeAlias from urllib.parse import urlsplit, urlunsplit import httpx @@ -17,6 +21,8 @@ NCCL_READY_MARKER = "NCCL_READY" ADMIN_TIMEOUT_S = 300.0 UPDATE_WEIGHTS_TIMEOUT_S = 720.0 +DISCOVERY_REQUEST_TIMEOUT_S = 10.0 +DISCOVERY_POLL_INTERVAL_S = 1.0 _REQUIRED_ROUTES = frozenset( { "init_weights_update_group", @@ -27,6 +33,63 @@ } ) +WorkerRole: TypeAlias = Literal["agg", "prefill", "decode"] + + +@dataclass(frozen=True, slots=True) +class DynamoWorker: + """Restart-safe identity and admin capabilities for one Dynamo worker.""" + + instance_id: int + component: str + role: WorkerRole + system_url: str + model: str + routes: frozenset[str] + + +@dataclass(frozen=True, slots=True) +class DynamoTopology: + """The exact worker shape Prime expects to administer.""" + + roles: tuple[WorkerRole, ...] + gpus_per_worker: int + + def __post_init__(self) -> None: + if not self.roles: + raise ValueError("Dynamo topology must contain at least one worker") + invalid_roles = set(self.roles) - {"agg", "prefill", "decode"} + if invalid_roles: + raise ValueError(f"Dynamo topology contains invalid roles: {sorted(invalid_roles)}") + if isinstance(self.gpus_per_worker, bool) or self.gpus_per_worker < 1: + raise ValueError("Dynamo topology gpus_per_worker must be at least one") + + def validate(self, workers: Sequence[DynamoWorker]) -> None: + expected = Counter(self.roles) + observed = Counter(worker.role for worker in workers) + if observed != expected: + raise ValueError( + "Dynamo worker topology does not match the configured roles: " + f"expected {dict(sorted(expected.items()))}, observed {dict(sorted(observed.items()))}" + ) + + def role_for_component(self, component: str) -> WorkerRole: + normalized = component.casefold() + if "prefill" in normalized: + return "prefill" + if "decode" in normalized: + return "decode" + if normalized in {"backend", "vllmworker", "agg", "aggregate", "aggregated", "vllmaggworker"}: + expected = set(self.roles) + if "decode" in expected and "agg" not in expected: + return "decode" + if "agg" in expected and "decode" not in expected: + return "agg" + raise ValueError( + f"Dynamo worker component {component!r} is ambiguous for configured roles {sorted(expected)}" + ) + raise ValueError(f"Dynamo worker component {component!r} has no recognized inference role") + def _root_url(url: str) -> str: return url.rstrip("/").removesuffix("/v1") @@ -52,62 +115,127 @@ def discovery_urls(config: ClientConfig) -> list[str]: return urls -async def discover_worker_urls( +def _worker_sort_key(worker: DynamoWorker) -> tuple[str, str, int, str]: + return (worker.role, worker.component, worker.instance_id, worker.system_url) + + +def _parse_worker(value: object, model_name: str, topology: DynamoTopology) -> DynamoWorker: + if not isinstance(value, Mapping): + raise ValueError("Dynamo worker discovery returned a non-object worker") + component = value.get("component") + instance_id = value.get("instance_id") + system_url = value.get("system_url") + model = value.get("model") + routes = value.get("routes") + if not isinstance(component, str) or not component: + raise ValueError("Dynamo worker discovery response is missing component") + if not isinstance(instance_id, int) or isinstance(instance_id, bool) or instance_id < 0: + raise ValueError(f"Dynamo worker {component!r} has an invalid instance_id") + if value.get("error"): + raise RuntimeError(f"Dynamo worker {component}[{instance_id}] is unhealthy: {value['error']}") + if not isinstance(system_url, str) or not system_url: + raise ValueError(f"Dynamo worker {component}[{instance_id}] is missing system_url") + if model != model_name: + raise ValueError( + f"Dynamo worker {component}[{instance_id}] at {system_url} serves {model!r}, expected {model_name!r}" + ) + if not isinstance(routes, list) or not all(isinstance(route, str) for route in routes): + raise ValueError(f"Dynamo worker {component}[{instance_id}] has invalid routes") + route_set = frozenset(routes) + missing = _REQUIRED_ROUTES - route_set + if missing: + raise ValueError(f"Dynamo worker {system_url} is missing RL routes: {sorted(missing)}") + return DynamoWorker( + instance_id=instance_id, + component=component, + role=topology.role_for_component(component), + system_url=_root_url(system_url), + model=model, + routes=route_set, + ) + + +def _parse_snapshot( + payload: object, + model_name: str, + topology: DynamoTopology, +) -> tuple[str, tuple[DynamoWorker, ...]]: + if not isinstance(payload, Mapping) or not isinstance(payload.get("workers"), list): + raise ValueError("Dynamo worker discovery returned an invalid response") + namespace = payload.get("namespace") + if not isinstance(namespace, str) or not namespace: + raise ValueError("Dynamo worker discovery response is missing namespace") + workers = tuple( + sorted( + (_parse_worker(value, model_name, topology) for value in payload["workers"]), + key=_worker_sort_key, + ) + ) + identities = {(worker.component, worker.instance_id) for worker in workers} + if len(identities) != len(workers): + raise ValueError("Dynamo worker discovery returned duplicate worker identities") + if len({worker.system_url for worker in workers}) != len(workers): + raise ValueError("Dynamo worker discovery returned duplicate system URLs") + return namespace, workers + + +async def discover_workers( discovery_clients: list[AsyncClient], - timeout: int, - model_name: str | None = None, -) -> list[str]: - """Wait for a complete Dynamo worker set and return stable system URLs.""" + timeout: float, + *, + model_name: str, + topology: DynamoTopology, +) -> tuple[DynamoWorker, ...]: + """Wait for all frontends to report one identical, complete worker set.""" + if not discovery_clients: + raise ValueError("Dynamo worker discovery requires at least one frontend") + if timeout <= 0: + raise TimeoutError("Dynamo worker discovery deadline has already expired") logger = get_logger() last_error: Exception | None = None - deadline = asyncio.get_running_loop().time() + timeout + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout - while asyncio.get_running_loop().time() < deadline: + while (remaining := deadline - loop.time()) > 0: try: - results = await asyncio.gather(*(client.get("/v1/rl/workers") for client in discovery_clients)) - workers: list[dict] = [] + request_timeout = min(remaining, DISCOVERY_REQUEST_TIMEOUT_S) + async with asyncio.timeout(remaining): + results = await asyncio.gather( + *( + client.get("/v1/rl/workers", timeout=httpx.Timeout(request_timeout)) + for client in discovery_clients + ) + ) + snapshots: list[tuple[str, tuple[DynamoWorker, ...]]] = [] for response in results: response.raise_for_status() - payload = response.json() - if not isinstance(payload, dict) or not isinstance(payload.get("workers"), list): - raise ValueError("Dynamo worker discovery returned an invalid response") - workers.extend(payload["workers"]) - - urls: list[str] = [] - for worker in workers: - if worker.get("error"): - raise RuntimeError( - f"Dynamo worker {worker.get('component', '')} is unhealthy: {worker['error']}" - ) - system_url = worker.get("system_url") - if not system_url: - raise ValueError("Dynamo worker discovery response is missing system_url") - worker_model = worker.get("model") - if model_name is not None and worker_model not in (None, model_name): - raise ValueError(f"Dynamo worker {system_url} serves {worker_model!r}, expected {model_name!r}") - missing = _REQUIRED_ROUTES - set(worker.get("routes", [])) - if missing: - raise ValueError(f"Dynamo worker {system_url} is missing RL routes: {sorted(missing)}") - urls.append(_root_url(system_url)) - - if urls: - if len(set(urls)) != len(urls): - raise ValueError("Dynamo worker discovery returned duplicate system URLs") - resolved = sorted(urls) - logger.info(f"Discovered {len(resolved)} Dynamo inference worker(s)") - return resolved - last_error = RuntimeError("Dynamo worker discovery returned no workers") + snapshots.append(_parse_snapshot(response.json(), model_name, topology)) + first = snapshots[0] + if any(snapshot != first for snapshot in snapshots[1:]): + raise ValueError("Dynamo discovery frontends returned inconsistent worker snapshots") + workers = first[1] + topology.validate(workers) + logger.info(f"Discovered {len(workers)} Dynamo inference worker(s)") + return workers except Exception as exc: last_error = exc - await asyncio.sleep(1) + remaining = deadline - loop.time() + if remaining > 0: + await asyncio.sleep(min(DISCOVERY_POLL_INTERVAL_S, remaining)) - raise TimeoutError(f"Dynamo workers were not ready after {timeout} seconds: {last_error}") + raise TimeoutError(f"Dynamo workers were not ready after {timeout} seconds: {last_error!r}") -def validate_worker_membership(expected: tuple[str, ...], discovered: list[str]) -> None: - if tuple(discovered) != expected: +def validate_worker_membership( + expected: Sequence[DynamoWorker], + discovered: Sequence[DynamoWorker], +) -> None: + expected_workers = tuple(sorted(expected, key=_worker_sort_key)) + discovered_workers = tuple(sorted(discovered, key=_worker_sort_key)) + if discovered_workers != expected_workers: raise RuntimeError( - f"Dynamo worker membership changed after initialization: expected {list(expected)}, discovered {discovered}" + "Dynamo worker membership changed after initialization: " + f"expected {expected_workers!r}, discovered {discovered_workers!r}" ) @@ -130,7 +258,25 @@ async def _post( body: dict | None = None, *, timeout_s: float = ADMIN_TIMEOUT_S, + retry_transient: bool = False, ) -> dict: + async def post_once() -> dict: + response = await client.post( + f"/engine/{method}", + json=body or {}, + timeout=httpx.Timeout(connect=10.0, read=timeout_s, write=60.0, pool=10.0), + ) + response.raise_for_status() + payload = response.json() + if not isinstance(payload, dict): + raise ValueError(f"Dynamo /engine/{method} returned a non-object response") + if payload.get("status") != "ok": + raise RuntimeError(payload.get("message", f"Dynamo /engine/{method} failed")) + return payload + + if not retry_transient: + return await post_once() + async for attempt in AsyncRetrying( retry=retry_if_exception(self._retryable), stop=stop_after_delay(2 * timeout_s) | stop_after_attempt(10), @@ -138,18 +284,7 @@ async def _post( reraise=True, ): with attempt: - response = await client.post( - f"/engine/{method}", - json=body or {}, - timeout=httpx.Timeout(connect=10.0, read=timeout_s, write=60.0, pool=10.0), - ) - response.raise_for_status() - payload = response.json() - if not isinstance(payload, dict): - raise ValueError(f"Dynamo /engine/{method} returned a non-object response") - if payload.get("status") != "ok": - raise RuntimeError(payload.get("message", f"Dynamo /engine/{method} failed")) - return payload + return await post_once() raise AssertionError("unreachable") async def initialize_nccl( @@ -160,12 +295,20 @@ async def initialize_nccl( port: int, timeout: int, inference_world_size: int | None, + gpus_per_worker: int, quantize_in_weight_transfer: bool, ) -> None: - world_size = inference_world_size or len(clients) - if not clients or world_size % len(clients) != 0: - raise ValueError(f"inference_world_size={world_size} must be divisible by {len(clients)} Dynamo workers") - gpus_per_worker = world_size // len(clients) + if not clients: + raise ValueError("Cannot initialize NCCL without Dynamo workers") + if isinstance(gpus_per_worker, bool) or gpus_per_worker < 1: + raise ValueError("gpus_per_worker must be at least one") + expected_world_size = len(clients) * gpus_per_worker + world_size = expected_world_size if inference_world_size is None else inference_world_size + if world_size != expected_world_size: + raise ValueError( + f"inference_world_size={world_size} does not match {len(clients)} Dynamo workers " + f"with {gpus_per_worker} GPUs each ({expected_world_size})" + ) await asyncio.gather( *( self._post( @@ -190,7 +333,15 @@ async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | No primary_error: BaseException | None = None try: await asyncio.gather( - *(self._post(client, "pause_generation", {"mode": "keep", "clear_cache": False}) for client in clients) + *( + self._post( + client, + "pause_generation", + {"mode": "keep", "clear_cache": False}, + retry_transient=True, + ) + for client in clients + ) ) if weight_dir is not None: marker = weight_dir / NCCL_READY_MARKER @@ -222,7 +373,9 @@ async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | No raise finally: try: - await asyncio.gather(*(self._post(client, "resume_generation") for client in clients)) + await asyncio.gather( + *(self._post(client, "resume_generation", retry_transient=True) for client in clients) + ) except BaseException as exc: if primary_error is None: raise diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 091620c52c..483325fc3c 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -18,7 +18,9 @@ from prime_rl.configs.shared import ClientConfig from prime_rl.inference.dynamo_admin import ( DynamoAdminAPI, - discover_worker_urls, + DynamoTopology, + DynamoWorker, + discover_workers, discovery_urls, validate_worker_membership, ) @@ -35,6 +37,19 @@ def client_identity(client: vf.ClientConfig) -> ClientIdentity: return (client.base_url, client.headers.get("X-data-parallel-rank")) +def _readiness_deadline(timeout: float) -> float: + if timeout <= 0: + raise TimeoutError("Inference readiness timeout must be greater than zero") + return asyncio.get_running_loop().time() + timeout + + +def _remaining_readiness_timeout(deadline: float, phase: str) -> float: + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + raise TimeoutError(f"Inference readiness deadline expired before {phase}") + return remaining + + @runtime_checkable class InferencePool(Protocol): """Protocol for inference pools (static or elastic).""" @@ -213,14 +228,16 @@ def __init__( async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: ready_timeout = timeout if timeout is not None else self._wait_for_ready_timeout + deadline = _readiness_deadline(ready_timeout) await check_health( self._frontend_admin_clients, - timeout=ready_timeout, + timeout=_remaining_readiness_timeout(deadline, "frontend health"), ) await maybe_check_has_model( self._frontend_admin_clients, model_name, skip_model_check=self._skip_model_check, + timeout=_remaining_readiness_timeout(deadline, "model registration"), ) async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: @@ -261,6 +278,15 @@ def __init__( renderer_config: RendererConfig | None = None, pool_size: int | None = None, ): + if client_config.dynamo_worker_roles is None or client_config.dynamo_gpus_per_worker is None: + raise ValueError( + "Dynamo clients require dynamo_worker_roles and dynamo_gpus_per_worker; " + "local RL configs derive them from the inference topology" + ) + topology = DynamoTopology( + roles=client_config.dynamo_worker_roles, + gpus_per_worker=client_config.dynamo_gpus_per_worker, + ) super().__init__( client_config, model_name, @@ -272,25 +298,52 @@ def __init__( self._frontend_admin_clients = setup_admin_clients(client_config, urls=client_config.base_url) self._discovery_clients = setup_admin_clients(client_config, urls=discovery_urls(client_config)) self._admin_clients: list[AsyncClient] = [] - self._worker_urls: tuple[str, ...] = () + self._topology = topology + self._workers: tuple[DynamoWorker, ...] = () self._admin = DynamoAdminAPI() async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: ready_timeout = timeout if timeout is not None else self._wait_for_ready_timeout - await check_health(self._frontend_admin_clients, timeout=ready_timeout) + deadline = _readiness_deadline(ready_timeout) + await check_health( + self._frontend_admin_clients, + timeout=_remaining_readiness_timeout(deadline, "frontend health"), + ) await maybe_check_has_model( self._frontend_admin_clients, model_name, skip_model_check=self._skip_model_check, + timeout=_remaining_readiness_timeout(deadline, "model registration"), ) - worker_urls = await discover_worker_urls(self._discovery_clients, ready_timeout, model_name=model_name) - self._worker_urls = tuple(worker_urls) - self._admin_clients = setup_admin_clients(self._client_config, urls=worker_urls) - await check_health(self._admin_clients, timeout=ready_timeout, strict=True) + workers = await discover_workers( + self._discovery_clients, + _remaining_readiness_timeout(deadline, "worker discovery"), + model_name=model_name, + topology=self._topology, + ) + admin_clients = setup_admin_clients(self._client_config, urls=[worker.system_url for worker in workers]) + try: + await check_health( + admin_clients, + timeout=_remaining_readiness_timeout(deadline, "worker health"), + strict=True, + ) + except BaseException: + await asyncio.gather(*(client.aclose() for client in admin_clients)) + raise + previous_clients = self._admin_clients + self._workers = workers + self._admin_clients = admin_clients + await asyncio.gather(*(client.aclose() for client in previous_clients)) async def _validate_membership(self) -> None: - discovered = await discover_worker_urls(self._discovery_clients, 30, model_name=self.model_name) - validate_worker_membership(self._worker_urls, discovered) + discovered = await discover_workers( + self._discovery_clients, + 30, + model_name=self.model_name, + topology=self._topology, + ) + validate_worker_membership(self._workers, discovered) async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: if lora_name is not None: @@ -313,6 +366,7 @@ async def init_nccl_broadcast( port=port, timeout=timeout, inference_world_size=inference_world_size, + gpus_per_worker=self._topology.gpus_per_worker, quantize_in_weight_transfer=quantize_in_weight_transfer, ) @@ -422,14 +476,21 @@ def _setup_admin_client(base_url: str) -> httpx.AsyncClient: async def maybe_check_has_model( - admin_clients: list[AsyncClient], model_name: str, skip_model_check: bool = False + admin_clients: list[AsyncClient], + model_name: str, + skip_model_check: bool = False, + timeout: float = 1800, ) -> None: if skip_model_check: return logger = get_logger() logger.debug(f"Checking if model {model_name} is in the inference pool") - results = await asyncio.gather(*[admin_client.get("/v1/models") for admin_client in admin_clients]) + async with asyncio.timeout(timeout): + results = await asyncio.gather( + *(admin_client.get("/v1/models", timeout=httpx.Timeout(timeout)) for admin_client in admin_clients) + ) for admin_client, result in zip(admin_clients, results): + result.raise_for_status() models = result.json()["data"] if not any(model["id"] == model_name for model in models): raise ValueError(f"Model {model_name} was not found in the inference pool on {admin_client.base_url}") @@ -438,34 +499,49 @@ async def maybe_check_has_model( async def check_health( admin_clients: list[AsyncClient], - interval: int = 1, - log_interval: int = 10, - timeout: int = 1800, + interval: float = 1, + log_interval: float = 10, + timeout: float = 1800, strict: bool = False, ) -> None: logger = get_logger() async def _check_health(admin_client: AsyncClient) -> None: - wait_time = 0 + loop = asyncio.get_running_loop() + started = loop.time() + deadline = started + timeout + next_log_at = log_interval + last_error: Exception | None = None logger.debug("Starting pinging /health to check health") - while wait_time < timeout: + while (remaining := deadline - loop.time()) > 0: try: - response = await admin_client.get("/health") + response = await admin_client.get( + "/health", + timeout=httpx.Timeout(min(remaining, 10.0)), + ) if strict: response.raise_for_status() - logger.debug(f"Inference pool is ready after {wait_time} seconds") + elapsed = loop.time() - started + logger.debug(f"Inference pool is ready after {elapsed:.1f} seconds") return except NotFoundError: logger.warning("The route /health does not exist. Skipping health check.") return except Exception as e: - if wait_time % log_interval == 0 and wait_time > 0: + last_error = e + elapsed = loop.time() - started + if elapsed >= next_log_at: logger.warning( - f"Inference server was not reached after {wait_time} seconds (Error: {e}) on {admin_client.base_url}" + f"Inference server was not reached after {elapsed:.1f} seconds " + f"(Error: {e}) on {admin_client.base_url}" ) - await asyncio.sleep(interval) - wait_time += interval - msg = f"Inference server is not ready after {wait_time} (>{timeout}) seconds. Aborting..." + next_log_at += log_interval + remaining = deadline - loop.time() + if remaining > 0: + await asyncio.sleep(min(interval, remaining)) + msg = ( + f"Inference server {admin_client.base_url} is not ready after {timeout} seconds; last error: {last_error!r}" + ) logger.error(msg) raise TimeoutError(msg) diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index 7bbc9ab287..e45f66cf0d 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -4,6 +4,7 @@ import pytest from prime_rl.configs.inference import InferenceConfig +from prime_rl.entrypoints import inference as inference_entrypoint from prime_rl.inference import dynamo from prime_rl.inference.dynamo import ( build_engine_config, @@ -131,6 +132,82 @@ def test_process_specs_own_canonical_commands_and_environment(tmp_path: Path): assert prefill.environment()["VLLM_NIXL_SIDE_CHANNEL_PORT"] == "20100" +def test_process_specs_preserve_custom_chat_template_and_parsers(tmp_path: Path): + source = tmp_path / "source-template.jinja" + source.write_text("{{ messages | length }}") + config = disaggregated_config( + model={ + "chat_template": str(source), + "tool_call_parser": "hermes", + "reasoning_parser": "qwen3", + } + ) + + frontend = build_frontend_process(config, output_dir=tmp_path / "generated") + worker = build_worker_process( + config, + "decode", + tmp_path / "decode.json", + nixl_host=None, + nixl_port=20100, + ) + + template_path = Path(frontend.arguments[frontend.arguments.index("--chat-template") + 1]) + assert template_path == tmp_path / "generated" / "chat-template.jinja" + assert template_path.read_text() == source.read_text() + assert frontend.arguments[-4:-2] == ("--dyn-chat-processor", "vllm") + assert worker.arguments[-4:] == ( + "--dyn-tool-call-parser", + "hermes", + "--dyn-reasoning-parser", + "qwen3", + ) + + +@pytest.mark.parametrize( + ("role", "component"), + [("prefill", "prefill"), ("decode", "backend"), ("agg", "backend")], +) +def test_worker_process_uses_deterministic_role_endpoint(tmp_path: Path, role: str, component: str): + process = build_worker_process( + disaggregated_config(env_vars={"DYN_NAMESPACE": "prime-test"}), + role, + tmp_path / f"{role}.json", + nixl_host=None, + nixl_port=20100, + ) + + endpoint = f"dyn://prime-test.{component}.generate" + assert process.environment()["DYN_NAMESPACE"] == "prime-test" + assert process.environment()["DYN_COMPONENT"] == component + assert process.environment()["DYN_ENDPOINT"] == endpoint + assert process.arguments[2:4] == ("--endpoint", endpoint) + + +def test_inline_chat_template_is_materialized_verbatim(tmp_path: Path): + config = disaggregated_config(model={"chat_template": "{{ messages }}"}) + + frontend = build_frontend_process(config, output_dir=tmp_path) + + template_path = Path(frontend.arguments[-1]) + assert template_path.read_text() == "{{ messages }}" + + +def test_frontend_runtime_chat_template_path_does_not_materialize_host_file(tmp_path: Path): + config = disaggregated_config(model={"chat_template": "{{ messages }}"}) + output_dir = tmp_path / "render-host" + runtime_path = Path("/etc/prime-rl/dynamo/chat-template.jinja") + + frontend = build_frontend_process( + config, + output_dir=output_dir, + runtime_chat_template_path=runtime_path, + ) + + assert frontend.arguments[-1] == str(runtime_path) + assert not output_dir.exists() + + def test_worker_environment_applies_only_matching_role_overrides(tmp_path: Path): config = disaggregated_config( deployment={ @@ -154,9 +231,49 @@ def test_worker_environment_applies_only_matching_role_overrides(tmp_path: Path) assert decode_env["CUDA_VISIBLE_DEVICES"] == "3" assert prefill_env["CUDA_VISIBLE_DEVICES"] == "7" assert decode_env["VLLM_PLUGINS"] == "prime_rl" + assert decode_env["DYN_COMPONENT"] == "backend" + assert prefill_env["DYN_COMPONENT"] == "prefill" + + +def test_aggregated_worker_uses_canonical_component_name(tmp_path: Path): + config = InferenceConfig.model_validate({"backend": {"type": "dynamo"}}) + spec = build_local_worker_specs(config, tmp_path, gpu_ids=["0"])[0] + + assert build_worker_environment(spec, {})["DYN_COMPONENT"] == "backend" + + +def test_dynamo_dry_run_uses_symbolic_gpu_slots(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + config = disaggregated_config(output_dir=tmp_path, dry_run=True) + captured_gpu_ids = [] + original_build_specs = dynamo.build_local_worker_specs + + class FakeLogger: + def info(self, _message): + pass + + def success(self, _message): + pass + + def build_specs(config, output_dir=None, gpu_ids=None, namespace=None): + captured_gpu_ids.extend(gpu_ids or []) + return original_build_specs(config, output_dir=output_dir, gpu_ids=gpu_ids, namespace=namespace) + + monkeypatch.setattr(dynamo, "_visible_gpu_ids", lambda: pytest.fail("dry-run queried physical GPUs")) + monkeypatch.setattr(dynamo, "build_local_worker_specs", build_specs) + monkeypatch.setattr(inference_entrypoint, "setup_logger", lambda *_args, **_kwargs: FakeLogger()) + + inference_entrypoint.inference_local(config) + + assert captured_gpu_ids == ["", "", "", ""] -def test_child_failure_tears_down_complete_process_group(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize(("child_code", "supervisor_code"), [(7, 7), (0, 1)]) +def test_child_exit_tears_down_complete_process_group( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + child_code: int, + supervisor_code: int, +): config = disaggregated_config(output_dir=tmp_path) processes = [] terminated = [] @@ -170,7 +287,7 @@ def poll(self): return self.returncode def popen(*_args, **_kwargs): - process = FakeProcess(7 if not processes else None) + process = FakeProcess(child_code if not processes else None) processes.append(process) return process @@ -182,6 +299,6 @@ def popen(*_args, **_kwargs): with pytest.raises(SystemExit) as exc: dynamo.run_dynamo_local(config) - assert exc.value.code == 7 + assert exc.value.code == supervisor_code assert len(processes) == 5 assert terminated == list(reversed(processes)) diff --git a/tests/unit/inference/test_dynamo_admin.py b/tests/unit/inference/test_dynamo_admin.py index 3a4e9dd743..48df7972da 100644 --- a/tests/unit/inference/test_dynamo_admin.py +++ b/tests/unit/inference/test_dynamo_admin.py @@ -7,11 +7,40 @@ from prime_rl.configs.shared import ClientConfig from prime_rl.inference.dynamo_admin import ( DynamoAdminAPI, - discover_worker_urls, + DynamoTopology, + DynamoWorker, + discover_workers, discovery_urls, validate_worker_membership, ) +ROUTES = frozenset( + { + "init_weights_update_group", + "pause_generation", + "resume_generation", + "update_weights_from_disk", + "update_weights_from_distributed", + } +) + + +def worker( + instance_id: int, + *, + component: str = "backend", + system_url: str | None = None, + model: str = "test-model", + routes: frozenset[str] = ROUTES, +) -> dict: + return { + "component": component, + "instance_id": instance_id, + "system_url": system_url or f"http://worker-{instance_id}:8081", + "model": model, + "routes": sorted(routes), + } + def async_client(handler, base_url: str = "http://worker:8081") -> httpx.AsyncClient: return httpx.AsyncClient(transport=httpx.MockTransport(handler), base_url=base_url) @@ -26,52 +55,161 @@ def test_discovery_url_defaults_to_rl_listener_port(monkeypatch: pytest.MonkeyPa @pytest.mark.asyncio async def test_worker_discovery_validates_and_sorts_system_urls(): - routes = [ - "init_weights_update_group", - "pause_generation", - "resume_generation", - "update_weights_from_disk", - "update_weights_from_distributed", - ] - def handler(_request: httpx.Request) -> httpx.Response: return httpx.Response( 200, json={ "namespace": "test", "workers": [ - {"system_url": "http://worker-b:8081", "routes": routes}, - {"system_url": "http://worker-a:8081", "routes": routes}, + worker(2, system_url="http://worker-b:8081"), + worker(1, system_url="http://worker-a:8081"), ], }, ) client = async_client(handler, "http://frontend:8001") try: - assert await discover_worker_urls([client], timeout=1) == [ - "http://worker-a:8081", - "http://worker-b:8081", - ] + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg", "agg"), gpus_per_worker=2), + ) finally: await client.aclose() + assert [item.system_url for item in discovered] == ["http://worker-a:8081", "http://worker-b:8081"] + assert discovered[0] == DynamoWorker( + instance_id=1, + component="backend", + role="agg", + system_url="http://worker-a:8081", + model="test-model", + routes=ROUTES, + ) + + +@pytest.mark.asyncio +async def test_worker_discovery_waits_for_exact_staggered_topology(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + workers = [worker(1, component="prefill")] + if calls > 1: + # Dynamo advertises decode as `backend`; the expected topology + # disambiguates it from an aggregated `backend` worker. + workers.append(worker(2, component="backend")) + return httpx.Response(200, json={"namespace": "test", "workers": workers}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("prefill", "decode"), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert {item.role for item in discovered} == {"prefill", "decode"} + + +@pytest.mark.asyncio +async def test_worker_discovery_deduplicates_consistent_frontend_snapshots(): + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"namespace": "test", "workers": [worker(7)]}) + + clients = [async_client(handler, f"http://frontend-{index}:8001") for index in range(2)] + try: + discovered = await discover_workers( + clients, + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + for client in clients: + await client.aclose() + + assert len(discovered) == 1 + assert discovered[0].instance_id == 7 + + +@pytest.mark.asyncio +async def test_worker_discovery_rejects_inconsistent_frontend_snapshots(): + def first(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"namespace": "test", "workers": [worker(7)]}) + + def restarted(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"namespace": "test", "workers": [worker(8)]}) + + clients = [async_client(first, "http://frontend-a:8001"), async_client(restarted, "http://frontend-b:8001")] + try: + with pytest.raises(TimeoutError, match="inconsistent worker snapshots"): + await discover_workers( + clients, + timeout=0.01, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + for client in clients: + await client.aclose() + + +@pytest.mark.asyncio +async def test_worker_discovery_bounds_each_get_by_remaining_deadline(): + request_timeouts: list[dict[str, float]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + request_timeouts.append(request.extensions["timeout"]) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1)]}) + + client = async_client(handler, "http://frontend:8001") + try: + await discover_workers( + [client], + timeout=0.5, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert len(request_timeouts) == 1 + assert 0 < request_timeouts[0]["read"] <= 0.5 + @pytest.mark.asyncio async def test_worker_discovery_rejects_incomplete_admin_surface(): def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(200, json={"workers": [{"system_url": "http://worker:8081", "routes": []}]}) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1, routes=frozenset())]}) client = async_client(handler, "http://frontend:8001") try: with pytest.raises(TimeoutError, match="missing RL routes"): - await discover_worker_urls([client], timeout=1) + await discover_workers( + [client], + timeout=0.01, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) finally: await client.aclose() def test_worker_membership_change_is_rejected(): + expected = (DynamoWorker(1, "backend", "agg", "http://worker:8081", "test-model", ROUTES),) + restarted = [ + DynamoWorker(2, "backend", "agg", "http://worker:8081", "test-model", ROUTES), + ] with pytest.raises(RuntimeError, match="membership changed"): - validate_worker_membership(("http://worker-a:8081",), ["http://worker-b:8081"]) + validate_worker_membership(expected, restarted) @pytest.mark.asyncio @@ -91,6 +229,7 @@ def handler(request: httpx.Request) -> httpx.Response: port=29511, timeout=12000, inference_world_size=4, + gpus_per_worker=2, quantize_in_weight_transfer=False, ) await admin.update_weights(clients, tmp_path / "step_1", step=1) @@ -113,6 +252,83 @@ def handler(request: httpx.Request) -> httpx.Response: assert paths.count("/engine/resume_generation") == 2 +@pytest.mark.asyncio +async def test_nccl_initialization_does_not_replay_ambiguous_timeout(): + paths: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + raise httpx.ReadTimeout("response lost after collective may have started", request=request) + + client = async_client(handler) + try: + with pytest.raises(httpx.ReadTimeout): + await DynamoAdminAPI().initialize_nccl( + [client], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=1, + gpus_per_worker=1, + quantize_in_weight_transfer=False, + ) + finally: + await client.aclose() + + assert paths == ["/engine/init_weights_update_group"] + + +@pytest.mark.asyncio +async def test_nccl_initialization_rejects_world_size_that_conflicts_with_topology(): + requests = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal requests + requests += 1 + return httpx.Response(200, json={"status": "ok"}) + + client = async_client(handler) + try: + with pytest.raises(ValueError, match="does not match"): + await DynamoAdminAPI().initialize_nccl( + [client], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=1, + gpus_per_worker=2, + quantize_in_weight_transfer=False, + ) + finally: + await client.aclose() + + assert requests == 0 + + +@pytest.mark.asyncio +async def test_weight_collective_does_not_replay_ambiguous_timeout(tmp_path: Path): + paths: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + if request.url.path.endswith("update_weights_from_disk"): + raise httpx.ReadTimeout("response lost after update may have committed", request=request) + return httpx.Response(200, json={"status": "ok"}) + + client = async_client(handler) + try: + with pytest.raises(httpx.ReadTimeout): + await DynamoAdminAPI().update_weights([client], tmp_path / "weights", step=1) + finally: + await client.aclose() + + assert paths == [ + "/engine/pause_generation", + "/engine/update_weights_from_disk", + "/engine/resume_generation", + ] + + @pytest.mark.asyncio async def test_engine_status_error_is_not_accepted(): paths = [] diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 4204fa3539..9abbde7b7e 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -251,6 +251,28 @@ def test_dynamo_disaggregated_config_rejects_disabled_prefix_caching(): ) +def test_dynamo_backend_rejects_lora(): + with pytest.raises(ValidationError, match="does not support LoRA"): + InferenceConfig.model_validate({"backend": {"type": "dynamo"}, "enable_lora": True}) + + +def test_rl_config_rejects_lora_auto_setup_for_dynamo_backend(): + with pytest.raises(ValidationError, match="does not support LoRA"): + RLConfig.model_validate( + { + "trainer": {"model": {"lora": {}}}, + "orchestrator": {}, + "inference": {"backend": {"type": "dynamo"}}, + "deployment": { + "type": "single_node", + "gpus_per_node": 2, + "num_train_gpus": 1, + "num_infer_gpus": 1, + }, + } + ) + + def test_native_disaggregated_config_still_requires_slurm(): with pytest.raises(ValidationError, match="Must use SLURM"): InferenceConfig.model_validate({"deployment": {"type": "disaggregated"}}) @@ -287,10 +309,56 @@ def test_single_node_dynamo_disaggregated_keeps_per_worker_dp_and_sets_nccl_worl assert config.inference.parallel.dp == 1 assert config.orchestrator.model.client.admin_api == "dynamo" assert config.orchestrator.model.client.dp_rank_count == 1 + assert config.orchestrator.model.client.dynamo_worker_roles == ("prefill", "prefill", "decode", "decode") + assert config.orchestrator.model.client.dynamo_gpus_per_worker == 1 assert config.trainer.weight_broadcast.inference_world_size == 4 assert config.orchestrator.weight_broadcast.inference_world_size == 4 +def test_single_node_dynamo_aggregated_derives_one_multi_gpu_worker(): + config = RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": {}, + "inference": {"backend": {"type": "dynamo"}, "parallel": {"tp": 1}}, + "deployment": { + "type": "single_node", + "gpus_per_node": 4, + "num_train_gpus": 2, + "num_infer_gpus": 2, + }, + } + ) + + assert config.inference is not None + assert config.orchestrator.model.client.dynamo_worker_roles == ("agg",) + assert config.orchestrator.model.client.dynamo_gpus_per_worker == 2 + + +def test_dynamo_topology_metadata_cannot_conflict_with_inference_config(): + with pytest.raises(ValidationError, match="dynamo_worker_roles conflicts"): + RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": { + "model": { + "client": { + "dynamo_worker_roles": ["decode"], + "dynamo_gpus_per_worker": 1, + } + } + }, + "inference": {"backend": {"type": "dynamo"}, "parallel": {"tp": 1}}, + "deployment": { + "type": "single_node", + "gpus_per_node": 2, + "num_train_gpus": 1, + "num_infer_gpus": 1, + }, + } + ) + + def test_multi_node_auto_inference_client_dp_rank_count_uses_router_url(): config = RLConfig.model_validate( { diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index d615715633..1f8ee514da 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -3,9 +3,11 @@ from unittest.mock import AsyncMock, MagicMock import httpx +import pytest from verifiers.v1.clients.config import EvalClientConfig from prime_rl.configs.shared import ClientConfig +from prime_rl.inference.dynamo_admin import DynamoTopology, DynamoWorker from prime_rl.utils.client import ( DynamoInferencePool, StaticInferencePool, @@ -137,6 +139,8 @@ def test_setup_inference_pool_selects_dynamo_pool_once(): ClientConfig( base_url=["http://frontend:8000/v1"], admin_api="dynamo", + dynamo_worker_roles=("agg",), + dynamo_gpus_per_worker=1, ), model_name="test-model", ) @@ -144,3 +148,60 @@ def test_setup_inference_pool_selects_dynamo_pool_once(): assert isinstance(pool, DynamoInferencePool) asyncio.run(pool.stop()) + + +@pytest.mark.asyncio +async def test_dynamo_readiness_shares_one_monotonic_deadline(monkeypatch: pytest.MonkeyPatch): + pool = DynamoInferencePool( + ClientConfig( + base_url=["http://frontend:8000/v1"], + admin_api="dynamo", + dynamo_worker_roles=("agg",), + dynamo_gpus_per_worker=2, + ), + model_name="test-model", + ) + observed_timeouts: list[float] = [] + worker = DynamoWorker( + instance_id=11, + component="backend", + role="agg", + system_url="http://worker:8081", + model="test-model", + routes=frozenset( + { + "init_weights_update_group", + "pause_generation", + "resume_generation", + "update_weights_from_disk", + "update_weights_from_distributed", + } + ), + ) + + async def health(_clients, *, timeout, **_kwargs): + observed_timeouts.append(timeout) + await asyncio.sleep(0.01) + + async def models(_clients, _model_name, *, skip_model_check, timeout): + assert skip_model_check is False + observed_timeouts.append(timeout) + await asyncio.sleep(0.01) + + async def discover(_clients, timeout, *, model_name, topology): + assert model_name == "test-model" + assert topology == DynamoTopology(roles=("agg",), gpus_per_worker=2) + observed_timeouts.append(timeout) + await asyncio.sleep(0.01) + return (worker,) + + monkeypatch.setattr("prime_rl.utils.client.check_health", health) + monkeypatch.setattr("prime_rl.utils.client.maybe_check_has_model", models) + monkeypatch.setattr("prime_rl.utils.client.discover_workers", discover) + try: + await pool.wait_for_ready("test-model", timeout=1) + finally: + await pool.stop() + + assert len(observed_timeouts) == 4 + assert all(later < earlier for earlier, later in zip(observed_timeouts, observed_timeouts[1:])) From 0d22c2b74a772a6aa38530ef97f2f92a69d8d94a Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 16:33:08 -0700 Subject: [PATCH 13/18] fix(dynamo): close topology and readiness gaps Signed-off-by: Biswa Panda --- .../src/prime_rl/configs/inference.py | 54 ++++++++++++-- .../src/prime_rl/configs/rl.py | 14 +--- src/prime_rl/inference/dynamo.py | 21 +++--- src/prime_rl/utils/client.py | 42 ++++++++--- tests/unit/inference/test_dynamo.py | 40 +++++++++-- tests/unit/test_configs.py | 70 +++++++++++++++++++ tests/unit/utils/test_client.py | 60 +++++++++++++++- 7 files changed, 260 insertions(+), 41 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index 0577fdd9b1..26e12f528f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -24,7 +24,7 @@ class ServerConfig(BaseConfig): class ParallelConfig(BaseConfig): - tp: int = 1 + tp: int = Field(1, ge=1) """Tensor parallel size. Forwarded to vLLM as ``--tensor-parallel-size``.""" dp: int = Field(1, ge=1) @@ -251,7 +251,7 @@ def validate_scorers(self): class BaseInferenceDeploymentConfig(BaseConfig): - gpus_per_node: int = 8 + gpus_per_node: int = Field(8, ge=1) """GPUs per node.""" backend_port: int = 8100 @@ -443,6 +443,29 @@ class InferenceConfig(BaseConfig): dry_run: bool = False """Only validate and dump resolved configs, then exit early.""" + @property + def dynamo_worker_roles(self) -> tuple[Literal["agg", "prefill", "decode"], ...]: + """Canonical admin and launch order for Dynamo worker groups.""" + if self.deployment.type == "disaggregated": + return ("prefill",) * self.deployment.num_prefill_replicas + ( + "decode", + ) * self.deployment.num_decode_replicas + return ("agg",) + + @property + def dynamo_gpus_per_worker(self) -> int: + """GPU allocation owned by one independently administered Dynamo worker.""" + if self.deployment.type == "disaggregated": + return self.deployment.gpus_per_node + return self.parallel.tp * self.parallel.dp + + @property + def dynamo_local_dp(self) -> int: + """vLLM data-parallel ranks inside one Dynamo worker.""" + if self.deployment.type == "disaggregated": + return self.deployment.gpus_per_node // self.parallel.tp + return self.parallel.dp + @model_validator(mode="after") def validate_multi_node_requires_slurm(self): if self.deployment.type == "multi_node" and self.slurm is None: @@ -451,6 +474,29 @@ def validate_multi_node_requires_slurm(self): raise ValueError("Must use SLURM for multi-node / disaggregated deployment.") return self + @model_validator(mode="after") + def validate_disaggregated_topology(self): + if self.deployment.type != "disaggregated": + return self + if self.deployment.gpus_per_node % self.parallel.tp != 0: + raise ValueError( + "inference.deployment.gpus_per_node must be divisible by inference.parallel.tp " + "so every worker contains whole tensor-parallel groups." + ) + if self.backend.type == "dynamo": + local_dp = self.deployment.gpus_per_node // self.parallel.tp + if "dp" in self.parallel.model_fields_set and self.parallel.dp != local_dp: + raise ValueError( + "inference.parallel.dp must equal inference.deployment.gpus_per_node / " + "inference.parallel.tp for a Dynamo disaggregated worker." + ) + if self.data_parallel_size_local is not None and self.data_parallel_size_local != local_dp: + raise ValueError( + "inference.data_parallel_size_local must equal inference.deployment.gpus_per_node / " + "inference.parallel.tp for a Dynamo disaggregated worker." + ) + return self + @model_validator(mode="after") def validate_dynamo_backend(self): if self.backend.type != "dynamo": @@ -503,9 +549,7 @@ def auto_setup_disaggregated(self): self.enable_expert_parallel = True if "enable_eplb" not in self.model_fields_set: self.enable_eplb = False - gpus_per_node = self.deployment.gpus_per_node - tp = self.parallel.tp - dp_per_node = gpus_per_node // tp + dp_per_node = self.dynamo_local_dp if self.data_parallel_size_local is None: self.data_parallel_size_local = dp_per_node if self.parallel.dp == 1: 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 6abafc72ba..0a50e17e38 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -710,19 +710,9 @@ def auto_setup_inference_client(self): else: client_updates["dp_rank_count"] = self.inference.data_parallel_size_local or self.inference.parallel.dp if self.inference.backend.type == "dynamo": - if self.inference.deployment.type == "disaggregated": - deployment = self.inference.deployment - expected_roles = ("prefill",) * deployment.num_prefill_replicas + ( - "decode", - ) * deployment.num_decode_replicas - expected_gpus_per_worker = deployment.gpus_per_node - else: - expected_roles = ("agg",) - expected_gpus_per_worker = self.inference.parallel.tp * self.inference.parallel.dp - expected_topology = { - "dynamo_worker_roles": expected_roles, - "dynamo_gpus_per_worker": expected_gpus_per_worker, + "dynamo_worker_roles": self.inference.dynamo_worker_roles, + "dynamo_gpus_per_worker": self.inference.dynamo_gpus_per_worker, } for field, expected in expected_topology.items(): if field in client.model_fields_set and getattr(client, field) != expected: diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py index 764e324e39..fc6ca3fc6f 100644 --- a/src/prime_rl/inference/dynamo.py +++ b/src/prime_rl/inference/dynamo.py @@ -39,10 +39,16 @@ ) _RESERVED_ENGINE_KEYS = frozenset( { + "data_parallel_rpc_port", + "data_parallel_size", + "data_parallel_size_local", "disaggregation_mode", + "enable_prefix_caching", "enable_rl", "kv_events_config", "kv_transfer_config", + "pipeline_parallel_size", + "tensor_parallel_size", "worker_extension_cls", } ) @@ -278,7 +284,7 @@ def build_engine_config( if config.deployment.type == "disaggregated": # Each generated worker is an independent vLLM server. Preserve local # DP within a worker, but never turn the P/D worker count into vLLM DP. - local_dp = config.deployment.gpus_per_node // config.parallel.tp + local_dp = config.dynamo_local_dp values["data_parallel_size"] = local_dp if local_dp == 1: values.pop("data_parallel_size_local", None) @@ -356,11 +362,11 @@ def build_local_worker_specs( raise ValueError("Local Dynamo requires one prefill node per prefill replica") if deployment.num_decode_nodes != deployment.num_decode_replicas: raise ValueError("Local Dynamo requires one decode node per decode replica") - roles: list[Role] = ["decode"] * deployment.num_decode_replicas + ["prefill"] * deployment.num_prefill_replicas - gpus_per_worker = deployment.gpus_per_node + roles = list(config.dynamo_worker_roles) + gpus_per_worker = config.dynamo_gpus_per_worker else: - roles = ["agg"] - gpus_per_worker = config.parallel.tp * config.parallel.dp + roles = list(config.dynamo_worker_roles) + gpus_per_worker = config.dynamo_gpus_per_worker required = len(roles) * gpus_per_worker if len(available) < required: @@ -404,10 +410,9 @@ def build_dry_run_worker_specs( ) -> list[DynamoWorkerSpec]: """Build local specs without consulting host GPU hardware.""" if config.deployment.type == "disaggregated": - worker_count = config.deployment.num_prefill_replicas + config.deployment.num_decode_replicas - gpu_count = worker_count * config.deployment.gpus_per_node + gpu_count = len(config.dynamo_worker_roles) * config.dynamo_gpus_per_worker else: - gpu_count = config.parallel.tp * config.parallel.dp + gpu_count = config.dynamo_gpus_per_worker return build_local_worker_specs( config, output_dir=output_dir, diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 483325fc3c..300d4aaea8 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -308,6 +308,7 @@ async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> N await check_health( self._frontend_admin_clients, timeout=_remaining_readiness_timeout(deadline, "frontend health"), + strict=True, ) await maybe_check_has_model( self._frontend_admin_clients, @@ -385,6 +386,9 @@ async def setup_inference_pool( pool_size: int | None = None, ) -> InferencePool: """Create an inference pool from config (static or elastic).""" + if client_config.admin_api == "dynamo" and client_config.is_elastic: + raise ValueError("Dynamo admin API does not support elastic inference pools") + if client_config.is_elastic: from prime_rl.utils.elastic import ElasticInferencePool @@ -480,20 +484,42 @@ async def maybe_check_has_model( model_name: str, skip_model_check: bool = False, timeout: float = 1800, + interval: float = 1, ) -> None: if skip_model_check: return logger = get_logger() + deadline = _readiness_deadline(timeout) logger.debug(f"Checking if model {model_name} is in the inference pool") - async with asyncio.timeout(timeout): - results = await asyncio.gather( - *(admin_client.get("/v1/models", timeout=httpx.Timeout(timeout)) for admin_client in admin_clients) + + async def _check_has_model(admin_client: AsyncClient) -> None: + last_error: Exception | None = None + loop = asyncio.get_running_loop() + while (remaining := deadline - loop.time()) > 0: + try: + async with asyncio.timeout(remaining): + result = await admin_client.get( + "/v1/models", + timeout=httpx.Timeout(min(remaining, 10.0)), + ) + result.raise_for_status() + models = result.json()["data"] + if not any(model["id"] == model_name for model in models): + raise ValueError(f"Model {model_name} is not registered") + return + except Exception as error: + last_error = error + remaining = deadline - loop.time() + if remaining > 0: + await asyncio.sleep(min(interval, remaining)) + + message = ( + f"Model {model_name} was not registered on {admin_client.base_url} " + f"before the {timeout}-second readiness deadline; last error: {last_error!r}" ) - for admin_client, result in zip(admin_clients, results): - result.raise_for_status() - models = result.json()["data"] - if not any(model["id"] == model_name for model in models): - raise ValueError(f"Model {model_name} was not found in the inference pool on {admin_client.base_url}") + raise TimeoutError(message) from last_error + + await asyncio.gather(*(_check_has_model(admin_client) for admin_client in admin_clients)) logger.debug(f"Model {model_name} was found in the inference pool") diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index e45f66cf0d..5240713c37 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -69,17 +69,47 @@ def test_role_overrides_are_isolated(): assert "max_num_seqs" not in prefill -@pytest.mark.parametrize("key", ["kv_transfer_config", "kv_events_config", "worker_extension_cls"]) +@pytest.mark.parametrize( + "key", + [ + "data_parallel_rpc_port", + "data_parallel_size", + "data_parallel_size_local", + "disaggregation_mode", + "enable_prefix_caching", + "enable_rl", + "kv_transfer_config", + "kv_events_config", + "pipeline_parallel_size", + "tensor_parallel_size", + "worker_extension_cls", + ], +) def test_reserved_engine_override_is_rejected(key: str): config = disaggregated_config(vllm_extra={key: {}}) with pytest.raises(ValueError, match="Dynamo-managed"): build_engine_config(config, "prefill", kv_events_port=20080) +def test_reserved_role_engine_override_is_rejected(): + config = disaggregated_config( + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + "prefill_vllm_overrides": {"tensor_parallel_size": 2}, + } + ) + + with pytest.raises(ValueError, match="prefill_vllm_overrides.*tensor_parallel_size"): + build_engine_config(config, "prefill", kv_events_port=20080) + + def test_local_specs_allocate_four_workers_and_unique_ports(tmp_path: Path): specs = build_local_worker_specs(disaggregated_config(), tmp_path, gpu_ids=["4", "5", "6", "7"]) - assert [spec.role for spec in specs] == ["decode", "decode", "prefill", "prefill"] + assert [spec.role for spec in specs] == list(disaggregated_config().dynamo_worker_roles) assert [spec.gpu_ids for spec in specs] == [("4",), ("5",), ("6",), ("7",)] assert len({spec.system_port for spec in specs}) == 4 assert len({spec.process.environment()["VLLM_NIXL_SIDE_CHANNEL_PORT"] for spec in specs}) == 4 @@ -221,15 +251,15 @@ def test_worker_environment_applies_only_matching_role_overrides(tmp_path: Path) "decode_env_vars": {"ROLE_SETTING": "decode"}, } ) - decode, prefill = build_local_worker_specs(config, tmp_path, gpu_ids=["3", "7"]) + prefill, decode = build_local_worker_specs(config, tmp_path, gpu_ids=["3", "7"]) decode_env = build_worker_environment(decode, {"COMMON": "value"}) prefill_env = build_worker_environment(prefill, {"COMMON": "value"}) assert decode_env["ROLE_SETTING"] == "decode" assert prefill_env["ROLE_SETTING"] == "prefill" - assert decode_env["CUDA_VISIBLE_DEVICES"] == "3" - assert prefill_env["CUDA_VISIBLE_DEVICES"] == "7" + assert prefill_env["CUDA_VISIBLE_DEVICES"] == "3" + assert decode_env["CUDA_VISIBLE_DEVICES"] == "7" assert decode_env["VLLM_PLUGINS"] == "prime_rl" assert decode_env["DYN_COMPONENT"] == "backend" assert prefill_env["DYN_COMPONENT"] == "prefill" diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 9abbde7b7e..54c46d2075 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -240,6 +240,76 @@ def test_dynamo_disaggregated_config_is_local_and_enables_prefix_caching(): assert config.use_pd_kv_transfer is True +@pytest.mark.parametrize( + "inference", + [ + {"parallel": {"tp": 0}}, + {"deployment": {"type": "single_node", "gpus_per_node": 0}}, + ], +) +def test_inference_topology_rejects_non_positive_gpu_dimensions(inference: dict): + with pytest.raises(ValidationError, match="greater than or equal to 1"): + InferenceConfig.model_validate(inference) + + +def test_dynamo_disaggregated_topology_requires_whole_tp_groups(): + with pytest.raises(ValidationError, match="gpus_per_node must be divisible"): + InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "parallel": {"tp": 2}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 3, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + }, + } + ) + + +@pytest.mark.parametrize( + "inference_override", + [ + {"parallel": {"tp": 2, "dp": 3}}, + {"parallel": {"tp": 2}, "data_parallel_size_local": 1}, + ], +) +def test_dynamo_disaggregated_topology_rejects_conflicting_dp(inference_override: dict): + inference = { + "backend": {"type": "dynamo"}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 4, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + }, + **inference_override, + } + + with pytest.raises(ValidationError, match="must equal.*gpus_per_node / inference.parallel.tp"): + InferenceConfig.model_validate(inference) + + +def test_dynamo_topology_is_derived_once_from_inference_config(): + config = InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "parallel": {"tp": 2}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 4, + "num_prefill_replicas": 2, + "num_decode_replicas": 1, + }, + } + ) + + assert config.dynamo_worker_roles == ("prefill", "prefill", "decode") + assert config.dynamo_gpus_per_worker == 4 + assert config.dynamo_local_dp == 2 + + def test_dynamo_disaggregated_config_rejects_disabled_prefix_caching(): with pytest.raises(ValidationError, match="requires prefix caching"): InferenceConfig.model_validate( diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 1f8ee514da..36301c8d46 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -1,18 +1,19 @@ import asyncio from pathlib import Path -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from verifiers.v1.clients.config import EvalClientConfig -from prime_rl.configs.shared import ClientConfig +from prime_rl.configs.shared import ClientConfig, ElasticConfig from prime_rl.inference.dynamo_admin import DynamoTopology, DynamoWorker from prime_rl.utils.client import ( DynamoInferencePool, StaticInferencePool, _is_retryable_lora_error, load_lora_adapter, + maybe_check_has_model, setup_clients, setup_inference_pool, ) @@ -150,6 +151,58 @@ def test_setup_inference_pool_selects_dynamo_pool_once(): asyncio.run(pool.stop()) +@pytest.mark.asyncio +async def test_setup_inference_pool_rejects_dynamo_elastic_before_pool_selection(): + client_config = ClientConfig( + base_url=["http://frontend:8000/v1"], + admin_api="dynamo", + dynamo_worker_roles=("agg",), + dynamo_gpus_per_worker=1, + elastic=ElasticConfig(hostname="inference.example"), + ) + original_config = client_config.model_copy(deep=True) + + with patch("prime_rl.utils.elastic.ElasticInferencePool.from_config", new=AsyncMock()) as from_config: + with pytest.raises(ValueError, match="Dynamo admin API does not support elastic inference pools"): + await setup_inference_pool(client_config, model_name="test-model") + + from_config.assert_not_awaited() + assert client_config == original_config + + +@pytest.mark.asyncio +async def test_model_registration_retries_transient_status_and_empty_models(): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + request = httpx.Request("GET", "http://frontend:8000/v1/models") + unavailable = httpx.Response( + 503, + request=request, + ) + empty = httpx.Response(200, json={"data": []}, request=request) + ready = httpx.Response(200, json={"data": [{"id": "test-model"}]}, request=request) + client.get.side_effect = [unavailable, empty, ready] + + await maybe_check_has_model([client], "test-model", timeout=1, interval=0) + + assert client.get.await_count == 3 + + +@pytest.mark.asyncio +async def test_model_registration_timeout_reports_last_error(): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + client.get.return_value = httpx.Response( + 503, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ) + + with pytest.raises(TimeoutError, match=r"test-model.*frontend:8000.*503 Service Unavailable"): + await maybe_check_has_model([client], "test-model", timeout=0.02, interval=0.001) + + assert client.get.await_count > 1 + + @pytest.mark.asyncio async def test_dynamo_readiness_shares_one_monotonic_deadline(monkeypatch: pytest.MonkeyPatch): pool = DynamoInferencePool( @@ -179,7 +232,8 @@ async def test_dynamo_readiness_shares_one_monotonic_deadline(monkeypatch: pytes ), ) - async def health(_clients, *, timeout, **_kwargs): + async def health(_clients, *, timeout, strict=False, **_kwargs): + assert strict is True observed_timeouts.append(timeout) await asyncio.sleep(0.01) From 4a8a185d5981f625c9c40bbc1865acb574396cbd Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 16:58:23 -0700 Subject: [PATCH 14/18] fix(dynamo): classify readiness and fixed pools --- src/prime_rl/orchestrator/algo/opd.py | 8 +- src/prime_rl/utils/client.py | 83 +++++++++++++++----- tests/unit/orchestrator/test_algorithms.py | 28 ++++++- tests/unit/utils/test_client.py | 89 +++++++++++++++++++++- 4 files changed, 183 insertions(+), 25 deletions(-) diff --git a/src/prime_rl/orchestrator/algo/opd.py b/src/prime_rl/orchestrator/algo/opd.py index 3135f2a9b2..e35444cc49 100644 --- a/src/prime_rl/orchestrator/algo/opd.py +++ b/src/prime_rl/orchestrator/algo/opd.py @@ -5,7 +5,7 @@ from prime_rl.configs.algorithm import OPDAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm -from prime_rl.utils.client import StaticInferencePool +from prime_rl.utils.client import FixedInferencePool if TYPE_CHECKING: from prime_rl.orchestrator.types import Rollout @@ -29,12 +29,12 @@ class OPDAlgorithm(Algorithm): def __init__(self, config: OPDAlgoConfig, policy_pool: InferencePool): super().__init__(config, policy_pool) self.teacher = config.teacher - self.teacher_pool: StaticInferencePool | None = None # static teacher endpoint, connected in setup() + self.teacher_pool: FixedInferencePool | None = None # fixed teacher endpoint, connected in setup() async def setup(self) -> None: pool = await self.connect(self.teacher) - if not isinstance(pool, StaticInferencePool): - raise TypeError("opd teacher must be a static endpoint — prefill scoring needs fixed endpoints") + if not isinstance(pool, FixedInferencePool): + raise TypeError("opd teacher must be a fixed endpoint — prefill scoring needs fixed endpoints") self.teacher_pool = pool async def score_rollout(self, rollout: Rollout) -> None: diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 300d4aaea8..16c956cef7 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -10,7 +10,7 @@ import httpx import verifiers.v1 as vf from httpx import AsyncClient -from openai import AsyncOpenAI, NotFoundError +from openai import AsyncOpenAI from renderers import RendererConfig from tenacity import AsyncRetrying, retry, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential from verifiers.v1.clients.config import EvalClientConfig, TrainClientConfig @@ -147,8 +147,8 @@ async def aclose(self) -> None: await asyncio.gather(*(c.close() for c in self._clients.values())) -class _StaticClientPool: - """Shared request-client behavior for fixed inference pools.""" +class FixedInferencePool: + """Base capability for pools whose endpoint set is fixed for their lifetime.""" def __init__( self, @@ -203,7 +203,7 @@ async def score(self, token_ids: list[int]) -> list[float]: return await self._scorer.score(self.train_clients, self.model_name, token_ids) -class StaticInferencePool(_StaticClientPool): +class StaticInferencePool(FixedInferencePool): """Static native-vLLM inference pool with fixed clients.""" def __init__( @@ -266,7 +266,7 @@ async def stop(self) -> None: await asyncio.gather(*(client.aclose() for client in clients)) -class DynamoInferencePool(_StaticClientPool): +class DynamoInferencePool(FixedInferencePool): """Static Dynamo pool with worker discovery and Dynamo administration.""" def __init__( @@ -479,6 +479,43 @@ def _setup_admin_client(base_url: str) -> httpx.AsyncClient: return [_setup_admin_client(base_url) for base_url in urls] +_RETRYABLE_READINESS_STATUS_CODES = frozenset({408, 409, 429}) + + +def _is_retryable_readiness_error(error: BaseException) -> bool: + """Return whether a readiness request can plausibly succeed unchanged. + + HTTP 408 and 429 are request/proxy backpressure. HTTP 409 is also transient + here because inference servers can report a state conflict while their model + lifecycle is transitioning. Other 4xx responses require a caller or routing + change and must fail immediately; server failures and transports are retried. + """ + if isinstance(error, httpx.HTTPStatusError): + status_code = error.response.status_code + return status_code in _RETRYABLE_READINESS_STATUS_CODES or 500 <= status_code < 600 + return isinstance(error, (httpx.TransportError, TimeoutError)) + + +def _model_ids(response: httpx.Response) -> frozenset[str]: + """Validate an OpenAI models response before interpreting model absence.""" + try: + payload = response.json() + except ValueError as error: + raise ValueError("Invalid /v1/models response: expected valid JSON") from error + if not isinstance(payload, dict): + raise ValueError("Invalid /v1/models response: expected a JSON object") + models = payload.get("data") + if not isinstance(models, list): + raise ValueError("Invalid /v1/models response: 'data' must be a list") + + model_ids: set[str] = set() + for index, model in enumerate(models): + if not isinstance(model, dict) or not isinstance(model.get("id"), str): + raise ValueError(f"Invalid /v1/models response: data[{index}].id must be a string") + model_ids.add(model["id"]) + return frozenset(model_ids) + + async def maybe_check_has_model( admin_clients: list[AsyncClient], model_name: str, @@ -503,15 +540,20 @@ async def _check_has_model(admin_client: AsyncClient) -> None: timeout=httpx.Timeout(min(remaining, 10.0)), ) result.raise_for_status() - models = result.json()["data"] - if not any(model["id"] == model_name for model in models): - raise ValueError(f"Model {model_name} is not registered") - return except Exception as error: + if not _is_retryable_readiness_error(error): + raise last_error = error - remaining = deadline - loop.time() - if remaining > 0: - await asyncio.sleep(min(interval, remaining)) + else: + if model_name in _model_ids(result): + return + # A valid 200 response without the requested model is expected + # while registration is still in progress, so it is retryable. + last_error = RuntimeError(f"Model {model_name} is not registered") + + remaining = deadline - loop.time() + if remaining > 0: + await asyncio.sleep(min(interval, remaining)) message = ( f"Model {model_name} was not registered on {admin_client.base_url} " @@ -530,6 +572,12 @@ async def check_health( timeout: float = 1800, strict: bool = False, ) -> None: + """Wait for healthy endpoints, retrying only transient request failures. + + Native endpoints may omit ``/health``; ``strict=True`` requires the route + and is used for Dynamo frontends and workers. All other non-success statuses + are classified by :func:`_is_retryable_readiness_error`. + """ logger = get_logger() async def _check_health(admin_client: AsyncClient) -> None: @@ -545,15 +593,16 @@ async def _check_health(admin_client: AsyncClient) -> None: "/health", timeout=httpx.Timeout(min(remaining, 10.0)), ) - if strict: - response.raise_for_status() + if not strict and response.status_code == 404: + logger.warning("The route /health does not exist. Skipping health check.") + return + response.raise_for_status() elapsed = loop.time() - started logger.debug(f"Inference pool is ready after {elapsed:.1f} seconds") return - except NotFoundError: - logger.warning("The route /health does not exist. Skipping health check.") - return except Exception as e: + if not _is_retryable_readiness_error(e): + raise last_error = e elapsed = loop.time() - started if elapsed >= next_log_at: diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index 7e7a427112..af5ebe2326 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -1,5 +1,5 @@ import asyncio -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pydantic import pytest @@ -8,10 +8,12 @@ from verifiers.v1.types import AssistantMessage, ToolMessage, UserMessage from prime_rl.configs.algorithm import AlgoConfig, FrozenModelConfig -from prime_rl.orchestrator.algo import EchoAlgorithm, stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.algo import EchoAlgorithm, OPDAlgorithm, stamp_advantages, stamp_loss_routing from prime_rl.orchestrator.trajectories import trace_to_samples from prime_rl.orchestrator.types import Rollout from prime_rl.transport.types import TrainingSample +from prime_rl.utils.client import DynamoInferencePool, StaticInferencePool +from prime_rl.utils.elastic import ElasticInferencePool FROZEN = {"name": "org/ref-model", "base_url": ["http://ref:8001/v1"]} @@ -78,6 +80,28 @@ def test_opd_teacher_must_be_a_frozen_endpoint(): _build(type="opd", teacher="policy") +@pytest.mark.asyncio +@pytest.mark.parametrize("pool_type", [StaticInferencePool, DynamoInferencePool]) +async def test_opd_accepts_fixed_teacher_pools(pool_type): + pool = object.__new__(pool_type) + algo = OPDAlgorithm(_build(type="opd", teacher=FROZEN), MagicMock()) + algo.connect = AsyncMock(return_value=pool) + + await algo.setup() + + assert algo.teacher_pool is pool + + +@pytest.mark.asyncio +async def test_opd_rejects_elastic_teacher_pool(): + pool = object.__new__(ElasticInferencePool) + algo = OPDAlgorithm(_build(type="opd", teacher=FROZEN), MagicMock()) + algo.connect = AsyncMock(return_value=pool) + + with pytest.raises(TypeError, match="fixed endpoint"): + await algo.setup() + + def test_sft_requires_teacher(): with pytest.raises(ValueError, match="needs a teacher to sample rollouts from"): _build(type="sft") diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 36301c8d46..cecc91b24b 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -12,6 +12,7 @@ DynamoInferencePool, StaticInferencePool, _is_retryable_lora_error, + check_health, load_lora_adapter, maybe_check_has_model, setup_clients, @@ -175,17 +176,101 @@ async def test_model_registration_retries_transient_status_and_empty_models(): client = AsyncMock() client.base_url = httpx.URL("http://frontend:8000") request = httpx.Request("GET", "http://frontend:8000/v1/models") + conflict = httpx.Response(409, request=request) unavailable = httpx.Response( 503, request=request, ) empty = httpx.Response(200, json={"data": []}, request=request) ready = httpx.Response(200, json={"data": [{"id": "test-model"}]}, request=request) - client.get.side_effect = [unavailable, empty, ready] + client.get.side_effect = [httpx.ConnectError("not listening", request=request), conflict, unavailable, empty, ready] await maybe_check_has_model([client], "test-model", timeout=1, interval=0) - assert client.get.await_count == 3 + assert client.get.await_count == 5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [401, 403, 404]) +async def test_model_registration_fails_permanent_http_status_immediately(status_code: int): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + client.get.return_value = httpx.Response( + status_code, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ) + + with pytest.raises(httpx.HTTPStatusError): + await maybe_check_has_model([client], "test-model", timeout=1, interval=0) + + client.get.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + httpx.Response( + 200, + content=b"not-json", + headers={"content-type": "application/json"}, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ), + httpx.Response( + 200, + json={"data": {}}, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ), + httpx.Response( + 200, + json={"data": [{"object": "model"}]}, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ), + ], + ids=["invalid-json", "data-not-list", "model-id-missing"], +) +async def test_model_registration_fails_invalid_response_immediately(response: httpx.Response): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + client.get.return_value = response + + with pytest.raises(ValueError, match=r"Invalid /v1/models response"): + await maybe_check_has_model([client], "test-model", timeout=1, interval=0) + + client.get.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_strict_health_retries_only_transient_failures(): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + request = httpx.Request("GET", "http://frontend:8000/health") + client.get.side_effect = [ + httpx.ConnectError("not listening", request=request), + httpx.Response(429, request=request), + httpx.Response(503, request=request), + httpx.Response(200, request=request), + ] + + await check_health([client], timeout=1, interval=0, strict=True) + + assert client.get.await_count == 4 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [401, 403, 404]) +async def test_strict_health_fails_permanent_http_status_immediately(status_code: int): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + client.get.return_value = httpx.Response( + status_code, + request=httpx.Request("GET", "http://frontend:8000/health"), + ) + + with pytest.raises(httpx.HTTPStatusError): + await check_health([client], timeout=1, interval=0, strict=True) + + client.get.assert_awaited_once() @pytest.mark.asyncio From f5a9d56c642a18b0928c64cbe2a46238190e31b4 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 17:21:45 -0700 Subject: [PATCH 15/18] fix(dynamo): fail fast on discovery contracts --- src/prime_rl/inference/dynamo_admin.py | 44 ++++- tests/unit/inference/test_dynamo_admin.py | 222 +++++++++++++++++++++- 2 files changed, 259 insertions(+), 7 deletions(-) diff --git a/src/prime_rl/inference/dynamo_admin.py b/src/prime_rl/inference/dynamo_admin.py index 14fc9211a3..57e8dfb309 100644 --- a/src/prime_rl/inference/dynamo_admin.py +++ b/src/prime_rl/inference/dynamo_admin.py @@ -23,6 +23,7 @@ UPDATE_WEIGHTS_TIMEOUT_S = 720.0 DISCOVERY_REQUEST_TIMEOUT_S = 10.0 DISCOVERY_POLL_INTERVAL_S = 1.0 +_RETRYABLE_DISCOVERY_HTTP_STATUS_CODES = frozenset({408, 409, 429}) _REQUIRED_ROUTES = frozenset( { "init_weights_update_group", @@ -36,6 +37,10 @@ WorkerRole: TypeAlias = Literal["agg", "prefill", "decode"] +class _DiscoveryConvergenceError(RuntimeError): + """A structurally valid discovery state that may converge during startup.""" + + @dataclass(frozen=True, slots=True) class DynamoWorker: """Restart-safe identity and admin capabilities for one Dynamo worker.""" @@ -131,8 +136,11 @@ def _parse_worker(value: object, model_name: str, topology: DynamoTopology) -> D raise ValueError("Dynamo worker discovery response is missing component") if not isinstance(instance_id, int) or isinstance(instance_id, bool) or instance_id < 0: raise ValueError(f"Dynamo worker {component!r} has an invalid instance_id") - if value.get("error"): - raise RuntimeError(f"Dynamo worker {component}[{instance_id}] is unhealthy: {value['error']}") + error = value.get("error") + if error is not None: + if not isinstance(error, str) or not error: + raise ValueError(f"Dynamo worker {component}[{instance_id}] has an invalid error") + raise _DiscoveryConvergenceError(f"Dynamo worker {component}[{instance_id}] is unhealthy: {error}") if not isinstance(system_url, str) or not system_url: raise ValueError(f"Dynamo worker {component}[{instance_id}] is missing system_url") if model != model_name: @@ -179,6 +187,21 @@ def _parse_snapshot( return namespace, workers +def _retryable_discovery_error(error: Exception) -> bool: + if isinstance(error, httpx.HTTPStatusError): + status_code = error.response.status_code + return status_code in _RETRYABLE_DISCOVERY_HTTP_STATUS_CODES or status_code >= 500 + return isinstance( + error, + ( + _DiscoveryConvergenceError, + httpx.TimeoutException, + httpx.TransportError, + TimeoutError, + ), + ) + + async def discover_workers( discovery_clients: list[AsyncClient], timeout: float, @@ -186,7 +209,13 @@ async def discover_workers( model_name: str, topology: DynamoTopology, ) -> tuple[DynamoWorker, ...]: - """Wait for all frontends to report one identical, complete worker set.""" + """Wait for all frontends to report one identical, complete worker set. + + Incomplete membership, inconsistent valid snapshots, and worker probe errors + are startup convergence states. A model mismatch or missing RL route is an + incompatible deployment contract and fails immediately, as do malformed + payloads and permanent HTTP errors. + """ if not discovery_clients: raise ValueError("Dynamo worker discovery requires at least one frontend") if timeout <= 0: @@ -212,12 +241,17 @@ async def discover_workers( snapshots.append(_parse_snapshot(response.json(), model_name, topology)) first = snapshots[0] if any(snapshot != first for snapshot in snapshots[1:]): - raise ValueError("Dynamo discovery frontends returned inconsistent worker snapshots") + raise _DiscoveryConvergenceError("Dynamo discovery frontends returned inconsistent worker snapshots") workers = first[1] - topology.validate(workers) + try: + topology.validate(workers) + except ValueError as exc: + raise _DiscoveryConvergenceError(str(exc)) from exc logger.info(f"Discovered {len(workers)} Dynamo inference worker(s)") return workers except Exception as exc: + if not _retryable_discovery_error(exc): + raise last_error = exc remaining = deadline - loop.time() if remaining > 0: diff --git a/tests/unit/inference/test_dynamo_admin.py b/tests/unit/inference/test_dynamo_admin.py index 48df7972da..a4f9d11c2a 100644 --- a/tests/unit/inference/test_dynamo_admin.py +++ b/tests/unit/inference/test_dynamo_admin.py @@ -187,21 +187,239 @@ def handler(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_worker_discovery_rejects_incomplete_admin_surface(): + calls = 0 + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 return httpx.Response(200, json={"namespace": "test", "workers": [worker(1, routes=frozenset())]}) client = async_client(handler, "http://frontend:8001") try: - with pytest.raises(TimeoutError, match="missing RL routes"): + with pytest.raises(ValueError, match="missing RL routes"): await discover_workers( [client], - timeout=0.01, + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 1 + + +@pytest.mark.asyncio +async def test_worker_discovery_rejects_model_mismatch_without_retry(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1, model="other-model")]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + with pytest.raises(ValueError, match="serves 'other-model', expected 'test-model'"): + await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [400, 401, 403, 404]) +async def test_worker_discovery_rejects_permanent_http_errors_without_retry( + monkeypatch: pytest.MonkeyPatch, + status_code: int, +): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(status_code) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + with pytest.raises(httpx.HTTPStatusError): + await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [408, 409, 429, 500]) +async def test_worker_discovery_retries_transient_http_errors( + monkeypatch: pytest.MonkeyPatch, + status_code: int, +): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + return httpx.Response(status_code) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1)]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert discovered[0].instance_id == 1 + + +@pytest.mark.asyncio +async def test_worker_discovery_retries_transport_errors(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + raise httpx.ConnectError("frontend is starting", request=request) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1)]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert discovered[0].instance_id == 1 + + +@pytest.mark.asyncio +async def test_worker_discovery_retries_request_timeouts(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + raise httpx.ReadTimeout("frontend response timed out", request=request) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1)]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert discovered[0].instance_id == 1 + + +@pytest.mark.asyncio +async def test_worker_discovery_retries_transient_worker_probe_errors(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + discovered_worker = worker(1) + if calls == 1: + discovered_worker["error"] = "worker endpoint has not converged" + return httpx.Response(200, json={"namespace": "test", "workers": [discovered_worker]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert discovered[0].instance_id == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "error_type", "match"), + [ + (httpx.Response(200, content=b"{"), json.JSONDecodeError, "Expecting property name"), + (httpx.Response(200, json={"namespace": "test"}), ValueError, "invalid response"), + ( + httpx.Response( + 200, + json={"namespace": "test", "workers": [{**worker(1), "error": 123}]}, + ), + ValueError, + "invalid error", + ), + ], +) +async def test_worker_discovery_rejects_invalid_payload_without_retry( + monkeypatch: pytest.MonkeyPatch, + response: httpx.Response, + error_type: type[Exception], + match: str, +): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return response + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + with pytest.raises(error_type, match=match): + await discover_workers( + [client], + timeout=1, model_name="test-model", topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), ) finally: await client.aclose() + assert calls == 1 + def test_worker_membership_change_is_rejected(): expected = (DynamoWorker(1, "backend", "agg", "http://worker:8081", "test-model", ROUTES),) From 66c40e420c0a4a1defde08f032c2a60fb048fcd9 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 18:46:40 -0700 Subject: [PATCH 16/18] fix(dynamo): make policy updates fail closed --- docs/training.md | 2 +- .../src/prime_rl/configs/orchestrator.py | 4 +- src/prime_rl/inference/dynamo.py | 3 + src/prime_rl/inference/dynamo_admin.py | 99 ++- src/prime_rl/orchestrator/algo/__init__.py | 12 +- src/prime_rl/orchestrator/algo/base.py | 18 +- src/prime_rl/orchestrator/algo/opsd.py | 17 +- .../orchestrator/component_supervision.py | 20 + src/prime_rl/orchestrator/dispatcher.py | 305 ++++++--- .../orchestrator/dispatcher_metrics.py | 58 ++ src/prime_rl/orchestrator/envs.py | 4 +- src/prime_rl/orchestrator/orchestrator.py | 20 +- src/prime_rl/orchestrator/policy_gate.py | 97 +++ src/prime_rl/orchestrator/pool_identity.py | 42 ++ src/prime_rl/orchestrator/types.py | 6 +- src/prime_rl/orchestrator/watcher.py | 101 ++- src/prime_rl/utils/async_utils.py | 18 + tests/unit/inference/test_dynamo.py | 24 + tests/unit/inference/test_dynamo_admin.py | 2 + .../inference/test_dynamo_admin_barriers.py | 263 +++++++ tests/unit/orchestrator/test_algorithms.py | 26 +- .../orchestrator/test_orchestrator_setup.py | 15 + tests/unit/orchestrator/test_pool_identity.py | 52 ++ .../test_weight_update_barrier.py | 643 ++++++++++++++++++ 24 files changed, 1678 insertions(+), 173 deletions(-) create mode 100644 src/prime_rl/orchestrator/component_supervision.py create mode 100644 src/prime_rl/orchestrator/dispatcher_metrics.py create mode 100644 src/prime_rl/orchestrator/policy_gate.py create mode 100644 src/prime_rl/orchestrator/pool_identity.py create mode 100644 tests/unit/inference/test_dynamo_admin_barriers.py create mode 100644 tests/unit/orchestrator/test_pool_identity.py create mode 100644 tests/unit/orchestrator/test_weight_update_barrier.py diff --git a/docs/training.md b/docs/training.md index 0372fd68c1..0f253ca0d0 100644 --- a/docs/training.md +++ b/docs/training.md @@ -58,7 +58,7 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli |---|---| | `orchestrator.batch_size` | Tasks per trainer step. | | `orchestrator.group_size` | Rollouts generated per task. | -| `orchestrator.max_off_policy_steps` | How many distinct policies may have contributed to one rollout before it's discarded (default 8). The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. | +| `orchestrator.max_off_policy_steps` | On the vLLM admin backend, how many distinct policies may contribute to one rollout before it is discarded (default 8). Dynamo instead drains all live-policy/eval requests before mutating weights, so this setting does not apply to Dynamo runs. | | `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). | | `[[orchestrator.train.env]]` | Training environments. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Environments](configuration.md#environments-orchestratortrainenv). | | `[[orchestrator.eval.env]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). | 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 c50f5af42c..a708dc6f12 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -529,7 +529,9 @@ class OrchestratorConfig(BaseConfig): """Maximum training steps. If None, runs indefinitely.""" max_off_policy_steps: int = Field(8, ge=0) - """Maximum policies allowed to generate a single rollout. Rollouts generated more than ``max_off_policy_steps`` ahead of training are discarded. Higher values yield better throughput at the cost of off-policy noise.""" + """Maximum policies allowed to generate one rollout on the vLLM admin backend. Dynamo uses a strict + application drain before every weight mutation, so live-policy and eval requests never span versions and + this tolerance does not apply there.""" bench: bool = False """Benchmark mode. Sets ``max_steps`` to 5 and disables W&B.""" diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py index fc6ca3fc6f..2ae26d2f73 100644 --- a/src/prime_rl/inference/dynamo.py +++ b/src/prime_rl/inference/dynamo.py @@ -107,6 +107,9 @@ def _validate_overrides(source: str, values: dict[str, Any]) -> None: conflicts = sorted(_RESERVED_ENGINE_KEYS & values.keys()) if conflicts: raise ValueError(f"{source} cannot override Dynamo-managed engine keys: {conflicts}") + wrapper_only = sorted(_ENGINE_CONFIG_EXCLUDED & values.keys()) + if wrapper_only: + raise ValueError(f"{source} keys {wrapper_only} are wrapper/server-only and cannot enter a vLLM engine config") def _environment_items(values: dict[str, str]) -> tuple[tuple[str, str], ...]: diff --git a/src/prime_rl/inference/dynamo_admin.py b/src/prime_rl/inference/dynamo_admin.py index 57e8dfb309..410ecb4ec7 100644 --- a/src/prime_rl/inference/dynamo_admin.py +++ b/src/prime_rl/inference/dynamo_admin.py @@ -5,7 +5,7 @@ import asyncio import os from collections import Counter -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Iterable, Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import Literal, TypeAlias @@ -16,6 +16,7 @@ from tenacity import AsyncRetrying, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.async_utils import gather_shielded from prime_rl.utils.logger import get_logger NCCL_READY_MARKER = "NCCL_READY" @@ -278,6 +279,20 @@ class DynamoAdminAPI: def __init__(self) -> None: self._distributed_updates = False + self._distributed_initialization_indeterminate = False + self._weight_update_indeterminate = False + + def _require_unambiguous_admin_state(self) -> None: + if self._distributed_initialization_indeterminate: + raise RuntimeError( + "Dynamo distributed weight-group initialization is indeterminate after a prior failure or " + "cancellation; refusing further admin mutation" + ) + if self._weight_update_indeterminate: + raise RuntimeError( + "Dynamo worker weight state is indeterminate after a prior update or resume failure; refusing " + "further admin mutation" + ) @staticmethod def _retryable(exception: BaseException) -> bool: @@ -321,6 +336,23 @@ async def post_once() -> dict: return await post_once() raise AssertionError("unreachable") + @staticmethod + async def _settle_fanout(awaitables: Iterable[Awaitable[dict]], operation: str) -> None: + """Await every sibling even when one fails or the caller is cancelled.""" + tasks = [asyncio.create_task(awaitable) for awaitable in awaitables] + if not tasks: + return + results, cancellation = await gather_shielded(*tasks) + + failures = [result for result in results if isinstance(result, BaseException)] + primary: BaseException | None = cancellation or (failures[0] if failures else None) + if primary is None: + return + siblings = failures if cancellation is not None else failures[1:] + for sibling in siblings: + primary.add_note(f"Dynamo {operation} sibling also failed: {sibling!r}") + raise primary + async def initialize_nccl( self, clients: list[AsyncClient], @@ -332,6 +364,7 @@ async def initialize_nccl( gpus_per_worker: int, quantize_in_weight_transfer: bool, ) -> None: + self._require_unambiguous_admin_state() if not clients: raise ValueError("Cannot initialize NCCL without Dynamo workers") if isinstance(gpus_per_worker, bool) or gpus_per_worker < 1: @@ -343,39 +376,49 @@ async def initialize_nccl( f"inference_world_size={world_size} does not match {len(clients)} Dynamo workers " f"with {gpus_per_worker} GPUs each ({expected_world_size})" ) - await asyncio.gather( - *( - self._post( - client, - "init_weights_update_group", - { - "host": host, - "port": port, - "rank_offset": index * gpus_per_worker, - "inference_world_size": world_size, - "timeout": timeout, - "quantize_in_weight_transfer": quantize_in_weight_transfer, - "engine_rpc": "init_broadcaster", - }, - ) - for index, client in enumerate(clients) + try: + await self._settle_fanout( + ( + self._post( + client, + "init_weights_update_group", + { + "host": host, + "port": port, + "rank_offset": index * gpus_per_worker, + "inference_world_size": world_size, + "timeout": timeout, + "quantize_in_weight_transfer": quantize_in_weight_transfer, + "engine_rpc": "init_broadcaster", + }, + ) + for index, client in enumerate(clients) + ), + "init_weights_update_group", ) - ) + except BaseException: + # A lost response or one failed sibling cannot prove whether every + # rank committed the collective group. This object must not choose + # the filesystem path or retry into that ambiguous state. + self._distributed_initialization_indeterminate = True + raise self._distributed_updates = True async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | None, step: int) -> None: + self._require_unambiguous_admin_state() primary_error: BaseException | None = None try: - await asyncio.gather( - *( + await self._settle_fanout( + ( self._post( client, "pause_generation", - {"mode": "keep", "clear_cache": False}, + {"mode": "wait", "clear_cache": False}, retry_transient=True, ) for client in clients - ) + ), + "pause_generation", ) if weight_dir is not None: marker = weight_dir / NCCL_READY_MARKER @@ -399,18 +442,22 @@ async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | No } method = "update_weights_from_disk" - await asyncio.gather( - *(self._post(client, method, body, timeout_s=UPDATE_WEIGHTS_TIMEOUT_S) for client in clients) + await self._settle_fanout( + (self._post(client, method, body, timeout_s=UPDATE_WEIGHTS_TIMEOUT_S) for client in clients), + method, ) except BaseException as exc: primary_error = exc + self._weight_update_indeterminate = True raise finally: try: - await asyncio.gather( - *(self._post(client, "resume_generation", retry_transient=True) for client in clients) + await self._settle_fanout( + (self._post(client, "resume_generation", retry_transient=True) for client in clients), + "resume_generation", ) except BaseException as exc: + self._weight_update_indeterminate = True if primary_error is None: raise primary_error.add_note(f"Dynamo resume_generation cleanup also failed: {exc!r}") diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index 8d1baa60a3..c79090d410 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/src/prime_rl/orchestrator/algo/__init__.py @@ -40,6 +40,7 @@ if TYPE_CHECKING: from prime_rl.configs.algorithm import AlgoConfig + from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.utils.client import InferencePool # Runtime dispatch is keyed on ``algo.type`` — it names the algorithm, and @@ -54,7 +55,12 @@ } -def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm: +def build_algorithm( + config: AlgoConfig, + policy_pool: InferencePool, + *, + policy_gate: MutablePolicyGate | None = None, +) -> Algorithm: cls = ALGORITHM_CLASSES[config.type] assert cls.action_loss_type == config.action_loss_type # config and runtime declare in two places # The Algorithm is the runtime of the algorithm config's training signal @@ -62,7 +68,9 @@ def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm # handed the live policy pool — opsd self-distills against it, others may # judge against it or ignore it. Other models (a frozen teacher, a hint # renderer) are built from the algorithm's own config in setup(). - return cls(config, policy_pool) + algorithm = cls(config, policy_pool) + algorithm.policy_gate = policy_gate + return algorithm __all__ = [ diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index b12d17df62..db7a1de24f 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -44,11 +44,13 @@ from prime_rl.configs.algorithm import ActionLossType, AlgoConfig, FrozenModelConfig from prime_rl.orchestrator.algo.routing import stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.policy_gate import PolicyRequestRejected from prime_rl.utils.logger import get_logger if TYPE_CHECKING: from renderers import RendererConfig + from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.orchestrator.types import Rollout from prime_rl.utils.client import InferencePool @@ -119,8 +121,15 @@ class Algorithm: action_loss_type: ClassVar[ActionLossType] = "rl" - def __init__(self, config: AlgoConfig, policy_pool: InferencePool): + def __init__( + self, + config: AlgoConfig, + policy_pool: InferencePool, + *, + policy_gate: MutablePolicyGate | None = None, + ): self.policy_pool = policy_pool + self.policy_gate = policy_gate self.connected_pools: list[InferencePool] = [] # frozen pools connected in setup(); closed at shutdown async def setup(self) -> None: @@ -152,7 +161,12 @@ async def finalize_rollout(self, rollout: Rollout) -> None: """Arrival phase (non-virtual): rollout-local scoring as each rollout is tokenized.""" if rollout.samples: - await self.score_rollout(rollout) + try: + await self.score_rollout(rollout) + except PolicyRequestRejected as exc: + # Preserve the sink's one-arrival-per-group accounting while + # dropping scoring that cannot use the generating version. + rollout.capture_error(exc) async def finalize_group(self, rollouts: list[Rollout]) -> None: """Group phase (non-virtual): group-relative scoring, then stamp each diff --git a/src/prime_rl/orchestrator/algo/opsd.py b/src/prime_rl/orchestrator/algo/opsd.py index 737666bea9..1a3e9e7d94 100644 --- a/src/prime_rl/orchestrator/algo/opsd.py +++ b/src/prime_rl/orchestrator/algo/opsd.py @@ -9,6 +9,7 @@ if TYPE_CHECKING: from renderers.base import Renderer + from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.orchestrator.types import Rollout from prime_rl.transport import TrainingSample from prime_rl.utils.client import InferencePool @@ -30,8 +31,14 @@ class OPSDAlgorithm(Algorithm): action_loss_type = "ref_kl" - def __init__(self, config: OPSDAlgoConfig, policy_pool: InferencePool): - super().__init__(config, policy_pool) + def __init__( + self, + config: OPSDAlgoConfig, + policy_pool: InferencePool, + *, + policy_gate: MutablePolicyGate | None = None, + ): + super().__init__(config, policy_pool, policy_gate=policy_gate) self.demo_key = config.demo_key self.template = config.template self.renderer_config = config.renderer @@ -74,4 +81,8 @@ async def score_sample(sample: TrainingSample) -> None: # sample.token_ids (demo-conditioned, the trainer's ref_kl target). sample.ref_logprobs = full_logprobs[len(hint_block) :] - await asyncio.gather(*(score_sample(sample) for sample in rollout.samples)) + if self.policy_gate is None: + await asyncio.gather(*(score_sample(sample) for sample in rollout.samples)) + return + async with self.policy_gate.request(expected_version=rollout.policy_version): + await asyncio.gather(*(score_sample(sample) for sample in rollout.samples)) diff --git a/src/prime_rl/orchestrator/component_supervision.py b/src/prime_rl/orchestrator/component_supervision.py new file mode 100644 index 0000000000..7b7df0f14a --- /dev/null +++ b/src/prime_rl/orchestrator/component_supervision.py @@ -0,0 +1,20 @@ +"""Failure propagation for orchestrator background components.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Sequence + + +def raise_if_component_failed(tasks: Sequence[asyncio.Task]) -> None: + """Raise when a background loop exits while the main loop is still live.""" + for task in tasks: + if not task.done(): + continue + name = task.get_name() + if task.cancelled(): + raise RuntimeError(f"Orchestrator component {name!r} was cancelled unexpectedly") + error = task.exception() + if error is not None: + raise error + raise RuntimeError(f"Orchestrator component {name!r} stopped unexpectedly") diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 11a1ab4d34..7d5d42020d 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -11,13 +11,8 @@ schedule next. Transitions are level-triggered (driven by the eval source's emptiness), so in-flight rollouts of the opposite kind drain naturally on either side of an eval boundary. -- ``on_version_pending`` (called by the watcher before the engines pause for - the weight update) bumps ``off_policy_steps`` on in-flight train rollouts and - drops groups past ``max_off_policy_steps``. - Eval rollouts are measurements for the policy version they started with, - so they are allowed to finish even if training advances. Train rollouts - sampled from a frozen model never age — their sampler doesn't change - with policy updates. +- Dynamo fences and settles eval/live-policy work before worker pause; other + backends retain the off-policy window. Frozen-pool rollouts survive. Cancellations surface as synthetic ``Cancelled`` markers so the sink's count-to-``group_size`` finalization still fires. """ @@ -27,15 +22,16 @@ import asyncio import uuid from collections import Counter, defaultdict -from dataclasses import dataclass, field from enum import Enum, auto -from typing import Literal import verifiers.v1 as vf from aiolimiter import AsyncLimiter +from prime_rl.orchestrator.dispatcher_metrics import DispatcherMetrics from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_source import EvalSource +from prime_rl.orchestrator.policy_gate import MutablePolicyGate +from prime_rl.orchestrator.pool_identity import pools_may_alias from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( GroupState, @@ -44,7 +40,7 @@ Rollout, RolloutKind, ) -from prime_rl.utils.async_utils import safe_cancel, safe_cancel_all +from prime_rl.utils.async_utils import gather_shielded, safe_cancel, safe_cancel_all from prime_rl.utils.client import InferencePool, client_identity from prime_rl.utils.logger import get_logger @@ -56,69 +52,12 @@ class DispatcherMode(Enum): PREFER_EVAL = auto() -@dataclass -class DispatcherMetrics: - """Per-tick drain counters for the orchestrator's periodic log. - ``drained()`` returns the current values and clears them; point-in-time - gauges live on ``RolloutDispatcher.gauges`` instead.""" - - cancelled_by_kind_env: dict[tuple[Literal["train", "eval"], str], int] = field( - default_factory=lambda: defaultdict(int) - ) - errored_by_kind_env: dict[tuple[Literal["train", "eval"], str], int] = field( - default_factory=lambda: defaultdict(int) - ) - - def record_cancellation(self, *, kind: Literal["train", "eval"], env_name: str, n: int = 1) -> None: - self.cancelled_by_kind_env[(kind, env_name)] += n - - def record_error(self, *, kind: Literal["train", "eval"], env_name: str) -> None: - self.errored_by_kind_env[(kind, env_name)] += 1 - - def drained(self, *, train_envs: set[str], eval_envs: set[str]) -> dict[str, float]: - """Return per-tick counters and clear them. Emits the full pre- - registered key set every tick (zero when no activity) so the wandb - time axis stays dense and ``define_metric`` lines up.""" - out: dict[str, float] = {} - for kind in ("train", "eval"): - envs = train_envs if kind == "train" else eval_envs - cancelled_total = sum(self.cancelled_by_kind_env.get((kind, e), 0) for e in envs) - errored_total = sum(self.errored_by_kind_env.get((kind, e), 0) for e in envs) - out[f"dispatcher/cancelled/{kind}"] = float(cancelled_total) - out[f"dispatcher/errored/{kind}"] = float(errored_total) - for env in train_envs | eval_envs: - out[f"dispatcher/cancelled/{env}"] = float( - self.cancelled_by_kind_env.get(("train", env), 0) + self.cancelled_by_kind_env.get(("eval", env), 0) - ) - out[f"dispatcher/errored/{env}"] = float( - self.errored_by_kind_env.get(("train", env), 0) + self.errored_by_kind_env.get(("eval", env), 0) - ) - self.cancelled_by_kind_env.clear() - self.errored_by_kind_env.clear() - return out - - @staticmethod - def drain_keys(*, train_envs: set[str], eval_envs: set[str]) -> list[str]: - """Full set of keys ``drained`` may emit; used by the periodic - logger for ``wandb.define_metric``.""" - keys = [ - "dispatcher/cancelled/train", - "dispatcher/cancelled/eval", - "dispatcher/errored/train", - "dispatcher/errored/eval", - ] - for env in train_envs | eval_envs: - keys.append(f"dispatcher/cancelled/{env}") - keys.append(f"dispatcher/errored/{env}") - return keys - - class RolloutDispatcher: """``await dispatcher.start()`` runs the dispatch loop until ``stop()``. Pulls examples from ``TrainSource`` / ``EvalSource``, schedules rollouts under shared capacity, and emits ``Rollout``\\ s to - ``out_q``. The watcher drives ``on_version_pending`` for off-policy - cancellation; the orchestrator triggers eval epochs.""" + ``out_q``. The watcher drives ``on_version_pending`` for the policy-update + barrier; the orchestrator triggers eval epochs.""" def __init__( self, @@ -132,6 +71,8 @@ def __init__( max_inflight_rollouts: int, tasks_per_minute: float | None, max_off_policy_steps: int, + enforce_policy_update_barrier: bool, + policy_gate: MutablePolicyGate | None = None, ) -> None: self.policy = policy self.train_envs = train_envs @@ -142,6 +83,8 @@ def __init__( self.train_source = train_source self.eval_source = eval_source self.max_off_policy_steps = max_off_policy_steps + self.enforce_policy_update_barrier = enforce_policy_update_barrier + self.policy_gate = policy_gate or MutablePolicyGate(policy, enabled=enforce_policy_update_barrier) self.max_inflight = max_inflight_rollouts self.inflight_permits = 0 @@ -156,14 +99,11 @@ def __init__( self.out_q: asyncio.Queue[Rollout] = asyncio.Queue(maxsize=max(8, self.max_inflight)) self.mode: DispatcherMode = DispatcherMode.PREFER_TRAIN - # Set by the orchestrator after the final train step; pipeline then - # winds down without scheduling new train rollouts + # Set after the final train step to wind down new scheduling. self.train_scheduling_disabled: bool = False self.metrics = DispatcherMetrics() - # Orchestrator-owned gate. When clear, ``fill_inflight`` returns - # without scheduling new groups. The dispatcher itself doesn't know - # *why* — the orchestrator toggles this based on step / policy lead. + # Orchestrator-owned step/policy-lead gate. self.dispatch_allowed = asyncio.Event() self.dispatch_allowed.set() @@ -267,40 +207,181 @@ async def stop(self) -> None: self.task = None async def on_version_pending(self, step: int) -> None: - """Bump off-policy counters and drop groups past - ``max_off_policy_steps`` (drop_group emits ``Cancelled`` markers so - the sink still finalizes the partial group). Eval rollouts are not - aged because they are tied to their start-time policy version. - - Runs *before* the inference engines are paused for the weight update so - the resulting aborts are processed while the engine is still stepping — - otherwise the orphaned KV transfers crash the decode engine on resume - (see ``WeightWatcher.apply_policy_update``).""" + """Prepare for mutation: Dynamo fences/drains mutable-policy requests + before worker pause; other backends retain their off-policy window. + Pre-pause cancellation lets P/D abort and connector cleanup settle.""" + if not self.enforce_policy_update_barrier: + await self._advance_off_policy_window() + return + + await self.policy_gate.begin_update(step=step) + claimed_groups, claimed_tasks = self._claim_mutable_policy_work() + results, cancellation = await gather_shielded( + self._settle_policy_requests(claimed_groups, claimed_tasks), + self.policy_gate.wait_idle(), + ) + failures = [result for result in results if isinstance(result, BaseException)] + primary: BaseException | None = cancellation or (failures[0] if failures else None) + if primary is not None: + siblings = failures if cancellation is not None else failures[1:] + for sibling in siblings: + primary.add_note(f"Another policy barrier drain failed: {sibling!r}") + raise primary + + async def on_new_version(self, step: int) -> None: + """Reopen admission after the new policy is live.""" + if self.enforce_policy_update_barrier: + await self._reopen_policy_admission() + + async def on_version_update_failed(self, step: int, error: BaseException) -> None: + """Roll back only the transition fence; the policy stays unchanged.""" + if self.enforce_policy_update_barrier: + await self._reopen_policy_admission() + + @property + def policy_update_pending(self) -> bool: + return self.policy_gate.pending + + async def _reopen_policy_admission(self) -> None: + await self.policy_gate.finish_update() + + async def _advance_off_policy_window(self) -> None: + """Retain the established non-Dynamo in-flight tolerance policy.""" stale_groups: set[uuid.UUID] = set() - cancelled = 0 for meta in self.inflight.values(): if meta.kind != "train": continue - # Frozen-sourced rollouts never go stale — their sampler doesn't - # change with policy updates. if not self.train_envs.get(meta.env_name).sampler.samples_from_live_policy: continue meta.off_policy_steps += 1 if meta.off_policy_steps > self.max_off_policy_steps: stale_groups.add(meta.group_id) - for gid in stale_groups: - removed = await self.drop_group(gid) - cancelled += removed - + cancelled = 0 + for group_id in stale_groups: + cancelled += await self.drop_group(group_id) if cancelled: get_logger().warning( f"Cancelled {cancelled} train rollouts past max_off_policy_steps={self.max_off_policy_steps}. " "Consider increasing it to avoid this." ) - async def on_new_version(self, step: int) -> None: - """No-op: the dispatcher drains in ``on_version_pending`` (pre-pause).""" + def _uses_mutable_policy(self, kind: RolloutKind, env_name: str) -> bool: + if kind == "eval": + return True + pool, _model_name, samples_from_live_policy = self._train_pool_for(env_name) + # A separately-constructed "frozen" pool is safe only when its model + # and request/admin endpoints do not alias the mutable policy service. + return samples_from_live_policy or pools_may_alias(pool, self.policy_pool) + + def _claim_mutable_policy_work( + self, + ) -> tuple[dict[uuid.UUID, GroupState], list[tuple[asyncio.Task, InflightRollout]]]: + group_ids = { + group_id for group_id, group in self.groups.items() if self._uses_mutable_policy(group.kind, group.env_name) + } + group_ids.update( + meta.group_id for meta in self.inflight.values() if self._uses_mutable_policy(meta.kind, meta.env_name) + ) + + claimed_groups = { + group_id: group for group_id in group_ids if (group := self.groups.pop(group_id, None)) is not None + } + claimed_tasks: list[tuple[asyncio.Task, InflightRollout]] = [] + for task, meta in list(self.inflight.items()): + if meta.group_id not in group_ids: + continue + del self.inflight[task] + self.release(meta.rollout_count) + claimed_tasks.append((task, meta)) + return claimed_groups, claimed_tasks + + async def _settle_policy_requests( + self, + groups: dict[uuid.UUID, GroupState], + claimed: list[tuple[asyncio.Task, InflightRollout]], + ) -> None: + tasks = [task for task, _meta in claimed] + already_settled = {task for task in tasks if task.done()} + for task in tasks: + task.cancel() + + results: list[object] = [] + cancellation: asyncio.CancelledError | None = None + if tasks: + results, cancellation = await gather_shielded(*tasks) + + metadata_by_group: dict[uuid.UUID, InflightRollout] = {} + for _task, meta in claimed: + metadata_by_group.setdefault(meta.group_id, meta) + + marker_results, marker_cancellation = await gather_shielded( + self._emit_policy_cancellation_markers(groups, metadata_by_group) + ) + + failures = [ + result + for task, result in zip(tasks, results, strict=True) + if task not in already_settled + and isinstance(result, BaseException) + and not isinstance(result, asyncio.CancelledError) + ] + failures.extend(result for result in marker_results if isinstance(result, BaseException)) + primary: BaseException | None = cancellation or marker_cancellation or (failures[0] if failures else None) + if primary is not None: + siblings = failures if cancellation is not None or marker_cancellation is not None else failures[1:] + for sibling in siblings: + primary.add_note(f"Another policy barrier cleanup failed: {sibling!r}") + raise primary + + async def _emit_policy_cancellation_markers( + self, + groups: dict[uuid.UUID, GroupState], + metadata_by_group: dict[uuid.UUID, InflightRollout], + ) -> None: + """Finish every claimed group; cancellation cannot interrupt this transaction.""" + cancelled = 0 + for group_id, group in groups.items(): + owed = max(0, group.target_rollouts - group.emitted) + if owed == 0: + continue + meta = metadata_by_group.get(group_id) or InflightRollout( + kind=group.kind, + env_name=group.env_name, + group_id=group_id, + policy_version=group.policy_version_at_start, + rollout_count=1, + eval_step=group.eval_step, + ) + for _ in range(owed): + trace = Rollout( + task=vf.Task(idx=group.task_idx, prompt=None), + errors=[vf.Error(type="Cancelled", message="Policy update barrier")], + stop_condition="error", + ) + await self._emit_claimed_rollout(meta, group, trace) + self.metrics.record_cancellation(kind=meta.kind, env_name=meta.env_name, n=owed) + cancelled += owed + if cancelled: + get_logger().debug(f"Policy update barrier cancelled {cancelled} mutable-policy rollout(s)") + + async def _emit_claimed_rollout( + self, + meta: InflightRollout, + group: GroupState, + rollout: Rollout, + ) -> None: + """Enqueue a barrier-owned marker before committing group accounting.""" + rollout.kind = meta.kind + rollout.env_name = meta.env_name + rollout.group_id = meta.group_id + rollout.policy_version = group.policy_version_at_start + rollout.off_policy_steps = meta.off_policy_steps + if meta.kind == "eval": + assert group.eval_step is not None, "eval rollout missing eval_step" + rollout.eval_step = group.eval_step + await self.out_q.put(rollout) + group.emitted += 1 async def fill_inflight(self) -> None: """Schedule new rollouts up to ``max_inflight``, honoring @@ -309,29 +390,27 @@ async def fill_inflight(self) -> None: respects it. When ``PREFER_EVAL``'s source exhausts we flip back to ``PREFER_TRAIN`` so the eval tail drains alongside fresh train.""" while True: - if self.available_permits <= 0: - return - - if self.mode == DispatcherMode.PREFER_EVAL: - # PREFER_EVAL is only entered when the orchestrator triggers - # eval, which requires ``eval_source`` to be configured - assert self.eval_source is not None - if not self.eval_has_work: - # Eval source + all eval groups fully dispatched. Flip - # to PREFER_TRAIN so any remaining permits go to train - # while the in-flight eval tail completes naturally - self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained") - continue - scheduled = await self.try_schedule("eval") - if not scheduled: - return - else: # PREFER_TRAIN — respects the orchestrator's dispatch gate - if not self.dispatch_allowed.is_set(): - return - scheduled = await self.try_schedule("train") - if not scheduled: + async with self.policy_gate.scheduling_admission() as admitted: + if not admitted or self.available_permits <= 0: return + if self.mode == DispatcherMode.PREFER_EVAL: + # PREFER_EVAL implies a configured eval source. + assert self.eval_source is not None + if not self.eval_has_work: + # Fill remaining permits with train while eval drains. + self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained") + continue + scheduled = await self.try_schedule("eval") + if not scheduled: + return + else: # PREFER_TRAIN — respects the orchestrator's dispatch gate + if not self.dispatch_allowed.is_set(): + return + scheduled = await self.try_schedule("train") + if not scheduled: + return + def switch_mode(self, new_mode: DispatcherMode, *, reason: str) -> None: if new_mode == self.mode: return diff --git a/src/prime_rl/orchestrator/dispatcher_metrics.py b/src/prime_rl/orchestrator/dispatcher_metrics.py new file mode 100644 index 0000000000..cf4e5608ea --- /dev/null +++ b/src/prime_rl/orchestrator/dispatcher_metrics.py @@ -0,0 +1,58 @@ +"""Drain counters owned by the rollout dispatcher.""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Literal + + +@dataclass +class DispatcherMetrics: + """Per-tick cancellation and error counters for pipeline logging.""" + + cancelled_by_kind_env: dict[tuple[Literal["train", "eval"], str], int] = field( + default_factory=lambda: defaultdict(int) + ) + errored_by_kind_env: dict[tuple[Literal["train", "eval"], str], int] = field( + default_factory=lambda: defaultdict(int) + ) + + def record_cancellation(self, *, kind: Literal["train", "eval"], env_name: str, n: int = 1) -> None: + self.cancelled_by_kind_env[(kind, env_name)] += n + + def record_error(self, *, kind: Literal["train", "eval"], env_name: str) -> None: + self.errored_by_kind_env[(kind, env_name)] += 1 + + def drained(self, *, train_envs: set[str], eval_envs: set[str]) -> dict[str, float]: + """Return the dense counter set for this tick and clear it.""" + out: dict[str, float] = {} + for kind in ("train", "eval"): + envs = train_envs if kind == "train" else eval_envs + out[f"dispatcher/cancelled/{kind}"] = float( + sum(self.cancelled_by_kind_env.get((kind, env), 0) for env in envs) + ) + out[f"dispatcher/errored/{kind}"] = float(sum(self.errored_by_kind_env.get((kind, env), 0) for env in envs)) + for env in train_envs | eval_envs: + out[f"dispatcher/cancelled/{env}"] = float( + self.cancelled_by_kind_env.get(("train", env), 0) + self.cancelled_by_kind_env.get(("eval", env), 0) + ) + out[f"dispatcher/errored/{env}"] = float( + self.errored_by_kind_env.get(("train", env), 0) + self.errored_by_kind_env.get(("eval", env), 0) + ) + self.cancelled_by_kind_env.clear() + self.errored_by_kind_env.clear() + return out + + @staticmethod + def drain_keys(*, train_envs: set[str], eval_envs: set[str]) -> list[str]: + """Return every key :meth:`drained` may emit.""" + keys = [ + "dispatcher/cancelled/train", + "dispatcher/cancelled/eval", + "dispatcher/errored/train", + "dispatcher/errored/eval", + ] + for env in train_envs | eval_envs: + keys.extend((f"dispatcher/cancelled/{env}", f"dispatcher/errored/{env}")) + return keys diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 75e1a55d8b..439c02af87 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -280,14 +280,14 @@ class TrainEnvs(Envs[TrainEnv]): :class:`Sampler` and runtime :class:`Algorithm`, built from the env's resolved algorithm config.""" - def __init__(self, configs: Sequence[TrainEnvConfig], *, policy_pool, renderer_config=None): + def __init__(self, configs: Sequence[TrainEnvConfig], *, policy_pool, renderer_config=None, policy_gate=None): self._envs: dict[str, TrainEnv] = {} for config in configs: assert config.algo is not None, "TrainEnvConfig.algo must be resolved before env construction" env = TrainEnv( config, Sampler(config.algo.sampling, policy_pool, renderer_config), - build_algorithm(config.algo, policy_pool), + build_algorithm(config.algo, policy_pool, policy_gate=policy_gate), ) self._envs[env.name] = env diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 4cad1fcbe4..8a36e03025 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -39,7 +39,9 @@ import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive imports from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.ckpt import setup_ckpt_manager -from prime_rl.orchestrator.dispatcher import DispatcherMetrics, DispatcherMode, RolloutDispatcher +from prime_rl.orchestrator.component_supervision import raise_if_component_failed +from prime_rl.orchestrator.dispatcher import DispatcherMode, RolloutDispatcher +from prime_rl.orchestrator.dispatcher_metrics import DispatcherMetrics from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_sink import EvalSink from prime_rl.orchestrator.eval_source import EvalSource @@ -50,6 +52,7 @@ monkey_patch_oai_iterable_types, ) from prime_rl.orchestrator.periodic_logger import PeriodicLogger +from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.orchestrator.train_sink import TrainSink from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( @@ -154,6 +157,10 @@ def __init__(self, config: OrchestratorConfig) -> None: self.progress = Progress() self.ckpt_manager = setup_ckpt_manager(config.output_dir, config.ckpt) self.policy = Policy(version=0, model_name="") + self.policy_gate = MutablePolicyGate( + self.policy, + enabled=config.model.client.admin_api == "dynamo", + ) self.stopped = asyncio.Event() # True after the final train step ships — pipeline winds down without # scheduling new train rollouts @@ -241,7 +248,10 @@ async def setup(self) -> None: get_logger().info("Loading training environments") self.train_envs = TrainEnvs( - config.train.env, policy_pool=self.policy_inference, renderer_config=config.renderer + config.train.env, + policy_pool=self.policy_inference, + renderer_config=config.renderer, + policy_gate=self.policy_gate, ) get_logger().debug( f"Loaded {len(self.train_envs)} training environment(s) ({', '.join(self.train_envs.names)})" @@ -349,6 +359,8 @@ async def setup(self) -> None: max_inflight_rollouts=config.max_inflight_rollouts, tasks_per_minute=config.tasks_per_minute, max_off_policy_steps=config.max_off_policy_steps, + enforce_policy_update_barrier=config.model.client.admin_api == "dynamo", + policy_gate=self.policy_gate, ) self.train_sink = TrainSink( config, @@ -453,6 +465,7 @@ async def main_loop(self) -> None: to the train / eval sink. Both sinks return a finalized batch (or ``None``) from ``add()``; we just dispatch on the result.""" while not self.stopped.is_set(): + raise_if_component_failed(self.component_tasks) if self.draining and self.dispatcher.is_idle: get_logger().info("Pipeline drained, exiting main loop") self.stopped.set() @@ -847,6 +860,9 @@ async def on_new_version(self, step: int) -> None: re-evaluate the dispatch gate (may resume if the trainer caught up).""" self.update_dispatch_gate() + async def on_version_update_failed(self, step: int, error: BaseException) -> None: + """No transition state is owned here; the policy version is unchanged.""" + async def stop(self) -> None: """Bounded best-effort teardown of all components. Has a global timeout so a wedged peer can't keep the process alive forever — diff --git a/src/prime_rl/orchestrator/policy_gate.py b/src/prime_rl/orchestrator/policy_gate.py new file mode 100644 index 0000000000..e5e8227004 --- /dev/null +++ b/src/prime_rl/orchestrator/policy_gate.py @@ -0,0 +1,97 @@ +"""Admission and in-flight accounting for mutable-policy inference calls.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +from prime_rl.orchestrator.types import Policy + + +class PolicyRequestRejected(RuntimeError): + """A request cannot safely run against the mutable policy version.""" + + +class MutablePolicyGate: + """Serialize policy mutation with every request that depends on its weights. + + Dispatcher scheduling holds :meth:`scheduling_admission` until a rollout + task is registered. Other policy I/O, such as OPSD prefill scoring, uses + :meth:`request` for its full lifetime. Closing the gate prevents new work; + callers can then cancel dispatcher-owned tasks and await :meth:`wait_idle` + before pausing the engine. + """ + + def __init__(self, policy: Policy, *, enabled: bool) -> None: + self.policy = policy + self.enabled = enabled + self._admission_lock = asyncio.Lock() + self._pending = False + self._active_requests = 0 + self._idle = asyncio.Event() + self._idle.set() + + @property + def pending(self) -> bool: + return self.enabled and self._pending + + @asynccontextmanager + async def scheduling_admission(self) -> AsyncIterator[bool]: + """Hold the admission serialization point through scheduling commit.""" + if not self.enabled: + yield True + return + async with self._admission_lock: + yield not self._pending + + @asynccontextmanager + async def request(self, *, expected_version: int) -> AsyncIterator[None]: + """Register one non-dispatcher policy call for its complete lifetime.""" + if not self.enabled: + yield + return + + async with self._admission_lock: + if self._pending: + raise PolicyRequestRejected( + f"Mutable-policy request rejected because a policy update is pending (expected version " + f"{expected_version})" + ) + if expected_version != self.policy.version: + raise PolicyRequestRejected( + f"Mutable-policy request expected policy version {expected_version}, but current version is " + f"{self.policy.version}" + ) + self._active_requests += 1 + self._idle.clear() + + try: + yield + finally: + # No await here: even repeated cancellation must not strand the + # active count and deadlock the mutation barrier. + self._active_requests -= 1 + if self._active_requests == 0: + self._idle.set() + + async def begin_update(self, *, step: int) -> None: + """Close admission after every in-progress scheduling commit.""" + if not self.enabled: + return + async with self._admission_lock: + if self._pending: + raise RuntimeError(f"A policy update is already pending while preparing step {step}") + self._pending = True + + async def finish_update(self) -> None: + """Reopen admission after a successful or proven pre-mutation failure.""" + if not self.enabled: + return + async with self._admission_lock: + self._pending = False + + async def wait_idle(self) -> None: + """Wait until every already-admitted non-dispatcher request settles.""" + if self.enabled: + await self._idle.wait() diff --git a/src/prime_rl/orchestrator/pool_identity.py b/src/prime_rl/orchestrator/pool_identity.py new file mode 100644 index 0000000000..fc1838daed --- /dev/null +++ b/src/prime_rl/orchestrator/pool_identity.py @@ -0,0 +1,42 @@ +"""Serving-resource identity used by the mutable-policy barrier.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from prime_rl.utils.client import InferencePool + + +def _normalized_url(value: object) -> str: + return str(value).rstrip("/").removesuffix("/v1") + + +def serving_identity(pool: InferencePool) -> tuple[str, frozenset[str], frozenset[str]]: + """Return model, request endpoints, and admin endpoints for one pool.""" + request_endpoints = frozenset(_normalized_url(client.base_url) for client in getattr(pool, "train_clients", ())) + admin_endpoints = frozenset(_normalized_url(client.base_url) for client in getattr(pool, "admin_clients", ())) + return str(pool.model_name), request_endpoints, admin_endpoints + + +def pools_may_alias(left: InferencePool, right: InferencePool) -> bool: + """Conservatively detect two pool objects backed by one mutable model. + + Inline frozen references construct a separate Python pool, so object + identity is only the fast path. Equal model names plus an overlapping + request or admin endpoint identify an alias. Missing endpoint information + is ambiguous and therefore treated as mutable. + """ + if left is right: + return True + left_model, left_requests, left_admin = serving_identity(left) + right_model, right_requests, right_admin = serving_identity(right) + if left_model != right_model: + return False + if not left_requests or not right_requests: + return True + if left_requests & right_requests: + return True + if not left_admin or not right_admin: + return True + return bool(left_admin & right_admin) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index e3d0e93198..0b3406dd65 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -172,8 +172,12 @@ class VersionObserver(Protocol): ``on_version_pending`` fires *before* the inference engines are paused for the weight update; ``on_new_version`` fires *after* the new weights are live - and ``Policy`` has been mutated.""" + and ``Policy`` has been mutated. ``on_version_update_failed`` rolls back + transition-only state only when failure is known to precede engine + mutation; indeterminate engine-update failures remain fenced.""" async def on_version_pending(self, step: int) -> None: ... async def on_new_version(self, step: int) -> None: ... + + async def on_version_update_failed(self, step: int, error: BaseException) -> None: ... diff --git a/src/prime_rl/orchestrator/watcher.py b/src/prime_rl/orchestrator/watcher.py index c01d349f40..8230938ac8 100644 --- a/src/prime_rl/orchestrator/watcher.py +++ b/src/prime_rl/orchestrator/watcher.py @@ -1,6 +1,6 @@ -"""WeightWatcher: polls the broadcast dir, advances ``Policy``, notifies -observers (dispatcher → off-policy cancel). Standalone async task; the -orchestrator's barrier bounds the in-flight lead.""" +"""WeightWatcher: polls broadcasts and applies policy updates behind the +dispatcher admission/drain barrier. Standalone async task; the orchestrator's +lead gate separately bounds sampling ahead of the trainer.""" from __future__ import annotations @@ -9,7 +9,7 @@ from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.types import Policy, VersionObserver -from prime_rl.utils.async_utils import safe_cancel +from prime_rl.utils.async_utils import gather_shielded, safe_cancel from prime_rl.utils.client import InferencePool from prime_rl.utils.logger import format_time, get_logger from prime_rl.utils.pathing import get_broadcast_dir, get_step_path, wait_for_path @@ -92,8 +92,11 @@ async def apply_policy_update(self, next_step: int) -> None: f"Orchestrator resumed: checkpoint {next_step} ready (after {format_time(self.last_wait_for_ckpt_time)})" ) - # Drain off-policy rollouts BEFORE pausing the inference engines. - # Aborting a rollout triggers vLLM's KV-connector cleanup (NIXL's + # Establish the backend-specific application transition BEFORE + # pausing inference engines. Dynamo fences new dispatch and drains + # every mutable-policy request; the vLLM admin path retains its + # configured off-policy cancellation window. Aborting a rollout + # triggers vLLM's KV-connector cleanup (NIXL's # ``_reqs_not_processed``), which is only propagated to the workers # while the engine is stepping. If we drain after resume instead, # the aborts race with the flush of KV transfers that completed @@ -102,17 +105,32 @@ async def apply_policy_update(self, next_step: int) -> None: # the engine and cascading to every DP rank. Draining first lets the # aborts settle under normal stepping. ``on_new_version`` (below) # still runs post-update for observers that need the live version. - for observer in self.observers: - try: + entered_observers: list[VersionObserver] = [] + try: + for observer in self.observers: + # Include the observer before entry so a partially-applied + # barrier gets its failure rollback hook. + entered_observers.append(observer) await observer.on_version_pending(next_step) - except Exception as exc: - get_logger().warning( - f"Observer {type(observer).__name__}.on_version_pending({next_step}) raised: {exc!r}" - ) + except BaseException as exc: + await self._notify_update_failed(entered_observers, next_step, exc) + raise get_logger().debug(f"Updating weights to step {next_step}") t1 = time.perf_counter() - await self.inference.update_weights(weights_path, lora_name=self.lora_name, step=next_step) + try: + await self.inference.update_weights(weights_path, lora_name=self.lora_name, step=next_step) + except BaseException as exc: + # Once an admin update starts, an error cannot prove that no + # worker committed new weights (a resume failure is even later). + # Keep every transition fence closed and let the component + # failure terminate the run; reopening would stamp old-version + # requests onto mixed or fully-updated workers. + exc.add_note( + f"Policy update {next_step} may have mutated inference workers; mutable-policy admission " + "remains fail-closed" + ) + raise self.last_update_weights_time = time.perf_counter() - t1 self.update_count += 1 get_logger().debug(f"Updated weights to step {next_step} in {format_time(self.last_update_weights_time)}") @@ -123,13 +141,60 @@ async def apply_policy_update(self, next_step: int) -> None: self.inference.update_model_name(self.lora_name) self.policy.model_name = self.lora_name - for observer in self.observers: + await self._notify_update_succeeded(entered_observers, next_step) + + @staticmethod + async def _notify_update_succeeded(observers: list[VersionObserver], step: int) -> None: + """Notify every observer, preserving cancellation until fences reopen.""" + fatal_error: BaseException | None = None + # Complete observers in reverse order: the orchestrator first + # re-evaluates its long-lived lead gate while the dispatcher's short + # transition fence remains closed, then the dispatcher reopens + # admission on the new version. + for observer in reversed(observers): + try: + await observer.on_new_version(step) + except BaseException as exc: + # A success callback may have partially applied transition + # state before failing. Its failure hook is also the idempotent + # release fallback; continue through the remaining observers. try: - await observer.on_new_version(next_step) - except Exception as exc: - get_logger().warning( - f"Observer {type(observer).__name__}.on_new_version({next_step}) raised: {exc!r}" + await observer.on_version_update_failed(step, exc) + except BaseException as cleanup_error: + exc.add_note( + f"Observer {type(observer).__name__}.on_version_update_failed({step}) also failed: " + f"{cleanup_error!r}" + ) + if not isinstance(cleanup_error, Exception) and fatal_error is None: + fatal_error = cleanup_error + if isinstance(exc, Exception): + get_logger().warning(f"Observer {type(observer).__name__}.on_new_version({step}) raised: {exc!r}") + elif fatal_error is None: + fatal_error = exc + else: + fatal_error.add_note( + f"Observer {type(observer).__name__}.on_new_version({step}) also failed: {exc!r}" ) + if fatal_error is not None: + await WeightWatcher._notify_update_failed(observers, step, fatal_error) + raise fatal_error + + @staticmethod + async def _notify_update_failed( + observers: list[VersionObserver], + step: int, + primary_error: BaseException, + ) -> None: + for observer in reversed(observers): + results, cancellation = await gather_shielded(observer.on_version_update_failed(step, primary_error)) + failures = [result for result in results if isinstance(result, BaseException)] + if cancellation is not None: + failures.append(cancellation) + for cleanup_error in failures: + primary_error.add_note( + f"Observer {type(observer).__name__}.on_version_update_failed({step}) also failed: " + f"{cleanup_error!r}" + ) def gauges(self) -> dict[str, float]: return { diff --git a/src/prime_rl/utils/async_utils.py b/src/prime_rl/utils/async_utils.py index 49e03500e2..79a80d1bf1 100644 --- a/src/prime_rl/utils/async_utils.py +++ b/src/prime_rl/utils/async_utils.py @@ -2,12 +2,30 @@ import asyncio from collections import deque +from collections.abc import Awaitable from time import perf_counter import numpy as np from pydantic import BaseModel +async def gather_shielded( + *awaitables: Awaitable[object], +) -> tuple[list[object], asyncio.CancelledError | None]: + """Settle all awaitables despite repeated cancellation of the caller.""" + settling = asyncio.gather(*awaitables, return_exceptions=True) + cancellation: asyncio.CancelledError | None = None + while not settling.done(): + try: + await asyncio.shield(settling) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + else: + cancellation.add_note("Caller cancelled again while bounded siblings were settling") + return list(settling.result()), cancellation + + async def safe_cancel(task: asyncio.Task) -> None: """Safely cancels and awaits an asyncio.Task.""" task.cancel() diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index 5240713c37..a08cad8fd6 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -106,6 +106,30 @@ def test_reserved_role_engine_override_is_rejected(): build_engine_config(config, "prefill", kv_events_port=20080) +@pytest.mark.parametrize("key", sorted(dynamo._ENGINE_CONFIG_EXCLUDED)) +def test_wrapper_only_global_engine_override_is_rejected(key: str): + config = disaggregated_config(vllm_extra={key: "invalid"}) + + with pytest.raises(ValueError, match=rf"vllm_extra.*{key}.*wrapper/server-only"): + build_engine_config(config, "prefill", kv_events_port=20080) + + +@pytest.mark.parametrize("key", sorted(dynamo._ENGINE_CONFIG_EXCLUDED)) +def test_wrapper_only_role_engine_override_is_rejected(key: str): + config = disaggregated_config( + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + "decode_vllm_overrides": {key: "invalid"}, + } + ) + + with pytest.raises(ValueError, match=rf"decode_vllm_overrides.*{key}.*wrapper/server-only"): + build_engine_config(config, "decode") + + def test_local_specs_allocate_four_workers_and_unique_ports(tmp_path: Path): specs = build_local_worker_specs(disaggregated_config(), tmp_path, gpu_ids=["4", "5", "6", "7"]) diff --git a/tests/unit/inference/test_dynamo_admin.py b/tests/unit/inference/test_dynamo_admin.py index a4f9d11c2a..e3b4c5cf2e 100644 --- a/tests/unit/inference/test_dynamo_admin.py +++ b/tests/unit/inference/test_dynamo_admin.py @@ -468,6 +468,8 @@ def handler(request: httpx.Request) -> httpx.Response: paths = [path for path, _body in requests] assert paths.count("/engine/pause_generation") == 2 assert paths.count("/engine/resume_generation") == 2 + pause_bodies = [body for path, body in requests if path.endswith("/pause_generation")] + assert pause_bodies == [{"mode": "wait", "clear_cache": False}] * 2 @pytest.mark.asyncio diff --git a/tests/unit/inference/test_dynamo_admin_barriers.py b/tests/unit/inference/test_dynamo_admin_barriers.py new file mode 100644 index 0000000000..83bb3d96e7 --- /dev/null +++ b/tests/unit/inference/test_dynamo_admin_barriers.py @@ -0,0 +1,263 @@ +"""Cancellation and partial-failure contracts for stateful Dynamo admin fanouts.""" + +import asyncio +from pathlib import Path + +import pytest + +from prime_rl.inference.dynamo_admin import DynamoAdminAPI + + +@pytest.mark.asyncio +async def test_nccl_initialization_settles_siblings_and_fails_closed_after_partial_error( + monkeypatch: pytest.MonkeyPatch, +): + admin = DynamoAdminAPI() + delayed_started = asyncio.Event() + release_delayed = asyncio.Event() + delayed_finished = asyncio.Event() + calls = 0 + + async def post(client, method, *_args, **_kwargs): + nonlocal calls + calls += 1 + assert method == "init_weights_update_group" + if client == "failed": + raise RuntimeError("init failed") + delayed_started.set() + await release_delayed.wait() + delayed_finished.set() + return {} + + monkeypatch.setattr(admin, "_post", post) + initialize = asyncio.create_task( + admin.initialize_nccl( + ["failed", "delayed"], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=2, + gpus_per_worker=1, + quantize_in_weight_transfer=False, + ) + ) + await delayed_started.wait() + await asyncio.sleep(0) + assert not initialize.done() + + release_delayed.set() + with pytest.raises(RuntimeError, match="init failed"): + await initialize + assert delayed_finished.is_set() + + with pytest.raises(RuntimeError, match="indeterminate"): + await admin.initialize_nccl( + ["failed", "delayed"], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=2, + gpus_per_worker=1, + quantize_in_weight_transfer=False, + ) + with pytest.raises(RuntimeError, match="indeterminate"): + await admin.update_weights(["failed", "delayed"], Path("weights"), step=1) + assert calls == 2 + + +@pytest.mark.asyncio +async def test_nccl_initialization_settles_before_propagating_repeated_cancellation( + monkeypatch: pytest.MonkeyPatch, +): + admin = DynamoAdminAPI() + init_started = asyncio.Event() + release_init = asyncio.Event() + init_finished = asyncio.Event() + + async def post(_client, method, *_args, **_kwargs): + assert method == "init_weights_update_group" + init_started.set() + await release_init.wait() + init_finished.set() + return {} + + monkeypatch.setattr(admin, "_post", post) + initialize = asyncio.create_task( + admin.initialize_nccl( + ["worker"], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=1, + gpus_per_worker=1, + quantize_in_weight_transfer=False, + ) + ) + await init_started.wait() + initialize.cancel() + await asyncio.sleep(0) + initialize.cancel() + await asyncio.sleep(0) + assert not initialize.done() + + release_init.set() + with pytest.raises(asyncio.CancelledError): + await initialize + assert init_finished.is_set() + + with pytest.raises(RuntimeError, match="indeterminate"): + await admin.update_weights(["worker"], Path("weights"), step=1) + + +@pytest.mark.asyncio +async def test_pause_fanout_settles_delayed_sibling_before_resume(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + delayed_started = asyncio.Event() + release_delayed = asyncio.Event() + delayed_finished = asyncio.Event() + resume_started = asyncio.Event() + + async def post(client, method, *_args, **_kwargs): + if method == "pause_generation": + if client == "fast-failure": + raise RuntimeError("pause failed") + delayed_started.set() + await release_delayed.wait() + delayed_finished.set() + return {} + if method == "resume_generation": + resume_started.set() + assert delayed_finished.is_set() + return {} + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + update = asyncio.create_task(admin.update_weights(["fast-failure", "delayed"], Path("weights"), step=1)) + await delayed_started.wait() + await asyncio.sleep(0) + assert not resume_started.is_set() + + release_delayed.set() + with pytest.raises(RuntimeError, match="pause failed"): + await update + + assert resume_started.is_set() + + +@pytest.mark.asyncio +async def test_collective_update_settles_delayed_sibling_before_resume(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + delayed_started = asyncio.Event() + release_delayed = asyncio.Event() + delayed_finished = asyncio.Event() + resume_started = asyncio.Event() + + async def post(client, method, *_args, **_kwargs): + if method == "pause_generation": + return {} + if method == "update_weights_from_disk": + if client == "fast-failure": + raise RuntimeError("collective failed") + delayed_started.set() + await release_delayed.wait() + delayed_finished.set() + return {} + if method == "resume_generation": + resume_started.set() + assert delayed_finished.is_set() + return {} + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + update = asyncio.create_task(admin.update_weights(["fast-failure", "delayed"], Path("weights"), step=1)) + await delayed_started.wait() + await asyncio.sleep(0) + assert not resume_started.is_set() + + release_delayed.set() + with pytest.raises(RuntimeError, match="collective failed"): + await update + + assert resume_started.is_set() + + +@pytest.mark.asyncio +async def test_collective_update_is_settled_before_propagating_cancellation(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + update_started = asyncio.Event() + release_update = asyncio.Event() + update_finished = asyncio.Event() + resume_started = asyncio.Event() + + async def post(_client, method, *_args, **_kwargs): + if method == "pause_generation": + return {} + if method == "update_weights_from_disk": + update_started.set() + await release_update.wait() + update_finished.set() + return {} + if method == "resume_generation": + resume_started.set() + assert update_finished.is_set() + return {} + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + update = asyncio.create_task(admin.update_weights(["worker"], Path("weights"), step=1)) + await update_started.wait() + update.cancel() + await asyncio.sleep(0) + update.cancel() + await asyncio.sleep(0) + + assert not update.done() + assert not resume_started.is_set() + + release_update.set() + with pytest.raises(asyncio.CancelledError): + await update + + assert update_finished.is_set() + assert resume_started.is_set() + + +@pytest.mark.asyncio +async def test_resume_failure_after_successful_mutation_is_reported(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + operations: list[str] = [] + + async def post(_client, method, *_args, **_kwargs): + operations.append(method) + if method == "resume_generation": + raise RuntimeError("resume failed after mutation") + return {} + + monkeypatch.setattr(admin, "_post", post) + + with pytest.raises(RuntimeError, match="resume failed after mutation"): + await admin.update_weights(["worker"], Path("weights"), step=1) + + assert operations == ["pause_generation", "update_weights_from_disk", "resume_generation"] + with pytest.raises(RuntimeError, match="weight state is indeterminate"): + await admin.update_weights(["worker"], Path("weights"), step=2) + assert operations == ["pause_generation", "update_weights_from_disk", "resume_generation"] + + +@pytest.mark.asyncio +async def test_fanout_preserves_primary_error_and_annotates_siblings(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + + async def post(client, method, *_args, **_kwargs): + if method == "pause_generation": + raise RuntimeError(f"{client} pause failed") + if method == "resume_generation": + return {} + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + + with pytest.raises(RuntimeError, match="primary pause failed") as exc: + await admin.update_weights(["primary", "sibling"], Path("weights"), step=1) + + assert exc.value.__notes__ == ["Dynamo pause_generation sibling also failed: RuntimeError('sibling pause failed')"] diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index af5ebe2326..9d79b5eb50 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -8,9 +8,10 @@ from verifiers.v1.types import AssistantMessage, ToolMessage, UserMessage from prime_rl.configs.algorithm import AlgoConfig, FrozenModelConfig -from prime_rl.orchestrator.algo import EchoAlgorithm, OPDAlgorithm, stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.algo import EchoAlgorithm, OPDAlgorithm, OPSDAlgorithm, stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import Policy, Rollout from prime_rl.transport.types import TrainingSample from prime_rl.utils.client import DynamoInferencePool, StaticInferencePool from prime_rl.utils.elastic import ElasticInferencePool @@ -102,6 +103,27 @@ async def test_opd_rejects_elastic_teacher_pool(): await algo.setup() +@pytest.mark.asyncio +async def test_opsd_rejects_late_live_policy_score_as_rollout_error(): + policy = Policy(version=0, model_name="policy") + gate = MutablePolicyGate(policy, enabled=True) + pool = MagicMock(model_name="policy") + pool.score = AsyncMock(return_value=[0.0, -0.1]) + algo = OPSDAlgorithm(_build(type="opsd"), pool, policy_gate=gate) + algo.renderer = MagicMock() + algo.renderer.render_ids.return_value = [999] + rollout = _make_rollout([_make_sample()]) + rollout.info["demonstration"] = "expert answer" + rollout.policy_version = 0 + + await gate.begin_update(step=1) + await algo.finalize_rollout(rollout) + + assert rollout.has_error + assert rollout.error.type == "PolicyRequestRejected" + pool.score.assert_not_awaited() + + def test_sft_requires_teacher(): with pytest.raises(ValueError, match="needs a teacher to sample rollouts from"): _build(type="sft") diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index 73768a7698..bf75d3bb4f 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -2,11 +2,26 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, patch +import pytest from renderers import Qwen3VLRendererConfig +from prime_rl.orchestrator.component_supervision import raise_if_component_failed from prime_rl.orchestrator.utils import setup_policy_inference_pool +def test_component_failure_is_raised_into_main_loop(): + async def run() -> None: + async def failed_watcher() -> None: + raise RuntimeError("indeterminate policy update") + + watcher = asyncio.create_task(failed_watcher(), name="watcher") + await asyncio.wait({watcher}) + with pytest.raises(RuntimeError, match="indeterminate policy update"): + raise_if_component_failed([watcher]) + + asyncio.run(run()) + + def test_setup_policy_inference_pool_uses_renderer_when_enabled(): async def run() -> None: tokenizer = object() diff --git a/tests/unit/orchestrator/test_pool_identity.py b/tests/unit/orchestrator/test_pool_identity.py new file mode 100644 index 0000000000..a1e090a1c0 --- /dev/null +++ b/tests/unit/orchestrator/test_pool_identity.py @@ -0,0 +1,52 @@ +from types import SimpleNamespace + +import pytest + +from prime_rl.orchestrator.pool_identity import pools_may_alias + + +def _pool(model: str, request: str | None, admin: str | None): + return SimpleNamespace( + model_name=model, + train_clients=[] if request is None else [SimpleNamespace(base_url=request)], + admin_clients=[] if admin is None else [SimpleNamespace(base_url=admin)], + ) + + +@pytest.mark.parametrize( + ("left", "right", "aliases"), + [ + ( + _pool("policy", "http://frontend/v1", "http://worker:8081"), + _pool("policy", "http://frontend", "http://worker:8081/"), + True, + ), + ( + _pool("policy", "http://policy/v1", "http://policy-worker:8081"), + _pool("policy", "http://frozen/v1", "http://frozen-worker:8081"), + False, + ), + ( + _pool("policy", "http://router-a/v1", "http://shared-worker:8081"), + _pool("policy", "http://router-b/v1", "http://shared-worker:8081"), + True, + ), + ( + _pool("policy", "http://frontend/v1", "http://worker:8081"), + _pool("other-model", "http://frontend/v1", "http://worker:8081"), + False, + ), + ( + _pool("policy", None, None), + _pool("policy", "http://frozen/v1", "http://frozen-worker:8081"), + True, + ), + ], +) +def test_pool_aliasing_uses_model_request_and_admin_identity(left, right, aliases: bool): + assert pools_may_alias(left, right) is aliases + + +def test_pool_aliasing_accepts_object_identity(): + pool = _pool("policy", "http://frontend/v1", "http://worker:8081") + assert pools_may_alias(pool, pool) diff --git a/tests/unit/orchestrator/test_weight_update_barrier.py b/tests/unit/orchestrator/test_weight_update_barrier.py new file mode 100644 index 0000000000..5d391ce65e --- /dev/null +++ b/tests/unit/orchestrator/test_weight_update_barrier.py @@ -0,0 +1,643 @@ +import asyncio +import uuid +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from prime_rl.orchestrator.dispatcher import DispatcherMode, RolloutDispatcher +from prime_rl.orchestrator.policy_gate import MutablePolicyGate, PolicyRequestRejected +from prime_rl.orchestrator.types import GroupState, InflightRollout, Policy +from prime_rl.orchestrator.watcher import WeightWatcher +from prime_rl.utils.pathing import get_broadcast_dir, get_step_path + + +class _TrainEnvs: + def __init__(self, *, live: bool, pool: object) -> None: + self._live = live + self._pool = pool + + def get(self, _name: str): + return SimpleNamespace(sampler=SimpleNamespace(samples_from_live_policy=self._live, pool=self._pool)) + + +def _dispatcher( + *, + max_inflight: int = 1, + live_train: bool = True, + frozen_uses_policy_pool: bool = False, + frozen_aliases_policy_pool: bool = False, + enforce_policy_update_barrier: bool = True, + max_off_policy_steps: int = 8, + policy_gate: MutablePolicyGate | None = None, +) -> RolloutDispatcher: + def pool(model_name: str, request_url: str, admin_url: str): + return SimpleNamespace( + model_name=model_name, + train_clients=[SimpleNamespace(base_url=request_url, headers={})], + admin_clients=[SimpleNamespace(base_url=admin_url)], + ) + + policy_pool = pool("policy", "http://policy/v1", "http://policy-worker:8081") + if live_train or frozen_uses_policy_pool: + train_pool = policy_pool + elif frozen_aliases_policy_pool: + # A separately-constructed pool can still address the exact mutable + # serving resource. Object identity is not a topology identity. + train_pool = pool("policy", "http://policy/v1", "http://policy-worker:8081") + else: + train_pool = pool("frozen", "http://frozen/v1", "http://frozen-worker:8081") + return RolloutDispatcher( + train_envs=_TrainEnvs(live=live_train, pool=train_pool), + eval_envs=object(), + train_source=object(), + eval_source=object(), + policy_pool=policy_pool, + policy=Policy(version=0, model_name="policy"), + max_inflight_rollouts=max_inflight, + tasks_per_minute=None, + max_off_policy_steps=max_off_policy_steps, + enforce_policy_update_barrier=enforce_policy_update_barrier, + policy_gate=policy_gate, + ) + + +@pytest.mark.asyncio +async def test_policy_barrier_serializes_with_suspended_eval_scheduling(): + dispatcher = _dispatcher() + dispatcher.mode = DispatcherMode.PREFER_EVAL + scheduling_started = asyncio.Event() + allow_schedule_to_commit = asyncio.Event() + request_cleanup_finished = asyncio.Event() + + async def active_request() -> None: + try: + await asyncio.Future() + finally: + # Represents the client/connector abort cleanup performed while + # unwinding a cancelled rollout request. + await asyncio.sleep(0) + request_cleanup_finished.set() + + async def schedule_one(_kind: str) -> bool: + scheduling_started.set() + await allow_schedule_to_commit.wait() + group_id = uuid.uuid4() + request = asyncio.create_task(active_request()) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[request] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits += 1 + await asyncio.sleep(0) + return True + + dispatcher.try_schedule = schedule_one # type: ignore[method-assign] + + fill = asyncio.create_task(dispatcher.fill_inflight()) + await scheduling_started.wait() + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await asyncio.sleep(0) + + # The barrier must serialize behind the scheduling commit. Otherwise it + # can take an empty snapshot and let an old-policy request escape. + assert not barrier.done() + + allow_schedule_to_commit.set() + await fill + await barrier + + assert request_cleanup_finished.is_set() + assert not dispatcher.inflight + assert dispatcher.policy_update_pending + + # The transition fence remains closed until the watcher reports either + # success or failure. + await dispatcher.fill_inflight() + assert not dispatcher.inflight + + await dispatcher.on_new_version(1) + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("frozen_uses_policy_pool", "frozen_aliases_policy_pool", "cancelled_by_barrier"), + [(False, False, False), (True, False, True), (False, True, True)], +) +async def test_policy_barrier_only_preserves_frozen_requests_on_a_distinct_pool( + frozen_uses_policy_pool: bool, + frozen_aliases_policy_pool: bool, + cancelled_by_barrier: bool, +): + dispatcher = _dispatcher( + live_train=False, + frozen_uses_policy_pool=frozen_uses_policy_pool, + frozen_aliases_policy_pool=frozen_aliases_policy_pool, + ) + group_id = uuid.uuid4() + request_started = asyncio.Event() + + async def request() -> None: + request_started.set() + await asyncio.Future() + + task = asyncio.create_task(request()) + await request_started.wait() + dispatcher.groups[group_id] = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + policy_version_at_start=0, + ) + dispatcher.inflight[task] = InflightRollout( + kind="train", + env_name="train", + group_id=group_id, + policy_version=0, + rollout_count=1, + ) + dispatcher.inflight_permits = 1 + + await dispatcher.on_version_pending(1) + + assert task.cancelled() is cancelled_by_barrier + assert (task in dispatcher.inflight) is not cancelled_by_barrier + + await dispatcher.on_new_version(1) + await dispatcher.cancel_inflight_rollouts() + + +@pytest.mark.asyncio +async def test_failed_policy_barrier_skips_engine_mutation_and_reopens_admission(tmp_path: Path): + dispatcher = _dispatcher() + group_id = uuid.uuid4() + cleanup_attempted = asyncio.Event() + + async def request_with_failed_cleanup() -> None: + try: + await asyncio.Future() + except asyncio.CancelledError: + cleanup_attempted.set() + raise RuntimeError("connector cleanup failed") + + request = asyncio.create_task(request_with_failed_cleanup()) + await asyncio.sleep(0) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[request] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits = 1 + + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _Inference: + def __init__(self) -> None: + self.update_calls = 0 + + async def update_weights(self, *_args, **_kwargs) -> None: + self.update_calls += 1 + + inference = _Inference() + policy = dispatcher.policy + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=policy, + inference=inference, + observers=[dispatcher], + lora_name=None, + ) + + with pytest.raises(RuntimeError, match="connector cleanup failed"): + await watcher.apply_policy_update(1) + + assert cleanup_attempted.is_set() + assert inference.update_calls == 0 + assert watcher.ckpt_step == 0 + assert policy.version == 0 + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_policy_barrier_accepts_request_error_that_was_already_settled(): + dispatcher = _dispatcher() + group_id = uuid.uuid4() + + async def failed_request() -> None: + raise RuntimeError("ordinary rollout failure") + + request = asyncio.create_task(failed_request()) + await asyncio.wait({request}) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[request] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits = 1 + + await dispatcher.on_version_pending(1) + + assert dispatcher.policy_update_pending + assert not dispatcher.inflight + await dispatcher.on_new_version(1) + + +@pytest.mark.asyncio +async def test_policy_barrier_settles_cleanup_before_propagating_cancellation(): + dispatcher = _dispatcher() + group_id = uuid.uuid4() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + cleanup_finished = asyncio.Event() + + async def request() -> None: + try: + await asyncio.Future() + finally: + cleanup_started.set() + await release_cleanup.wait() + cleanup_finished.set() + + task = asyncio.create_task(request()) + await asyncio.sleep(0) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[task] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits = 1 + + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await cleanup_started.wait() + barrier.cancel() + await asyncio.sleep(0) + barrier.cancel() + await asyncio.sleep(0) + + assert not barrier.done() + assert not cleanup_finished.is_set() + + release_cleanup.set() + with pytest.raises(asyncio.CancelledError): + await barrier + + assert cleanup_finished.is_set() + await dispatcher.on_version_update_failed(1, asyncio.CancelledError()) + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_policy_barrier_shields_marker_enqueue_and_commits_accounting_after_put(): + dispatcher = _dispatcher() + group_id = uuid.uuid4() + request_started = asyncio.Event() + marker_put_started = asyncio.Event() + + class _ObservedQueue(asyncio.Queue): + async def put(self, item) -> None: + marker_put_started.set() + await super().put(item) + + dispatcher.out_q = _ObservedQueue(maxsize=1) + dispatcher.out_q.put_nowait(object()) + + async def request() -> None: + request_started.set() + await asyncio.Future() + + request_task = asyncio.create_task(request()) + await request_started.wait() + group = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.groups[group_id] = group + dispatcher.inflight[request_task] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits = 1 + + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await marker_put_started.wait() + assert group.emitted == 0 + + barrier.cancel() + await asyncio.sleep(0) + barrier.cancel() + await asyncio.sleep(0) + assert not barrier.done() + assert group.emitted == 0 + + dispatcher.out_q.get_nowait() + with pytest.raises(asyncio.CancelledError): + await barrier + + marker = dispatcher.out_q.get_nowait() + assert marker.error.type == "Cancelled" + assert group.emitted == 1 + assert not dispatcher.groups + assert not dispatcher.inflight + assert dispatcher.inflight_permits == 0 + + await dispatcher.on_version_update_failed(1, asyncio.CancelledError()) + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_policy_barrier_drains_active_policy_calls_and_rejects_late_or_stale_calls(): + policy = Policy(version=0, model_name="policy") + policy_gate = MutablePolicyGate(policy, enabled=True) + dispatcher = _dispatcher(policy_gate=policy_gate) + # The helper constructs its own Policy, so make both components share the + # exact version object as production does. + dispatcher.policy = policy + + call_started = asyncio.Event() + release_call = asyncio.Event() + + async def active_policy_call() -> None: + async with policy_gate.request(expected_version=0): + call_started.set() + await release_call.wait() + + score = asyncio.create_task(active_policy_call()) + await call_started.wait() + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await asyncio.sleep(0) + assert dispatcher.policy_update_pending + assert not barrier.done() + + with pytest.raises(PolicyRequestRejected, match="update is pending"): + async with policy_gate.request(expected_version=0): + pass + + release_call.set() + await score + await barrier + + policy.version = 1 + await dispatcher.on_new_version(1) + with pytest.raises(PolicyRequestRejected, match="expected policy version 0.*current version is 1"): + async with policy_gate.request(expected_version=0): + pass + + +@pytest.mark.asyncio +async def test_non_dynamo_dispatcher_retains_configured_off_policy_window(): + dispatcher = _dispatcher(enforce_policy_update_barrier=False, max_off_policy_steps=1) + group_id = uuid.uuid4() + request_started = asyncio.Event() + + async def request() -> None: + request_started.set() + await asyncio.Future() + + task = asyncio.create_task(request()) + await request_started.wait() + dispatcher.groups[group_id] = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + policy_version_at_start=0, + ) + dispatcher.inflight[task] = InflightRollout( + kind="train", + env_name="train", + group_id=group_id, + policy_version=0, + rollout_count=1, + ) + dispatcher.inflight_permits = 1 + + await dispatcher.on_version_pending(1) + assert dispatcher.inflight[task].off_policy_steps == 1 + assert not dispatcher.policy_update_pending + + await dispatcher.on_version_pending(2) + assert task.cancelled() + assert not dispatcher.inflight + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["collective update failed", "resume_generation failed"]) +async def test_indeterminate_engine_failure_keeps_admission_fail_closed_without_advancing_version( + tmp_path: Path, + failure: str, +): + dispatcher = _dispatcher() + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _Inference: + async def update_weights(self, *_args, **_kwargs) -> None: + raise RuntimeError(failure) + + policy = dispatcher.policy + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=policy, + inference=_Inference(), + observers=[dispatcher], + lora_name=None, + ) + + with pytest.raises(RuntimeError, match=failure) as exc: + await watcher.apply_policy_update(1) + + assert watcher.ckpt_step == 0 + assert policy.version == 0 + assert dispatcher.policy_update_pending + assert exc.value.__notes__ == [ + "Policy update 1 may have mutated inference workers; mutable-policy admission remains fail-closed" + ] + + +@pytest.mark.asyncio +async def test_success_updates_lead_gate_before_reopening_transition_fence(tmp_path: Path): + events: list[str] = [] + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _Observer: + def __init__(self, name: str) -> None: + self.name = name + + async def on_version_pending(self, _step: int) -> None: + events.append(f"{self.name}:pending") + + async def on_new_version(self, _step: int) -> None: + events.append(f"{self.name}:new") + + async def on_version_update_failed(self, _step: int, _error: BaseException) -> None: + events.append(f"{self.name}:failed") + + class _Inference: + async def update_weights(self, *_args, **_kwargs) -> None: + events.append("engine:update") + + policy = Policy(version=0, model_name="policy") + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=policy, + inference=_Inference(), + observers=[_Observer("fence"), _Observer("lead-gate")], + lora_name=None, + ) + + await watcher.apply_policy_update(1) + + assert events == [ + "fence:pending", + "lead-gate:pending", + "engine:update", + "lead-gate:new", + "fence:new", + ] + assert watcher.ckpt_step == 1 + assert policy.version == 1 + + +@pytest.mark.asyncio +async def test_success_callback_cancellation_still_reopens_transition_fence(tmp_path: Path): + dispatcher = _dispatcher() + callback_started = asyncio.Event() + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _BlockingLeadGate: + async def on_version_pending(self, _step: int) -> None: + return + + async def on_new_version(self, _step: int) -> None: + callback_started.set() + await asyncio.Future() + + async def on_version_update_failed(self, _step: int, _error: BaseException) -> None: + return + + class _Inference: + async def update_weights(self, *_args, **_kwargs) -> None: + return + + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=dispatcher.policy, + inference=_Inference(), + observers=[dispatcher, _BlockingLeadGate()], + lora_name=None, + ) + + update = asyncio.create_task(watcher.apply_policy_update(1)) + await callback_started.wait() + update.cancel() + with pytest.raises(asyncio.CancelledError): + await update + + assert watcher.ckpt_step == 1 + assert dispatcher.policy.version == 1 + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_success_callback_error_still_reopens_transition_fence(tmp_path: Path): + dispatcher = _dispatcher() + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _BrokenLeadGate: + async def on_version_pending(self, _step: int) -> None: + return + + async def on_new_version(self, _step: int) -> None: + raise RuntimeError("lead gate callback failed") + + async def on_version_update_failed(self, _step: int, _error: BaseException) -> None: + return + + class _Inference: + async def update_weights(self, *_args, **_kwargs) -> None: + return + + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=dispatcher.policy, + inference=_Inference(), + observers=[dispatcher, _BrokenLeadGate()], + lora_name=None, + ) + + await watcher.apply_policy_update(1) + + assert watcher.ckpt_step == 1 + assert dispatcher.policy.version == 1 + assert not dispatcher.policy_update_pending From 3ece154a0d97e8f5692d5083f0e45bd7d0027720 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 20:19:34 -0700 Subject: [PATCH 17/18] fix(dynamo): harden policy update transactions --- .../src/prime_rl/configs/shared.py | 3 +- src/prime_rl/inference/dynamo.py | 45 ++- src/prime_rl/inference/dynamo_admin.py | 54 ++- src/prime_rl/orchestrator/algo/opsd.py | 16 +- .../orchestrator/component_supervision.py | 89 ++++- src/prime_rl/orchestrator/dispatcher.py | 360 ++++++++++-------- .../orchestrator/dispatcher_transactions.py | 128 +++++++ src/prime_rl/orchestrator/orchestrator.py | 211 ++-------- src/prime_rl/orchestrator/policy_gate.py | 70 +++- src/prime_rl/orchestrator/pool_identity.py | 18 +- .../orchestrator/train_finalization.py | 242 ++++++++++++ src/prime_rl/orchestrator/types.py | 6 +- src/prime_rl/orchestrator/utils.py | 3 +- src/prime_rl/orchestrator/watcher.py | 36 +- src/prime_rl/utils/client.py | 11 +- src/prime_rl/utils/elastic.py | 4 + src/prime_rl/utils/policy_client_config.py | 50 +++ tests/unit/inference/test_dynamo.py | 35 +- tests/unit/inference/test_dynamo_admin.py | 1 - .../inference/test_dynamo_admin_barriers.py | 12 +- tests/unit/orchestrator/test_algorithms.py | 42 ++ .../orchestrator/test_orchestrator_setup.py | 257 ++++++++++++- .../test_policy_gate_cancellation.py | 68 ++++ tests/unit/orchestrator/test_pool_identity.py | 4 +- .../orchestrator/test_train_finalization.py | 150 ++++++++ .../test_weight_update_barrier.py | 201 +++++++--- .../test_weight_update_barrier_mutability.py | 211 ++++++++++ tests/unit/test_transport_config.py | 21 + tests/unit/utils/test_client.py | 31 ++ 29 files changed, 1902 insertions(+), 477 deletions(-) create mode 100644 src/prime_rl/orchestrator/dispatcher_transactions.py create mode 100644 src/prime_rl/orchestrator/train_finalization.py create mode 100644 src/prime_rl/utils/policy_client_config.py create mode 100644 tests/unit/orchestrator/test_policy_gate_cancellation.py create mode 100644 tests/unit/orchestrator/test_train_finalization.py create mode 100644 tests/unit/orchestrator/test_weight_update_barrier_mutability.py create mode 100644 tests/unit/test_transport_config.py 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 fbc3e2f4c2..225d544013 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -267,7 +267,8 @@ class MetricsServerConfig(BaseConfig): class BaseTransportConfig(BaseConfig): - pass + send_timeout_seconds: float = Field(300.0, gt=0, allow_inf_nan=False) + """Maximum time to ship one rollout batch before the orchestrator fails closed.""" class FileSystemTransportConfig(BaseTransportConfig): diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py index 2ae26d2f73..ca9f9e0533 100644 --- a/src/prime_rl/inference/dynamo.py +++ b/src/prime_rl/inference/dynamo.py @@ -85,6 +85,34 @@ class DynamoWorkerSpec: process: DynamoProcessSpec +@dataclass(frozen=True) +class DynamoWorkerPorts: + """Host-local ports reserved by one worker process.""" + + system: int + nixl: int + data_parallel_rpc: int + kv_events: int + + +_LOCAL_WORKER_PORT_BASE = 18_000 +_LOCAL_WORKER_PORT_STRIDE = 4 + + +def _allocate_local_worker_ports(worker_index: int) -> DynamoWorkerPorts: + """Allocate one non-overlapping port block for a same-host worker.""" + base = _LOCAL_WORKER_PORT_BASE + worker_index * _LOCAL_WORKER_PORT_STRIDE + ports = DynamoWorkerPorts( + system=base, + nixl=base + 1, + data_parallel_rpc=base + 2, + kv_events=base + 3, + ) + if ports.kv_events > 65_535: + raise ValueError(f"Local Dynamo worker {worker_index} exceeds the available TCP port range") + return ports + + def _json_default(value: Any) -> Any: if isinstance(value, Path): return str(value) @@ -272,6 +300,7 @@ def build_engine_config( role: Role, *, kv_events_port: int | None = None, + data_parallel_rpc_port: int | None = None, ) -> dict[str, Any]: """Build one deterministic vLLM ``AsyncEngineArgs`` object.""" _validate_overrides("vllm_extra", config.vllm_extra) @@ -294,6 +323,8 @@ def build_engine_config( values.pop("data_parallel_rpc_port", None) else: values["data_parallel_size_local"] = local_dp + if data_parallel_rpc_port is not None: + values["data_parallel_rpc_port"] = data_parallel_rpc_port if role in ("prefill", "agg") and kv_events_port is not None: values["kv_events_config"] = { @@ -382,24 +413,30 @@ def build_local_worker_specs( role_indexes[role] += 1 start = worker_index * gpus_per_worker worker_gpus = tuple(available[start : start + gpus_per_worker]) - kv_events_port = 20080 + role_index if role in ("prefill", "agg") else None + ports = _allocate_local_worker_ports(worker_index) + kv_events_port = ports.kv_events if role in ("prefill", "agg") else None name = f"{role}-{role_index}" engine_path = _write_json( config_dir / f"{name}-engine.json", - build_engine_config(config, role, kv_events_port=kv_events_port), + build_engine_config( + config, + role, + kv_events_port=kv_events_port, + data_parallel_rpc_port=ports.data_parallel_rpc, + ), ) specs.append( DynamoWorkerSpec( name=name, role=role, gpu_ids=worker_gpus, - system_port=8081 + worker_index, + system_port=ports.system, process=build_worker_process( config, role, engine_path, nixl_host="127.0.0.1", - nixl_port=20100 + worker_index, + nixl_port=ports.nixl, namespace=resolved_namespace, ), ) diff --git a/src/prime_rl/inference/dynamo_admin.py b/src/prime_rl/inference/dynamo_admin.py index 410ecb4ec7..02fcaf8bec 100644 --- a/src/prime_rl/inference/dynamo_admin.py +++ b/src/prime_rl/inference/dynamo_admin.py @@ -406,7 +406,6 @@ async def initialize_nccl( async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | None, step: int) -> None: self._require_unambiguous_admin_state() - primary_error: BaseException | None = None try: await self._settle_fanout( ( @@ -420,6 +419,13 @@ async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | No ), "pause_generation", ) + except BaseException as exc: + # Pausing cannot change weights. Settle the fanout, undo any + # successful pauses, and preserve the original pre-mutation error. + await self._resume_after_pre_mutation_failure(clients, exc) + raise + + try: if weight_dir is not None: marker = weight_dir / NCCL_READY_MARKER marker.parent.mkdir(parents=True, exist_ok=True) @@ -441,23 +447,43 @@ async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | No "engine_rpc": "update_weights_from_path", } method = "update_weights_from_disk" + except BaseException as exc: + # Validation/filesystem preparation failed before any update RPC + # was issued, so generation can safely resume on the old weights. + await self._resume_after_pre_mutation_failure(clients, exc) + raise + try: await self._settle_fanout( (self._post(client, method, body, timeout_s=UPDATE_WEIGHTS_TIMEOUT_S) for client in clients), method, ) - except BaseException as exc: - primary_error = exc + except BaseException: + # At least one mutation RPC was issued. A failure or cancellation + # cannot prove which workers committed, so keep generation paused + # and permanently poison this admin object. self._weight_update_indeterminate = True raise - finally: - try: - await self._settle_fanout( - (self._post(client, "resume_generation", retry_transient=True) for client in clients), - "resume_generation", - ) - except BaseException as exc: - self._weight_update_indeterminate = True - if primary_error is None: - raise - primary_error.add_note(f"Dynamo resume_generation cleanup also failed: {exc!r}") + + try: + await self._resume_generation(clients) + except BaseException: + self._weight_update_indeterminate = True + raise + + async def _resume_generation(self, clients: list[AsyncClient]) -> None: + await self._settle_fanout( + (self._post(client, "resume_generation", retry_transient=True) for client in clients), + "resume_generation", + ) + + async def _resume_after_pre_mutation_failure( + self, + clients: list[AsyncClient], + primary_error: BaseException, + ) -> None: + try: + await self._resume_generation(clients) + except BaseException as resume_error: + self._weight_update_indeterminate = True + primary_error.add_note(f"Dynamo resume_generation cleanup also failed: {resume_error!r}") diff --git a/src/prime_rl/orchestrator/algo/opsd.py b/src/prime_rl/orchestrator/algo/opsd.py index 1a3e9e7d94..a8c5791218 100644 --- a/src/prime_rl/orchestrator/algo/opsd.py +++ b/src/prime_rl/orchestrator/algo/opsd.py @@ -1,10 +1,10 @@ from __future__ import annotations -import asyncio from typing import TYPE_CHECKING from prime_rl.configs.algorithm import OPSDAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.utils.async_utils import gather_shielded if TYPE_CHECKING: from renderers.base import Renderer @@ -81,8 +81,18 @@ async def score_sample(sample: TrainingSample) -> None: # sample.token_ids (demo-conditioned, the trainer's ref_kl target). sample.ref_logprobs = full_logprobs[len(hint_block) :] + async def settle_scores() -> None: + results, cancellation = await gather_shielded(*(score_sample(sample) for sample in rollout.samples)) + failures = [result for result in results if isinstance(result, BaseException)] + primary: BaseException | None = cancellation or (failures[0] if failures else None) + if primary is not None: + siblings = failures if cancellation is not None else failures[1:] + for sibling in siblings: + primary.add_note(f"Another OPSD score failed: {sibling!r}") + raise primary + if self.policy_gate is None: - await asyncio.gather(*(score_sample(sample) for sample in rollout.samples)) + await settle_scores() return async with self.policy_gate.request(expected_version=rollout.policy_version): - await asyncio.gather(*(score_sample(sample) for sample in rollout.samples)) + await settle_scores() diff --git a/src/prime_rl/orchestrator/component_supervision.py b/src/prime_rl/orchestrator/component_supervision.py index 7b7df0f14a..5d1cff41ab 100644 --- a/src/prime_rl/orchestrator/component_supervision.py +++ b/src/prime_rl/orchestrator/component_supervision.py @@ -3,7 +3,55 @@ from __future__ import annotations import asyncio -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Sequence +from typing import TypeVar + +T = TypeVar("T") + +# Cleanup must not extend a user-facing operation timeout indefinitely. A task +# that suppresses cancellation is retained below and observed when it settles. +SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS = 1.0 +_ORPHANED_OPERATIONS: set[asyncio.Task] = set() + + +def _observe_operation(task: asyncio.Task) -> None: + _ORPHANED_OPERATIONS.discard(task) + try: + task.exception() + except BaseException: + pass + + +async def _cancel_operation_with_grace(task: asyncio.Task) -> asyncio.CancelledError | None: + task.cancel() + loop = asyncio.get_running_loop() + grace_elapsed = loop.create_future() + caller_cancellation: asyncio.CancelledError | None = None + + def finish_grace() -> None: + if not grace_elapsed.done(): + grace_elapsed.set_result(None) + + timer = loop.call_later(SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS, finish_grace) + try: + while not task.done() and not grace_elapsed.done(): + try: + await asyncio.wait((task, grace_elapsed), return_when=asyncio.FIRST_COMPLETED) + except asyncio.CancelledError as error: + # Finish bounded cleanup before propagating caller cancellation. + # Repeated cancellation must not make cleanup unbounded. + if caller_cancellation is None: + caller_cancellation = error + continue + finally: + timer.cancel() + + if task.done(): + _observe_operation(task) + else: + _ORPHANED_OPERATIONS.add(task) + task.add_done_callback(_observe_operation) + return caller_cancellation def raise_if_component_failed(tasks: Sequence[asyncio.Task]) -> None: @@ -18,3 +66,42 @@ def raise_if_component_failed(tasks: Sequence[asyncio.Task]) -> None: if error is not None: raise error raise RuntimeError(f"Orchestrator component {name!r} stopped unexpectedly") + + +async def run_with_component_supervision( + operation: Callable[[], Awaitable[T]], + component_tasks: Sequence[asyncio.Task], + *, + timeout: float | None = None, + timeout_description: str | None = None, +) -> T | None: + """Run one operation while racing every supervised component. + + A component failure wins when both sides complete in the same event-loop + turn. ``None`` denotes timeout unless ``timeout_description`` is supplied, + in which case timeout raises after cancelling the operation. + """ + raise_if_component_failed(component_tasks) + task = asyncio.create_task(operation()) + component_failure_selected = False + try: + await asyncio.wait( + [task, *component_tasks], + return_when=asyncio.FIRST_COMPLETED, + timeout=timeout, + ) + try: + raise_if_component_failed(component_tasks) + except BaseException: + component_failure_selected = True + raise + if task.done(): + return await task + if timeout_description is not None: + raise TimeoutError(f"{timeout_description} timed out after {timeout} seconds") + return None + finally: + if not task.done(): + caller_cancellation = await _cancel_operation_with_grace(task) + if caller_cancellation is not None and not component_failure_selected: + raise caller_cancellation diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 7d5d42020d..c384e86686 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -2,11 +2,8 @@ - Capacity (``max_inflight_rollouts``) is shared across train + eval. A group-scoring task that runs N rollouts in one call reserves N permits. -- Optional rate limiting via ``AsyncLimiter(tasks_per_minute, 60)``. -- Emit-everything invariant: every dispatched rollout eventually reaches - ``out_q`` exactly once as a ``Rollout``. Failures - (env error, empty trajectory, task exception, off-policy cancel) carry - ``trace.error`` set; sinks decide drop / partial-train policy. +- Emit-everything invariant: every dispatched rollout reaches ``out_q`` once. + Failures carry ``trace.error``; sinks decide drop / partial-train policy. - ``DispatcherMode.PREFER_TRAIN`` / ``PREFER_EVAL`` controls which kind to schedule next. Transitions are level-triggered (driven by the eval source's emptiness), so in-flight rollouts of the opposite kind drain @@ -28,10 +25,16 @@ from aiolimiter import AsyncLimiter from prime_rl.orchestrator.dispatcher_metrics import DispatcherMetrics +from prime_rl.orchestrator.dispatcher_transactions import ( + EmissionRecord, + EmissionTracker, + emit_policy_cancellation_markers, + settle_transaction_cleanup, +) from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_source import EvalSource -from prime_rl.orchestrator.policy_gate import MutablePolicyGate -from prime_rl.orchestrator.pool_identity import pools_may_alias +from prime_rl.orchestrator.policy_gate import MutablePolicyGate, PolicyUpdateToken, SchedulingEpoch +from prime_rl.orchestrator.pool_identity import client_may_alias_pool, pools_may_alias from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( GroupState, @@ -85,6 +88,7 @@ def __init__( self.max_off_policy_steps = max_off_policy_steps self.enforce_policy_update_barrier = enforce_policy_update_barrier self.policy_gate = policy_gate or MutablePolicyGate(policy, enabled=enforce_policy_update_barrier) + self._policy_update: tuple[int, PolicyUpdateToken] | None = None self.max_inflight = max_inflight_rollouts self.inflight_permits = 0 @@ -94,6 +98,7 @@ def __init__( self.inflight: dict[asyncio.Task, InflightRollout] = {} self.groups: dict[uuid.UUID, GroupState] = {} + self._emissions = EmissionTracker() # Bounded so the dispatcher backpressures on a slow sink self.out_q: asyncio.Queue[Rollout] = asyncio.Queue(maxsize=max(8, self.max_inflight)) @@ -214,36 +219,49 @@ async def on_version_pending(self, step: int) -> None: await self._advance_off_policy_window() return - await self.policy_gate.begin_update(step=step) - claimed_groups, claimed_tasks = self._claim_mutable_policy_work() - results, cancellation = await gather_shielded( - self._settle_policy_requests(claimed_groups, claimed_tasks), - self.policy_gate.wait_idle(), - ) - failures = [result for result in results if isinstance(result, BaseException)] - primary: BaseException | None = cancellation or (failures[0] if failures else None) - if primary is not None: - siblings = failures if cancellation is not None else failures[1:] - for sibling in siblings: - primary.add_note(f"Another policy barrier drain failed: {sibling!r}") - raise primary + token = await self.policy_gate.begin_update(step=step) + self._policy_update = (step, token) + try: + claimed_groups, claimed_tasks, claimed_emissions = self._claim_mutable_policy_work() + results, cancellation = await gather_shielded( + self._settle_policy_requests(claimed_groups, claimed_tasks, claimed_emissions), + self.policy_gate.wait_idle(), + ) + failures = [result for result in results if isinstance(result, BaseException)] + primary: BaseException | None = cancellation or (failures[0] if failures else None) + if primary is not None: + siblings = failures if cancellation is not None else failures[1:] + for sibling in siblings: + primary.add_note(f"Another policy barrier drain failed: {sibling!r}") + raise primary + except BaseException as primary_error: + await settle_transaction_cleanup( + self._reopen_policy_admission(step, token), + primary_error, + "Policy transition rollback", + ) + raise async def on_new_version(self, step: int) -> None: """Reopen admission after the new policy is live.""" if self.enforce_policy_update_barrier: - await self._reopen_policy_admission() + await self._reopen_policy_admission(step) async def on_version_update_failed(self, step: int, error: BaseException) -> None: """Roll back only the transition fence; the policy stays unchanged.""" if self.enforce_policy_update_barrier: - await self._reopen_policy_admission() + await self._reopen_policy_admission(step) @property def policy_update_pending(self) -> bool: return self.policy_gate.pending - async def _reopen_policy_admission(self) -> None: - await self.policy_gate.finish_update() + async def _reopen_policy_admission(self, step: int, token: PolicyUpdateToken | None = None) -> None: + owned = self._policy_update + if owned is None or owned[0] != step or (token is not None and owned[1] is not token): + raise RuntimeError(f"Dispatcher does not own the pending policy transition for step {step}") + await self.policy_gate.finish_update(owned[1]) + self._policy_update = None async def _advance_off_policy_window(self) -> None: """Retain the established non-Dynamo in-flight tolerance policy.""" @@ -251,7 +269,8 @@ async def _advance_off_policy_window(self) -> None: for meta in self.inflight.values(): if meta.kind != "train": continue - if not self.train_envs.get(meta.env_name).sampler.samples_from_live_policy: + meta.uses_mutable_policy = meta.uses_mutable_policy or self._uses_mutable_policy(meta.kind, meta.env_name) + if not meta.uses_mutable_policy: continue meta.off_policy_steps += 1 if meta.off_policy_steps > self.max_off_policy_steps: @@ -276,13 +295,22 @@ def _uses_mutable_policy(self, kind: RolloutKind, env_name: str) -> bool: def _claim_mutable_policy_work( self, - ) -> tuple[dict[uuid.UUID, GroupState], list[tuple[asyncio.Task, InflightRollout]]]: - group_ids = { - group_id for group_id, group in self.groups.items() if self._uses_mutable_policy(group.kind, group.env_name) - } - group_ids.update( - meta.group_id for meta in self.inflight.values() if self._uses_mutable_policy(meta.kind, meta.env_name) - ) + ) -> tuple[ + dict[uuid.UUID, GroupState], + list[tuple[asyncio.Task, InflightRollout]], + list[EmissionRecord], + ]: + group_ids: set[uuid.UUID] = set() + for group_id, group in self.groups.items(): + group.uses_mutable_policy = group.uses_mutable_policy or self._uses_mutable_policy( + group.kind, group.env_name + ) + if group.uses_mutable_policy: + group_ids.add(group_id) + for meta in self.inflight.values(): + meta.uses_mutable_policy = meta.uses_mutable_policy or self._uses_mutable_policy(meta.kind, meta.env_name) + if meta.uses_mutable_policy: + group_ids.add(meta.group_id) claimed_groups = { group_id: group for group_id in group_ids if (group := self.groups.pop(group_id, None)) is not None @@ -294,12 +322,14 @@ def _claim_mutable_policy_work( del self.inflight[task] self.release(meta.rollout_count) claimed_tasks.append((task, meta)) - return claimed_groups, claimed_tasks + claimed_emissions = self._emissions.claim(group_ids) + return claimed_groups, claimed_tasks, claimed_emissions async def _settle_policy_requests( self, groups: dict[uuid.UUID, GroupState], claimed: list[tuple[asyncio.Task, InflightRollout]], + emissions: list[EmissionRecord], ) -> None: tasks = [task for task, _meta in claimed] already_settled = {task for task in tasks if task.done()} @@ -311,12 +341,25 @@ async def _settle_policy_requests( if tasks: results, cancellation = await gather_shielded(*tasks) + emission_results: list[object] = [] + emission_cancellation: asyncio.CancelledError | None = None + if emissions: + emission_results, emission_cancellation = await gather_shielded( + *(record.done.wait() for record in emissions) + ) + metadata_by_group: dict[uuid.UUID, InflightRollout] = {} for _task, meta in claimed: metadata_by_group.setdefault(meta.group_id, meta) marker_results, marker_cancellation = await gather_shielded( - self._emit_policy_cancellation_markers(groups, metadata_by_group) + emit_policy_cancellation_markers( + groups, + metadata_by_group, + out_q=self.out_q, + stopped=self.stopped, + metrics=self.metrics, + ) ) failures = [ @@ -327,62 +370,21 @@ async def _settle_policy_requests( and not isinstance(result, asyncio.CancelledError) ] failures.extend(result for result in marker_results if isinstance(result, BaseException)) - primary: BaseException | None = cancellation or marker_cancellation or (failures[0] if failures else None) + failures.extend(record.error for record in emissions if record.error is not None) + failures.extend(result for result in emission_results if isinstance(result, BaseException)) + primary: BaseException | None = ( + cancellation or emission_cancellation or marker_cancellation or (failures[0] if failures else None) + ) if primary is not None: - siblings = failures if cancellation is not None or marker_cancellation is not None else failures[1:] + siblings = ( + failures + if cancellation is not None or emission_cancellation is not None or marker_cancellation is not None + else failures[1:] + ) for sibling in siblings: primary.add_note(f"Another policy barrier cleanup failed: {sibling!r}") raise primary - async def _emit_policy_cancellation_markers( - self, - groups: dict[uuid.UUID, GroupState], - metadata_by_group: dict[uuid.UUID, InflightRollout], - ) -> None: - """Finish every claimed group; cancellation cannot interrupt this transaction.""" - cancelled = 0 - for group_id, group in groups.items(): - owed = max(0, group.target_rollouts - group.emitted) - if owed == 0: - continue - meta = metadata_by_group.get(group_id) or InflightRollout( - kind=group.kind, - env_name=group.env_name, - group_id=group_id, - policy_version=group.policy_version_at_start, - rollout_count=1, - eval_step=group.eval_step, - ) - for _ in range(owed): - trace = Rollout( - task=vf.Task(idx=group.task_idx, prompt=None), - errors=[vf.Error(type="Cancelled", message="Policy update barrier")], - stop_condition="error", - ) - await self._emit_claimed_rollout(meta, group, trace) - self.metrics.record_cancellation(kind=meta.kind, env_name=meta.env_name, n=owed) - cancelled += owed - if cancelled: - get_logger().debug(f"Policy update barrier cancelled {cancelled} mutable-policy rollout(s)") - - async def _emit_claimed_rollout( - self, - meta: InflightRollout, - group: GroupState, - rollout: Rollout, - ) -> None: - """Enqueue a barrier-owned marker before committing group accounting.""" - rollout.kind = meta.kind - rollout.env_name = meta.env_name - rollout.group_id = meta.group_id - rollout.policy_version = group.policy_version_at_start - rollout.off_policy_steps = meta.off_policy_steps - if meta.kind == "eval": - assert group.eval_step is not None, "eval rollout missing eval_step" - rollout.eval_step = group.eval_step - await self.out_q.put(rollout) - group.emitted += 1 - async def fill_inflight(self) -> None: """Schedule new rollouts up to ``max_inflight``, honoring ``self.mode``. Eval scheduling ignores the orchestrator's dispatch @@ -390,26 +392,26 @@ async def fill_inflight(self) -> None: respects it. When ``PREFER_EVAL``'s source exhausts we flip back to ``PREFER_TRAIN`` so the eval tail drains alongside fresh train.""" while True: - async with self.policy_gate.scheduling_admission() as admitted: - if not admitted or self.available_permits <= 0: + epoch = await self.policy_gate.scheduling_epoch() + if epoch is None or self.available_permits <= 0: + return + + if self.mode == DispatcherMode.PREFER_EVAL: + # PREFER_EVAL implies a configured eval source. + assert self.eval_source is not None + if not self.eval_has_work: + # Fill remaining permits with train while eval drains. + self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained") + continue + scheduled = await self.try_schedule("eval", epoch=epoch) + if not scheduled: + return + else: # PREFER_TRAIN — respects the orchestrator's dispatch gate + if not self.dispatch_allowed.is_set(): + return + scheduled = await self.try_schedule("train", epoch=epoch) + if not scheduled: return - - if self.mode == DispatcherMode.PREFER_EVAL: - # PREFER_EVAL implies a configured eval source. - assert self.eval_source is not None - if not self.eval_has_work: - # Fill remaining permits with train while eval drains. - self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained") - continue - scheduled = await self.try_schedule("eval") - if not scheduled: - return - else: # PREFER_TRAIN — respects the orchestrator's dispatch gate - if not self.dispatch_allowed.is_set(): - return - scheduled = await self.try_schedule("train") - if not scheduled: - return def switch_mode(self, new_mode: DispatcherMode, *, reason: str) -> None: if new_mode == self.mode: @@ -418,7 +420,7 @@ def switch_mode(self, new_mode: DispatcherMode, *, reason: str) -> None: get_logger().info(f"Switching dispatcher mode to prefer {prefer} rollouts because {reason}") self.mode = new_mode - async def try_schedule(self, kind: RolloutKind) -> bool: + async def try_schedule(self, kind: RolloutKind, *, epoch: SchedulingEpoch) -> bool: """Schedule one rollout of ``kind``: prefer continuing an existing group (keeps prefix-cache hits); otherwise open a fresh group from the corresponding source. Returns False if nothing could be @@ -435,14 +437,20 @@ async def try_schedule(self, kind: RolloutKind) -> bool: env = envs.get(group.env_name) cost = group.rollouts_to_schedule if env.requires_group_scoring else 1 if cost <= self.available_permits: - return await self.schedule_group_rollout(gid, group) + return await self.schedule_group_rollout(gid, group, epoch=epoch) - fresh = self.next_fresh_group(kind, envs) - if fresh is None: - return False - gid = uuid.uuid4() - self.groups[gid] = fresh - return await self.schedule_group_rollout(gid, fresh) + # Pop and publish a fresh group inside the short scheduling commit. + # An update either sees the group in its claim snapshot or invalidates + # this epoch before the source is consumed. + async with self.policy_gate.scheduling_commit(epoch) as admitted: + if not admitted or self.available_permits <= 0: + return False + fresh = self.next_fresh_group(kind, envs) + if fresh is None: + return False + gid = uuid.uuid4() + self.groups[gid] = fresh + return await self.schedule_group_rollout(gid, fresh, epoch=epoch) def next_fresh_group(self, kind: RolloutKind, envs) -> GroupState | None: """Pop the next example from the corresponding source and wrap it in @@ -469,9 +477,16 @@ def next_fresh_group(self, kind: RolloutKind, envs) -> GroupState | None: target_rollouts=group_size, eval_step=eval_step, policy_version_at_start=self.policy.version, + uses_mutable_policy=self._uses_mutable_policy(kind, env_name), ) - async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) -> bool: + async def schedule_group_rollout( + self, + group_id: uuid.UUID, + group: GroupState, + *, + epoch: SchedulingEpoch, + ) -> bool: """Dispatch one ``run_rollout`` / ``run_group`` task for this group. Returns False only if we couldn't even schedule one rollout (no clients @@ -488,7 +503,8 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - else: pool, model_name, live_sourced = self._train_pool_for(group.env_name) - # Pin a single client per group to keep prefix-cache hits + # Resolve a client and rate-limit outside the gate. Both operations may + # wait indefinitely for elastic discovery or quota replenishment. if group.pinned_client is None: if group.kind == "eval": client = await pool.get_eval_client() @@ -499,7 +515,6 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - client = await pool.select_train_client(load) if group_id not in self.groups: return False - group.pinned_client = client else: client = group.pinned_client @@ -507,58 +522,70 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - if env_collection is None: return False env = env_collection.get(group.env_name) - # Frozen-sourced train rollouts hit a frozen pool; salting per policy - # version would invalidate its prefix cache every weight update for - # no reason. - if live_sourced: - cache_salt = str(group.policy_version_at_start) - else: - cache_salt = None - if env.requires_group_scoring: permits = group.rollouts_to_schedule - group.rollouts_to_schedule = 0 - await self.acquire(permits) - task: asyncio.Task = asyncio.create_task( - env.run_group( - client=client, - task_idx=group.task_idx, - model_name=model_name, - group_size=permits, - cache_salt=cache_salt, - ) - ) else: permits = 1 - group.rollouts_to_schedule -= 1 - await self.acquire(permits) - task = asyncio.create_task( - env.run_rollout( - client=client, - task_idx=group.task_idx, - model_name=model_name, - cache_salt=cache_salt, + # Snapshot selected identity before elastic churn can occur at the next await. + group.uses_mutable_policy = ( + group.uses_mutable_policy + or live_sourced + or pools_may_alias(pool, self.policy_pool) + or client_may_alias_pool(client, self.policy_pool) + ) + await self._wait_for_rate_limit(permits) + + async with self.policy_gate.scheduling_commit(epoch) as admitted: + if ( + not admitted + or self.groups.get(group_id) is not group + or permits > self.available_permits + or group.rollouts_to_schedule < permits + or (group.kind == "train" and (self.train_scheduling_disabled or not self.dispatch_allowed.is_set())) + ): + return False + + group.pinned_client = client + cache_salt = str(group.policy_version_at_start) if group.uses_mutable_policy else None + if env.requires_group_scoring: + group.rollouts_to_schedule = 0 + task: asyncio.Task = asyncio.create_task( + env.run_group( + client=client, + task_idx=group.task_idx, + model_name=model_name, + group_size=permits, + cache_salt=cache_salt, + ) + ) + else: + group.rollouts_to_schedule -= 1 + task = asyncio.create_task( + env.run_rollout( + client=client, + task_idx=group.task_idx, + model_name=model_name, + cache_salt=cache_salt, + ) ) - ) - self.inflight[task] = InflightRollout( - kind=group.kind, - env_name=group.env_name, - group_id=group_id, - policy_version=group.policy_version_at_start, - rollout_count=permits, - client_config=client, - eval_step=group.eval_step, - ) - return True + self.inflight_permits += permits + self.inflight[task] = InflightRollout( + kind=group.kind, + env_name=group.env_name, + group_id=group_id, + policy_version=group.policy_version_at_start, + rollout_count=permits, + client_config=client, + eval_step=group.eval_step, + uses_mutable_policy=group.uses_mutable_policy, + ) + return True - async def acquire(self, n: int) -> None: - """Reserve ``n`` permits + rate-limit each one. Caller must precheck - ``available_permits >= n``; this is not a blocking acquire.""" + async def _wait_for_rate_limit(self, n: int) -> None: for _ in range(n): if self.rate_limiter is not None: await self.rate_limiter.acquire() - self.inflight_permits += 1 def release(self, n: int) -> None: self.inflight_permits -= n @@ -575,6 +602,16 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: return # already handled by drop_group / cancel_inflight_rollouts self.release(meta.rollout_count) group = self.groups.get(meta.group_id) + async with self._emissions.track(meta.group_id): + await self._emit_completed_task(task, meta, group) + + async def _emit_completed_task( + self, + task: asyncio.Task, + meta: InflightRollout, + group: GroupState | None, + ) -> None: + """Convert one settled task into its complete output transaction.""" is_synth_exception = False try: @@ -615,9 +652,6 @@ async def emit_rollout(self, meta: InflightRollout, group: GroupState | None, ro if group is not None: eval_step = group.eval_step policy_version = group.policy_version_at_start - group.emitted += 1 - if group.emitted >= group.target_rollouts: - self.groups.pop(meta.group_id, None) rollout.kind = meta.kind rollout.env_name = meta.env_name @@ -628,6 +662,10 @@ async def emit_rollout(self, meta: InflightRollout, group: GroupState | None, ro assert eval_step is not None, "eval rollout missing eval_step" rollout.eval_step = eval_step await self.out_q.put(rollout) + if group is not None: + group.emitted += 1 + if group.emitted >= group.target_rollouts: + self.groups.pop(meta.group_id, None) async def drop_group(self, group_id: uuid.UUID) -> int: """Cancel remaining in-flight tasks for this group and emit a @@ -678,6 +716,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: policy_version=group.policy_version_at_start, rollout_count=1, eval_step=group.eval_step, + uses_mutable_policy=group.uses_mutable_policy, ) unscheduled_cancelled = group.rollouts_to_schedule for _ in range(unscheduled_cancelled): @@ -698,6 +737,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: policy_version=group.policy_version_at_start if group else 0, rollout_count=1, eval_step=group.eval_step, + uses_mutable_policy=group.uses_mutable_policy, ) if group is not None else None diff --git a/src/prime_rl/orchestrator/dispatcher_transactions.py b/src/prime_rl/orchestrator/dispatcher_transactions.py new file mode 100644 index 0000000000..d9e4c1cee7 --- /dev/null +++ b/src/prime_rl/orchestrator/dispatcher_transactions.py @@ -0,0 +1,128 @@ +"""Small transactional helpers for dispatcher output ownership.""" + +from __future__ import annotations + +import asyncio +import uuid +from collections import defaultdict +from collections.abc import AsyncIterator, Awaitable, Iterable +from contextlib import asynccontextmanager +from dataclasses import dataclass, field + +import verifiers.v1 as vf + +from prime_rl.orchestrator.dispatcher_metrics import DispatcherMetrics +from prime_rl.orchestrator.types import GroupState, InflightRollout, Rollout +from prime_rl.utils.async_utils import gather_shielded, safe_cancel +from prime_rl.utils.logger import get_logger + + +@dataclass +class EmissionRecord: + """A completion handler that still owns group-output accounting.""" + + done: asyncio.Event = field(default_factory=asyncio.Event) + error: BaseException | None = None + + +class EmissionTracker: + """Make completion handlers visible after their inflight task is popped.""" + + def __init__(self) -> None: + self._by_group: dict[uuid.UUID, list[EmissionRecord]] = defaultdict(list) + + @asynccontextmanager + async def track(self, group_id: uuid.UUID) -> AsyncIterator[None]: + record = EmissionRecord() + self._by_group[group_id].append(record) + try: + yield + except BaseException as exc: + record.error = exc + raise + finally: + record.done.set() + records = self._by_group[group_id] + records.remove(record) + if not records: + self._by_group.pop(group_id, None) + + def claim(self, group_ids: Iterable[uuid.UUID]) -> list[EmissionRecord]: + return [record for group_id in group_ids for record in self._by_group.get(group_id, ())] + + +async def settle_transaction_cleanup( + cleanup: Awaitable[object], + primary_error: BaseException, + description: str, +) -> None: + """Settle cleanup despite repeated cancellation, preserving the primary error.""" + results, cancellation = await gather_shielded(cleanup) + failures = [result for result in results if isinstance(result, BaseException)] + for cleanup_error in failures: + primary_error.add_note(f"{description} also failed: {cleanup_error!r}") + if cancellation is not None: + primary_error.add_note(f"{description} was cancelled again but settled before propagation") + + +async def _put_before_stop( + out_q: asyncio.Queue[Rollout], + stopped_event: asyncio.Event, + rollout: Rollout, +) -> None: + put = asyncio.create_task(out_q.put(rollout)) + stopped = asyncio.create_task(stopped_event.wait()) + try: + await asyncio.wait((put, stopped), return_when=asyncio.FIRST_COMPLETED) + if put.done(): + put.result() + else: + raise RuntimeError("Dispatcher stopped while emitting policy barrier cancellation markers") + finally: + if not put.done(): + await safe_cancel(put) + if not stopped.done(): + await safe_cancel(stopped) + + +async def emit_policy_cancellation_markers( + groups: dict[uuid.UUID, GroupState], + metadata_by_group: dict[uuid.UUID, InflightRollout], + *, + out_q: asyncio.Queue[Rollout], + stopped: asyncio.Event, + metrics: DispatcherMetrics, +) -> None: + """Finish every claimed group without hanging after dispatcher stop.""" + cancelled = 0 + for group_id, group in groups.items(): + owed = max(0, group.target_rollouts - group.emitted) + if owed == 0: + continue + meta = metadata_by_group.get(group_id) or InflightRollout( + kind=group.kind, + env_name=group.env_name, + group_id=group_id, + policy_version=group.policy_version_at_start, + rollout_count=1, + eval_step=group.eval_step, + uses_mutable_policy=group.uses_mutable_policy, + ) + for _ in range(owed): + rollout = Rollout( + task=vf.Task(idx=group.task_idx, prompt=None), + errors=[vf.Error(type="Cancelled", message="Policy update barrier")], + stop_condition="error", + kind=meta.kind, + env_name=meta.env_name, + group_id=meta.group_id, + policy_version=group.policy_version_at_start, + off_policy_steps=meta.off_policy_steps, + eval_step=group.eval_step if meta.kind == "eval" else None, + ) + await _put_before_stop(out_q, stopped, rollout) + group.emitted += 1 + metrics.record_cancellation(kind=meta.kind, env_name=meta.env_name, n=owed) + cancelled += owed + if cancelled: + get_logger().debug(f"Policy update barrier cancelled {cancelled} mutable-policy rollout(s)") diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 8a36e03025..f883b1bcd7 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -39,7 +39,7 @@ import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive imports from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.ckpt import setup_ckpt_manager -from prime_rl.orchestrator.component_supervision import raise_if_component_failed +from prime_rl.orchestrator.component_supervision import raise_if_component_failed, run_with_component_supervision from prime_rl.orchestrator.dispatcher import DispatcherMode, RolloutDispatcher from prime_rl.orchestrator.dispatcher_metrics import DispatcherMetrics from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs @@ -53,6 +53,7 @@ ) from prime_rl.orchestrator.periodic_logger import PeriodicLogger from prime_rl.orchestrator.policy_gate import MutablePolicyGate +from prime_rl.orchestrator.train_finalization import finalize_train_batch as finalize_train_batch_step from prime_rl.orchestrator.train_sink import TrainSink from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( @@ -92,11 +93,6 @@ # shutdown wedges (env-server ZMQ recv, vLLM admin aclose, etc) SHUTDOWN_TIMEOUT_S = 300 -# Abort after this many consecutive train batches drop all rollouts to -# post-batch filters — usually a misconfigured filter or homogeneous-reward -# dataset; fail loudly instead of spinning -MAX_CONSECUTIVE_EMPTY_BATCHES = 10 - # Maximum batches the orchestrator may run ahead of the trainer. The # dispatcher is paused via ``update_dispatch_gate`` once this is exceeded; # resumed when the watcher advances ``policy.version``. @@ -125,6 +121,7 @@ class Orchestrator: train_source: TrainSource train_sink: TrainSink dispatcher: RolloutDispatcher + policy_gate: MutablePolicyGate watcher: WeightWatcher lag_monitor: EventLoopLagMonitor periodic_logger: PeriodicLogger @@ -157,10 +154,6 @@ def __init__(self, config: OrchestratorConfig) -> None: self.progress = Progress() self.ckpt_manager = setup_ckpt_manager(config.output_dir, config.ckpt) self.policy = Policy(version=0, model_name="") - self.policy_gate = MutablePolicyGate( - self.policy, - enabled=config.model.client.admin_api == "dynamo", - ) self.stopped = asyncio.Event() # True after the final train step ships — pipeline winds down without # scheduling new train rollouts @@ -216,6 +209,12 @@ async def setup(self) -> None: self.renderer, self.policy_inference = await setup_policy_inference_pool( config=config, tokenizer=self.tokenizer ) + # The deployment boundary may resolve the client from DYN_RL_TOPOLOGY; + # derive the strict gate from the pool that was actually constructed. + self.policy_gate = MutablePolicyGate( + self.policy, + enabled=self.policy_inference.admin_api == "dynamo", + ) self.mm_token_type_ids_mapping = ( getattr(self.renderer, "mm_token_type_id_map", None) if self.renderer is not None else None ) @@ -359,7 +358,7 @@ async def setup(self) -> None: max_inflight_rollouts=config.max_inflight_rollouts, tasks_per_minute=config.tasks_per_minute, max_off_policy_steps=config.max_off_policy_steps, - enforce_policy_update_barrier=config.model.client.admin_api == "dynamo", + enforce_policy_update_barrier=self.policy_gate.enabled, policy_gate=self.policy_gate, ) self.train_sink = TrainSink( @@ -471,9 +470,8 @@ async def main_loop(self) -> None: self.stopped.set() break - try: - rollout: Rollout = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5) - except asyncio.TimeoutError: + rollout = await self._next_rollout() + if rollout is None: continue # Every completed rollout — errored, filtered, or never batched — lands in the @@ -502,136 +500,25 @@ async def main_loop(self) -> None: await self.finalize_train_batch(train_batch) async def finalize_train_batch(self, batch: TrainBatch) -> None: - """Ship one ``TrainBatch`` out to the trainer and handle the I/O - side-effects (ckpt, save_rollouts, reference scoring, sender.send, - metrics, heartbeat, progress, eval trigger). The sink has already - done all data-transformation work.""" - config = self.config - step = self.progress.step - - # Sink-to-sink cycle time — the actual time between batches, not - # including the orchestrator's ship I/O (overlapped with the - # dispatcher producing the next batch) - now = time.perf_counter() - step_time = (now - self.last_batch_at) if self.last_batch_at is not None else 0.0 - self.last_batch_at = now - - if config.max_steps is not None and step > config.max_steps: - self.draining = True - self.dispatcher.disable_train_scheduling() - n_cancelled = await self.dispatcher.cancel_inflight_train_rollouts() - get_logger().info( - f"Draining pipeline (cancelled {n_cancelled} in-flight train rollout(s); " - f"any in-flight evals will complete)" - ) - return - - if not batch.samples: - self.consecutive_empty_batches += 1 - get_logger().warning( - f"Step {step}: empty train batch (0 of {len(batch.rollouts)} generated rollouts shipped — " - f"all errored or filtered out) " - f"(consecutive empty batches: {self.consecutive_empty_batches}/{MAX_CONSECUTIVE_EMPTY_BATCHES})" - ) - if self.consecutive_empty_batches >= MAX_CONSECUTIVE_EMPTY_BATCHES: - raise RuntimeError( - f"{self.consecutive_empty_batches} consecutive empty train batches — " - "check filter config (pre_batch_filters / post_batch_filters) or task difficulty." - ) - return - self.consecutive_empty_batches = 0 - n_trainable = sum(1 for r in batch.rollouts if r.is_trainable) - if n_trainable / len(batch.rollouts) <= 0.1: - get_logger().warning( - f"Only {n_trainable}/{len(batch.rollouts)} generated rollouts are trainable " - f"({n_trainable / len(batch.rollouts):.1%}) — consider reviewing task difficulty / filter config" - ) - - # The effective (clean, trained-on) subset lands in the per-step ``effective`` trace file - # at ship time; the full arrival window already streamed into ``all`` on arrival. - # to_record drops the per-node training tensors — they're for training, not the rollout - # record, and can't round-trip json (raw numpy bytes). - effective = batch.rollouts.effective - records = [r.to_record() for r in effective] - await asyncio.to_thread(save_rollouts, records, get_trace_path(config.output_dir, step, "train", "effective")) - - await self.sender.send(TrainingBatch(examples=batch.samples, step=step)) - self.progress.step += 1 - self.update_dispatch_gate() - # Checkpoint the step we just shipped (resume point: continue at step + 1). - save_ckpt_time = await self.maybe_save_ckpt(step) - trim_process_memory() - - # Rollout metrics over the {agg,} × {all,effective} matrix. ``batch.rollouts`` is the - # full arrival window (errored + filtered included); ``.effective`` is the clean subset. - metrics: dict[str, float] = {} - for subset, pool in (("all", batch.rollouts), ("effective", effective)): - metrics |= pool.metrics.to_wandb(prefix="train/agg", subset=subset) - for env_name, env_pool in pool.by_env().items(): - metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) - - # Progress / timing / env-share / pre-filter accounting (assembled here, not in the metrics - # objects). ``num_tokens`` is over the full arrival window; the input/output breakdown is over - # the effective (shipped) subset, summing the same ``vf.Trace`` token properties the metric - # matrix reports. - num_tokens = sum(r.num_total_tokens for r in batch.rollouts) - num_input = sum(r.num_input_tokens for r in effective) - num_output = sum(r.num_output_tokens for r in effective) - num_rollouts = len(batch.rollouts) - num_unique_examples = len({r.group_id for r in batch.rollouts}) - metrics |= { - "progress/tokens": num_tokens, - "progress/input_tokens": num_input, - "progress/output_tokens": num_output, - "progress/rollouts": num_rollouts, - "progress/tasks": num_unique_examples, - "progress/total_tokens": self.progress.total_tokens, - "progress/total_rollouts": self.progress.total_samples, - "progress/total_tasks": self.progress.total_problems, - "time/step": step_time, - "time/save_ckpt": save_ckpt_time, - "time/wait_for_policy": self.wait_for_policy_time, - "step": step, - } - for env_name, env_pool in batch.rollouts.by_env().items(): - metrics[f"batch/{env_name}"] = len(env_pool) / len(batch.rollouts) - if self.train_sink.pre_filter_seen > 0: - metrics["pre_filters/all/dropped_rate"] = ( - self.train_sink.pre_filter_dropped / self.train_sink.pre_filter_seen - ) - for name, count in self.train_sink.pre_filter_dropped_by_name.items(): - metrics[f"pre_filters/all/{name}/rate"] = count / self.train_sink.pre_filter_seen - self.monitor.log(metrics, step=step) - self.wait_for_policy_time = 0.0 - self.monitor.log_samples(effective.rollouts, step=step) - self.monitor.log_distributions( - distributions={ - "rewards": [r.reward for r in effective], - "advantages": [a for r in effective if (a := r.scalar_advantage()) is not None], - }, - step=step, + """Persist, ship, checkpoint, and report one finalized train batch.""" + await finalize_train_batch_step(self, batch) + + async def _next_rollout(self) -> Rollout | None: + """Wait for output while supervising fatal background components.""" + return await run_with_component_supervision( + self.dispatcher.out_q.get, + self.component_tasks, + timeout=0.5, ) - if self.usage_reporter is not None: - run_id = os.getenv("RUN_ID", "") - if run_id: - self.usage_reporter.report_training_usage( - run_id=run_id, - step=step, - tokens=num_input + num_output, - ) - if self.heart is not None: - self.heart.beat() - - self.progress.total_tokens += num_tokens - self.progress.total_samples += num_rollouts - self.progress.total_problems += num_unique_examples - - self.log_train_batch(batch, step=step, step_time=step_time) - - self.train_sink.reset_pre_filter_stats() - self.maybe_trigger_eval(self.progress.step) - trim_process_memory() + async def _send_to_trainer(self, batch: TrainingBatch) -> None: + """Race shipment against fatal components, preferring any failure.""" + await run_with_component_supervision( + lambda: self.sender.send(batch), + self.component_tasks, + timeout=self.config.rollout_transport.send_timeout_seconds, + timeout_description=f"Training batch {batch.step} send", + ) def maybe_trigger_eval(self, step: int) -> None: """Fire eligible eval epochs and flip to ``PREFER_EVAL`` if anything @@ -716,46 +603,6 @@ def collect_pipeline_view(self) -> tuple[str, dict[str, float]]: payload["event_loop_lag/n"] = float(lag_stats.n) return body, payload - def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> None: - """Per-step ``Step …`` success line. Multi-env runs append an indented ``╰─`` line per env. - ``Error`` is the sink-level rate (errored arrivals / total arrivals, over the full window); - the quality metrics are over the effective (clean, trained-on) subset; ``Trainable`` is - relative to all generated rollouts.""" - rollouts = batch.rollouts - effective = rollouts.effective - eff = effective.metrics - n_generated = len(rollouts) - n_trainable = sum(1 for r in rollouts if r.is_trainable) - trainable_rate = (n_trainable / n_generated) if n_generated else 0.0 - max_off_policy = max((r.off_policy_steps for r in effective), default=0) - - head = ( - f"Step {step} | {format_time(step_time):>7} | Reward {eff.reward.mean():.4f} | " - f"Trainable {n_trainable}/{n_generated} ({trainable_rate:.1%}) | " - f"Turns {eff.num_turns.mean():.1f} | Branches {eff.num_branches.mean():.1f} | " - f"Max Off-Policy {max_off_policy} | " - f"Error {rollouts.metrics.has_error.mean():.1%} | Truncation {eff.is_truncated.mean():.1%}" - ) - if len(self.train_envs) <= 1: - get_logger().success(head) - return - - by_env = rollouts.by_env() - name_width = max((len(n) for n in by_env), default=0) - lines = [head] - for env_name in sorted(by_env): - pool = by_env[env_name] - env_eff_pool = pool.effective - env_eff = env_eff_pool.metrics - ratio = (len(pool) / n_generated) if n_generated else 0.0 - lines.append( - f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | Reward {env_eff.reward.mean():.4f} | " - f"Turns {env_eff.num_turns.mean():.1f} | Branches {env_eff.num_branches.mean():.1f} | " - f"Max Off-Policy {max((r.off_policy_steps for r in env_eff_pool), default=0)} | " - f"Error {pool.metrics.has_error.mean():.1%} | Truncation {env_eff.is_truncated.mean():.1%}" - ) - get_logger().success("\n\t\t ".join(lines)) - async def finalize_eval_batch(self, batch: EvalBatch) -> None: """Persist + log one completed eval epoch (save_rollouts, monitor.log_eval_samples, monitor.log).""" diff --git a/src/prime_rl/orchestrator/policy_gate.py b/src/prime_rl/orchestrator/policy_gate.py index e5e8227004..8c10b2de5b 100644 --- a/src/prime_rl/orchestrator/policy_gate.py +++ b/src/prime_rl/orchestrator/policy_gate.py @@ -5,6 +5,7 @@ import asyncio from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from dataclasses import dataclass from prime_rl.orchestrator.types import Policy @@ -13,37 +14,67 @@ class PolicyRequestRejected(RuntimeError): """A request cannot safely run against the mutable policy version.""" +@dataclass(frozen=True) +class SchedulingEpoch: + """Snapshot revalidated at the dispatcher's non-yielding commit point.""" + + value: int + + +@dataclass(frozen=True) +class PolicyUpdateToken: + """Unique ownership of one closed-gate transition.""" + + step: int + epoch: int + + class MutablePolicyGate: """Serialize policy mutation with every request that depends on its weights. - Dispatcher scheduling holds :meth:`scheduling_admission` until a rollout - task is registered. Other policy I/O, such as OPSD prefill scoring, uses - :meth:`request` for its full lifetime. Closing the gate prevents new work; - callers can then cancel dispatcher-owned tasks and await :meth:`wait_idle` - before pausing the engine. + Dispatcher scheduling prepares outside the lock, then revalidates a + :class:`SchedulingEpoch` at its short commit point. Other policy I/O, such + as OPSD prefill scoring, uses :meth:`request` for its full lifetime. + Closing the gate prevents new work; callers can then cancel + dispatcher-owned tasks and await :meth:`wait_idle` before pausing. """ def __init__(self, policy: Policy, *, enabled: bool) -> None: self.policy = policy self.enabled = enabled self._admission_lock = asyncio.Lock() - self._pending = False + self._epoch = 0 + self._pending_token: PolicyUpdateToken | None = None self._active_requests = 0 self._idle = asyncio.Event() self._idle.set() @property def pending(self) -> bool: - return self.enabled and self._pending + return self.enabled and self._pending_token is not None + + async def scheduling_epoch(self) -> SchedulingEpoch | None: + """Take a short-lived admission snapshot before slow preparation.""" + if not self.enabled: + return SchedulingEpoch(self._epoch) + async with self._admission_lock: + if self._pending_token is not None: + return None + return SchedulingEpoch(self._epoch) @asynccontextmanager - async def scheduling_admission(self) -> AsyncIterator[bool]: - """Hold the admission serialization point through scheduling commit.""" + async def scheduling_commit(self, epoch: SchedulingEpoch) -> AsyncIterator[bool]: + """Revalidate and serialize only the scheduling state commit. + + The caller must not await inside the admitted branch. Client discovery + and rate limiting belong before this context so an update can close the + gate promptly. + """ if not self.enabled: yield True return async with self._admission_lock: - yield not self._pending + yield self._pending_token is None and epoch.value == self._epoch @asynccontextmanager async def request(self, *, expected_version: int) -> AsyncIterator[None]: @@ -53,7 +84,7 @@ async def request(self, *, expected_version: int) -> AsyncIterator[None]: return async with self._admission_lock: - if self._pending: + if self._pending_token is not None: raise PolicyRequestRejected( f"Mutable-policy request rejected because a policy update is pending (expected version " f"{expected_version})" @@ -75,21 +106,26 @@ async def request(self, *, expected_version: int) -> AsyncIterator[None]: if self._active_requests == 0: self._idle.set() - async def begin_update(self, *, step: int) -> None: + async def begin_update(self, *, step: int) -> PolicyUpdateToken: """Close admission after every in-progress scheduling commit.""" if not self.enabled: - return + return PolicyUpdateToken(step=step, epoch=self._epoch) async with self._admission_lock: - if self._pending: + if self._pending_token is not None: raise RuntimeError(f"A policy update is already pending while preparing step {step}") - self._pending = True + self._epoch += 1 + token = PolicyUpdateToken(step=step, epoch=self._epoch) + self._pending_token = token + return token - async def finish_update(self) -> None: + async def finish_update(self, token: PolicyUpdateToken) -> None: """Reopen admission after a successful or proven pre-mutation failure.""" if not self.enabled: return async with self._admission_lock: - self._pending = False + if self._pending_token is not token: + raise RuntimeError(f"Policy update token for step {token.step} does not own the pending transition") + self._pending_token = None async def wait_idle(self) -> None: """Wait until every already-admitted non-dispatcher request settles.""" diff --git a/src/prime_rl/orchestrator/pool_identity.py b/src/prime_rl/orchestrator/pool_identity.py index fc1838daed..6429bb6732 100644 --- a/src/prime_rl/orchestrator/pool_identity.py +++ b/src/prime_rl/orchestrator/pool_identity.py @@ -5,6 +5,8 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: + import verifiers.v1 as vf + from prime_rl.utils.client import InferencePool @@ -14,11 +16,17 @@ def _normalized_url(value: object) -> str: def serving_identity(pool: InferencePool) -> tuple[str, frozenset[str], frozenset[str]]: """Return model, request endpoints, and admin endpoints for one pool.""" - request_endpoints = frozenset(_normalized_url(client.base_url) for client in getattr(pool, "train_clients", ())) - admin_endpoints = frozenset(_normalized_url(client.base_url) for client in getattr(pool, "admin_clients", ())) + request_endpoints = frozenset(_normalized_url(client.base_url) for client in pool.train_clients) + admin_endpoints = frozenset(_normalized_url(client.base_url) for client in pool.admin_clients) return str(pool.model_name), request_endpoints, admin_endpoints +def client_may_alias_pool(client: vf.ClientConfig, pool: InferencePool) -> bool: + """Whether one selected request client addresses a pool request endpoint.""" + selected_endpoint = _normalized_url(client.base_url) + return any(selected_endpoint == _normalized_url(candidate.base_url) for candidate in pool.train_clients) + + def pools_may_alias(left: InferencePool, right: InferencePool) -> bool: """Conservatively detect two pool objects backed by one mutable model. @@ -31,6 +39,10 @@ def pools_may_alias(left: InferencePool, right: InferencePool) -> bool: return True left_model, left_requests, left_admin = serving_identity(left) right_model, right_requests, right_admin = serving_identity(right) + # The worker-admin endpoint identifies the mutable engine itself. It + # aliases even when two logical model names or frontend routes differ. + if left_admin and right_admin and left_admin & right_admin: + return True if left_model != right_model: return False if not left_requests or not right_requests: @@ -39,4 +51,4 @@ def pools_may_alias(left: InferencePool, right: InferencePool) -> bool: return True if not left_admin or not right_admin: return True - return bool(left_admin & right_admin) + return False diff --git a/src/prime_rl/orchestrator/train_finalization.py b/src/prime_rl/orchestrator/train_finalization.py new file mode 100644 index 0000000000..00932ab901 --- /dev/null +++ b/src/prime_rl/orchestrator/train_finalization.py @@ -0,0 +1,242 @@ +"""Ordered persistence, shipment, and reporting for one train batch.""" + +from __future__ import annotations + +import asyncio +import os +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +from prime_rl.orchestrator.utils import save_rollouts, trim_process_memory +from prime_rl.transport import TrainingBatch +from prime_rl.utils.logger import format_time, get_logger +from prime_rl.utils.pathing import get_trace_path + +if TYPE_CHECKING: + from prime_rl.configs.orchestrator import OrchestratorConfig + from prime_rl.orchestrator.dispatcher import RolloutDispatcher + from prime_rl.orchestrator.envs import TrainEnvs + from prime_rl.orchestrator.metrics import TrainRollouts + from prime_rl.orchestrator.train_sink import TrainSink + from prime_rl.orchestrator.types import Progress, TrainBatch + from prime_rl.utils.heartbeat import Heartbeat + from prime_rl.utils.monitor.base import Monitor + from prime_rl.utils.usage_reporter import UsageReporter + +MAX_CONSECUTIVE_EMPTY_BATCHES = 10 + + +class TrainFinalizationHost(Protocol): + config: OrchestratorConfig + progress: Progress + dispatcher: RolloutDispatcher + train_envs: TrainEnvs + train_sink: TrainSink + monitor: Monitor + usage_reporter: UsageReporter | None + heart: Heartbeat | None + last_batch_at: float | None + consecutive_empty_batches: int + draining: bool + wait_for_policy_time: float + + async def _send_to_trainer(self, batch: TrainingBatch) -> None: ... + + def update_dispatch_gate(self) -> None: ... + + async def maybe_save_ckpt(self, step: int) -> float: ... + + def maybe_trigger_eval(self, step: int) -> None: ... + + +@dataclass(frozen=True) +class TrainStepReport: + metrics: dict[str, float] + num_tokens: int + num_input: int + num_output: int + num_rollouts: int + num_unique_examples: int + + +async def finalize_train_batch(host: TrainFinalizationHost, batch: TrainBatch) -> None: + """Preserve the step's persist → send → checkpoint → report transaction.""" + step = host.progress.step + step_time = _start_step_clock(host) + if await _skip_unshippable_batch(host, batch, step): + return + effective, save_ckpt_time = await _persist_and_ship(host, batch, step) + report = _build_train_step_report(host, batch, effective, step, step_time, save_ckpt_time) + _publish_train_step_report(host, batch, effective, report, step, step_time) + + +def _start_step_clock(host: TrainFinalizationHost) -> float: + now = time.perf_counter() + step_time = (now - host.last_batch_at) if host.last_batch_at is not None else 0.0 + host.last_batch_at = now + return step_time + + +async def _skip_unshippable_batch(host: TrainFinalizationHost, batch: TrainBatch, step: int) -> bool: + if host.config.max_steps is not None and step > host.config.max_steps: + host.draining = True + host.dispatcher.disable_train_scheduling() + n_cancelled = await host.dispatcher.cancel_inflight_train_rollouts() + get_logger().info( + f"Draining pipeline (cancelled {n_cancelled} in-flight train rollout(s); any in-flight evals will complete)" + ) + return True + if not batch.samples: + host.consecutive_empty_batches += 1 + get_logger().warning( + f"Step {step}: empty train batch (0 of {len(batch.rollouts)} generated rollouts shipped — " + f"all errored or filtered out) (consecutive empty batches: " + f"{host.consecutive_empty_batches}/{MAX_CONSECUTIVE_EMPTY_BATCHES})" + ) + if host.consecutive_empty_batches >= MAX_CONSECUTIVE_EMPTY_BATCHES: + raise RuntimeError( + f"{host.consecutive_empty_batches} consecutive empty train batches — " + "check filter config (pre_batch_filters / post_batch_filters) or task difficulty." + ) + return True + host.consecutive_empty_batches = 0 + n_trainable = sum(1 for rollout in batch.rollouts if rollout.is_trainable) + if n_trainable / len(batch.rollouts) <= 0.1: + get_logger().warning( + f"Only {n_trainable}/{len(batch.rollouts)} generated rollouts are trainable " + f"({n_trainable / len(batch.rollouts):.1%}) — consider reviewing task difficulty / filter config" + ) + return False + + +async def _persist_and_ship( + host: TrainFinalizationHost, + batch: TrainBatch, + step: int, +) -> tuple[TrainRollouts, float]: + effective = batch.rollouts.effective + records = [rollout.to_record() for rollout in effective] + await asyncio.to_thread( + save_rollouts, + records, + get_trace_path(host.config.output_dir, step, "train", "effective"), + ) + await host._send_to_trainer(TrainingBatch(examples=batch.samples, step=step)) + host.progress.step += 1 + host.update_dispatch_gate() + save_ckpt_time = await host.maybe_save_ckpt(step) + trim_process_memory() + return effective, save_ckpt_time + + +def _build_train_step_report( + host: TrainFinalizationHost, + batch: TrainBatch, + effective: TrainRollouts, + step: int, + step_time: float, + save_ckpt_time: float, +) -> TrainStepReport: + metrics: dict[str, float] = {} + for subset, pool in (("all", batch.rollouts), ("effective", effective)): + metrics |= pool.metrics.to_wandb(prefix="train/agg", subset=subset) + for env_name, env_pool in pool.by_env().items(): + metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) + num_tokens = sum(rollout.num_total_tokens for rollout in batch.rollouts) + num_input = sum(rollout.num_input_tokens for rollout in effective) + num_output = sum(rollout.num_output_tokens for rollout in effective) + num_rollouts = len(batch.rollouts) + num_unique_examples = len({rollout.group_id for rollout in batch.rollouts}) + metrics |= { + "progress/tokens": num_tokens, + "progress/input_tokens": num_input, + "progress/output_tokens": num_output, + "progress/rollouts": num_rollouts, + "progress/tasks": num_unique_examples, + "progress/total_tokens": host.progress.total_tokens, + "progress/total_rollouts": host.progress.total_samples, + "progress/total_tasks": host.progress.total_problems, + "time/step": step_time, + "time/save_ckpt": save_ckpt_time, + "time/wait_for_policy": host.wait_for_policy_time, + "step": step, + } + for env_name, env_pool in batch.rollouts.by_env().items(): + metrics[f"batch/{env_name}"] = len(env_pool) / num_rollouts + if host.train_sink.pre_filter_seen > 0: + metrics["pre_filters/all/dropped_rate"] = host.train_sink.pre_filter_dropped / host.train_sink.pre_filter_seen + for name, count in host.train_sink.pre_filter_dropped_by_name.items(): + metrics[f"pre_filters/all/{name}/rate"] = count / host.train_sink.pre_filter_seen + return TrainStepReport(metrics, num_tokens, num_input, num_output, num_rollouts, num_unique_examples) + + +def _publish_train_step_report( + host: TrainFinalizationHost, + batch: TrainBatch, + effective: TrainRollouts, + report: TrainStepReport, + step: int, + step_time: float, +) -> None: + host.monitor.log(report.metrics, step=step) + host.wait_for_policy_time = 0.0 + host.monitor.log_samples(effective.rollouts, step=step) + host.monitor.log_distributions( + distributions={ + "rewards": [rollout.reward for rollout in effective], + "advantages": [value for rollout in effective if (value := rollout.scalar_advantage()) is not None], + }, + step=step, + ) + if host.usage_reporter is not None and (run_id := os.getenv("RUN_ID", "")): + host.usage_reporter.report_training_usage( + run_id=run_id, + step=step, + tokens=report.num_input + report.num_output, + ) + if host.heart is not None: + host.heart.beat() + host.progress.total_tokens += report.num_tokens + host.progress.total_samples += report.num_rollouts + host.progress.total_problems += report.num_unique_examples + _log_train_batch(host, batch, step=step, step_time=step_time) + host.train_sink.reset_pre_filter_stats() + host.maybe_trigger_eval(host.progress.step) + trim_process_memory() + + +def _log_train_batch(host: TrainFinalizationHost, batch: TrainBatch, *, step: int, step_time: float) -> None: + rollouts = batch.rollouts + effective = rollouts.effective + metrics = effective.metrics + n_generated = len(rollouts) + n_trainable = sum(1 for rollout in rollouts if rollout.is_trainable) + trainable_rate = (n_trainable / n_generated) if n_generated else 0.0 + max_off_policy = max((rollout.off_policy_steps for rollout in effective), default=0) + head = ( + f"Step {step} | {format_time(step_time):>7} | Reward {metrics.reward.mean():.4f} | " + f"Trainable {n_trainable}/{n_generated} ({trainable_rate:.1%}) | " + f"Turns {metrics.num_turns.mean():.1f} | Branches {metrics.num_branches.mean():.1f} | " + f"Max Off-Policy {max_off_policy} | Error {rollouts.metrics.has_error.mean():.1%} | " + f"Truncation {metrics.is_truncated.mean():.1%}" + ) + if len(host.train_envs) <= 1: + get_logger().success(head) + return + by_env = rollouts.by_env() + name_width = max((len(name) for name in by_env), default=0) + lines = [head] + for env_name in sorted(by_env): + pool = by_env[env_name] + env_effective = pool.effective + env_metrics = env_effective.metrics + ratio = (len(pool) / n_generated) if n_generated else 0.0 + lines.append( + f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | " + f"Reward {env_metrics.reward.mean():.4f} | Turns {env_metrics.num_turns.mean():.1f} | " + f"Branches {env_metrics.num_branches.mean():.1f} | " + f"Max Off-Policy {max((r.off_policy_steps for r in env_effective), default=0)} | " + f"Error {pool.metrics.has_error.mean():.1%} | Truncation {env_metrics.is_truncated.mean():.1%}" + ) + get_logger().success("\n\t\t ".join(lines)) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 0b3406dd65..0db6e4219d 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -51,6 +51,7 @@ class InflightRollout: client_config: vf.ClientConfig | None = None off_policy_steps: int = 0 eval_step: int | None = None + uses_mutable_policy: bool = False @dataclass @@ -67,6 +68,7 @@ class GroupState: eval_step: int | None = None pinned_client: vf.ClientConfig | None = None policy_version_at_start: int = 0 + uses_mutable_policy: bool = False class Rollout(vf.Trace[DataT], Generic[DataT]): @@ -174,7 +176,9 @@ class VersionObserver(Protocol): the weight update; ``on_new_version`` fires *after* the new weights are live and ``Policy`` has been mutated. ``on_version_update_failed`` rolls back transition-only state only when failure is known to precede engine - mutation; indeterminate engine-update failures remain fenced.""" + mutation; indeterminate engine-update failures remain fenced. A pending + hook that raises must roll back state it partially entered before raising; + only hooks that return successfully receive a later failure callback.""" async def on_version_pending(self, step: int) -> None: ... diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 2b23369045..a844d5fe5e 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -12,6 +12,7 @@ from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.utils.client import setup_inference_pool from prime_rl.utils.logger import InterceptHandler, get_logger, setup_logger +from prime_rl.utils.policy_client_config import policy_client_config_from_environment from prime_rl.utils.utils import ( get_broadcast_dir, get_ckpt_dir, @@ -31,7 +32,7 @@ async def setup_policy_inference_pool(*, config: OrchestratorConfig, tokenizer): use plain chat-completions.""" from renderers.base import create_renderer - client_config = config.model.client + client_config = policy_client_config_from_environment(config.model.client) model_name = config.model.name renderer = create_renderer(tokenizer, config.renderer) get_logger().info(f"Initialized {type(renderer).__name__} for {model_name}") diff --git a/src/prime_rl/orchestrator/watcher.py b/src/prime_rl/orchestrator/watcher.py index 8230938ac8..966d69dfdb 100644 --- a/src/prime_rl/orchestrator/watcher.py +++ b/src/prime_rl/orchestrator/watcher.py @@ -108,10 +108,10 @@ async def apply_policy_update(self, next_step: int) -> None: entered_observers: list[VersionObserver] = [] try: for observer in self.observers: - # Include the observer before entry so a partially-applied - # barrier gets its failure rollback hook. - entered_observers.append(observer) await observer.on_version_pending(next_step) + # A hook owns its own partial-entry rollback. Only a fully + # entered observer may receive a later transition cleanup. + entered_observers.append(observer) except BaseException as exc: await self._notify_update_failed(entered_observers, next_step, exc) raise @@ -145,39 +145,13 @@ async def apply_policy_update(self, next_step: int) -> None: @staticmethod async def _notify_update_succeeded(observers: list[VersionObserver], step: int) -> None: - """Notify every observer, preserving cancellation until fences reopen.""" - fatal_error: BaseException | None = None + """Notify observers in dependency order and fail closed on any error.""" # Complete observers in reverse order: the orchestrator first # re-evaluates its long-lived lead gate while the dispatcher's short # transition fence remains closed, then the dispatcher reopens # admission on the new version. for observer in reversed(observers): - try: - await observer.on_new_version(step) - except BaseException as exc: - # A success callback may have partially applied transition - # state before failing. Its failure hook is also the idempotent - # release fallback; continue through the remaining observers. - try: - await observer.on_version_update_failed(step, exc) - except BaseException as cleanup_error: - exc.add_note( - f"Observer {type(observer).__name__}.on_version_update_failed({step}) also failed: " - f"{cleanup_error!r}" - ) - if not isinstance(cleanup_error, Exception) and fatal_error is None: - fatal_error = cleanup_error - if isinstance(exc, Exception): - get_logger().warning(f"Observer {type(observer).__name__}.on_new_version({step}) raised: {exc!r}") - elif fatal_error is None: - fatal_error = exc - else: - fatal_error.add_note( - f"Observer {type(observer).__name__}.on_new_version({step}) also failed: {exc!r}" - ) - if fatal_error is not None: - await WeightWatcher._notify_update_failed(observers, step, fatal_error) - raise fatal_error + await observer.on_new_version(step) @staticmethod async def _notify_update_failed( diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 16c956cef7..19b0e6443a 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 Literal, Protocol, runtime_checkable import httpx import verifiers.v1 as vf @@ -59,6 +59,11 @@ def model_name(self) -> str: """Get current model name for inference requests.""" ... + @property + def admin_api(self) -> Literal["vllm", "dynamo"]: + """Administration protocol used by the resolved pool.""" + ... + @property def train_clients(self) -> list[vf.ClientConfig]: """Get inference clients.""" @@ -179,6 +184,10 @@ def __init__( def train_clients(self) -> list[vf.ClientConfig]: return self._train_clients + @property + def admin_api(self) -> Literal["vllm", "dynamo"]: + return self._client_config.admin_api + @property def admin_clients(self) -> list[AsyncClient]: return self._admin_clients diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index 4266beffe4..208f9ebcd6 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -225,6 +225,10 @@ def train_clients(self) -> list[vf.ClientConfig]: self._rebuild_clients() return self._train_clients + @property + def admin_api(self) -> Literal["vllm", "dynamo"]: + return self.client_config.admin_api + @property def eval_clients(self) -> list[vf.ClientConfig]: self._rebuild_clients() diff --git a/src/prime_rl/utils/policy_client_config.py b/src/prime_rl/utils/policy_client_config.py new file mode 100644 index 0000000000..8250713776 --- /dev/null +++ b/src/prime_rl/utils/policy_client_config.py @@ -0,0 +1,50 @@ +"""Resolve the generated DGD policy-client deployment boundary.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + +from prime_rl.configs.shared import ClientConfig + +DYNAMO_TOPOLOGY_ENV = "DYN_RL_TOPOLOGY" + + +class _DynamoClientTopologyEnvironment(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + schema_version: Literal[1] + admin_api: Literal["dynamo"] + base_url: tuple[str, ...] = Field(min_length=1) + rl_base_url: tuple[str, ...] = Field(min_length=1) + dynamo_worker_roles: tuple[Literal["agg", "prefill", "decode"], ...] = Field(min_length=1) + dynamo_gpus_per_worker: int = Field(ge=1) + + +def policy_client_config_from_environment( + client_config: ClientConfig, + environment: Mapping[str, str] | None = None, +) -> ClientConfig: + """Apply the generated policy topology without mutating user config.""" + source = environment if environment is not None else os.environ + serialized = source.get(DYNAMO_TOPOLOGY_ENV) + if serialized is None: + return client_config + topology = _DynamoClientTopologyEnvironment.model_validate_json(serialized) + updates = { + "admin_api": topology.admin_api, + "base_url": list(topology.base_url), + "rl_base_url": list(topology.rl_base_url), + "dynamo_worker_roles": topology.dynamo_worker_roles, + "dynamo_gpus_per_worker": topology.dynamo_gpus_per_worker, + } + for field, expected in updates.items(): + if field in client_config.model_fields_set and getattr(client_config, field) != expected: + raise ValueError( + f"orchestrator.model.client.{field} conflicts with the generated " + f"{DYNAMO_TOPOLOGY_ENV} deployment boundary" + ) + return client_config.model_copy(update=updates) diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py index a08cad8fd6..b08ead2a6a 100644 --- a/tests/unit/inference/test_dynamo.py +++ b/tests/unit/inference/test_dynamo.py @@ -140,13 +140,40 @@ def test_local_specs_allocate_four_workers_and_unique_ports(tmp_path: Path): prefill_configs = [ json.loads(Path(spec.process.arguments[1]).read_text()) for spec in specs if spec.role == "prefill" ] - assert [config["kv_events_config"]["endpoint"] for config in prefill_configs] == [ - "tcp://*:20080", - "tcp://*:20081", - ] + assert len({config["kv_events_config"]["endpoint"] for config in prefill_configs}) == 2 assert all("--enable-rl" in spec.process.command() for spec in specs) +def test_local_multi_gpu_workers_allocate_globally_unique_coordinator_ports(tmp_path: Path): + config = disaggregated_config( + parallel={"tp": 1}, + deployment={ + "type": "disaggregated", + "gpus_per_node": 2, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + }, + ) + + specs = build_local_worker_specs(config, tmp_path, gpu_ids=["0", "1", "2", "3"]) + engine_configs = [json.loads(Path(spec.process.arguments[1]).read_text()) for spec in specs] + allocated_ports = [ + *(spec.system_port for spec in specs), + *(int(spec.process.environment()["VLLM_NIXL_SIDE_CHANNEL_PORT"]) for spec in specs), + *(engine["data_parallel_rpc_port"] for engine in engine_configs), + *( + int(engine["kv_events_config"]["endpoint"].rsplit(":", 1)[1]) + for engine in engine_configs + if "kv_events_config" in engine + ), + ] + + assert len({engine["data_parallel_rpc_port"] for engine in engine_configs}) == len(specs) + assert len(allocated_ports) == len(set(allocated_ports)) + + def test_wrapper_options_are_not_written_to_engine_json(): engine = build_engine_config(disaggregated_config(), "prefill", kv_events_port=20080) assert "disaggregation_mode" not in engine diff --git a/tests/unit/inference/test_dynamo_admin.py b/tests/unit/inference/test_dynamo_admin.py index e3b4c5cf2e..408e08cc17 100644 --- a/tests/unit/inference/test_dynamo_admin.py +++ b/tests/unit/inference/test_dynamo_admin.py @@ -545,7 +545,6 @@ def handler(request: httpx.Request) -> httpx.Response: assert paths == [ "/engine/pause_generation", "/engine/update_weights_from_disk", - "/engine/resume_generation", ] diff --git a/tests/unit/inference/test_dynamo_admin_barriers.py b/tests/unit/inference/test_dynamo_admin_barriers.py index 83bb3d96e7..d3486af182 100644 --- a/tests/unit/inference/test_dynamo_admin_barriers.py +++ b/tests/unit/inference/test_dynamo_admin_barriers.py @@ -145,7 +145,9 @@ async def post(client, method, *_args, **_kwargs): @pytest.mark.asyncio -async def test_collective_update_settles_delayed_sibling_before_resume(monkeypatch: pytest.MonkeyPatch): +async def test_collective_update_settles_delayed_sibling_without_resuming_indeterminate_workers( + monkeypatch: pytest.MonkeyPatch, +): admin = DynamoAdminAPI() delayed_started = asyncio.Event() release_delayed = asyncio.Event() @@ -178,7 +180,9 @@ async def post(client, method, *_args, **_kwargs): with pytest.raises(RuntimeError, match="collective failed"): await update - assert resume_started.is_set() + assert not resume_started.is_set() + with pytest.raises(RuntimeError, match="weight state is indeterminate"): + await admin.update_weights(["fast-failure", "delayed"], Path("weights"), step=2) @pytest.mark.asyncio @@ -219,7 +223,9 @@ async def post(_client, method, *_args, **_kwargs): await update assert update_finished.is_set() - assert resume_started.is_set() + assert not resume_started.is_set() + with pytest.raises(RuntimeError, match="weight state is indeterminate"): + await admin.update_weights(["worker"], Path("weights"), step=2) @pytest.mark.asyncio diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index 9d79b5eb50..62b6de4f87 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -124,6 +124,48 @@ async def test_opsd_rejects_late_live_policy_score_as_rollout_error(): pool.score.assert_not_awaited() +@pytest.mark.asyncio +async def test_opsd_settles_every_score_sibling_before_releasing_policy_gate(): + policy = Policy(version=0, model_name="policy") + gate = MutablePolicyGate(policy, enabled=True) + sibling_started = asyncio.Event() + release_sibling = asyncio.Event() + calls = 0 + + async def score(_token_ids: list[int]) -> list[float]: + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("score failed") + sibling_started.set() + await release_sibling.wait() + return [0.0] * 7 + + pool = MagicMock(model_name="policy") + pool.score = score + algo = OPSDAlgorithm(_build(type="opsd"), pool, policy_gate=gate) + algo.renderer = MagicMock() + algo.renderer.render_ids.return_value = [999] + rollout = _make_rollout([_make_sample(), _make_sample()]) + rollout.info["demonstration"] = "expert answer" + rollout.policy_version = 0 + + scoring = asyncio.create_task(algo.score_rollout(rollout)) + await sibling_started.wait() + update_token = await gate.begin_update(step=1) + idle = asyncio.create_task(gate.wait_idle()) + await asyncio.sleep(0) + + assert not scoring.done() + assert not idle.done() + + release_sibling.set() + with pytest.raises(RuntimeError, match="score failed"): + await scoring + await idle + await gate.finish_update(update_token) + + def test_sft_requires_teacher(): with pytest.raises(ValueError, match="needs a teacher to sample rollouts from"): _build(type="sft") diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index bf75d3bb4f..811480a293 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -1,11 +1,14 @@ import asyncio +from contextlib import suppress from types import SimpleNamespace from unittest.mock import AsyncMock, patch import pytest from renderers import Qwen3VLRendererConfig -from prime_rl.orchestrator.component_supervision import raise_if_component_failed +from prime_rl.configs.shared import ClientConfig +from prime_rl.orchestrator import component_supervision +from prime_rl.orchestrator.component_supervision import raise_if_component_failed, run_with_component_supervision from prime_rl.orchestrator.utils import setup_policy_inference_pool @@ -22,6 +25,228 @@ async def failed_watcher() -> None: asyncio.run(run()) +def test_component_failure_wins_over_queued_output_and_prevents_trainer_send(): + async def run() -> None: + async def failed_watcher() -> None: + raise RuntimeError("indeterminate policy update") + + watcher = asyncio.create_task(failed_watcher(), name="watcher") + await asyncio.wait({watcher}) + operation = AsyncMock(return_value=object()) + + with pytest.raises(RuntimeError, match="indeterminate policy update"): + await run_with_component_supervision(operation, [watcher]) + + operation.assert_not_called() + + asyncio.run(run()) + + +def test_component_failure_cancels_in_progress_external_operation(): + async def run() -> None: + operation_started = asyncio.Event() + operation_cancelled = asyncio.Event() + fail_component = asyncio.Event() + + async def operation() -> None: + operation_started.set() + try: + await asyncio.Future() + finally: + operation_cancelled.set() + + async def watcher() -> None: + await fail_component.wait() + raise RuntimeError("watcher failed during send") + + watcher_task = asyncio.create_task(watcher(), name="watcher") + supervised = asyncio.create_task( + run_with_component_supervision( + operation, + [watcher_task], + timeout=1.0, + timeout_description="training batch send", + ) + ) + await operation_started.wait() + fail_component.set() + + with pytest.raises(RuntimeError, match="watcher failed during send"): + await supervised + assert operation_cancelled.is_set() + + asyncio.run(run()) + + +def test_supervised_operation_timeout_cancels_stuck_send_and_raises(): + async def run() -> None: + operation_cancelled = asyncio.Event() + + async def stuck_send() -> None: + try: + await asyncio.Future() + finally: + operation_cancelled.set() + + with pytest.raises(TimeoutError, match=r"training batch send timed out after 0\.01 seconds"): + await run_with_component_supervision( + stuck_send, + [], + timeout=0.01, + timeout_description="training batch send", + ) + + assert operation_cancelled.is_set() + + asyncio.run(run()) + + +def test_timeout_cleanup_is_bounded_when_operation_suppresses_cancellation(monkeypatch: pytest.MonkeyPatch): + async def run() -> None: + operation_started = asyncio.Event() + cancellation_seen = asyncio.Event() + release_orphan = asyncio.Event() + + async def cancellation_suppressing_send() -> None: + operation_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_seen.set() + await release_orphan.wait() + raise RuntimeError("late orphan failure") + + monkeypatch.setattr( + component_supervision, + "SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS", + 0.01, + raising=False, + ) + supervised = asyncio.create_task( + run_with_component_supervision( + cancellation_suppressing_send, + [], + timeout=0.01, + timeout_description="training batch send", + ) + ) + try: + with pytest.raises(TimeoutError, match=r"training batch send timed out after 0\.01 seconds"): + await asyncio.wait_for(asyncio.shield(supervised), timeout=0.1) + assert cancellation_seen.is_set() + assert len(component_supervision._ORPHANED_OPERATIONS) == 1 + finally: + release_orphan.set() + if not supervised.done(): + with suppress(BaseException): + await supervised + + for _ in range(10): + if not component_supervision._ORPHANED_OPERATIONS: + break + await asyncio.sleep(0) + assert not component_supervision._ORPHANED_OPERATIONS + + asyncio.run(run()) + + +def test_caller_cancellation_during_timeout_cleanup_is_not_swallowed(monkeypatch: pytest.MonkeyPatch): + async def run() -> None: + cancellation_seen = asyncio.Event() + release_orphan = asyncio.Event() + + async def cancellation_suppressing_poll() -> None: + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_seen.set() + await release_orphan.wait() + + monkeypatch.setattr(component_supervision, "SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS", 0.01) + supervised = asyncio.create_task( + run_with_component_supervision(cancellation_suppressing_poll, [], timeout=0.01) + ) + await cancellation_seen.wait() + supervised.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(supervised), timeout=0.1) + assert len(component_supervision._ORPHANED_OPERATIONS) == 1 + + release_orphan.set() + for _ in range(10): + if not component_supervision._ORPHANED_OPERATIONS: + break + await asyncio.sleep(0) + assert not component_supervision._ORPHANED_OPERATIONS + + asyncio.run(run()) + + +def test_component_failure_precedes_caller_cancellation_during_cleanup(monkeypatch: pytest.MonkeyPatch): + async def run() -> None: + operation_started = asyncio.Event() + cancellation_seen = asyncio.Event() + release_orphan = asyncio.Event() + fail_watcher = asyncio.Event() + + async def cancellation_suppressing_send() -> None: + operation_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_seen.set() + await release_orphan.wait() + + async def failed_watcher() -> None: + await fail_watcher.wait() + raise RuntimeError("component failed before cleanup cancellation") + + monkeypatch.setattr(component_supervision, "SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS", 0.01) + watcher = asyncio.create_task(failed_watcher(), name="watcher") + supervised = asyncio.create_task(run_with_component_supervision(cancellation_suppressing_send, [watcher])) + await operation_started.wait() + fail_watcher.set() + await cancellation_seen.wait() + supervised.cancel() + + with pytest.raises(RuntimeError, match="component failed before cleanup cancellation"): + await asyncio.wait_for(asyncio.shield(supervised), timeout=0.1) + assert len(component_supervision._ORPHANED_OPERATIONS) == 1 + + release_orphan.set() + for _ in range(10): + if not component_supervision._ORPHANED_OPERATIONS: + break + await asyncio.sleep(0) + assert not component_supervision._ORPHANED_OPERATIONS + + asyncio.run(run()) + + +def test_component_failure_wins_when_operation_completes_in_same_loop_turn(): + async def run() -> None: + release = asyncio.Event() + + async def operation() -> str: + await release.wait() + return "sent" + + async def watcher() -> None: + await release.wait() + raise RuntimeError("simultaneous watcher failure") + + watcher_task = asyncio.create_task(watcher(), name="watcher") + supervised = asyncio.create_task(run_with_component_supervision(operation, [watcher_task])) + await asyncio.sleep(0) + release.set() + + with pytest.raises(RuntimeError, match="simultaneous watcher failure"): + await supervised + + asyncio.run(run()) + + def test_setup_policy_inference_pool_uses_renderer_when_enabled(): async def run() -> None: tokenizer = object() @@ -111,3 +336,33 @@ async def run() -> None: ) asyncio.run(run()) + + +def test_setup_policy_inference_pool_uses_environment_resolved_client(monkeypatch: pytest.MonkeyPatch): + async def run() -> None: + monkeypatch.setenv( + "DYN_RL_TOPOLOGY", + """{"schema_version":1,"admin_api":"dynamo","base_url":["http://frontend:8000/v1"],"rl_base_url":["http://frontend-rl:8001"],"dynamo_worker_roles":["agg"],"dynamo_gpus_per_worker":1}""", + ) + config = SimpleNamespace( + model=SimpleNamespace(client=ClientConfig(), name="policy-model"), + renderer=Qwen3VLRendererConfig(), + pool_size=None, + any_policy_sourced=True, + ) + + with ( + patch("renderers.base.create_renderer", return_value=object()), + patch( + "prime_rl.orchestrator.utils.setup_inference_pool", + new=AsyncMock(return_value=object()), + ) as setup_pool, + ): + await setup_policy_inference_pool(config=config, tokenizer=object()) + + resolved = setup_pool.await_args.args[0] + assert resolved.admin_api == "dynamo" + assert resolved.base_url == ["http://frontend:8000/v1"] + assert config.model.client.admin_api == "vllm" + + asyncio.run(run()) diff --git a/tests/unit/orchestrator/test_policy_gate_cancellation.py b/tests/unit/orchestrator/test_policy_gate_cancellation.py new file mode 100644 index 0000000000..0738a806ad --- /dev/null +++ b/tests/unit/orchestrator/test_policy_gate_cancellation.py @@ -0,0 +1,68 @@ +import asyncio + +import pytest + +from prime_rl.orchestrator.dispatcher import RolloutDispatcher +from prime_rl.orchestrator.types import Policy + + +def _dispatcher() -> RolloutDispatcher: + pool = type( + "Pool", + (), + { + "model_name": "policy", + "train_clients": [], + "admin_clients": [], + }, + )() + return RolloutDispatcher( + train_envs=object(), + eval_envs=None, + train_source=object(), + eval_source=None, + policy_pool=pool, + policy=Policy(version=0, model_name="policy"), + max_inflight_rollouts=1, + tasks_per_minute=None, + max_off_policy_steps=0, + enforce_policy_update_barrier=True, + ) + + +@pytest.mark.asyncio +async def test_repeated_cancellation_cannot_interrupt_partial_entry_rollback(): + dispatcher = _dispatcher() + settle_started = asyncio.Event() + release_settle = asyncio.Event() + rollback_started = asyncio.Event() + + async def settle_policy_requests(*_args) -> None: + settle_started.set() + await release_settle.wait() + + dispatcher._settle_policy_requests = settle_policy_requests # type: ignore[method-assign] + finish_update = dispatcher.policy_gate.finish_update + + async def observed_finish_update(token) -> None: + rollback_started.set() + await finish_update(token) + + dispatcher.policy_gate.finish_update = observed_finish_update # type: ignore[method-assign] + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await settle_started.wait() + + await dispatcher.policy_gate._admission_lock.acquire() + barrier.cancel() + release_settle.set() + await rollback_started.wait() + barrier.cancel() + await asyncio.sleep(0) + rollback_was_interrupted = barrier.done() + dispatcher.policy_gate._admission_lock.release() + + assert not rollback_was_interrupted + with pytest.raises(asyncio.CancelledError) as exc: + await barrier + assert not dispatcher.policy_update_pending + assert exc.value.__notes__ == ["Policy transition rollback was cancelled again but settled before propagation"] diff --git a/tests/unit/orchestrator/test_pool_identity.py b/tests/unit/orchestrator/test_pool_identity.py index a1e090a1c0..e17c4e02df 100644 --- a/tests/unit/orchestrator/test_pool_identity.py +++ b/tests/unit/orchestrator/test_pool_identity.py @@ -33,8 +33,8 @@ def _pool(model: str, request: str | None, admin: str | None): ), ( _pool("policy", "http://frontend/v1", "http://worker:8081"), - _pool("other-model", "http://frontend/v1", "http://worker:8081"), - False, + _pool("other-model", "http://other-frontend/v1", "http://worker:8081"), + True, ), ( _pool("policy", None, None), diff --git a/tests/unit/orchestrator/test_train_finalization.py b/tests/unit/orchestrator/test_train_finalization.py new file mode 100644 index 0000000000..6df8627123 --- /dev/null +++ b/tests/unit/orchestrator/test_train_finalization.py @@ -0,0 +1,150 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from prime_rl.orchestrator.train_finalization import finalize_train_batch +from prime_rl.orchestrator.types import Progress + + +class _Stat: + def mean(self) -> float: + return 1.0 + + +class _Metrics: + reward = _Stat() + num_turns = _Stat() + num_branches = _Stat() + is_truncated = _Stat() + has_error = _Stat() + + def to_wandb(self, *, prefix: str, subset: str) -> dict[str, float]: + return {f"{prefix}/{subset}/metric": 1.0} + + +class _Rollout: + is_trainable = True + off_policy_steps = 0 + num_total_tokens = 3 + num_input_tokens = 1 + num_output_tokens = 2 + group_id = "group" + reward = 1.0 + + def to_record(self) -> dict: + return {"group": self.group_id} + + def scalar_advantage(self) -> float: + return 0.5 + + +class _Rollouts(list): + metrics = _Metrics() + + @property + def effective(self): + return self + + @property + def rollouts(self): + return list(self) + + def by_env(self) -> dict: + return {} + + +class _Monitor: + def __init__(self, events: list[str]) -> None: + self.events = events + self.logged_metrics: dict[str, float] = {} + + def log(self, metrics: dict[str, float], *, step: int) -> None: + self.events.append("monitor:log") + self.logged_metrics = metrics + + def log_samples(self, _rollouts, *, step: int) -> None: + self.events.append("monitor:samples") + + def log_distributions(self, *, distributions, step: int) -> None: + self.events.append("monitor:distributions") + + +class _Host: + def __init__(self, tmp_path: Path, events: list[str]) -> None: + self.config = SimpleNamespace(output_dir=tmp_path, max_steps=None) + self.progress = Progress() + self.last_batch_at = 1.0 + self.draining = False + self.consecutive_empty_batches = 0 + self.wait_for_policy_time = 0.25 + self.monitor = _Monitor(events) + self.usage_reporter = None + self.heart = None + self.train_envs = [] + self.train_sink = SimpleNamespace( + pre_filter_seen=0, + pre_filter_dropped=0, + pre_filter_dropped_by_name={}, + reset_pre_filter_stats=lambda: events.append("reset-filters"), + ) + self.events = events + + async def _send_to_trainer(self, _batch) -> None: + self.events.append("send") + + def update_dispatch_gate(self) -> None: + assert self.progress.step == 2 + self.events.append("dispatch-gate") + + async def maybe_save_ckpt(self, step: int) -> float: + assert step == 1 + self.events.append("checkpoint") + return 0.5 + + def maybe_trigger_eval(self, step: int) -> None: + assert step == 2 + self.events.append("eval") + + +@pytest.mark.asyncio +async def test_finalize_train_batch_preserves_ship_and_reporting_order(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + events: list[str] = [] + host = _Host(tmp_path, events) + batch = SimpleNamespace(samples=[object()], rollouts=_Rollouts([_Rollout()])) + + monkeypatch.setattr( + "prime_rl.orchestrator.train_finalization.save_rollouts", + lambda *_args: events.append("save-rollouts"), + ) + monkeypatch.setattr( + "prime_rl.orchestrator.train_finalization.trim_process_memory", + lambda: events.append("trim"), + ) + monkeypatch.setattr( + "prime_rl.orchestrator.train_finalization.get_logger", + lambda: SimpleNamespace(success=lambda _message: events.append("success")), + ) + + await finalize_train_batch(host, batch) + + assert events == [ + "save-rollouts", + "send", + "dispatch-gate", + "checkpoint", + "trim", + "monitor:log", + "monitor:samples", + "monitor:distributions", + "success", + "reset-filters", + "eval", + "trim", + ] + assert host.progress.step == 2 + assert host.progress.total_tokens == 3 + assert host.progress.total_samples == 1 + assert host.progress.total_problems == 1 + assert host.monitor.logged_metrics["progress/total_tokens"] == 0 + assert host.wait_for_policy_time == 0.0 diff --git a/tests/unit/orchestrator/test_weight_update_barrier.py b/tests/unit/orchestrator/test_weight_update_barrier.py index 5d391ce65e..7ab0bc145a 100644 --- a/tests/unit/orchestrator/test_weight_update_barrier.py +++ b/tests/unit/orchestrator/test_weight_update_barrier.py @@ -4,10 +4,11 @@ from types import SimpleNamespace import pytest +import verifiers.v1 as vf from prime_rl.orchestrator.dispatcher import DispatcherMode, RolloutDispatcher from prime_rl.orchestrator.policy_gate import MutablePolicyGate, PolicyRequestRejected -from prime_rl.orchestrator.types import GroupState, InflightRollout, Policy +from prime_rl.orchestrator.types import GroupState, InflightRollout, Policy, Rollout from prime_rl.orchestrator.watcher import WeightWatcher from prime_rl.utils.pathing import get_broadcast_dir, get_step_path @@ -63,47 +64,44 @@ def pool(model_name: str, request_url: str, admin_url: str): @pytest.mark.asyncio -async def test_policy_barrier_serializes_with_suspended_eval_scheduling(): +async def test_policy_barrier_invalidates_slow_scheduling_before_short_commit(): dispatcher = _dispatcher() dispatcher.mode = DispatcherMode.PREFER_EVAL scheduling_started = asyncio.Event() allow_schedule_to_commit = asyncio.Event() - request_cleanup_finished = asyncio.Event() + request_started = asyncio.Event() async def active_request() -> None: - try: - await asyncio.Future() - finally: - # Represents the client/connector abort cleanup performed while - # unwinding a cancelled rollout request. - await asyncio.sleep(0) - request_cleanup_finished.set() + request_started.set() + await asyncio.Future() - async def schedule_one(_kind: str) -> bool: + async def schedule_one(_kind: str, *, epoch) -> bool: scheduling_started.set() await allow_schedule_to_commit.wait() - group_id = uuid.uuid4() - request = asyncio.create_task(active_request()) - dispatcher.groups[group_id] = GroupState( - kind="eval", - env_name="eval", - task_idx=0, - rollouts_to_schedule=0, - target_rollouts=1, - eval_step=1, - policy_version_at_start=0, - ) - dispatcher.inflight[request] = InflightRollout( - kind="eval", - env_name="eval", - group_id=group_id, - policy_version=0, - rollout_count=1, - eval_step=1, - ) - dispatcher.inflight_permits += 1 - await asyncio.sleep(0) - return True + async with dispatcher.policy_gate.scheduling_commit(epoch) as admitted: + if not admitted: + return False + group_id = uuid.uuid4() + request = asyncio.create_task(active_request()) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[request] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits += 1 + return True dispatcher.try_schedule = schedule_one # type: ignore[method-assign] @@ -112,17 +110,13 @@ async def schedule_one(_kind: str) -> bool: barrier = asyncio.create_task(dispatcher.on_version_pending(1)) await asyncio.sleep(0) - # The barrier must serialize behind the scheduling commit. Otherwise it - # can take an empty snapshot and let an old-policy request escape. - assert not barrier.done() + await barrier + assert dispatcher.policy_update_pending allow_schedule_to_commit.set() await fill - await barrier - - assert request_cleanup_finished.is_set() + assert not request_started.is_set() assert not dispatcher.inflight - assert dispatcher.policy_update_pending # The transition fence remains closed until the watcher reports either # success or failure. @@ -336,7 +330,6 @@ async def request() -> None: await barrier assert cleanup_finished.is_set() - await dispatcher.on_version_update_failed(1, asyncio.CancelledError()) assert not dispatcher.policy_update_pending @@ -403,7 +396,113 @@ async def request() -> None: assert not dispatcher.inflight assert dispatcher.inflight_permits == 0 - await dispatcher.on_version_update_failed(1, asyncio.CancelledError()) + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_policy_barrier_waits_for_completed_handler_before_computing_markers(): + dispatcher = _dispatcher(max_inflight=2) + dispatcher.out_q = asyncio.Queue(maxsize=1) + group_id = uuid.uuid4() + + async def completed_group() -> list[Rollout]: + return [ + Rollout( + task=vf.Task(idx=0, prompt=None), + errors=[vf.Error(type="Existing", message="result")], + stop_condition="error", + ) + for _ in range(2) + ] + + task = asyncio.create_task(completed_group()) + await task + group = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=2, + eval_step=1, + policy_version_at_start=0, + uses_mutable_policy=True, + ) + dispatcher.groups[group_id] = group + dispatcher.inflight[task] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=2, + eval_step=1, + uses_mutable_policy=True, + ) + dispatcher.inflight_permits = 2 + + handler = asyncio.create_task(dispatcher.handle_completed_rollout(task)) + while dispatcher.out_q.qsize() != 1: + await asyncio.sleep(0) + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await asyncio.sleep(0) + + assert not barrier.done() + + first = dispatcher.out_q.get_nowait() + await handler + await barrier + second = dispatcher.out_q.get_nowait() + + assert [first.error.type, second.error.type] == ["Existing", "Existing"] + assert group.emitted == 2 + await dispatcher.on_new_version(1) + + +@pytest.mark.asyncio +async def test_policy_barrier_marker_enqueue_aborts_promptly_on_dispatcher_stop(): + dispatcher = _dispatcher() + marker_put_started = asyncio.Event() + + class _ObservedQueue(asyncio.Queue): + async def put(self, item) -> None: + marker_put_started.set() + await super().put(item) + + dispatcher.out_q = _ObservedQueue(maxsize=1) + dispatcher.out_q.put_nowait(object()) + group_id = uuid.uuid4() + + async def request() -> None: + await asyncio.Future() + + task = asyncio.create_task(request()) + await asyncio.sleep(0) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + uses_mutable_policy=True, + ) + dispatcher.inflight[task] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + uses_mutable_policy=True, + ) + dispatcher.inflight_permits = 1 + + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await marker_put_started.wait() + await dispatcher.stop() + + with pytest.raises(RuntimeError, match="stopped while emitting policy barrier"): + await asyncio.wait_for(barrier, timeout=0.5) assert not dispatcher.policy_update_pending @@ -518,6 +617,13 @@ async def update_weights(self, *_args, **_kwargs) -> None: "Policy update 1 may have mutated inference workers; mutable-policy admission remains fail-closed" ] + second_weight_path = get_step_path(get_broadcast_dir(tmp_path), 2) + second_weight_path.mkdir(parents=True) + (second_weight_path / "STABLE").touch() + with pytest.raises(RuntimeError, match="already pending"): + await watcher.apply_policy_update(2) + assert dispatcher.policy_update_pending + @pytest.mark.asyncio async def test_success_updates_lead_gate_before_reopening_transition_fence(tmp_path: Path): @@ -566,7 +672,7 @@ async def update_weights(self, *_args, **_kwargs) -> None: @pytest.mark.asyncio -async def test_success_callback_cancellation_still_reopens_transition_fence(tmp_path: Path): +async def test_success_callback_cancellation_propagates_and_keeps_transition_fence_closed(tmp_path: Path): dispatcher = _dispatcher() callback_started = asyncio.Event() weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) @@ -604,11 +710,11 @@ async def update_weights(self, *_args, **_kwargs) -> None: assert watcher.ckpt_step == 1 assert dispatcher.policy.version == 1 - assert not dispatcher.policy_update_pending + assert dispatcher.policy_update_pending @pytest.mark.asyncio -async def test_success_callback_error_still_reopens_transition_fence(tmp_path: Path): +async def test_success_callback_error_propagates_and_keeps_transition_fence_closed(tmp_path: Path): dispatcher = _dispatcher() weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) weight_path.mkdir(parents=True) @@ -636,8 +742,9 @@ async def update_weights(self, *_args, **_kwargs) -> None: lora_name=None, ) - await watcher.apply_policy_update(1) + with pytest.raises(RuntimeError, match="lead gate callback failed"): + await watcher.apply_policy_update(1) assert watcher.ckpt_step == 1 assert dispatcher.policy.version == 1 - assert not dispatcher.policy_update_pending + assert dispatcher.policy_update_pending diff --git a/tests/unit/orchestrator/test_weight_update_barrier_mutability.py b/tests/unit/orchestrator/test_weight_update_barrier_mutability.py new file mode 100644 index 0000000000..7f0988a84c --- /dev/null +++ b/tests/unit/orchestrator/test_weight_update_barrier_mutability.py @@ -0,0 +1,211 @@ +import asyncio +import uuid +from types import SimpleNamespace + +import pytest + +from prime_rl.orchestrator.dispatcher import RolloutDispatcher +from prime_rl.orchestrator.types import GroupState, InflightRollout, Policy + + +def _pool(model: str, request: str, admin: str): + return SimpleNamespace( + model_name=model, + train_clients=[SimpleNamespace(base_url=request, headers={})], + admin_clients=[SimpleNamespace(base_url=admin)], + ) + + +class _EnvCollection: + def __init__(self, pool, *, run_rollout=None) -> None: + self.env = SimpleNamespace( + sampler=SimpleNamespace(samples_from_live_policy=False, pool=pool), + requires_group_scoring=False, + config=SimpleNamespace(group_size=1), + run_rollout=run_rollout, + ) + + def get(self, _name: str): + return self.env + + +def _dispatcher(*, enforce_barrier: bool, max_off_policy_steps: int = 0, run_rollout=None): + policy_pool = _pool("policy", "http://policy/v1", "http://policy-worker:8081") + distinct_pool = _pool("frozen", "http://frozen/v1", "http://frozen-worker:8081") + train_envs = _EnvCollection(distinct_pool, run_rollout=run_rollout) + dispatcher = RolloutDispatcher( + train_envs=train_envs, + eval_envs=None, + train_source=object(), + eval_source=None, + policy_pool=policy_pool, + policy=Policy(version=4, model_name="policy"), + max_inflight_rollouts=1, + tasks_per_minute=None, + max_off_policy_steps=max_off_policy_steps, + enforce_policy_update_barrier=enforce_barrier, + ) + return dispatcher + + +@pytest.mark.asyncio +async def test_group_mutability_is_monotonic_for_cache_salt_after_pool_identity_churn(): + called: dict[str, object] = {} + + async def run_rollout(**kwargs): + called.update(kwargs) + await asyncio.Future() + + dispatcher = _dispatcher(enforce_barrier=True, run_rollout=run_rollout) + group_id = uuid.uuid4() + formerly_aliased_client = SimpleNamespace(base_url="http://policy/v1", headers={}) + group = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=1, + target_rollouts=1, + pinned_client=formerly_aliased_client, + policy_version_at_start=4, + uses_mutable_policy=False, + ) + dispatcher.groups[group_id] = group + epoch = await dispatcher.policy_gate.scheduling_epoch() + assert epoch is not None + + assert await dispatcher.schedule_group_rollout(group_id, group, epoch=epoch) + await asyncio.sleep(0) + + assert called["cache_salt"] == "4" + assert next(iter(dispatcher.inflight.values())).uses_mutable_policy + await dispatcher.cancel_inflight_rollouts() + + +@pytest.mark.asyncio +async def test_group_mutability_is_monotonic_for_non_dynamo_off_policy_aging(): + dispatcher = _dispatcher(enforce_barrier=False, max_off_policy_steps=0) + group_id = uuid.uuid4() + + async def request() -> None: + await asyncio.Future() + + task = asyncio.create_task(request()) + await asyncio.sleep(0) + dispatcher.groups[group_id] = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + policy_version_at_start=4, + uses_mutable_policy=True, + ) + dispatcher.inflight[task] = InflightRollout( + kind="train", + env_name="train", + group_id=group_id, + policy_version=4, + rollout_count=1, + uses_mutable_policy=True, + ) + dispatcher.inflight_permits = 1 + + await dispatcher.on_version_pending(5) + + assert task.cancelled() + assert not dispatcher.inflight + + +@pytest.mark.asyncio +async def test_policy_update_does_not_wait_for_slow_client_selection(): + selection_started = asyncio.Event() + release_selection = asyncio.Event() + + class _BlockingPool: + model_name = "policy" + train_clients = [SimpleNamespace(base_url="http://policy/v1", headers={})] + admin_clients = [SimpleNamespace(base_url="http://policy-worker:8081")] + + async def select_train_client(self, _load): + selection_started.set() + await release_selection.wait() + return self.train_clients[0] + + dispatcher = _dispatcher(enforce_barrier=True, run_rollout=None) + dispatcher.train_envs.env.sampler.pool = _BlockingPool() + group_id = uuid.uuid4() + group = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=1, + target_rollouts=1, + policy_version_at_start=4, + uses_mutable_policy=True, + ) + dispatcher.groups[group_id] = group + epoch = await dispatcher.policy_gate.scheduling_epoch() + assert epoch is not None + + scheduling = asyncio.create_task(dispatcher.schedule_group_rollout(group_id, group, epoch=epoch)) + await selection_started.wait() + await asyncio.wait_for(dispatcher.on_version_pending(5), timeout=0.5) + + assert dispatcher.policy_update_pending + assert group_id not in dispatcher.groups + + release_selection.set() + assert not await scheduling + assert not dispatcher.inflight + await dispatcher.on_new_version(5) + + +@pytest.mark.asyncio +async def test_selected_mutable_endpoint_survives_pool_churn_during_rate_limit_wait(): + rate_limit_started = asyncio.Event() + release_rate_limit = asyncio.Event() + selected_client = SimpleNamespace(base_url="http://policy/v1", headers={}) + + class _ChurningPool: + model_name = "frozen" + train_clients = [selected_client] + admin_clients = [SimpleNamespace(base_url="http://frozen-worker:8081")] + + async def select_train_client(self, _load): + # The elastic snapshot loses the selected endpoint before the + # dispatcher reaches its next slow wait. + self.train_clients = [SimpleNamespace(base_url="http://frozen/v1", headers={})] + return selected_client + + dispatcher = _dispatcher(enforce_barrier=True, run_rollout=None) + dispatcher.train_envs.env.sampler.pool = _ChurningPool() + + async def wait_for_rate_limit(_permits: int) -> None: + rate_limit_started.set() + await release_rate_limit.wait() + + dispatcher._wait_for_rate_limit = wait_for_rate_limit # type: ignore[method-assign] + group_id = uuid.uuid4() + group = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=1, + target_rollouts=1, + policy_version_at_start=4, + uses_mutable_policy=False, + ) + dispatcher.groups[group_id] = group + epoch = await dispatcher.policy_gate.scheduling_epoch() + assert epoch is not None + + scheduling = asyncio.create_task(dispatcher.schedule_group_rollout(group_id, group, epoch=epoch)) + await rate_limit_started.wait() + + assert group.uses_mutable_policy + await asyncio.wait_for(dispatcher.on_version_pending(5), timeout=0.5) + assert group_id not in dispatcher.groups + + release_rate_limit.set() + assert not await scheduling + await dispatcher.on_new_version(5) diff --git a/tests/unit/test_transport_config.py b/tests/unit/test_transport_config.py new file mode 100644 index 0000000000..9a3f60fbaa --- /dev/null +++ b/tests/unit/test_transport_config.py @@ -0,0 +1,21 @@ +import math + +import pytest +from pydantic import ValidationError + +from prime_rl.configs.shared import FileSystemTransportConfig, ZMQTransportConfig + + +@pytest.mark.parametrize("config_type", [FileSystemTransportConfig, ZMQTransportConfig]) +def test_rollout_send_timeout_has_finite_positive_default(config_type): + timeout = config_type().send_timeout_seconds + + assert timeout == 300.0 + assert math.isfinite(timeout) + + +@pytest.mark.parametrize("config_type", [FileSystemTransportConfig, ZMQTransportConfig]) +@pytest.mark.parametrize("timeout", [0, -1, float("inf"), float("-inf"), float("nan")]) +def test_rollout_send_timeout_rejects_nonpositive_or_nonfinite_values(config_type, timeout: float): + with pytest.raises(ValidationError): + config_type(send_timeout_seconds=timeout) diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index cecc91b24b..7053bf69aa 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -18,6 +18,7 @@ setup_clients, setup_inference_pool, ) +from prime_rl.utils.policy_client_config import policy_client_config_from_environment def test_is_retryable_lora_error_returns_true_for_404(): @@ -152,6 +153,36 @@ def test_setup_inference_pool_selects_dynamo_pool_once(): asyncio.run(pool.stop()) +def test_generated_dgd_topology_selects_dynamo_pool_without_mutating_config(monkeypatch: pytest.MonkeyPatch): + client_config = ClientConfig() + original_config = client_config.model_copy(deep=True) + monkeypatch.setenv( + "DYN_RL_TOPOLOGY", + """{"schema_version":1,"admin_api":"dynamo","base_url":["http://frontend:8000/v1"],"rl_base_url":["http://frontend-rl:8001"],"dynamo_worker_roles":["prefill","decode"],"dynamo_gpus_per_worker":1}""", + ) + + resolved = policy_client_config_from_environment(client_config) + pool = asyncio.run(setup_inference_pool(resolved, model_name="test-model")) + + assert isinstance(pool, DynamoInferencePool) + assert pool.admin_api == "dynamo" + assert resolved.base_url == ["http://frontend:8000/v1"] + assert resolved.rl_base_url == ["http://frontend-rl:8001"] + assert client_config == original_config + asyncio.run(pool.stop()) + + +def test_generated_dgd_topology_rejects_explicit_client_conflict(monkeypatch: pytest.MonkeyPatch): + client_config = ClientConfig(admin_api="vllm") + monkeypatch.setenv( + "DYN_RL_TOPOLOGY", + """{"schema_version":1,"admin_api":"dynamo","base_url":["http://frontend:8000/v1"],"rl_base_url":["http://frontend-rl:8001"],"dynamo_worker_roles":["agg"],"dynamo_gpus_per_worker":1}""", + ) + + with pytest.raises(ValueError, match="admin_api.*conflicts"): + policy_client_config_from_environment(client_config) + + @pytest.mark.asyncio async def test_setup_inference_pool_rejects_dynamo_elastic_before_pool_selection(): client_config = ClientConfig( From f077fbf9620f144ec2cd5d74c4dc0eb147f5ceb2 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sat, 11 Jul 2026 03:43:44 -0700 Subject: [PATCH 18/18] fix(orchestrator): adapt policy markers to TraceTask --- src/prime_rl/orchestrator/dispatcher_transactions.py | 2 +- tests/unit/orchestrator/test_weight_update_barrier.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher_transactions.py b/src/prime_rl/orchestrator/dispatcher_transactions.py index d9e4c1cee7..fd92956c6c 100644 --- a/src/prime_rl/orchestrator/dispatcher_transactions.py +++ b/src/prime_rl/orchestrator/dispatcher_transactions.py @@ -110,7 +110,7 @@ async def emit_policy_cancellation_markers( ) for _ in range(owed): rollout = Rollout( - task=vf.Task(idx=group.task_idx, prompt=None), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=group.task_idx, prompt=None)), errors=[vf.Error(type="Cancelled", message="Policy update barrier")], stop_condition="error", kind=meta.kind, diff --git a/tests/unit/orchestrator/test_weight_update_barrier.py b/tests/unit/orchestrator/test_weight_update_barrier.py index 7ab0bc145a..637f3b988a 100644 --- a/tests/unit/orchestrator/test_weight_update_barrier.py +++ b/tests/unit/orchestrator/test_weight_update_barrier.py @@ -408,7 +408,7 @@ async def test_policy_barrier_waits_for_completed_handler_before_computing_marke async def completed_group() -> list[Rollout]: return [ Rollout( - task=vf.Task(idx=0, prompt=None), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), errors=[vf.Error(type="Existing", message="result")], stop_condition="error", )