From 3690dc064da218afdfcc8298e2e5b4b3725c8678 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 3 Jul 2026 12:24:02 -0700 Subject: [PATCH 01/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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", ) From d5e8bdf110dad826e87c7b908a6fb0d38b362e61 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sat, 4 Jul 2026 20:41:26 -0700 Subject: [PATCH 19/30] feat(k8s): render Dynamo graph deployments Signed-off-by: Biswa Panda --- k8s/prime-rl/templates/_helpers.tpl | 17 + k8s/prime-rl/templates/deployment.yaml | 32 +- .../templates/dynamo-engine-config.yaml | 18 + .../templates/dynamo-graph-deployment.yaml | 3 + k8s/prime-rl/templates/dynamo-rl-service.yaml | 23 ++ k8s/prime-rl/templates/service.yaml | 4 +- k8s/prime-rl/values.schema.json | 15 + k8s/prime-rl/values.yaml | 6 + pyproject.toml | 1 + src/prime_rl/inference/dgd.py | 317 ++++++++++++++++++ tests/unit/inference/test_dgd.py | 135 ++++++++ tests/unit/inference/test_helm_dgd.py | 112 +++++++ tests/unit/test_deployment_assets.py | 5 + 13 files changed, 665 insertions(+), 23 deletions(-) create mode 100644 k8s/prime-rl/templates/dynamo-engine-config.yaml create mode 100644 k8s/prime-rl/templates/dynamo-graph-deployment.yaml create mode 100644 k8s/prime-rl/templates/dynamo-rl-service.yaml create mode 100644 k8s/prime-rl/values.schema.json create mode 100644 src/prime_rl/inference/dgd.py create mode 100644 tests/unit/inference/test_dgd.py create mode 100644 tests/unit/inference/test_helm_dgd.py create mode 100644 tests/unit/test_deployment_assets.py diff --git a/k8s/prime-rl/templates/_helpers.tpl b/k8s/prime-rl/templates/_helpers.tpl index cd6d7183d5..fbefa6eac0 100644 --- a/k8s/prime-rl/templates/_helpers.tpl +++ b/k8s/prime-rl/templates/_helpers.tpl @@ -5,6 +5,23 @@ Expand the name of the chart. {{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }} {{- end }} +{{- define "prime-rl.inferenceUrls" -}} +{{- if eq .Values.inference.mode "dynamoGraph" -}} +{{- printf "http://%s-frontend.%s.svc.cluster.local:8000/v1" .Release.Name .Values.namespace -}} +{{- else -}} +{{- $releaseName := .Release.Name -}} +{{- $namespace := .Values.namespace -}} +{{- $port := int .Values.inference.service.port -}} +{{- $replicas := int .Values.inference.replicas -}} +{{- $urls := list -}} +{{- range $i := until $replicas -}} +{{- $url := printf "http://%s-inference-%d.%s-inference-headless.%s.svc.cluster.local:%d/v1" $releaseName $i $releaseName $namespace $port -}} +{{- $urls = append $urls $url -}} +{{- end -}} +{{- $urls | join "," -}} +{{- end -}} +{{- end }} + {{/* Create a default fully qualified app name. */}} diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index 1ed16b9f5e..9b269fc0a9 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -62,16 +62,11 @@ spec: - name: HEADLESS_SERVICE value: "{{ .Release.Name }}-orchestrator-headless.{{ .Values.namespace }}.svc.cluster.local" - name: INFERENCE_URL - {{- $releaseName := .Release.Name }} - {{- $namespace := .Values.namespace }} - {{- $port := int .Values.inference.service.port }} - {{- $replicas := int .Values.inference.replicas }} - {{- $urls := list }} - {{- range $i := until $replicas }} - {{- $url := printf "http://%s-inference-%d.%s-inference-headless.%s.svc.cluster.local:%d/v1" $releaseName $i $releaseName $namespace $port }} - {{- $urls = append $urls $url }} - {{- end }} - value: {{ $urls | join "," | quote }} + value: {{ include "prime-rl.inferenceUrls" . | quote }} + {{- if eq .Values.inference.mode "dynamoGraph" }} + - name: DYN_RL_DISCOVERY_URL + value: "http://{{ .Release.Name }}-frontend-rl.{{ .Values.namespace }}.svc.cluster.local:8001" + {{- end }} {{- with .Values.orchestrator.env }} {{- toYaml . | nindent 8 }} {{- end }} @@ -104,7 +99,7 @@ spec: {{- end }} {{- end }} --- -{{- if .Values.inference.enabled }} +{{- if and .Values.inference.enabled (eq .Values.inference.mode "statefulset") }} apiVersion: apps/v1 kind: StatefulSet metadata: @@ -301,16 +296,11 @@ spec: value: {{ .Values.trainer.pytorchCudaAllocConf | quote }} {{- end }} - name: INFERENCE_URL - {{- $releaseName := .Release.Name }} - {{- $namespace := .Values.namespace }} - {{- $port := int .Values.inference.service.port }} - {{- $replicas := int .Values.inference.replicas }} - {{- $urls := list }} - {{- range $i := until $replicas }} - {{- $url := printf "http://%s-inference-%d.%s-inference-headless.%s.svc.cluster.local:%d/v1" $releaseName $i $releaseName $namespace $port }} - {{- $urls = append $urls $url }} - {{- end }} - value: {{ $urls | join "," | quote }} + value: {{ include "prime-rl.inferenceUrls" . | quote }} + {{- if eq .Values.inference.mode "dynamoGraph" }} + - name: DYN_RL_DISCOVERY_URL + value: "http://{{ .Release.Name }}-frontend-rl.{{ .Values.namespace }}.svc.cluster.local:8001" + {{- end }} {{- with .Values.trainer.env }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/k8s/prime-rl/templates/dynamo-engine-config.yaml b/k8s/prime-rl/templates/dynamo-engine-config.yaml new file mode 100644 index 0000000000..3ac896bc57 --- /dev/null +++ b/k8s/prime-rl/templates/dynamo-engine-config.yaml @@ -0,0 +1,18 @@ +{{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ required "inference.dynamoGraph.engineConfig.name is required" .Values.inference.dynamoGraph.engineConfig.name }} + namespace: {{ .Values.namespace }} + labels: + {{- include "prime-rl.labels" . | nindent 4 }} + role: inference + annotations: + {{- toYaml .Values.inference.dynamoGraph.engineConfig.annotations | nindent 4 }} +immutable: true +data: + {{- range $name, $content := .Values.inference.dynamoGraph.engineConfig.data }} + {{ $name }}: | +{{ $content | nindent 4 }} + {{- end }} +{{- end }} diff --git a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml new file mode 100644 index 0000000000..a29e592e20 --- /dev/null +++ b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml @@ -0,0 +1,3 @@ +{{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} +{{- toYaml (required "inference.dynamoGraph.resource is required" .Values.inference.dynamoGraph.resource) }} +{{- end }} diff --git a/k8s/prime-rl/templates/dynamo-rl-service.yaml b/k8s/prime-rl/templates/dynamo-rl-service.yaml new file mode 100644 index 0000000000..371f5a3d08 --- /dev/null +++ b/k8s/prime-rl/templates/dynamo-rl-service.yaml @@ -0,0 +1,23 @@ +{{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} +apiVersion: v1 +kind: Service +metadata: + name: {{ .Release.Name }}-frontend-rl + namespace: {{ .Values.namespace }} + labels: + {{- include "prime-rl.labels" . | nindent 4 }} + role: inference + annotations: + {{- toYaml .Values.inference.dynamoGraph.resource.metadata.annotations | nindent 4 }} +spec: + type: ClusterIP + selector: + nvidia.com/dynamo-graph-deployment-name: {{ .Release.Name }} + nvidia.com/dynamo-component: Frontend + nvidia.com/dynamo-component-type: frontend + ports: + - name: rl + protocol: TCP + port: 8001 + targetPort: rl +{{- end }} diff --git a/k8s/prime-rl/templates/service.yaml b/k8s/prime-rl/templates/service.yaml index b783a7816d..6f7bf66be9 100644 --- a/k8s/prime-rl/templates/service.yaml +++ b/k8s/prime-rl/templates/service.yaml @@ -46,7 +46,7 @@ spec: name: nccl {{- end }} --- -{{- if .Values.inference.enabled }} +{{- if and .Values.inference.enabled (eq .Values.inference.mode "statefulset") }} {{- if .Values.inference.service.enabled }} apiVersion: v1 kind: Service @@ -98,7 +98,7 @@ spec: {{- end }} {{- end }} --- -{{- if .Values.inference.enabled }} +{{- if and .Values.inference.enabled (eq .Values.inference.mode "statefulset") }} apiVersion: v1 kind: Service metadata: diff --git a/k8s/prime-rl/values.schema.json b/k8s/prime-rl/values.schema.json new file mode 100644 index 0000000000..d7bbb88f22 --- /dev/null +++ b/k8s/prime-rl/values.schema.json @@ -0,0 +1,15 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "inference": { + "type": "object", + "properties": { + "mode": { + "type": "string", + "enum": ["statefulset", "dynamoGraph"] + } + } + } + } +} diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml index 7fc3b48989..56e52b974b 100644 --- a/k8s/prime-rl/values.yaml +++ b/k8s/prime-rl/values.yaml @@ -53,8 +53,14 @@ orchestrator: # Inference component inference: enabled: true + mode: statefulset replicas: 1 + # Generated by `dynamo-dgd` when mode is dynamoGraph. + dynamoGraph: + engineConfig: {} + resource: {} + # Auto-start configuration (set to false to use sleep infinity for debugging) autoStart: false command: "" # e.g., "uv run inference @ /app/examples/reverse_text/rl/infer.toml" diff --git a/pyproject.toml b/pyproject.toml index 0daa7259ab..440f0f8808 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ dependencies = [ rl = "prime_rl.entrypoints.rl:main" sft = "prime_rl.entrypoints.sft:main" inference = "prime_rl.entrypoints.inference:main" +dynamo-dgd = "prime_rl.inference.dgd:main" trainer = "prime_rl.entrypoints.trainer:main" orchestrator = "prime_rl.entrypoints.orchestrator:main" env-server = "prime_rl.orchestrator.env_server.env_server:main" diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py new file mode 100644 index 0000000000..b286c2d6e7 --- /dev/null +++ b/src/prime_rl/inference/dgd.py @@ -0,0 +1,317 @@ +"""Compile a Prime Dynamo inference config into Helm DGD values.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from prime_rl.configs.inference import DisaggregatedInferenceDeploymentConfig, InferenceConfig +from prime_rl.inference.dynamo import ( + DynamoProcessSpec, + build_frontend_process, + build_worker_process, + write_role_engine_configs, +) +from prime_rl.utils.config import cli + +ENGINE_MOUNT_PATH = "/etc/prime-rl/dynamo" + + +@dataclass(frozen=True) +class DynamoGraphRenderOptions: + release_name: str + namespace: str + image: str + output_dir: Path + prime_sha: str + dynamo_sha: str + image_digest: str + run_name: str + model_cache_pvc: str | None = None + shared_pvc: str | None = None + image_pull_secrets: tuple[str, ...] = () + hf_token_secret: str | None = None + + def __post_init__(self) -> None: + for name, value in (("prime_sha", self.prime_sha), ("dynamo_sha", self.dynamo_sha)): + if re.fullmatch(r"[0-9a-f]{40}", value) is None: + raise ValueError(f"{name} must be a full 40-character Git commit SHA") + if re.fullmatch(r"sha256:[0-9a-f]{64}", self.image_digest) is None: + raise ValueError("image_digest must be a full sha256 digest") + if not self.image.endswith(f"@{self.image_digest}"): + raise ValueError("DGD image must be pinned to image_digest") + image_tag = self.image.rsplit("@", 1)[0] + if self.prime_sha[:12] not in image_tag or self.dynamo_sha[:12] not in image_tag: + raise ValueError("DGD image tag must include the Prime and Dynamo commit suffixes") + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _canonical_json(value: Any) -> bytes: + return (json.dumps(value, indent=2, sort_keys=True) + "\n").encode() + + +def _worker_env(process: DynamoProcessSpec) -> list[dict[str, Any]]: + values: list[dict[str, Any]] = [ + {"name": name, "value": value} for name, value in sorted(process.environment().items()) + ] + values.append( + { + "name": "VLLM_NIXL_SIDE_CHANNEL_HOST", + "valueFrom": {"fieldRef": {"fieldPath": "status.podIP"}}, + } + ) + return values + + +def _apply_pod_credentials( + pod_spec: dict[str, Any], + container: dict[str, Any], + options: DynamoGraphRenderOptions, +) -> None: + if options.image_pull_secrets: + pod_spec["imagePullSecrets"] = [{"name": name} for name in options.image_pull_secrets] + if options.hf_token_secret and not any(item["name"] == "HF_TOKEN" for item in container.get("env", [])): + container.setdefault("env", []).append( + { + "name": "HF_TOKEN", + "valueFrom": { + "secretKeyRef": { + "name": options.hf_token_secret, + "key": "HF_TOKEN", + "optional": True, + } + }, + } + ) + + +def _worker_service( + config: InferenceConfig, + options: DynamoGraphRenderOptions, + *, + role: str, + replicas: int, + config_map_name: str, + engine_file: str, +) -> dict[str, Any]: + assert config.deployment.type == "disaggregated" + process = build_worker_process( + config, + role, + Path(ENGINE_MOUNT_PATH) / engine_file, + nixl_host=None, + nixl_port=20100, + ) + container = { + "image": options.image, + "imagePullPolicy": "IfNotPresent", + "command": ["python3", "-m", process.module], + "args": list(process.arguments), + "env": _worker_env(process), + "volumeMounts": [ + { + "name": "dynamo-engine-config", + "mountPath": ENGINE_MOUNT_PATH, + "readOnly": True, + } + ], + } + pod_spec = { + "runtimeClassName": "nvidia", + "tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists", "effect": "NoSchedule"}], + "volumes": [ + { + "name": "dynamo-engine-config", + "configMap": {"name": config_map_name}, + } + ], + "mainContainer": container, + } + _apply_pod_credentials(pod_spec, container, options) + return { + "componentType": "worker", + "subComponentType": role, + "replicas": replicas, + "sharedMemory": {"size": "64Gi"}, + "resources": { + "requests": {"gpu": str(config.deployment.gpus_per_node)}, + "limits": {"gpu": str(config.deployment.gpus_per_node)}, + }, + "extraPodSpec": pod_spec, + } + + +def _add_pvc(resource: dict[str, Any], service: dict[str, Any], name: str | None, mount_point: str) -> None: + if not name: + return + pvcs = resource["spec"].setdefault("pvcs", []) + if not any(pvc["name"] == name for pvc in pvcs): + pvcs.append({"name": name, "create": False}) + service.setdefault("volumeMounts", []).append({"name": name, "mountPoint": mount_point}) + + +def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) -> dict[str, Any]: + if config.backend.type != "dynamo" or config.deployment.type != "disaggregated": + raise ValueError("DGD rendering requires a Dynamo disaggregated inference config") + deployment: DisaggregatedInferenceDeploymentConfig = config.deployment + if deployment.num_prefill_nodes != deployment.num_prefill_replicas: + raise ValueError("DGD rendering currently requires one pod per prefill replica") + if deployment.num_decode_nodes != deployment.num_decode_replicas: + raise ValueError("DGD rendering currently requires one pod per decode replica") + + engine_paths = write_role_engine_configs(config, options.output_dir) + prefill_text = engine_paths["prefill"].read_text() + decode_text = engine_paths["decode"].read_text() + engine_hash = _sha256_bytes((prefill_text + decode_text).encode()) + config_map_name = f"{options.release_name}-dynamo-engine-{engine_hash[:12]}" + + annotations = { + "prime-rl.nvidia.com/config-sha256": engine_hash, + "prime-rl.nvidia.com/dynamo-sha": options.dynamo_sha, + "prime-rl.nvidia.com/image-digest": options.image_digest, + "prime-rl.nvidia.com/prime-sha": options.prime_sha, + "prime-rl.nvidia.com/run-name": options.run_name, + } + frontend_process = build_frontend_process(config, host="0.0.0.0", port=8000) + frontend_container = { + "image": options.image, + "imagePullPolicy": "IfNotPresent", + "command": ["python3", "-m", frontend_process.module], + "args": list(frontend_process.arguments), + "env": [{"name": name, "value": value} for name, value in sorted(frontend_process.environment().items())], + "ports": [{"containerPort": 8001, "name": "rl"}], + } + frontend_pod_spec = {"mainContainer": frontend_container} + _apply_pod_credentials(frontend_pod_spec, frontend_container, options) + frontend = { + "componentType": "frontend", + "replicas": 1, + "extraPodSpec": frontend_pod_spec, + } + prefill = _worker_service( + config, + options, + role="prefill", + replicas=deployment.num_prefill_replicas, + config_map_name=config_map_name, + engine_file="prefill-engine.json", + ) + decode = _worker_service( + config, + options, + role="decode", + replicas=deployment.num_decode_replicas, + config_map_name=config_map_name, + engine_file="decode-engine.json", + ) + resource: dict[str, Any] = { + "apiVersion": "nvidia.com/v1alpha1", + "kind": "DynamoGraphDeployment", + "metadata": { + "name": options.release_name, + "namespace": options.namespace, + "annotations": annotations, + }, + "spec": { + "backendFramework": "vllm", + "services": { + "Frontend": frontend, + "VllmDecodeWorker": decode, + "VllmPrefillWorker": prefill, + }, + }, + } + for service in (frontend, prefill, decode): + _add_pvc(resource, service, options.model_cache_pvc, "/model-cache") + _add_pvc(resource, service, options.shared_pvc, "/data") + + manifest_hash = _sha256_bytes(_canonical_json(resource)) + annotations["prime-rl.nvidia.com/manifest-sha256"] = manifest_hash + return { + "namespace": options.namespace, + "inference": { + "mode": "dynamoGraph", + "dynamoGraph": { + "engineConfig": { + "name": config_map_name, + "sha256": engine_hash, + "annotations": annotations, + "data": { + "prefill-engine.json": prefill_text, + "decode-engine.json": decode_text, + }, + }, + "resource": resource, + }, + }, + } + + +def write_dgd_artifacts(config: InferenceConfig, options: DynamoGraphRenderOptions) -> dict[str, Path]: + options.output_dir.mkdir(parents=True, exist_ok=True) + values = build_dgd_values(config, options) + resource = values["inference"]["dynamoGraph"]["resource"] + paths = { + "values": options.output_dir / "dynamo-helm-values.json", + "resource": options.output_dir / "dynamo-graph-deployment.json", + } + paths["values"].write_bytes(_canonical_json(values)) + paths["resource"].write_bytes(_canonical_json(resource)) + manifest_entries = [] + for path in sorted(options.output_dir.glob("*.json")): + manifest_entries.append(f"{_sha256_bytes(path.read_bytes())} {path.name}") + manifest = options.output_dir / "artifact-manifest.sha256" + manifest.write_text("\n".join(manifest_entries) + "\n") + paths["manifest"] = manifest + return paths + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("inference_config", type=Path) + parser.add_argument("--release-name", required=True) + parser.add_argument("--namespace", required=True) + parser.add_argument("--image", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--prime-sha", required=True) + parser.add_argument("--dynamo-sha", required=True) + parser.add_argument("--image-digest", required=True) + parser.add_argument("--run-name") + parser.add_argument("--model-cache-pvc") + parser.add_argument("--shared-pvc") + parser.add_argument("--image-pull-secret", action="append", default=[]) + parser.add_argument("--hf-token-secret") + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + config = cli(InferenceConfig, args=["@", str(args.inference_config)]) + options = DynamoGraphRenderOptions( + release_name=args.release_name, + namespace=args.namespace, + image=args.image, + output_dir=args.output_dir, + prime_sha=args.prime_sha, + dynamo_sha=args.dynamo_sha, + image_digest=args.image_digest, + run_name=args.run_name or args.release_name, + model_cache_pvc=args.model_cache_pvc, + shared_pvc=args.shared_pvc, + image_pull_secrets=tuple(args.image_pull_secret), + hf_token_secret=args.hf_token_secret, + ) + for path in write_dgd_artifacts(config, options).values(): + print(path) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py new file mode 100644 index 0000000000..e9a6421b8a --- /dev/null +++ b/tests/unit/inference/test_dgd.py @@ -0,0 +1,135 @@ +import hashlib +import json +from pathlib import Path + +import pytest + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference.dgd import DynamoGraphRenderOptions, build_dgd_values, write_dgd_artifacts +from prime_rl.inference.dynamo import build_frontend_process, build_worker_process + +PRIME_SHA = "1" * 40 +DYNAMO_SHA = "2" * 40 +IMAGE_DIGEST = f"sha256:{'3' * 64}" + + +def inference_config() -> InferenceConfig: + return InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "model": {"name": "Qwen/Qwen3-30B-A3B-Thinking-2507"}, + "parallel": {"tp": 1}, + "weight_broadcast": {"type": "nccl"}, + "env_vars": {"HF_HOME": "/model-cache", "HF_HUB_OFFLINE": "1"}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_nodes": 2, + "num_decode_nodes": 2, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + ) + + +def render_options(tmp_path: Path) -> DynamoGraphRenderOptions: + return DynamoGraphRenderOptions( + release_name="p4-math", + namespace="bis-vllm", + image=f"nvcr.io/example/prime:prime-{PRIME_SHA[:12]}-dynamo-{DYNAMO_SHA[:12]}@{IMAGE_DIGEST}", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + model_cache_pvc="model-cache", + shared_pvc="p4-shared-data", + image_pull_secrets=("nvcrimagepullsecret",), + hf_token_secret="hf-token-secret", + ) + + +def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): + values = build_dgd_values(inference_config(), render_options(tmp_path)) + graph = values["inference"]["dynamoGraph"] + resource = graph["resource"] + services = resource["spec"]["services"] + + assert resource["apiVersion"] == "nvidia.com/v1alpha1" + assert resource["metadata"]["namespace"] == "bis-vllm" + assert services["VllmPrefillWorker"]["replicas"] == 2 + assert services["VllmDecodeWorker"]["replicas"] == 2 + assert services["VllmPrefillWorker"]["resources"]["limits"]["gpu"] == "1" + assert services["VllmDecodeWorker"]["resources"]["limits"]["gpu"] == "1" + assert resource["spec"]["pvcs"] == [ + {"name": "model-cache", "create": False}, + {"name": "p4-shared-data", "create": False}, + ] + for service in services.values(): + pod_spec = service["extraPodSpec"] + assert pod_spec["imagePullSecrets"] == [{"name": "nvcrimagepullsecret"}] + assert any(item["name"] == "HF_TOKEN" for item in pod_spec["mainContainer"]["env"]) + env = {item["name"]: item.get("value") for item in pod_spec["mainContainer"]["env"]} + assert env["HF_HOME"] == "/model-cache" + assert env["HF_HUB_OFFLINE"] == "1" + + gpu_toleration = {"key": "nvidia.com/gpu", "operator": "Exists", "effect": "NoSchedule"} + assert gpu_toleration in services["VllmPrefillWorker"]["extraPodSpec"]["tolerations"] + assert gpu_toleration in services["VllmDecodeWorker"]["extraPodSpec"]["tolerations"] + + prefill_args = services["VllmPrefillWorker"]["extraPodSpec"]["mainContainer"]["args"] + decode_args = services["VllmDecodeWorker"]["extraPodSpec"]["mainContainer"]["args"] + prefill_process = build_worker_process( + inference_config(), + "prefill", + Path("/etc/prime-rl/dynamo/prefill-engine.json"), + nixl_host=None, + nixl_port=20100, + ) + decode_process = build_worker_process( + inference_config(), + "decode", + Path("/etc/prime-rl/dynamo/decode-engine.json"), + nixl_host=None, + nixl_port=20100, + ) + frontend_process = build_frontend_process(inference_config(), host="0.0.0.0", port=8000) + assert prefill_args == list(prefill_process.arguments) + assert decode_args == list(decode_process.arguments) + assert services["Frontend"]["extraPodSpec"]["mainContainer"]["args"] == list(frontend_process.arguments) + + prefill_env = { + item["name"]: item.get("value") + for item in services["VllmPrefillWorker"]["extraPodSpec"]["mainContainer"]["env"] + } + frontend_env = { + item["name"]: item.get("value") for item in services["Frontend"]["extraPodSpec"]["mainContainer"]["env"] + } + assert {key: prefill_env[key] for key in prefill_process.environment()} == prefill_process.environment() + assert {key: frontend_env[key] for key in frontend_process.environment()} == frontend_process.environment() + + prefill = json.loads(graph["engineConfig"]["data"]["prefill-engine.json"]) + decode = json.loads(graph["engineConfig"]["data"]["decode-engine.json"]) + assert prefill["kv_transfer_config"] == decode["kv_transfer_config"] + assert "kv_events_config" in prefill + assert "kv_events_config" not in decode + assert "disaggregation_mode" not in prefill + assert "enable_rl" not in decode + + +def test_dgd_artifacts_are_deterministic_and_manifest_verifies(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + first_values = paths["values"].read_bytes() + write_dgd_artifacts(inference_config(), render_options(tmp_path)) + assert paths["values"].read_bytes() == first_values + + for line in paths["manifest"].read_text().splitlines(): + expected, name = line.split(" ", 1) + assert hashlib.sha256((tmp_path / name).read_bytes()).hexdigest() == expected + + +def test_dgd_rejects_native_backend(tmp_path: Path): + config = InferenceConfig.model_validate({}) + with pytest.raises(ValueError, match="Dynamo disaggregated"): + build_dgd_values(config, render_options(tmp_path)) diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py new file mode 100644 index 0000000000..7d74e2b8bb --- /dev/null +++ b/tests/unit/inference/test_helm_dgd.py @@ -0,0 +1,112 @@ +import json +import shutil +import subprocess +from pathlib import Path + +import pytest + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference.dgd import DynamoGraphRenderOptions, write_dgd_artifacts + +HELM = shutil.which("helm") +CHART = Path(__file__).parents[3] / "k8s" / "prime-rl" +PRIME_SHA = "1" * 40 +DYNAMO_SHA = "2" * 40 +IMAGE_DIGEST = f"sha256:{'3' * 64}" + + +def inference_config() -> InferenceConfig: + return InferenceConfig.model_validate( + { + "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, + }, + } + ) + + +def render_options(tmp_path: Path) -> DynamoGraphRenderOptions: + return DynamoGraphRenderOptions( + release_name="p4-math", + namespace="bis-vllm", + image=f"nvcr.io/example/prime:prime-{PRIME_SHA[:12]}-dynamo-{DYNAMO_SHA[:12]}@{IMAGE_DIGEST}", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + ) + + +def helm_template(*args: str) -> str: + if HELM is None: + pytest.skip("helm is not installed") + return subprocess.run( + [HELM, "template", "p4-math", str(CHART), *args], + check=True, + capture_output=True, + text=True, + ).stdout + + +def test_native_chart_still_renders_inference_statefulset(): + rendered = helm_template() + assert "name: p4-math-inference\n" in rendered + assert "kind: DynamoGraphDeployment" not in rendered + assert "p4-math-inference-0.p4-math-inference-headless" in rendered + + +def test_chart_rejects_unknown_inference_mode(): + with pytest.raises(subprocess.CalledProcessError): + helm_template("--set", "inference.mode=typo") + + +def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + rendered = helm_template("-f", str(paths["values"])) + graph = json.loads(paths["resource"].read_text()) + + assert rendered.count("kind: DynamoGraphDeployment") == 1 + assert rendered.count("kind: ConfigMap") == 1 + assert "name: p4-math-inference\n" not in rendered + assert "name: p4-math-frontend-rl" in rendered + assert "http://p4-math-frontend.bis-vllm.svc.cluster.local:8000/v1" in rendered + assert "http://p4-math-frontend-rl.bis-vllm.svc.cluster.local:8001" in rendered + assert graph["spec"]["services"]["VllmPrefillWorker"]["replicas"] == 2 + assert graph["spec"]["services"]["VllmDecodeWorker"]["replicas"] == 2 + assert not any(kind in rendered for kind in ("kind: ClusterRole", "kind: CustomResourceDefinition")) + + +def test_dgd_rejects_image_without_matching_digest(tmp_path: Path): + with pytest.raises(ValueError, match="must be pinned"): + DynamoGraphRenderOptions( + release_name="p4-math", + namespace="bis-vllm", + image="nvcr.io/example/prime:p4", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + ) + + +def test_dgd_rejects_image_without_commit_suffixes(tmp_path: Path): + with pytest.raises(ValueError, match="commit suffixes"): + DynamoGraphRenderOptions( + release_name="p4-math", + namespace="bis-vllm", + image=f"nvcr.io/example/prime:p4@{IMAGE_DIGEST}", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + ) diff --git a/tests/unit/test_deployment_assets.py b/tests/unit/test_deployment_assets.py new file mode 100644 index 0000000000..5768f5bb67 --- /dev/null +++ b/tests/unit/test_deployment_assets.py @@ -0,0 +1,5 @@ +from pathlib import Path + + +def test_chart_does_not_duplicate_prime_runtime_image(): + assert not (Path(__file__).parents[2] / "Dockerfile.dynamo").exists() From 4d710d11dba98962ba768a8932691fd45a0713f2 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 15:24:20 -0700 Subject: [PATCH 20/30] test: update DynamoGraph topology fixtures --- tests/unit/inference/test_dgd.py | 4 ++-- tests/unit/inference/test_helm_dgd.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index e9a6421b8a..66370f9341 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -24,8 +24,8 @@ def inference_config() -> 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, }, diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index 7d74e2b8bb..3eb2c5a32d 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -23,8 +23,8 @@ def inference_config() -> 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, }, From 6940e64f261a102cf434072197feb94c59319698 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 15:54:16 -0700 Subject: [PATCH 21/30] fix(k8s): pin DynamoGraph runtime resources --- k8s/prime-rl/templates/_helpers.tpl | 19 ++++++ k8s/prime-rl/templates/deployment.yaml | 30 +++++++-- .../templates/dynamo-graph-deployment.yaml | 17 ++++- k8s/prime-rl/templates/pvc.yaml | 2 +- k8s/prime-rl/values.schema.json | 56 +++++++++++++++- k8s/prime-rl/values.yaml | 7 +- src/prime_rl/inference/dgd.py | 53 +++++++++++++-- tests/unit/inference/test_dgd.py | 64 ++++++++++++++++-- tests/unit/inference/test_helm_dgd.py | 65 +++++++++++++++++-- 9 files changed, 288 insertions(+), 25 deletions(-) diff --git a/k8s/prime-rl/templates/_helpers.tpl b/k8s/prime-rl/templates/_helpers.tpl index fbefa6eac0..738aed204d 100644 --- a/k8s/prime-rl/templates/_helpers.tpl +++ b/k8s/prime-rl/templates/_helpers.tpl @@ -5,6 +5,25 @@ Expand the name of the chart. {{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }} {{- end }} +{{/* +Resolve the immutable image reference generated for DGD, with the native chart +repository/tag remaining as the fallback for statefulset mode. +*/}} +{{- define "prime-rl.image" -}} +{{- if .Values.image.reference -}} +{{- .Values.image.reference -}} +{{- else -}} +{{- printf "%s:%s" .Values.image.repository .Values.image.tag -}} +{{- end -}} +{{- end }} + +{{/* +Reuse a supplied shared claim or derive the chart-managed claim name. +*/}} +{{- define "prime-rl.storageClaimName" -}} +{{- default (printf "%s-shared-data" .Release.Name) .Values.storage.existingClaim -}} +{{- end }} + {{- define "prime-rl.inferenceUrls" -}} {{- if eq .Values.inference.mode "dynamoGraph" -}} {{- printf "http://%s-frontend.%s.svc.cluster.local:8000/v1" .Release.Name .Values.namespace -}} diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index 9b269fc0a9..a4a127b8bc 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -22,13 +22,19 @@ spec: {{- include "prime-rl.componentLabels" . | nindent 8 }} role: orchestrator spec: + {{- with .Values.image.pullSecrets }} + imagePullSecrets: + {{- range . }} + - name: {{ . | quote }} + {{- end }} + {{- end }} {{- with .Values.orchestrator.nodeSelector }} nodeSelector: {{- toYaml . | nindent 8 }} {{- end }} containers: - name: prime-rl-orchestrator - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + image: {{ include "prime-rl.image" . | quote }} imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.orchestrator.autoStart }} command: ["/bin/bash", "-c"] @@ -95,7 +101,7 @@ spec: volumes: - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ include "prime-rl.storageClaimName" . }} {{- end }} {{- end }} --- @@ -124,12 +130,18 @@ spec: {{- include "prime-rl.componentLabels" . | nindent 8 }} role: inference spec: + {{- with .Values.image.pullSecrets }} + imagePullSecrets: + {{- range . }} + - name: {{ . | quote }} + {{- end }} + {{- end }} {{- if .Values.inference.runtimeClassName }} runtimeClassName: {{ .Values.inference.runtimeClassName }} {{- end }} containers: - name: prime-rl-inference - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + image: {{ include "prime-rl.image" . | quote }} imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.inference.autoStart }} command: ["/bin/bash", "-c"] @@ -224,7 +236,7 @@ spec: volumes: - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ include "prime-rl.storageClaimName" . }} {{- end }} {{- end }} --- @@ -253,12 +265,18 @@ spec: {{- include "prime-rl.componentLabels" . | nindent 8 }} role: trainer spec: + {{- with .Values.image.pullSecrets }} + imagePullSecrets: + {{- range . }} + - name: {{ . | quote }} + {{- end }} + {{- end }} {{- if .Values.trainer.runtimeClassName }} runtimeClassName: {{ .Values.trainer.runtimeClassName }} {{- end }} containers: - name: prime-rl-trainer - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + image: {{ include "prime-rl.image" . | quote }} imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.trainer.autoStart }} command: ["/bin/bash", "-c"] @@ -365,6 +383,6 @@ spec: volumes: - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ include "prime-rl.storageClaimName" . }} {{- end }} {{- end }} diff --git a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml index a29e592e20..6446a088eb 100644 --- a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml +++ b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml @@ -1,3 +1,18 @@ {{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} -{{- toYaml (required "inference.dynamoGraph.resource is required" .Values.inference.dynamoGraph.resource) }} +{{- $resource := required "inference.dynamoGraph.resource is required" .Values.inference.dynamoGraph.resource }} +{{- $resourceName := required "inference.dynamoGraph.resource.metadata.name is required" $resource.metadata.name }} +{{- if ne .Release.Name $resourceName }} +{{- fail (printf "Helm Release.Name %q must match embedded DynamoGraphDeployment metadata.name %q" .Release.Name $resourceName) }} +{{- end }} +{{- $image := required "image.reference is required in dynamoGraph mode" .Values.image.reference }} +{{- if not (regexMatch "^.+@sha256:[0-9a-f]{64}$" $image) }} +{{- fail "image.reference must be pinned to a full sha256 digest in dynamoGraph mode" }} +{{- end }} +{{- range $serviceName, $service := $resource.spec.services }} +{{- $serviceImage := required (printf "DynamoGraphDeployment service %s image is required" $serviceName) $service.extraPodSpec.mainContainer.image }} +{{- if ne $image $serviceImage }} +{{- fail (printf "DynamoGraphDeployment service %s must use the same image.reference as orchestrator and trainer" $serviceName) }} +{{- end }} +{{- end }} +{{- toYaml $resource }} {{- end }} diff --git a/k8s/prime-rl/templates/pvc.yaml b/k8s/prime-rl/templates/pvc.yaml index 7afd6ff386..5bd740f7f6 100644 --- a/k8s/prime-rl/templates/pvc.yaml +++ b/k8s/prime-rl/templates/pvc.yaml @@ -1,4 +1,4 @@ -{{- if .Values.storage.enabled }} +{{- if and .Values.storage.enabled (not .Values.storage.existingClaim) }} apiVersion: v1 kind: PersistentVolumeClaim metadata: diff --git a/k8s/prime-rl/values.schema.json b/k8s/prime-rl/values.schema.json index d7bbb88f22..b494fd78aa 100644 --- a/k8s/prime-rl/values.schema.json +++ b/k8s/prime-rl/values.schema.json @@ -2,6 +2,22 @@ "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { + "image": { + "type": "object", + "properties": { + "reference": { + "type": "string" + }, + "pullSecrets": { + "type": "array", + "items": { + "type": "string", + "minLength": 1 + }, + "uniqueItems": true + } + } + }, "inference": { "type": "object", "properties": { @@ -10,6 +26,44 @@ "enum": ["statefulset", "dynamoGraph"] } } + }, + "storage": { + "type": "object", + "properties": { + "existingClaim": { + "type": "string" + } + } + } + }, + "allOf": [ + { + "if": { + "properties": { + "inference": { + "properties": { + "mode": { + "const": "dynamoGraph" + } + }, + "required": ["mode"] + } + }, + "required": ["inference"] + }, + "then": { + "properties": { + "image": { + "properties": { + "reference": { + "pattern": "^.+@sha256:[0-9a-f]{64}$" + } + }, + "required": ["reference"] + } + }, + "required": ["image"] + } } - } + ] } diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml index 56e52b974b..4bee20c572 100644 --- a/k8s/prime-rl/values.yaml +++ b/k8s/prime-rl/values.yaml @@ -10,11 +10,16 @@ image: repository: primeintellect/prime-rl pullPolicy: IfNotPresent tag: "main" + # Generated DGD values set the full reviewed repository:tag@sha256 reference. + reference: "" + pullSecrets: [] # Shared storage configuration storage: enabled: true - # PVC name will be automatically set to {{ .Release.Name }}-shared-data + # Reuse an existing ReadWriteMany claim instead of creating a release-owned PVC. + existingClaim: "" + # A chart-managed PVC is named {{ .Release.Name }}-shared-data when existingClaim is empty. storageClassName: nfs accessModes: - ReadWriteMany diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index b286c2d6e7..87f5204c1a 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -20,6 +20,12 @@ from prime_rl.utils.config import cli ENGINE_MOUNT_PATH = "/etc/prime-rl/dynamo" +MANIFEST_HASH_ANNOTATION = "prime-rl.nvidia.com/manifest-sha256" +MANIFEST_HASH_SCOPE_ANNOTATION = "prime-rl.nvidia.com/manifest-sha256-scope" +MANIFEST_HASH_SCOPE = ( + "resource; json.dumps(sort_keys=true,indent=2)+newline; " + "exclude=/metadata/annotations/prime-rl.nvidia.com~1manifest-sha256" +) @dataclass(frozen=True) @@ -58,6 +64,19 @@ def _canonical_json(value: Any) -> bytes: return (json.dumps(value, indent=2, sort_keys=True) + "\n").encode() +def _resource_manifest_hash(resource: dict[str, Any]) -> str: + annotations = resource["metadata"]["annotations"] + scoped_annotations = {key: value for key, value in annotations.items() if key != MANIFEST_HASH_ANNOTATION} + scoped_resource = { + **resource, + "metadata": { + **resource["metadata"], + "annotations": scoped_annotations, + }, + } + return _sha256_bytes(_canonical_json(scoped_resource)) + + def _worker_env(process: DynamoProcessSpec) -> list[dict[str, Any]]: values: list[dict[str, Any]] = [ {"name": name, "value": value} for name, value in sorted(process.environment().items()) @@ -166,6 +185,8 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) raise ValueError("DGD rendering currently requires one pod per prefill replica") if deployment.num_decode_nodes != deployment.num_decode_replicas: raise ValueError("DGD rendering currently requires one pod per decode replica") + if config.weight_broadcast.type == "filesystem" and not options.shared_pvc: + raise ValueError("Dynamo filesystem weight broadcast requires a shared existing PVC") engine_paths = write_role_engine_configs(config, options.output_dir) prefill_text = engine_paths["prefill"].read_text() @@ -179,6 +200,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "prime-rl.nvidia.com/image-digest": options.image_digest, "prime-rl.nvidia.com/prime-sha": options.prime_sha, "prime-rl.nvidia.com/run-name": options.run_name, + MANIFEST_HASH_SCOPE_ANNOTATION: MANIFEST_HASH_SCOPE, } frontend_process = build_frontend_process(config, host="0.0.0.0", port=8000) frontend_container = { @@ -231,12 +253,26 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) } for service in (frontend, prefill, decode): _add_pvc(resource, service, options.model_cache_pvc, "/model-cache") - _add_pvc(resource, service, options.shared_pvc, "/data") - - manifest_hash = _sha256_bytes(_canonical_json(resource)) - annotations["prime-rl.nvidia.com/manifest-sha256"] = manifest_hash - return { + if config.weight_broadcast.type == "filesystem": + for service in (prefill, decode): + _add_pvc(resource, service, options.shared_pvc, "/data") + + manifest_hash = _resource_manifest_hash(resource) + annotations = {**annotations, MANIFEST_HASH_ANNOTATION: manifest_hash} + resource = { + **resource, + "metadata": { + **resource["metadata"], + "annotations": annotations, + }, + } + values: dict[str, Any] = { "namespace": options.namespace, + "image": { + "reference": options.image, + "pullPolicy": "IfNotPresent", + "pullSecrets": list(options.image_pull_secrets), + }, "inference": { "mode": "dynamoGraph", "dynamoGraph": { @@ -253,6 +289,13 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) }, }, } + if options.shared_pvc: + values["storage"] = { + "enabled": True, + "existingClaim": options.shared_pvc, + "mountPath": "/data", + } + return values def write_dgd_artifacts(config: InferenceConfig, options: DynamoGraphRenderOptions) -> dict[str, Path]: diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index 66370f9341..5b3c4a732e 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -1,5 +1,7 @@ import hashlib import json +from copy import deepcopy +from dataclasses import replace from pathlib import Path import pytest @@ -13,13 +15,13 @@ IMAGE_DIGEST = f"sha256:{'3' * 64}" -def inference_config() -> InferenceConfig: +def inference_config(weight_broadcast: str = "nccl") -> InferenceConfig: return InferenceConfig.model_validate( { "backend": {"type": "dynamo"}, "model": {"name": "Qwen/Qwen3-30B-A3B-Thinking-2507"}, "parallel": {"tp": 1}, - "weight_broadcast": {"type": "nccl"}, + "weight_broadcast": {"type": weight_broadcast}, "env_vars": {"HF_HOME": "/model-cache", "HF_HUB_OFFLINE": "1"}, "deployment": { "type": "disaggregated", @@ -51,11 +53,17 @@ def render_options(tmp_path: Path) -> DynamoGraphRenderOptions: def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): - values = build_dgd_values(inference_config(), render_options(tmp_path)) + options = render_options(tmp_path) + values = build_dgd_values(inference_config(), options) graph = values["inference"]["dynamoGraph"] resource = graph["resource"] services = resource["spec"]["services"] + assert values["image"] == { + "reference": options.image, + "pullPolicy": "IfNotPresent", + "pullSecrets": ["nvcrimagepullsecret"], + } assert resource["apiVersion"] == "nvidia.com/v1alpha1" assert resource["metadata"]["namespace"] == "bis-vllm" assert services["VllmPrefillWorker"]["replicas"] == 2 @@ -64,10 +72,15 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): assert services["VllmDecodeWorker"]["resources"]["limits"]["gpu"] == "1" assert resource["spec"]["pvcs"] == [ {"name": "model-cache", "create": False}, - {"name": "p4-shared-data", "create": False}, ] + assert values["storage"] == { + "enabled": True, + "existingClaim": "p4-shared-data", + "mountPath": "/data", + } for service in services.values(): pod_spec = service["extraPodSpec"] + assert pod_spec["mainContainer"]["image"] == options.image assert pod_spec["imagePullSecrets"] == [{"name": "nvcrimagepullsecret"}] assert any(item["name"] == "HF_TOKEN" for item in pod_spec["mainContainer"]["env"]) env = {item["name"]: item.get("value") for item in pod_spec["mainContainer"]["env"]} @@ -109,6 +122,10 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): assert {key: prefill_env[key] for key in prefill_process.environment()} == prefill_process.environment() assert {key: frontend_env[key] for key in frontend_process.environment()} == frontend_process.environment() + assert {mount["name"] for mount in services["Frontend"]["volumeMounts"]} == {"model-cache"} + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + assert not any(mount["name"] == "p4-shared-data" for mount in services[role]["volumeMounts"]) + prefill = json.loads(graph["engineConfig"]["data"]["prefill-engine.json"]) decode = json.loads(graph["engineConfig"]["data"]["decode-engine.json"]) assert prefill["kv_transfer_config"] == decode["kv_transfer_config"] @@ -128,6 +145,45 @@ def test_dgd_artifacts_are_deterministic_and_manifest_verifies(tmp_path: Path): expected, name = line.split(" ", 1) assert hashlib.sha256((tmp_path / name).read_bytes()).hexdigest() == expected + resource = json.loads(paths["resource"].read_text()) + annotations = resource["metadata"]["annotations"] + assert annotations["prime-rl.nvidia.com/manifest-sha256-scope"] == ( + "resource; json.dumps(sort_keys=true,indent=2)+newline; " + "exclude=/metadata/annotations/prime-rl.nvidia.com~1manifest-sha256" + ) + expected_manifest_hash = annotations["prime-rl.nvidia.com/manifest-sha256"] + unhashed_resource = deepcopy(resource) + del unhashed_resource["metadata"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"] + canonical_resource = (json.dumps(unhashed_resource, indent=2, sort_keys=True) + "\n").encode() + assert hashlib.sha256(canonical_resource).hexdigest() == expected_manifest_hash + + +def test_filesystem_broadcast_requires_one_shared_existing_claim(tmp_path: Path): + options = render_options(tmp_path) + values = build_dgd_values(inference_config("filesystem"), options) + resource = values["inference"]["dynamoGraph"]["resource"] + services = resource["spec"]["services"] + + assert values["storage"] == { + "enabled": True, + "existingClaim": "p4-shared-data", + "mountPath": "/data", + } + assert resource["spec"]["pvcs"] == [ + {"name": "model-cache", "create": False}, + {"name": "p4-shared-data", "create": False}, + ] + assert not any(mount["name"] == "p4-shared-data" for mount in services["Frontend"]["volumeMounts"]) + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + assert {"name": "p4-shared-data", "mountPoint": "/data"} in services[role]["volumeMounts"] + + +def test_filesystem_broadcast_rejects_missing_shared_claim(tmp_path: Path): + options = replace(render_options(tmp_path), shared_pvc=None) + + with pytest.raises(ValueError, match="shared existing PVC"): + build_dgd_values(inference_config("filesystem"), options) + def test_dgd_rejects_native_backend(tmp_path: Path): config = InferenceConfig.model_validate({}) diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index 3eb2c5a32d..cad05d1800 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -15,11 +15,11 @@ IMAGE_DIGEST = f"sha256:{'3' * 64}" -def inference_config() -> InferenceConfig: +def inference_config(weight_broadcast: str = "nccl") -> InferenceConfig: return InferenceConfig.model_validate( { "backend": {"type": "dynamo"}, - "weight_broadcast": {"type": "nccl"}, + "weight_broadcast": {"type": weight_broadcast}, "deployment": { "type": "disaggregated", "gpus_per_node": 1, @@ -32,7 +32,7 @@ def inference_config() -> InferenceConfig: ) -def render_options(tmp_path: Path) -> DynamoGraphRenderOptions: +def render_options(tmp_path: Path, *, shared_pvc: str | None = None) -> DynamoGraphRenderOptions: return DynamoGraphRenderOptions( release_name="p4-math", namespace="bis-vllm", @@ -42,14 +42,16 @@ def render_options(tmp_path: Path) -> DynamoGraphRenderOptions: dynamo_sha=DYNAMO_SHA, image_digest=IMAGE_DIGEST, run_name="p4-run", + shared_pvc=shared_pvc, + image_pull_secrets=("nvcrimagepullsecret",), ) -def helm_template(*args: str) -> str: +def helm_template(*args: str, release_name: str = "p4-math") -> str: if HELM is None: pytest.skip("helm is not installed") return subprocess.run( - [HELM, "template", "p4-math", str(CHART), *args], + [HELM, "template", release_name, str(CHART), *args], check=True, capture_output=True, text=True, @@ -69,7 +71,8 @@ def test_chart_rejects_unknown_inference_mode(): def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_path: Path): - paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + options = render_options(tmp_path) + paths = write_dgd_artifacts(inference_config(), options) rendered = helm_template("-f", str(paths["values"])) graph = json.loads(paths["resource"].read_text()) @@ -81,9 +84,59 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert "http://p4-math-frontend-rl.bis-vllm.svc.cluster.local:8001" in rendered assert graph["spec"]["services"]["VllmPrefillWorker"]["replicas"] == 2 assert graph["spec"]["services"]["VllmDecodeWorker"]["replicas"] == 2 + assert rendered.count(f'image: "{options.image}"') == 2 + assert rendered.count(f"image: {options.image}") == 3 + assert rendered.count("nvcrimagepullsecret") == 5 assert not any(kind in rendered for kind in ("kind: ClusterRole", "kind: CustomResourceDefinition")) +def test_dgd_chart_rejects_release_name_mismatch(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), release_name="other-release") + + assert "must match embedded DynamoGraphDeployment metadata.name" in error.value.stderr + + +def test_filesystem_broadcast_reuses_existing_claim_without_rendering_pvc(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config("filesystem"), + render_options(tmp_path, shared_pvc="p4-shared-data"), + ) + rendered = helm_template("-f", str(paths["values"])) + + assert "kind: PersistentVolumeClaim" not in rendered + assert rendered.count("claimName: p4-shared-data") == 2 + + +def test_dgd_chart_rejects_mutable_prime_runtime_image(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError): + helm_template( + "-f", + str(paths["values"]), + "--set", + "image.reference=nvcr.io/example/prime:latest", + ) + + +def test_dgd_chart_rejects_runtime_image_that_differs_from_workers(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + other_digest = f"sha256:{'4' * 64}" + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "-f", + str(paths["values"]), + "--set", + f"image.reference=nvcr.io/example/prime:reviewed@{other_digest}", + ) + + assert "must use the same image.reference" in error.value.stderr + + def test_dgd_rejects_image_without_matching_digest(tmp_path: Path): with pytest.raises(ValueError, match="must be pinned"): DynamoGraphRenderOptions( From c75e4e52a8bc4f1718f075e9a542523a2b90e429 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 16:07:16 -0700 Subject: [PATCH 22/30] fix(k8s): mount Dynamo chat template --- src/prime_rl/inference/dgd.py | 43 ++++++++++++++++++--- tests/unit/inference/test_dgd.py | 54 ++++++++++++++++++++++++++- tests/unit/inference/test_helm_dgd.py | 22 ++++++++++- 3 files changed, 110 insertions(+), 9 deletions(-) diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index 87f5204c1a..f5aaf31b54 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -12,9 +12,11 @@ from prime_rl.configs.inference import DisaggregatedInferenceDeploymentConfig, InferenceConfig from prime_rl.inference.dynamo import ( + CHAT_TEMPLATE_ASSET, DynamoProcessSpec, build_frontend_process, build_worker_process, + resolve_chat_template_content, write_role_engine_configs, ) from prime_rl.utils.config import cli @@ -191,7 +193,14 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) engine_paths = write_role_engine_configs(config, options.output_dir) prefill_text = engine_paths["prefill"].read_text() decode_text = engine_paths["decode"].read_text() - engine_hash = _sha256_bytes((prefill_text + decode_text).encode()) + engine_data = { + "prefill-engine.json": prefill_text, + "decode-engine.json": decode_text, + } + chat_template_content = resolve_chat_template_content(config) + if chat_template_content is not None: + engine_data[CHAT_TEMPLATE_ASSET] = chat_template_content + engine_hash = _sha256_bytes(_canonical_json(engine_data)) config_map_name = f"{options.release_name}-dynamo-engine-{engine_hash[:12]}" annotations = { @@ -202,7 +211,15 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "prime-rl.nvidia.com/run-name": options.run_name, MANIFEST_HASH_SCOPE_ANNOTATION: MANIFEST_HASH_SCOPE, } - frontend_process = build_frontend_process(config, host="0.0.0.0", port=8000) + runtime_chat_template_path = ( + Path(ENGINE_MOUNT_PATH) / CHAT_TEMPLATE_ASSET if chat_template_content is not None else None + ) + frontend_process = build_frontend_process( + config, + host="0.0.0.0", + port=8000, + runtime_chat_template_path=runtime_chat_template_path, + ) frontend_container = { "image": options.image, "imagePullPolicy": "IfNotPresent", @@ -212,6 +229,23 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "ports": [{"containerPort": 8001, "name": "rl"}], } frontend_pod_spec = {"mainContainer": frontend_container} + if chat_template_content is not None: + frontend_container["volumeMounts"] = [ + { + "name": "dynamo-chat-template", + "mountPath": ENGINE_MOUNT_PATH, + "readOnly": True, + } + ] + frontend_pod_spec["volumes"] = [ + { + "name": "dynamo-chat-template", + "configMap": { + "name": config_map_name, + "items": [{"key": CHAT_TEMPLATE_ASSET, "path": CHAT_TEMPLATE_ASSET}], + }, + } + ] _apply_pod_credentials(frontend_pod_spec, frontend_container, options) frontend = { "componentType": "frontend", @@ -280,10 +314,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "name": config_map_name, "sha256": engine_hash, "annotations": annotations, - "data": { - "prefill-engine.json": prefill_text, - "decode-engine.json": decode_text, - }, + "data": engine_data, }, "resource": resource, }, diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index 5b3c4a732e..fcde460efb 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -15,11 +15,18 @@ IMAGE_DIGEST = f"sha256:{'3' * 64}" -def inference_config(weight_broadcast: str = "nccl") -> InferenceConfig: +def inference_config( + weight_broadcast: str = "nccl", + *, + chat_template: str | None = None, +) -> InferenceConfig: + model = {"name": "Qwen/Qwen3-30B-A3B-Thinking-2507"} + if chat_template is not None: + model["chat_template"] = chat_template return InferenceConfig.model_validate( { "backend": {"type": "dynamo"}, - "model": {"name": "Qwen/Qwen3-30B-A3B-Thinking-2507"}, + "model": model, "parallel": {"tp": 1}, "weight_broadcast": {"type": weight_broadcast}, "env_vars": {"HF_HOME": "/model-cache", "HF_HUB_OFFLINE": "1"}, @@ -158,6 +165,49 @@ def test_dgd_artifacts_are_deterministic_and_manifest_verifies(tmp_path: Path): assert hashlib.sha256(canonical_resource).hexdigest() == expected_manifest_hash +def test_dgd_embeds_and_mounts_content_addressed_chat_template(tmp_path: Path): + first = build_dgd_values( + inference_config(chat_template="template-v1: {{ messages }}"), + render_options(tmp_path / "first"), + ) + graph = first["inference"]["dynamoGraph"] + engine_config = graph["engineConfig"] + frontend = graph["resource"]["spec"]["services"]["Frontend"]["extraPodSpec"] + + assert engine_config["data"]["chat-template.jinja"] == "template-v1: {{ messages }}" + expected_hash = hashlib.sha256( + (json.dumps(engine_config["data"], indent=2, sort_keys=True) + "\n").encode() + ).hexdigest() + assert engine_config["sha256"] == expected_hash + assert engine_config["name"].endswith(expected_hash[:12]) + assert frontend["mainContainer"]["args"][-1] == "/etc/prime-rl/dynamo/chat-template.jinja" + assert frontend["mainContainer"]["volumeMounts"] == [ + { + "name": "dynamo-chat-template", + "mountPath": "/etc/prime-rl/dynamo", + "readOnly": True, + } + ] + assert frontend["volumes"] == [ + { + "name": "dynamo-chat-template", + "configMap": { + "name": engine_config["name"], + "items": [{"key": "chat-template.jinja", "path": "chat-template.jinja"}], + }, + } + ] + assert not (tmp_path / "first" / "chat-template.jinja").exists() + + second = build_dgd_values( + inference_config(chat_template="template-v2: {{ messages }}"), + render_options(tmp_path / "second"), + ) + second_config = second["inference"]["dynamoGraph"]["engineConfig"] + assert second_config["name"] != engine_config["name"] + assert second_config["sha256"] != engine_config["sha256"] + + def test_filesystem_broadcast_requires_one_shared_existing_claim(tmp_path: Path): options = render_options(tmp_path) values = build_dgd_values(inference_config("filesystem"), options) diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index cad05d1800..689f3d6e6b 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -15,10 +15,16 @@ IMAGE_DIGEST = f"sha256:{'3' * 64}" -def inference_config(weight_broadcast: str = "nccl") -> InferenceConfig: +def inference_config( + weight_broadcast: str = "nccl", + *, + chat_template: str | None = None, +) -> InferenceConfig: + model = {"chat_template": chat_template} if chat_template is not None else {} return InferenceConfig.model_validate( { "backend": {"type": "dynamo"}, + "model": model, "weight_broadcast": {"type": weight_broadcast}, "deployment": { "type": "disaggregated", @@ -90,6 +96,20 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert not any(kind in rendered for kind in ("kind: ClusterRole", "kind: CustomResourceDefinition")) +def test_dgd_chart_renders_chat_template_configmap_and_frontend_mount(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config(chat_template="template-marker: {{ messages }}"), + render_options(tmp_path), + ) + rendered = helm_template("-f", str(paths["values"])) + + assert "chat-template.jinja: |" in rendered + assert "template-marker: {{ messages }}" in rendered + assert "/etc/prime-rl/dynamo/chat-template.jinja" in rendered + assert "name: dynamo-chat-template" in rendered + assert "key: chat-template.jinja" in rendered + + def test_dgd_chart_rejects_release_name_mismatch(tmp_path: Path): paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) From bdae79e357fc77611ebac696e231f9765ebe750b Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 16:40:23 -0700 Subject: [PATCH 23/30] fix(k8s): enforce Dynamo release contracts Signed-off-by: Biswa Panda --- k8s/prime-rl/templates/deployment.yaml | 102 ++++++++++++++++- k8s/prime-rl/values.schema.json | 131 +++++++++++++++++++--- k8s/prime-rl/values.yaml | 14 +++ src/prime_rl/inference/dgd.py | 145 ++++++++++++++++++++++++- tests/unit/inference/test_dgd.py | 81 +++++++++++++- tests/unit/inference/test_helm_dgd.py | 21 +++- 6 files changed, 463 insertions(+), 31 deletions(-) diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index a4a127b8bc..8b3e1ed694 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -69,9 +69,15 @@ spec: value: "{{ .Release.Name }}-orchestrator-headless.{{ .Values.namespace }}.svc.cluster.local" - name: INFERENCE_URL value: {{ include "prime-rl.inferenceUrls" . | quote }} + {{- if .Values.modelCache.enabled }} + - name: HF_HOME + value: {{ .Values.modelCache.mountPath | quote }} + {{- end }} {{- if eq .Values.inference.mode "dynamoGraph" }} - name: DYN_RL_DISCOVERY_URL value: "http://{{ .Release.Name }}-frontend-rl.{{ .Values.namespace }}.svc.cluster.local:8001" + - name: DYN_RL_TOPOLOGY + value: {{ required "inference.dynamoGraph.clientTopology is required" .Values.inference.dynamoGraph.clientTopology | toJson | quote }} {{- end }} {{- with .Values.orchestrator.env }} {{- toYaml . | nindent 8 }} @@ -83,6 +89,15 @@ spec: name: {{ .Values.config.secrets.name }} key: wandb-api-key optional: true + {{- end }} + {{- if .Values.huggingFace.tokenSecretName }} + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.huggingFace.tokenSecretName }} + key: {{ .Values.huggingFace.tokenSecretKey }} + optional: true + {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: secretKeyRef: @@ -92,17 +107,30 @@ spec: {{- end }} resources: {{- toYaml .Values.orchestrator.resources | nindent 10 }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + mountPath: {{ .Values.modelCache.mountPath }} + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: claimName: {{ include "prime-rl.storageClaimName" . }} {{- end }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + persistentVolumeClaim: + claimName: {{ required "modelCache.existingClaim is required when modelCache.enabled" .Values.modelCache.existingClaim }} + {{- end }} + {{- end }} {{- end }} --- {{- if and .Values.inference.enabled (eq .Values.inference.mode "statefulset") }} @@ -172,6 +200,10 @@ spec: value: "{{ .Values.inference.replicas }}" - name: HEADLESS_SERVICE value: "{{ .Release.Name }}-inference-headless.{{ .Values.namespace }}.svc.cluster.local" + {{- if .Values.modelCache.enabled }} + - name: HF_HOME + value: {{ .Values.modelCache.mountPath | quote }} + {{- end }} {{- with .Values.inference.env }} {{- toYaml . | nindent 8 }} {{- end }} @@ -182,6 +214,15 @@ spec: name: {{ .Values.config.secrets.name }} key: wandb-api-key optional: true + {{- end }} + {{- if .Values.huggingFace.tokenSecretName }} + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.huggingFace.tokenSecretName }} + key: {{ .Values.huggingFace.tokenSecretKey }} + optional: true + {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: secretKeyRef: @@ -227,17 +268,30 @@ spec: failureThreshold: {{ .Values.inference.probes.readiness.failureThreshold }} timeoutSeconds: {{ .Values.inference.probes.readiness.timeoutSeconds }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + mountPath: {{ .Values.modelCache.mountPath }} + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: claimName: {{ include "prime-rl.storageClaimName" . }} {{- end }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + persistentVolumeClaim: + claimName: {{ required "modelCache.existingClaim is required when modelCache.enabled" .Values.modelCache.existingClaim }} + {{- end }} + {{- end }} {{- end }} --- {{- if .Values.trainer.enabled }} @@ -274,6 +328,14 @@ spec: {{- if .Values.trainer.runtimeClassName }} runtimeClassName: {{ .Values.trainer.runtimeClassName }} {{- end }} + {{- with .Values.trainer.nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.trainer.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: prime-rl-trainer image: {{ include "prime-rl.image" . | quote }} @@ -315,9 +377,15 @@ spec: {{- end }} - name: INFERENCE_URL value: {{ include "prime-rl.inferenceUrls" . | quote }} + {{- if .Values.modelCache.enabled }} + - name: HF_HOME + value: {{ .Values.modelCache.mountPath | quote }} + {{- end }} {{- if eq .Values.inference.mode "dynamoGraph" }} - name: DYN_RL_DISCOVERY_URL value: "http://{{ .Release.Name }}-frontend-rl.{{ .Values.namespace }}.svc.cluster.local:8001" + - name: DYN_RL_TOPOLOGY + value: {{ required "inference.dynamoGraph.clientTopology is required" .Values.inference.dynamoGraph.clientTopology | toJson | quote }} {{- end }} {{- with .Values.trainer.env }} {{- toYaml . | nindent 8 }} @@ -329,6 +397,15 @@ spec: name: {{ .Values.config.secrets.name }} key: wandb-api-key optional: true + {{- end }} + {{- if .Values.huggingFace.tokenSecretName }} + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.huggingFace.tokenSecretName }} + key: {{ .Values.huggingFace.tokenSecretKey }} + optional: true + {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: secretKeyRef: @@ -374,15 +451,28 @@ spec: failureThreshold: {{ .Values.trainer.probes.readiness.failureThreshold }} timeoutSeconds: {{ .Values.trainer.probes.readiness.timeoutSeconds }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + mountPath: {{ .Values.modelCache.mountPath }} + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: claimName: {{ include "prime-rl.storageClaimName" . }} {{- end }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + persistentVolumeClaim: + claimName: {{ required "modelCache.existingClaim is required when modelCache.enabled" .Values.modelCache.existingClaim }} + {{- end }} + {{- end }} {{- end }} diff --git a/k8s/prime-rl/values.schema.json b/k8s/prime-rl/values.schema.json index b494fd78aa..af83f236ef 100644 --- a/k8s/prime-rl/values.schema.json +++ b/k8s/prime-rl/values.schema.json @@ -5,15 +5,10 @@ "image": { "type": "object", "properties": { - "reference": { - "type": "string" - }, + "reference": {"type": "string"}, "pullSecrets": { "type": "array", - "items": { - "type": "string", - "minLength": 1 - }, + "items": {"type": "string", "minLength": 1}, "uniqueItems": true } } @@ -24,14 +19,93 @@ "mode": { "type": "string", "enum": ["statefulset", "dynamoGraph"] + }, + "dynamoGraph": { + "type": "object", + "properties": { + "clientTopology": { + "type": "object", + "properties": { + "schema_version": {"const": 1}, + "admin_api": {"const": "dynamo"}, + "base_url": { + "type": "array", + "minItems": 1, + "items": {"type": "string", "minLength": 1} + }, + "rl_base_url": { + "type": "array", + "minItems": 1, + "items": {"type": "string", "minLength": 1} + }, + "dynamo_worker_roles": { + "type": "array", + "minItems": 1, + "items": {"type": "string", "enum": ["agg", "prefill", "decode"]} + }, + "dynamo_gpus_per_worker": {"type": "integer", "minimum": 1} + }, + "required": [ + "schema_version", + "admin_api", + "base_url", + "rl_base_url", + "dynamo_worker_roles", + "dynamo_gpus_per_worker" + ], + "additionalProperties": false + }, + "engineConfig": {"type": "object"}, + "resource": {"type": "object"} + } } } }, "storage": { "type": "object", "properties": { - "existingClaim": { - "type": "string" + "existingClaim": {"type": "string"} + } + }, + "modelCache": { + "type": "object", + "properties": { + "enabled": {"type": "boolean"}, + "existingClaim": {"type": "string"}, + "mountPath": {"type": "string", "minLength": 1} + } + }, + "huggingFace": { + "type": "object", + "properties": { + "tokenSecretName": {"type": "string"}, + "tokenSecretKey": {"type": "string", "minLength": 1} + } + }, + "trainer": { + "type": "object", + "properties": { + "runtimeClassName": {"type": "string"}, + "nodeSelector": { + "type": "object", + "additionalProperties": {"type": "string"} + }, + "tolerations": { + "type": "array", + "items": { + "type": "object", + "properties": { + "key": {"type": "string", "minLength": 1}, + "operator": {"type": "string", "enum": ["Exists", "Equal"]}, + "value": {"type": "string"}, + "effect": { + "type": "string", + "enum": ["NoSchedule", "PreferNoSchedule", "NoExecute"] + } + }, + "required": ["key", "operator"], + "additionalProperties": false + } } } } @@ -41,11 +115,7 @@ "if": { "properties": { "inference": { - "properties": { - "mode": { - "const": "dynamoGraph" - } - }, + "properties": {"mode": {"const": "dynamoGraph"}}, "required": ["mode"] } }, @@ -55,14 +125,39 @@ "properties": { "image": { "properties": { - "reference": { - "pattern": "^.+@sha256:[0-9a-f]{64}$" - } + "reference": {"pattern": "^.+@sha256:[0-9a-f]{64}$"} }, "required": ["reference"] + }, + "inference": { + "properties": { + "dynamoGraph": { + "required": ["clientTopology", "engineConfig", "resource"] + } + }, + "required": ["dynamoGraph"] } }, - "required": ["image"] + "required": ["image", "inference"] + } + }, + { + "if": { + "properties": { + "modelCache": { + "properties": {"enabled": {"const": true}}, + "required": ["enabled"] + } + }, + "required": ["modelCache"] + }, + "then": { + "properties": { + "modelCache": { + "properties": {"existingClaim": {"minLength": 1}}, + "required": ["existingClaim"] + } + } } } ] diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml index 4bee20c572..e55d262871 100644 --- a/k8s/prime-rl/values.yaml +++ b/k8s/prime-rl/values.yaml @@ -26,6 +26,17 @@ storage: size: 1Ti mountPath: /data +# Existing model cache mounted consistently into generated runtime pods. +modelCache: + enabled: false + existingClaim: "" + mountPath: /model-cache + +# Existing Hugging Face token secret. The secret value is never embedded in values. +huggingFace: + tokenSecretName: "" + tokenSecretKey: HF_TOKEN + # Orchestrator component orchestrator: enabled: true @@ -134,6 +145,9 @@ trainer: runtimeClassName: nvidia + nodeSelector: {} + tolerations: [] + # Health probes for trainer (requires metrics_server config) probes: enabled: false diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index f5aaf31b54..2c271eada1 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -8,7 +8,7 @@ import re from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, Literal from prime_rl.configs.inference import DisaggregatedInferenceDeploymentConfig, InferenceConfig from prime_rl.inference.dynamo import ( @@ -28,6 +28,85 @@ "resource; json.dumps(sort_keys=true,indent=2)+newline; " "exclude=/metadata/annotations/prime-rl.nvidia.com~1manifest-sha256" ) +_DGD_RESERVED_ENV_KEYS = frozenset( + { + "DYN_COMPONENT", + "DYN_DISCOVERY_BACKEND", + "DYN_ENABLE_RL", + "DYN_ENDPOINT", + "DYN_ETCD_ENDPOINTS", + "DYN_EVENT_PLANE", + "DYN_FILE_KV", + "DYN_NAMESPACE", + "DYN_RL_ENDPOINT", + "DYN_RL_PORT", + "DYN_SYSTEM_PORT", + "VLLM_NIXL_SIDE_CHANNEL_HOST", + "VLLM_NIXL_SIDE_CHANNEL_PORT", + } +) + + +@dataclass(frozen=True, slots=True) +class KubernetesToleration: + key: str + operator: Literal["Exists", "Equal"] = "Exists" + effect: Literal["NoSchedule", "PreferNoSchedule", "NoExecute"] | None = "NoSchedule" + value: str | None = None + + def __post_init__(self) -> None: + if not self.key: + raise ValueError("Kubernetes toleration key must not be empty") + if self.operator == "Exists" and self.value is not None: + raise ValueError("An Exists toleration cannot define a value") + if self.operator == "Equal" and self.value is None: + raise ValueError("An Equal toleration requires a value") + + def as_manifest(self) -> dict[str, str]: + return { + "key": self.key, + "operator": self.operator, + **({"value": self.value} if self.value is not None else {}), + **({"effect": self.effect} if self.effect is not None else {}), + } + + +@dataclass(frozen=True, slots=True) +class GPUSchedulingProfile: + """One image-compatible placement contract for every generated GPU pod.""" + + runtime_class_name: str + architecture: str + product: str + node_pool: str + node_pool_label: str = "cloud.google.com/gke-nodepool" + tolerations: tuple[KubernetesToleration, ...] = (KubernetesToleration(key="nvidia.com/gpu"),) + + def __post_init__(self) -> None: + required = { + "runtime_class_name": self.runtime_class_name, + "architecture": self.architecture, + "product": self.product, + "node_pool": self.node_pool, + "node_pool_label": self.node_pool_label, + } + empty = [name for name, value in required.items() if not value] + if empty: + raise ValueError(f"GPU scheduling fields must not be empty: {empty}") + if not self.tolerations: + raise ValueError("GPU scheduling requires at least one toleration") + + @property + def node_selector(self) -> dict[str, str]: + return { + "kubernetes.io/arch": self.architecture, + "nvidia.com/gpu.product": self.product, + self.node_pool_label: self.node_pool, + } + + @property + def toleration_manifests(self) -> list[dict[str, str]]: + return [toleration.as_manifest() for toleration in self.tolerations] @dataclass(frozen=True) @@ -40,6 +119,7 @@ class DynamoGraphRenderOptions: dynamo_sha: str image_digest: str run_name: str + gpu_scheduling: GPUSchedulingProfile model_cache_pvc: str | None = None shared_pvc: str | None = None image_pull_secrets: tuple[str, ...] = () @@ -99,6 +179,8 @@ def _apply_pod_credentials( ) -> None: if options.image_pull_secrets: pod_spec["imagePullSecrets"] = [{"name": name} for name in options.image_pull_secrets] + if options.model_cache_pvc and not any(item["name"] == "HF_HOME" for item in container.get("env", [])): + container.setdefault("env", []).append({"name": "HF_HOME", "value": "/model-cache"}) if options.hf_token_secret and not any(item["name"] == "HF_TOKEN" for item in container.get("env", [])): container.setdefault("env", []).append( { @@ -146,8 +228,9 @@ def _worker_service( ], } pod_spec = { - "runtimeClassName": "nvidia", - "tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists", "effect": "NoSchedule"}], + "runtimeClassName": options.gpu_scheduling.runtime_class_name, + "nodeSelector": options.gpu_scheduling.node_selector, + "tolerations": options.gpu_scheduling.toleration_manifests, "volumes": [ { "name": "dynamo-engine-config", @@ -179,6 +262,21 @@ def _add_pvc(resource: dict[str, Any], service: dict[str, Any], name: str | None service.setdefault("volumeMounts", []).append({"name": name, "mountPoint": mount_point}) +def _validate_dgd_environment(config: InferenceConfig) -> None: + environment_sources = [("global", config.env_vars)] + if config.deployment.type == "disaggregated": + environment_sources.extend( + [ + ("prefill", config.deployment.prefill_env_vars), + ("decode", config.deployment.decode_env_vars), + ] + ) + for source, environment in environment_sources: + conflicts = sorted(_DGD_RESERVED_ENV_KEYS & environment.keys()) + if conflicts: + raise ValueError(f"{source} env_vars contains {conflicts}; these DGD keys are operator-owned") + + def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) -> dict[str, Any]: if config.backend.type != "dynamo" or config.deployment.type != "disaggregated": raise ValueError("DGD rendering requires a Dynamo disaggregated inference config") @@ -189,6 +287,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) raise ValueError("DGD rendering currently requires one pod per decode replica") if config.weight_broadcast.type == "filesystem" and not options.shared_pvc: raise ValueError("Dynamo filesystem weight broadcast requires a shared existing PVC") + _validate_dgd_environment(config) engine_paths = write_role_engine_configs(config, options.output_dir) prefill_text = engine_paths["prefill"].read_text() @@ -310,6 +409,18 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "inference": { "mode": "dynamoGraph", "dynamoGraph": { + "clientTopology": { + "schema_version": 1, + "admin_api": "dynamo", + "base_url": [ + f"http://{options.release_name}-frontend.{options.namespace}.svc.cluster.local:8000/v1" + ], + "rl_base_url": [ + f"http://{options.release_name}-frontend-rl.{options.namespace}.svc.cluster.local:8001" + ], + "dynamo_worker_roles": list(config.dynamo_worker_roles), + "dynamo_gpus_per_worker": config.dynamo_gpus_per_worker, + }, "engineConfig": { "name": config_map_name, "sha256": engine_hash, @@ -319,7 +430,23 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "resource": resource, }, }, + "trainer": { + "runtimeClassName": options.gpu_scheduling.runtime_class_name, + "nodeSelector": options.gpu_scheduling.node_selector, + "tolerations": options.gpu_scheduling.toleration_manifests, + }, } + if options.model_cache_pvc: + values["modelCache"] = { + "enabled": True, + "existingClaim": options.model_cache_pvc, + "mountPath": "/model-cache", + } + if options.hf_token_secret: + values["huggingFace"] = { + "tokenSecretName": options.hf_token_secret, + "tokenSecretKey": "HF_TOKEN", + } if options.shared_pvc: values["storage"] = { "enabled": True, @@ -359,6 +486,11 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--dynamo-sha", required=True) parser.add_argument("--image-digest", required=True) parser.add_argument("--run-name") + parser.add_argument("--gpu-runtime-class", default="nvidia") + parser.add_argument("--gpu-architecture", required=True) + parser.add_argument("--gpu-product", required=True) + parser.add_argument("--gpu-node-pool", required=True) + parser.add_argument("--gpu-node-pool-label", default="cloud.google.com/gke-nodepool") parser.add_argument("--model-cache-pvc") parser.add_argument("--shared-pvc") parser.add_argument("--image-pull-secret", action="append", default=[]) @@ -378,6 +510,13 @@ def main() -> None: dynamo_sha=args.dynamo_sha, image_digest=args.image_digest, run_name=args.run_name or args.release_name, + gpu_scheduling=GPUSchedulingProfile( + runtime_class_name=args.gpu_runtime_class, + architecture=args.gpu_architecture, + product=args.gpu_product, + node_pool=args.gpu_node_pool, + node_pool_label=args.gpu_node_pool_label, + ), model_cache_pvc=args.model_cache_pvc, shared_pvc=args.shared_pvc, image_pull_secrets=tuple(args.image_pull_secret), diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index fcde460efb..648d765612 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -7,12 +7,23 @@ import pytest from prime_rl.configs.inference import InferenceConfig -from prime_rl.inference.dgd import DynamoGraphRenderOptions, build_dgd_values, write_dgd_artifacts +from prime_rl.inference.dgd import ( + DynamoGraphRenderOptions, + GPUSchedulingProfile, + build_dgd_values, + write_dgd_artifacts, +) from prime_rl.inference.dynamo import build_frontend_process, build_worker_process PRIME_SHA = "1" * 40 DYNAMO_SHA = "2" * 40 IMAGE_DIGEST = f"sha256:{'3' * 64}" +GPU_SCHEDULING = GPUSchedulingProfile( + runtime_class_name="nvidia", + architecture="arm64", + product="NVIDIA-GB200", + node_pool="customer-gpu-o7v", +) def inference_config( @@ -52,6 +63,7 @@ def render_options(tmp_path: Path) -> DynamoGraphRenderOptions: dynamo_sha=DYNAMO_SHA, image_digest=IMAGE_DIGEST, run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, model_cache_pvc="model-cache", shared_pvc="p4-shared-data", image_pull_secrets=("nvcrimagepullsecret",), @@ -85,6 +97,23 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): "existingClaim": "p4-shared-data", "mountPath": "/data", } + assert values["modelCache"] == { + "enabled": True, + "existingClaim": "model-cache", + "mountPath": "/model-cache", + } + assert values["huggingFace"] == { + "tokenSecretName": "hf-token-secret", + "tokenSecretKey": "HF_TOKEN", + } + assert values["inference"]["dynamoGraph"]["clientTopology"] == { + "schema_version": 1, + "admin_api": "dynamo", + "base_url": ["http://p4-math-frontend.bis-vllm.svc.cluster.local:8000/v1"], + "rl_base_url": ["http://p4-math-frontend-rl.bis-vllm.svc.cluster.local:8001"], + "dynamo_worker_roles": ["prefill", "prefill", "decode", "decode"], + "dynamo_gpus_per_worker": 1, + } for service in services.values(): pod_spec = service["extraPodSpec"] assert pod_spec["mainContainer"]["image"] == options.image @@ -95,8 +124,19 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): assert env["HF_HUB_OFFLINE"] == "1" gpu_toleration = {"key": "nvidia.com/gpu", "operator": "Exists", "effect": "NoSchedule"} - assert gpu_toleration in services["VllmPrefillWorker"]["extraPodSpec"]["tolerations"] - assert gpu_toleration in services["VllmDecodeWorker"]["extraPodSpec"]["tolerations"] + expected_selector = { + "kubernetes.io/arch": "arm64", + "nvidia.com/gpu.product": "NVIDIA-GB200", + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + } + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + worker_pod = services[role]["extraPodSpec"] + assert worker_pod["runtimeClassName"] == "nvidia" + assert worker_pod["nodeSelector"] == expected_selector + assert worker_pod["tolerations"] == [gpu_toleration] + assert values["trainer"]["runtimeClassName"] == "nvidia" + assert values["trainer"]["nodeSelector"] == expected_selector + assert values["trainer"]["tolerations"] == [gpu_toleration] prefill_args = services["VllmPrefillWorker"]["extraPodSpec"]["mainContainer"]["args"] decode_args = services["VllmDecodeWorker"]["extraPodSpec"]["mainContainer"]["args"] @@ -235,6 +275,41 @@ def test_filesystem_broadcast_rejects_missing_shared_claim(tmp_path: Path): build_dgd_values(inference_config("filesystem"), options) +@pytest.mark.parametrize( + ("scope", "key"), + [ + ("global", "DYN_DISCOVERY_BACKEND"), + ("global", "DYN_NAMESPACE"), + ("prefill", "DYN_SYSTEM_PORT"), + ("decode", "DYN_ENDPOINT"), + ], +) +def test_dgd_rejects_operator_owned_environment(scope: str, key: str, tmp_path: Path): + config_data = inference_config().model_dump(mode="python") + if scope == "global": + config_data["env_vars"] = {key: "override"} + else: + config_data["deployment"][f"{scope}_env_vars"] = {key: "override"} + config = InferenceConfig.model_validate(config_data) + + with pytest.raises(ValueError, match=rf"{scope}.*{key}.*operator-owned"): + build_dgd_values(config, render_options(tmp_path)) + + +def test_gpu_scheduling_changes_manifest_identity(tmp_path: Path): + first = build_dgd_values(inference_config(), render_options(tmp_path / "first")) + changed_profile = replace(GPU_SCHEDULING, node_pool="customer-gpu-alternate") + changed_options = replace(render_options(tmp_path / "second"), gpu_scheduling=changed_profile) + second = build_dgd_values(inference_config(), changed_options) + + first_resource = first["inference"]["dynamoGraph"]["resource"] + second_resource = second["inference"]["dynamoGraph"]["resource"] + assert ( + first_resource["metadata"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"] + != (second_resource["metadata"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"]) + ) + + def test_dgd_rejects_native_backend(tmp_path: Path): config = InferenceConfig.model_validate({}) with pytest.raises(ValueError, match="Dynamo disaggregated"): diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index 689f3d6e6b..5274051705 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -6,13 +6,19 @@ import pytest from prime_rl.configs.inference import InferenceConfig -from prime_rl.inference.dgd import DynamoGraphRenderOptions, write_dgd_artifacts +from prime_rl.inference.dgd import DynamoGraphRenderOptions, GPUSchedulingProfile, write_dgd_artifacts HELM = shutil.which("helm") CHART = Path(__file__).parents[3] / "k8s" / "prime-rl" PRIME_SHA = "1" * 40 DYNAMO_SHA = "2" * 40 IMAGE_DIGEST = f"sha256:{'3' * 64}" +GPU_SCHEDULING = GPUSchedulingProfile( + runtime_class_name="nvidia", + architecture="arm64", + product="NVIDIA-GB200", + node_pool="customer-gpu-o7v", +) def inference_config( @@ -48,6 +54,9 @@ def render_options(tmp_path: Path, *, shared_pvc: str | None = None) -> DynamoGr dynamo_sha=DYNAMO_SHA, image_digest=IMAGE_DIGEST, run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, + model_cache_pvc="model-cache", + hf_token_secret="hf-token-secret", shared_pvc=shared_pvc, image_pull_secrets=("nvcrimagepullsecret",), ) @@ -93,6 +102,14 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert rendered.count(f'image: "{options.image}"') == 2 assert rendered.count(f"image: {options.image}") == 3 assert rendered.count("nvcrimagepullsecret") == 5 + assert rendered.count("name: DYN_RL_TOPOLOGY") == 2 + assert rendered.count("claimName: model-cache") == 2 + assert rendered.count("name: HF_TOKEN") == 5 + assert rendered.count("name: HF_HOME") == 5 + assert rendered.count("cloud.google.com/gke-nodepool: customer-gpu-o7v") == 3 + assert rendered.count("kubernetes.io/arch: arm64") == 3 + assert rendered.count("nvidia.com/gpu.product: NVIDIA-GB200") == 3 + assert rendered.count("key: nvidia.com/gpu") == 3 assert not any(kind in rendered for kind in ("kind: ClusterRole", "kind: CustomResourceDefinition")) @@ -168,6 +185,7 @@ def test_dgd_rejects_image_without_matching_digest(tmp_path: Path): dynamo_sha=DYNAMO_SHA, image_digest=IMAGE_DIGEST, run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, ) @@ -182,4 +200,5 @@ def test_dgd_rejects_image_without_commit_suffixes(tmp_path: Path): dynamo_sha=DYNAMO_SHA, image_digest=IMAGE_DIGEST, run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, ) From 16a6ffd5606849a7b7b7e029a58de2ba1fc0ccc0 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 17:01:00 -0700 Subject: [PATCH 24/30] fix(k8s): preserve operator runtime contracts --- k8s/prime-rl/templates/deployment.yaml | 7 + .../templates/dynamo-engine-config.yaml | 3 +- .../templates/dynamo-graph-deployment.yaml | 21 +++ k8s/prime-rl/values.schema.json | 43 ++++- k8s/prime-rl/values.yaml | 2 + src/prime_rl/inference/dgd.py | 154 ++++++++++++++++- tests/unit/inference/test_dgd.py | 159 +++++++++++++++++- tests/unit/inference/test_helm_dgd.py | 84 ++++++++- 8 files changed, 451 insertions(+), 22 deletions(-) diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index 8b3e1ed694..89b4470a95 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -32,6 +32,13 @@ spec: nodeSelector: {{- toYaml . | nindent 8 }} {{- end }} + {{- if .Values.orchestrator.runtimeClassName }} + runtimeClassName: {{ .Values.orchestrator.runtimeClassName }} + {{- end }} + {{- with .Values.orchestrator.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: prime-rl-orchestrator image: {{ include "prime-rl.image" . | quote }} diff --git a/k8s/prime-rl/templates/dynamo-engine-config.yaml b/k8s/prime-rl/templates/dynamo-engine-config.yaml index 3ac896bc57..3daf657552 100644 --- a/k8s/prime-rl/templates/dynamo-engine-config.yaml +++ b/k8s/prime-rl/templates/dynamo-engine-config.yaml @@ -12,7 +12,6 @@ metadata: immutable: true data: {{- range $name, $content := .Values.inference.dynamoGraph.engineConfig.data }} - {{ $name }}: | -{{ $content | nindent 4 }} + {{ $name }}: {{ $content | toJson }} {{- end }} {{- end }} diff --git a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml index 6446a088eb..b5c98a3702 100644 --- a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml +++ b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml @@ -1,9 +1,30 @@ {{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} {{- $resource := required "inference.dynamoGraph.resource is required" .Values.inference.dynamoGraph.resource }} {{- $resourceName := required "inference.dynamoGraph.resource.metadata.name is required" $resource.metadata.name }} +{{- $resourceNamespace := required "inference.dynamoGraph.resource.metadata.namespace is required" $resource.metadata.namespace }} {{- if ne .Release.Name $resourceName }} {{- fail (printf "Helm Release.Name %q must match embedded DynamoGraphDeployment metadata.name %q" .Release.Name $resourceName) }} {{- end }} +{{- if ne .Values.namespace $resourceNamespace }} +{{- fail (printf "values namespace %q must match embedded DynamoGraphDeployment metadata.namespace %q" .Values.namespace $resourceNamespace) }} +{{- end }} +{{- $topology := required "inference.dynamoGraph.clientTopology is required" .Values.inference.dynamoGraph.clientTopology }} +{{- $baseURLs := required "inference.dynamoGraph.clientTopology.base_url is required" $topology.base_url }} +{{- $rlBaseURLs := required "inference.dynamoGraph.clientTopology.rl_base_url is required" $topology.rl_base_url }} +{{- if or (ne (len $baseURLs) 1) (ne (index $baseURLs 0) (printf "http://%s-frontend.%s.svc.cluster.local:8000/v1" .Release.Name .Values.namespace)) }} +{{- fail "inference.dynamoGraph.clientTopology.base_url must match the generated release frontend URL" }} +{{- end }} +{{- if or (ne (len $rlBaseURLs) 1) (ne (index $rlBaseURLs 0) (printf "http://%s-frontend-rl.%s.svc.cluster.local:8001" .Release.Name .Values.namespace)) }} +{{- fail "inference.dynamoGraph.clientTopology.rl_base_url must match the generated release RL discovery URL" }} +{{- end }} +{{- $protectedEnv := dict "DYN_RL_DISCOVERY_URL" true "DYN_RL_TOPOLOGY" true "INFERENCE_URL" true "HF_HOME" true "HF_TOKEN" true "HUGGING_FACE_HUB_TOKEN" true }} +{{- range $componentName, $component := dict "orchestrator" .Values.orchestrator "trainer" .Values.trainer }} +{{- range $entry := $component.env }} +{{- if hasKey $protectedEnv $entry.name }} +{{- fail (printf "%s.env cannot override generated %s" $componentName $entry.name) }} +{{- end }} +{{- end }} +{{- end }} {{- $image := required "image.reference is required in dynamoGraph mode" .Values.image.reference }} {{- if not (regexMatch "^.+@sha256:[0-9a-f]{64}$" $image) }} {{- fail "image.reference must be pinned to a full sha256 digest in dynamoGraph mode" }} diff --git a/k8s/prime-rl/values.schema.json b/k8s/prime-rl/values.schema.json index af83f236ef..df8f781fb6 100644 --- a/k8s/prime-rl/values.schema.json +++ b/k8s/prime-rl/values.schema.json @@ -82,6 +82,33 @@ "tokenSecretKey": {"type": "string", "minLength": 1} } }, + "orchestrator": { + "type": "object", + "properties": { + "runtimeClassName": {"type": "string"}, + "nodeSelector": { + "type": "object", + "additionalProperties": {"type": "string"} + }, + "tolerations": { + "type": "array", + "items": { + "type": "object", + "properties": { + "key": {"type": "string", "minLength": 1}, + "operator": {"type": "string", "enum": ["Exists", "Equal"]}, + "value": {"type": "string"}, + "effect": { + "type": "string", + "enum": ["NoSchedule", "PreferNoSchedule", "NoExecute"] + } + }, + "required": ["key", "operator"], + "additionalProperties": false + } + } + } + }, "trainer": { "type": "object", "properties": { @@ -136,9 +163,23 @@ } }, "required": ["dynamoGraph"] + }, + "orchestrator": { + "properties": { + "nodeSelector": {"minProperties": 2}, + "tolerations": {"minItems": 2} + }, + "required": ["nodeSelector", "tolerations"] + }, + "trainer": { + "properties": { + "nodeSelector": {"minProperties": 3}, + "tolerations": {"minItems": 3} + }, + "required": ["runtimeClassName", "nodeSelector", "tolerations"] } }, - "required": ["image", "inference"] + "required": ["image", "inference", "orchestrator", "trainer"] } }, { diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml index e55d262871..5ea9c848e4 100644 --- a/k8s/prime-rl/values.yaml +++ b/k8s/prime-rl/values.yaml @@ -65,6 +65,8 @@ orchestrator: nodeSelector: {} # nvidia.com/gpu.present: "true" # Orchestrator doesn't need GPUs + runtimeClassName: "" + tolerations: [] # Inference component inference: diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index 2c271eada1..a37357ed2a 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -30,21 +30,36 @@ ) _DGD_RESERVED_ENV_KEYS = frozenset( { + "CONTAINER_NAME", + "DYNAMO_PORT", "DYN_COMPONENT", "DYN_DISCOVERY_BACKEND", "DYN_ENABLE_RL", "DYN_ENDPOINT", + "DYN_ENDPOINT_TYPES", "DYN_ETCD_ENDPOINTS", "DYN_EVENT_PLANE", "DYN_FILE_KV", + "DYN_HEALTH_CHECK_ENABLED", + "DYN_HTTP_PORT", + "DYN_KUBE_DISCOVERY_MODE", "DYN_NAMESPACE", + "DYN_NAMESPACE_PREFIX", + "DYN_NAMESPACE_WORKER_SUFFIX", + "DYN_PARENT_DGD_K8S_NAME", + "DYN_PARENT_DGD_K8S_NAMESPACE", "DYN_RL_ENDPOINT", "DYN_RL_PORT", - "DYN_SYSTEM_PORT", + "POD_NAME", + "POD_NAMESPACE", + "POD_UID", "VLLM_NIXL_SIDE_CHANNEL_HOST", "VLLM_NIXL_SIDE_CHANNEL_PORT", } ) +_DGD_RESERVED_ENV_PREFIXES = ("DYN_HEALTH_CHECK_", "DYN_SYSTEM_") +_RAW_HUGGING_FACE_TOKEN_KEYS = frozenset({"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"}) +_MODEL_CACHE_MOUNT_PATH = "/model-cache" @dataclass(frozen=True, slots=True) @@ -57,6 +72,10 @@ class KubernetesToleration: def __post_init__(self) -> None: if not self.key: raise ValueError("Kubernetes toleration key must not be empty") + if self.operator not in ("Exists", "Equal"): + raise ValueError(f"Unsupported Kubernetes toleration operator: {self.operator!r}") + if self.effect not in (None, "NoSchedule", "PreferNoSchedule", "NoExecute"): + raise ValueError(f"Unsupported Kubernetes toleration effect: {self.effect!r}") if self.operator == "Exists" and self.value is not None: raise ValueError("An Exists toleration cannot define a value") if self.operator == "Equal" and self.value is None: @@ -71,9 +90,39 @@ def as_manifest(self) -> dict[str, str]: } +def _parse_kubernetes_toleration(value: str) -> KubernetesToleration: + """Parse one CLI toleration from JSON without evaluating shell-like input.""" + try: + payload = json.loads(value) + except json.JSONDecodeError as error: + raise ValueError(f"Kubernetes toleration must be a JSON object: {error.msg}") from error + if not isinstance(payload, dict): + raise ValueError("Kubernetes toleration must be a JSON object") + supported = {"key", "operator", "effect", "value"} + unknown = sorted(payload.keys() - supported) + if unknown: + raise ValueError(f"Kubernetes toleration has unsupported fields: {unknown}") + if "key" not in payload: + raise ValueError("Kubernetes toleration requires a key") + return KubernetesToleration(**payload) + + +def _unique_tolerations( + tolerations: tuple[KubernetesToleration, ...], +) -> tuple[KubernetesToleration, ...]: + unique: list[KubernetesToleration] = [] + seen: set[tuple[tuple[str, str], ...]] = set() + for toleration in tolerations: + identity = tuple(sorted(toleration.as_manifest().items())) + if identity not in seen: + unique.append(toleration) + seen.add(identity) + return tuple(unique) + + @dataclass(frozen=True, slots=True) class GPUSchedulingProfile: - """One image-compatible placement contract for every generated GPU pod.""" + """Image placement plus the stricter placement required by GPU consumers.""" runtime_class_name: str architecture: str @@ -81,6 +130,8 @@ class GPUSchedulingProfile: node_pool: str node_pool_label: str = "cloud.google.com/gke-nodepool" tolerations: tuple[KubernetesToleration, ...] = (KubernetesToleration(key="nvidia.com/gpu"),) + additional_image_tolerations: tuple[KubernetesToleration, ...] = () + additional_gpu_tolerations: tuple[KubernetesToleration, ...] = () def __post_init__(self) -> None: required = { @@ -95,18 +146,46 @@ def __post_init__(self) -> None: raise ValueError(f"GPU scheduling fields must not be empty: {empty}") if not self.tolerations: raise ValueError("GPU scheduling requires at least one toleration") + if not any(toleration.key == "nvidia.com/gpu" for toleration in self.tolerations): + raise ValueError("GPU scheduling must tolerate the nvidia.com/gpu taint") @property - def node_selector(self) -> dict[str, str]: + def image_node_selector(self) -> dict[str, str]: return { "kubernetes.io/arch": self.architecture, - "nvidia.com/gpu.product": self.product, self.node_pool_label: self.node_pool, } + @property + def node_selector(self) -> dict[str, str]: + return { + **self.image_node_selector, + "nvidia.com/gpu.product": self.product, + } + + @property + def image_tolerations(self) -> tuple[KubernetesToleration, ...]: + required = ( + KubernetesToleration( + key="kubernetes.io/arch", + operator="Equal", + value=self.architecture, + ), + KubernetesToleration(key="prime-rl", operator="Equal", value="true"), + ) + return _unique_tolerations((*required, *self.additional_image_tolerations)) + + @property + def image_toleration_manifests(self) -> list[dict[str, str]]: + return [toleration.as_manifest() for toleration in self.image_tolerations] + + @property + def gpu_tolerations(self) -> tuple[KubernetesToleration, ...]: + return _unique_tolerations((*self.image_tolerations, *self.tolerations, *self.additional_gpu_tolerations)) + @property def toleration_manifests(self) -> list[dict[str, str]]: - return [toleration.as_manifest() for toleration in self.tolerations] + return [toleration.as_manifest() for toleration in self.gpu_tolerations] @dataclass(frozen=True) @@ -272,11 +351,42 @@ def _validate_dgd_environment(config: InferenceConfig) -> None: ] ) for source, environment in environment_sources: - conflicts = sorted(_DGD_RESERVED_ENV_KEYS & environment.keys()) + conflicts = sorted( + key + for key in environment + if key in _DGD_RESERVED_ENV_KEYS or any(key.startswith(prefix) for prefix in _DGD_RESERVED_ENV_PREFIXES) + ) if conflicts: raise ValueError(f"{source} env_vars contains {conflicts}; these DGD keys are operator-owned") +def _validate_typed_credentials( + config: InferenceConfig, + options: DynamoGraphRenderOptions, +) -> None: + environment_sources = [("global", config.env_vars)] + if config.deployment.type == "disaggregated": + environment_sources.extend( + [ + ("prefill", config.deployment.prefill_env_vars), + ("decode", config.deployment.decode_env_vars), + ] + ) + for source, environment in environment_sources: + token_conflicts = sorted(_RAW_HUGGING_FACE_TOKEN_KEYS & environment.keys()) + if token_conflicts: + raise ValueError( + f"{source} env_vars contains raw Hugging Face credentials {token_conflicts}; " + "use DynamoGraphRenderOptions.hf_token_secret" + ) + hf_home = environment.get("HF_HOME") + if options.model_cache_pvc and hf_home is not None and hf_home != _MODEL_CACHE_MOUNT_PATH: + raise ValueError( + f"{source} env_vars sets HF_HOME={hf_home!r}, but the typed model cache mount " + f"requires {_MODEL_CACHE_MOUNT_PATH!r}" + ) + + def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) -> dict[str, Any]: if config.backend.type != "dynamo" or config.deployment.type != "disaggregated": raise ValueError("DGD rendering requires a Dynamo disaggregated inference config") @@ -288,6 +398,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) if config.weight_broadcast.type == "filesystem" and not options.shared_pvc: raise ValueError("Dynamo filesystem weight broadcast requires a shared existing PVC") _validate_dgd_environment(config) + _validate_typed_credentials(config, options) engine_paths = write_role_engine_configs(config, options.output_dir) prefill_text = engine_paths["prefill"].read_text() @@ -325,9 +436,16 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "command": ["python3", "-m", frontend_process.module], "args": list(frontend_process.arguments), "env": [{"name": name, "value": value} for name, value in sorted(frontend_process.environment().items())], - "ports": [{"containerPort": 8001, "name": "rl"}], + "ports": [ + {"containerPort": 8000, "name": "http"}, + {"containerPort": 8001, "name": "rl"}, + ], + } + frontend_pod_spec = { + "nodeSelector": options.gpu_scheduling.image_node_selector, + "tolerations": options.gpu_scheduling.image_toleration_manifests, + "mainContainer": frontend_container, } - frontend_pod_spec = {"mainContainer": frontend_container} if chat_template_content is not None: frontend_container["volumeMounts"] = [ { @@ -435,6 +553,10 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "nodeSelector": options.gpu_scheduling.node_selector, "tolerations": options.gpu_scheduling.toleration_manifests, }, + "orchestrator": { + "nodeSelector": options.gpu_scheduling.image_node_selector, + "tolerations": options.gpu_scheduling.image_toleration_manifests, + }, } if options.model_cache_pvc: values["modelCache"] = { @@ -491,6 +613,20 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--gpu-product", required=True) parser.add_argument("--gpu-node-pool", required=True) parser.add_argument("--gpu-node-pool-label", default="cloud.google.com/gke-nodepool") + parser.add_argument( + "--image-toleration", + action="append", + default=[], + type=_parse_kubernetes_toleration, + help='Additional image-pod toleration as JSON, e.g. \'{"key":"dedicated","operator":"Exists"}\'', + ) + parser.add_argument( + "--gpu-toleration", + action="append", + default=[], + type=_parse_kubernetes_toleration, + help='Additional GPU-pod toleration as JSON, e.g. \'{"key":"capacity","operator":"Exists"}\'', + ) parser.add_argument("--model-cache-pvc") parser.add_argument("--shared-pvc") parser.add_argument("--image-pull-secret", action="append", default=[]) @@ -516,6 +652,8 @@ def main() -> None: product=args.gpu_product, node_pool=args.gpu_node_pool, node_pool_label=args.gpu_node_pool_label, + additional_image_tolerations=tuple(args.image_toleration), + additional_gpu_tolerations=tuple(args.gpu_toleration), ), model_cache_pvc=args.model_cache_pvc, shared_pvc=args.shared_pvc, diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index 648d765612..9bdb2b9e37 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -1,5 +1,6 @@ import hashlib import json +import sys from copy import deepcopy from dataclasses import replace from pathlib import Path @@ -10,6 +11,9 @@ from prime_rl.inference.dgd import ( DynamoGraphRenderOptions, GPUSchedulingProfile, + KubernetesToleration, + _parse_args, + _parse_kubernetes_toleration, build_dgd_values, write_dgd_artifacts, ) @@ -123,20 +127,49 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): assert env["HF_HOME"] == "/model-cache" assert env["HF_HUB_OFFLINE"] == "1" - gpu_toleration = {"key": "nvidia.com/gpu", "operator": "Exists", "effect": "NoSchedule"} - expected_selector = { + image_tolerations = [ + { + "key": "kubernetes.io/arch", + "operator": "Equal", + "value": "arm64", + "effect": "NoSchedule", + }, + { + "key": "prime-rl", + "operator": "Equal", + "value": "true", + "effect": "NoSchedule", + }, + ] + gpu_tolerations = [ + *image_tolerations, + {"key": "nvidia.com/gpu", "operator": "Exists", "effect": "NoSchedule"}, + ] + image_selector = { "kubernetes.io/arch": "arm64", - "nvidia.com/gpu.product": "NVIDIA-GB200", "cloud.google.com/gke-nodepool": "customer-gpu-o7v", } + gpu_selector = {**image_selector, "nvidia.com/gpu.product": "NVIDIA-GB200"} + frontend_pod = services["Frontend"]["extraPodSpec"] + assert frontend_pod["nodeSelector"] == image_selector + assert frontend_pod["tolerations"] == image_tolerations + assert "runtimeClassName" not in frontend_pod + assert frontend_pod["mainContainer"]["ports"] == [ + {"containerPort": 8000, "name": "http"}, + {"containerPort": 8001, "name": "rl"}, + ] for role in ("VllmPrefillWorker", "VllmDecodeWorker"): worker_pod = services[role]["extraPodSpec"] assert worker_pod["runtimeClassName"] == "nvidia" - assert worker_pod["nodeSelector"] == expected_selector - assert worker_pod["tolerations"] == [gpu_toleration] + assert worker_pod["nodeSelector"] == gpu_selector + assert worker_pod["tolerations"] == gpu_tolerations + assert values["orchestrator"] == { + "nodeSelector": image_selector, + "tolerations": image_tolerations, + } assert values["trainer"]["runtimeClassName"] == "nvidia" - assert values["trainer"]["nodeSelector"] == expected_selector - assert values["trainer"]["tolerations"] == [gpu_toleration] + assert values["trainer"]["nodeSelector"] == gpu_selector + assert values["trainer"]["tolerations"] == gpu_tolerations prefill_args = services["VllmPrefillWorker"]["extraPodSpec"]["mainContainer"]["args"] decode_args = services["VllmDecodeWorker"]["extraPodSpec"]["mainContainer"]["args"] @@ -280,7 +313,25 @@ def test_filesystem_broadcast_rejects_missing_shared_claim(tmp_path: Path): [ ("global", "DYN_DISCOVERY_BACKEND"), ("global", "DYN_NAMESPACE"), + ("global", "DYN_NAMESPACE_PREFIX"), + ("global", "DYN_NAMESPACE_WORKER_SUFFIX"), + ("global", "DYN_PARENT_DGD_K8S_NAME"), + ("global", "DYN_PARENT_DGD_K8S_NAMESPACE"), + ("global", "DYN_KUBE_DISCOVERY_MODE"), + ("global", "DYN_ENDPOINT_TYPES"), + ("global", "DYN_SYSTEM_ENABLED"), + ("global", "DYN_SYSTEM_HOST"), + ("global", "DYN_SYSTEM_HEALTH_PATH"), + ("global", "DYN_SYSTEM_LIVE_PATH"), + ("global", "DYN_SYSTEM_STARTING_HEALTH_STATUS"), + ("global", "DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"), + ("global", "DYN_HEALTH_CHECK_ENABLED"), + ("global", "POD_NAME"), + ("global", "POD_NAMESPACE"), + ("global", "POD_UID"), + ("global", "CONTAINER_NAME"), ("prefill", "DYN_SYSTEM_PORT"), + ("prefill", "DYN_SYSTEM_PORT1"), ("decode", "DYN_ENDPOINT"), ], ) @@ -296,6 +347,100 @@ def test_dgd_rejects_operator_owned_environment(scope: str, key: str, tmp_path: build_dgd_values(config, render_options(tmp_path)) +@pytest.mark.parametrize("scope", ["global", "prefill", "decode"]) +@pytest.mark.parametrize("key", ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"]) +def test_dgd_rejects_raw_hugging_face_credentials(scope: str, key: str, tmp_path: Path): + config_data = inference_config().model_dump(mode="python") + if scope == "global": + config_data["env_vars"][key] = "plaintext-secret" + else: + config_data["deployment"][f"{scope}_env_vars"] = {key: "plaintext-secret"} + config = InferenceConfig.model_validate(config_data) + + with pytest.raises(ValueError, match=rf"{scope}.*{key}.*hf_token_secret"): + build_dgd_values(config, render_options(tmp_path)) + + +@pytest.mark.parametrize("scope", ["global", "prefill", "decode"]) +def test_dgd_rejects_hf_home_that_conflicts_with_typed_model_cache(scope: str, tmp_path: Path): + config_data = inference_config().model_dump(mode="python") + if scope == "global": + config_data["env_vars"]["HF_HOME"] = "/wrong-cache" + else: + config_data["deployment"][f"{scope}_env_vars"] = {"HF_HOME": "/wrong-cache"} + config = InferenceConfig.model_validate(config_data) + + with pytest.raises(ValueError, match=rf"{scope}.*HF_HOME.*/model-cache"): + build_dgd_values(config, render_options(tmp_path)) + + +def test_dgd_allows_hf_home_matching_typed_model_cache(tmp_path: Path): + values = build_dgd_values(inference_config(), render_options(tmp_path)) + services = values["inference"]["dynamoGraph"]["resource"]["spec"]["services"] + + for service in services.values(): + env = service["extraPodSpec"]["mainContainer"]["env"] + assert [item for item in env if item["name"] == "HF_HOME"] == [{"name": "HF_HOME", "value": "/model-cache"}] + + +def test_cli_toleration_parser_is_typed_and_rejects_unknown_fields(): + parsed = _parse_kubernetes_toleration( + '{"key":"dedicated","operator":"Equal","value":"prime","effect":"NoSchedule"}' + ) + assert parsed == KubernetesToleration( + key="dedicated", + operator="Equal", + value="prime", + effect="NoSchedule", + ) + + with pytest.raises(ValueError, match="unsupported fields"): + _parse_kubernetes_toleration('{"key":"dedicated","command":"touch /tmp/pwned"}') + + with pytest.raises(ValueError, match="operator"): + _parse_kubernetes_toleration('{"key":"dedicated","operator":"NotARealOperator"}') + + +def test_cli_accepts_typed_additional_image_and_gpu_tolerations(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + sys, + "argv", + [ + "dynamo-dgd", + "inference.toml", + "--release-name", + "p4-math", + "--namespace", + "bis-vllm", + "--image", + "registry/image@sha256:digest", + "--output-dir", + "/tmp/artifacts", + "--prime-sha", + PRIME_SHA, + "--dynamo-sha", + DYNAMO_SHA, + "--image-digest", + IMAGE_DIGEST, + "--gpu-architecture", + "arm64", + "--gpu-product", + "NVIDIA-GB200", + "--gpu-node-pool", + "customer-gpu-o7v", + "--image-toleration", + '{"key":"image-extra","operator":"Exists"}', + "--gpu-toleration", + '{"key":"gpu-extra","operator":"Equal","value":"true"}', + ], + ) + + args = _parse_args() + + assert args.image_toleration == [KubernetesToleration(key="image-extra")] + assert args.gpu_toleration == [KubernetesToleration(key="gpu-extra", operator="Equal", value="true")] + + def test_gpu_scheduling_changes_manifest_identity(tmp_path: Path): first = build_dgd_values(inference_config(), render_options(tmp_path / "first")) changed_profile = replace(GPU_SCHEDULING, node_pool="customer-gpu-alternate") diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index 5274051705..948afb45b0 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -1,9 +1,11 @@ +import hashlib import json import shutil import subprocess from pathlib import Path import pytest +import yaml from prime_rl.configs.inference import InferenceConfig from prime_rl.inference.dgd import DynamoGraphRenderOptions, GPUSchedulingProfile, write_dgd_artifacts @@ -73,6 +75,18 @@ def helm_template(*args: str, release_name: str = "p4-math") -> str: ).stdout +def rendered_documents(rendered: str) -> list[dict]: + return [document for document in yaml.safe_load_all(rendered) if document] + + +def rendered_resource(rendered: str, kind: str, name: str) -> dict: + return next( + document + for document in rendered_documents(rendered) + if document.get("kind") == kind and document.get("metadata", {}).get("name") == name + ) + + def test_native_chart_still_renders_inference_statefulset(): rendered = helm_template() assert "name: p4-math-inference\n" in rendered @@ -90,6 +104,7 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat paths = write_dgd_artifacts(inference_config(), options) rendered = helm_template("-f", str(paths["values"])) graph = json.loads(paths["resource"].read_text()) + rendered_graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") assert rendered.count("kind: DynamoGraphDeployment") == 1 assert rendered.count("kind: ConfigMap") == 1 @@ -99,6 +114,10 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert "http://p4-math-frontend-rl.bis-vllm.svc.cluster.local:8001" in rendered assert graph["spec"]["services"]["VllmPrefillWorker"]["replicas"] == 2 assert graph["spec"]["services"]["VllmDecodeWorker"]["replicas"] == 2 + assert rendered_graph["spec"]["services"]["Frontend"]["extraPodSpec"]["mainContainer"]["ports"] == [ + {"containerPort": 8000, "name": "http"}, + {"containerPort": 8001, "name": "rl"}, + ] assert rendered.count(f'image: "{options.image}"') == 2 assert rendered.count(f"image: {options.image}") == 3 assert rendered.count("nvcrimagepullsecret") == 5 @@ -106,10 +125,13 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert rendered.count("claimName: model-cache") == 2 assert rendered.count("name: HF_TOKEN") == 5 assert rendered.count("name: HF_HOME") == 5 - assert rendered.count("cloud.google.com/gke-nodepool: customer-gpu-o7v") == 3 - assert rendered.count("kubernetes.io/arch: arm64") == 3 + assert rendered.count("cloud.google.com/gke-nodepool: customer-gpu-o7v") == 5 + assert rendered.count("kubernetes.io/arch: arm64") == 5 assert rendered.count("nvidia.com/gpu.product: NVIDIA-GB200") == 3 + assert rendered.count("runtimeClassName: nvidia") == 3 + assert rendered.count("key: kubernetes.io/arch") == 5 assert rendered.count("key: nvidia.com/gpu") == 3 + assert rendered.count("key: prime-rl") == 5 assert not any(kind in rendered for kind in ("kind: ClusterRole", "kind: CustomResourceDefinition")) @@ -119,14 +141,36 @@ def test_dgd_chart_renders_chat_template_configmap_and_frontend_mount(tmp_path: render_options(tmp_path), ) rendered = helm_template("-f", str(paths["values"])) + values = json.loads(paths["values"].read_text()) + config_map = rendered_resource( + rendered, + "ConfigMap", + values["inference"]["dynamoGraph"]["engineConfig"]["name"], + ) - assert "chat-template.jinja: |" in rendered - assert "template-marker: {{ messages }}" in rendered + assert config_map["data"]["chat-template.jinja"] == "template-marker: {{ messages }}" assert "/etc/prime-rl/dynamo/chat-template.jinja" in rendered assert "name: dynamo-chat-template" in rendered assert "key: chat-template.jinja" in rendered +def test_dgd_chart_preserves_exact_content_addressed_configmap_bytes(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config(chat_template="EXACT: {{ messages }}"), + render_options(tmp_path), + ) + values = json.loads(paths["values"].read_text()) + rendered = helm_template("-f", str(paths["values"])) + config_map_name = values["inference"]["dynamoGraph"]["engineConfig"]["name"] + config_map = rendered_resource(rendered, "ConfigMap", config_map_name) + expected_data = values["inference"]["dynamoGraph"]["engineConfig"]["data"] + + assert config_map["data"] == expected_data + expected_hash = values["inference"]["dynamoGraph"]["engineConfig"]["sha256"] + canonical_data = (json.dumps(config_map["data"], indent=2, sort_keys=True) + "\n").encode() + assert hashlib.sha256(canonical_data).hexdigest() == expected_hash + + def test_dgd_chart_rejects_release_name_mismatch(tmp_path: Path): paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) @@ -136,6 +180,38 @@ def test_dgd_chart_rejects_release_name_mismatch(tmp_path: Path): assert "must match embedded DynamoGraphDeployment metadata.name" in error.value.stderr +def test_dgd_chart_rejects_namespace_mismatch(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), "--set", "namespace=other-namespace") + + assert "must match embedded DynamoGraphDeployment metadata.namespace" in error.value.stderr + + +@pytest.mark.parametrize( + ("component", "name"), + [ + ("orchestrator", "DYN_RL_TOPOLOGY"), + ("orchestrator", "HF_TOKEN"), + ("trainer", "HF_HOME"), + ], +) +def test_dgd_chart_rejects_raw_env_that_overrides_typed_contract( + component: str, + name: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + overlay = tmp_path / f"{component}-{name}.json" + overlay.write_text(json.dumps({component: {"env": [{"name": name, "value": "override"}]}})) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), "-f", str(overlay)) + + assert f"{component}.env cannot override generated {name}" in error.value.stderr + + def test_filesystem_broadcast_reuses_existing_claim_without_rendering_pvc(tmp_path: Path): paths = write_dgd_artifacts( inference_config("filesystem"), From 7a24d56ac2f1397be0ca2d3e8bf2f00568805650 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 17:24:28 -0700 Subject: [PATCH 25/30] fix(k8s): bind rendered Dynamo identities --- .../templates/dynamo-graph-deployment.yaml | 98 +++++++++++- k8s/prime-rl/values.schema.json | 76 +++++++++- src/prime_rl/inference/dgd.py | 80 ++++++---- tests/unit/inference/test_dgd.py | 21 +++ tests/unit/inference/test_helm_dgd.py | 139 ++++++++++++++++++ 5 files changed, 384 insertions(+), 30 deletions(-) diff --git a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml index b5c98a3702..321b470343 100644 --- a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml +++ b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml @@ -1,5 +1,6 @@ {{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} -{{- $resource := required "inference.dynamoGraph.resource is required" .Values.inference.dynamoGraph.resource }} +{{- $graph := required "inference.dynamoGraph is required" .Values.inference.dynamoGraph }} +{{- $resource := required "inference.dynamoGraph.resource is required" $graph.resource }} {{- $resourceName := required "inference.dynamoGraph.resource.metadata.name is required" $resource.metadata.name }} {{- $resourceNamespace := required "inference.dynamoGraph.resource.metadata.namespace is required" $resource.metadata.namespace }} {{- if ne .Release.Name $resourceName }} @@ -8,7 +9,7 @@ {{- if ne .Values.namespace $resourceNamespace }} {{- fail (printf "values namespace %q must match embedded DynamoGraphDeployment metadata.namespace %q" .Values.namespace $resourceNamespace) }} {{- end }} -{{- $topology := required "inference.dynamoGraph.clientTopology is required" .Values.inference.dynamoGraph.clientTopology }} +{{- $topology := required "inference.dynamoGraph.clientTopology is required" $graph.clientTopology }} {{- $baseURLs := required "inference.dynamoGraph.clientTopology.base_url is required" $topology.base_url }} {{- $rlBaseURLs := required "inference.dynamoGraph.clientTopology.rl_base_url is required" $topology.rl_base_url }} {{- if or (ne (len $baseURLs) 1) (ne (index $baseURLs 0) (printf "http://%s-frontend.%s.svc.cluster.local:8000/v1" .Release.Name .Values.namespace)) }} @@ -17,6 +18,99 @@ {{- if or (ne (len $rlBaseURLs) 1) (ne (index $rlBaseURLs 0) (printf "http://%s-frontend-rl.%s.svc.cluster.local:8001" .Release.Name .Values.namespace)) }} {{- fail "inference.dynamoGraph.clientTopology.rl_base_url must match the generated release RL discovery URL" }} {{- end }} +{{- $engineConfig := required "inference.dynamoGraph.engineConfig is required" $graph.engineConfig }} +{{- $engineCanonical := required "inference.dynamoGraph.engineConfig.canonicalData is required" $engineConfig.canonicalData }} +{{- $engineHash := required "inference.dynamoGraph.engineConfig.sha256 is required" $engineConfig.sha256 }} +{{- $computedEngineHash := sha256sum $engineCanonical }} +{{- if ne $engineHash $computedEngineHash }} +{{- fail "inference.dynamoGraph.engineConfig.sha256 must match its canonical payload" }} +{{- end }} +{{- $canonicalEngineData := mustFromJson $engineCanonical }} +{{- if ne (toJson $canonicalEngineData) (toJson $engineConfig.data) }} +{{- fail "inference.dynamoGraph.engineConfig.data must match its canonical payload" }} +{{- end }} +{{- $expectedConfigName := printf "%s-dynamo-engine-%s" .Release.Name (trunc 12 $computedEngineHash) }} +{{- if ne (required "inference.dynamoGraph.engineConfig.name is required" $engineConfig.name) $expectedConfigName }} +{{- fail "inference.dynamoGraph.engineConfig.name must be content-addressed by engineConfig.sha256" }} +{{- end }} +{{- $resourceAnnotations := required "DynamoGraphDeployment metadata.annotations are required" $resource.metadata.annotations }} +{{- if ne (required "DynamoGraphDeployment config-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/config-sha256")) $engineHash }} +{{- fail "DynamoGraphDeployment config-sha256 must match engineConfig.sha256" }} +{{- end }} +{{- $engineAnnotations := required "inference.dynamoGraph.engineConfig.annotations are required" $engineConfig.annotations }} +{{- if ne (required "engineConfig config-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/config-sha256")) $engineHash }} +{{- fail "engineConfig config-sha256 annotation must match engineConfig.sha256" }} +{{- end }} +{{- $topologyBinding := required "inference.dynamoGraph.topologyBinding is required" $graph.topologyBinding }} +{{- $topologyCanonical := required "inference.dynamoGraph.topologyBinding.canonical is required" $topologyBinding.canonical }} +{{- $topologyHash := required "inference.dynamoGraph.topologyBinding.sha256 is required" $topologyBinding.sha256 }} +{{- if ne (sha256sum $topologyCanonical) $topologyHash }} +{{- fail "topology binding sha256 must match its canonical payload" }} +{{- end }} +{{- if ne (required "DynamoGraphDeployment topology-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/topology-sha256")) $topologyHash }} +{{- fail "DynamoGraphDeployment topology binding annotation must match topologyBinding.sha256" }} +{{- end }} +{{- $canonicalTopology := mustFromJson $topologyCanonical }} +{{- if ne (toJson (index $canonicalTopology "clientTopology")) (toJson $topology) }} +{{- fail "clientTopology must match the renderer topology binding" }} +{{- end }} +{{- $services := required "DynamoGraphDeployment services are required" $resource.spec.services }} +{{- $boundWorkers := required "topology binding workerServices are required" (index $canonicalTopology "workerServices") }} +{{- if or (ne (len $services) 3) (ne (len $boundWorkers) 2) }} +{{- fail "DynamoGraphDeployment services must match the renderer topology binding" }} +{{- end }} +{{- range $serviceName, $expectedWorker := $boundWorkers }} +{{- $service := required (printf "topology binding service %s is required" $serviceName) (index $services $serviceName) }} +{{- if ne (required (printf "%s componentType is required" $serviceName) $service.componentType) "worker" }} +{{- fail (printf "DynamoGraphDeployment service %s must be a worker in the renderer topology binding" $serviceName) }} +{{- end }} +{{- $actualWorker := dict + "role" (required (printf "%s subComponentType is required" $serviceName) $service.subComponentType) + "replicas" (required (printf "%s replicas is required" $serviceName) $service.replicas) + "requestsGpu" (required (printf "%s GPU request is required" $serviceName) $service.resources.requests.gpu) + "limitsGpu" (required (printf "%s GPU limit is required" $serviceName) $service.resources.limits.gpu) }} +{{- if ne (toJson $expectedWorker) (toJson $actualWorker) }} +{{- fail (printf "DynamoGraphDeployment service %s must match the renderer topology binding" $serviceName) }} +{{- end }} +{{- end }} +{{- $prefill := required "DynamoGraphDeployment VllmPrefillWorker is required" (index $services "VllmPrefillWorker") }} +{{- $decode := required "DynamoGraphDeployment VllmDecodeWorker is required" (index $services "VllmDecodeWorker") }} +{{- $expectedRoles := list }} +{{- range $_ := until (int (required "VllmPrefillWorker replicas is required" $prefill.replicas)) }} +{{- $expectedRoles = append $expectedRoles "prefill" }} +{{- end }} +{{- range $_ := until (int (required "VllmDecodeWorker replicas is required" $decode.replicas)) }} +{{- $expectedRoles = append $expectedRoles "decode" }} +{{- end }} +{{- if ne (toJson $expectedRoles) (toJson (required "clientTopology.dynamo_worker_roles is required" $topology.dynamo_worker_roles)) }} +{{- fail "clientTopology.dynamo_worker_roles must match the DGD worker topology binding" }} +{{- end }} +{{- $gpusPerWorker := printf "%v" (required "clientTopology.dynamo_gpus_per_worker is required" $topology.dynamo_gpus_per_worker) }} +{{- range $serviceName, $service := dict "VllmDecodeWorker" $decode "VllmPrefillWorker" $prefill }} +{{- if or (ne (printf "%v" $service.resources.requests.gpu) $gpusPerWorker) (ne (printf "%v" $service.resources.limits.gpu) $gpusPerWorker) }} +{{- fail (printf "DynamoGraphDeployment service %s GPU resources must match the client topology binding" $serviceName) }} +{{- end }} +{{- end }} +{{- $manifestCanonical := required "inference.dynamoGraph.manifestCanonical is required" $graph.manifestCanonical }} +{{- $manifestHash := required "DynamoGraphDeployment manifest-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/manifest-sha256") }} +{{- if ne (sha256sum $manifestCanonical) $manifestHash }} +{{- fail "DynamoGraphDeployment manifest-sha256 must match its canonical payload" }} +{{- end }} +{{- if ne (required "engineConfig manifest-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/manifest-sha256")) $manifestHash }} +{{- fail "engineConfig manifest-sha256 annotation must match the DynamoGraphDeployment manifest" }} +{{- end }} +{{- if ne (required "engineConfig topology-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/topology-sha256")) $topologyHash }} +{{- fail "engineConfig topology-sha256 annotation must match topologyBinding.sha256" }} +{{- end }} +{{- if ne (toJson $engineAnnotations) (toJson $resourceAnnotations) }} +{{- fail "engineConfig annotations must match DynamoGraphDeployment identity annotations" }} +{{- end }} +{{- $canonicalResource := mustFromJson $manifestCanonical }} +{{- $scopedResource := deepCopy $resource }} +{{- $_ := unset $scopedResource.metadata.annotations "prime-rl.nvidia.com/manifest-sha256" }} +{{- if ne (toJson $canonicalResource) (toJson $scopedResource) }} +{{- fail "DynamoGraphDeployment resource must match its canonical manifest payload" }} +{{- end }} {{- $protectedEnv := dict "DYN_RL_DISCOVERY_URL" true "DYN_RL_TOPOLOGY" true "INFERENCE_URL" true "HF_HOME" true "HF_TOKEN" true "HUGGING_FACE_HUB_TOKEN" true }} {{- range $componentName, $component := dict "orchestrator" .Values.orchestrator "trainer" .Values.trainer }} {{- range $entry := $component.env }} diff --git a/k8s/prime-rl/values.schema.json b/k8s/prime-rl/values.schema.json index df8f781fb6..5afebd8fdc 100644 --- a/k8s/prime-rl/values.schema.json +++ b/k8s/prime-rl/values.schema.json @@ -55,9 +55,69 @@ ], "additionalProperties": false }, - "engineConfig": {"type": "object"}, + "engineConfig": { + "type": "object", + "properties": { + "name": { + "type": "string", + "pattern": "^.+-dynamo-engine-[0-9a-f]{12}$" + }, + "sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "canonicalData": {"type": "string", "minLength": 1}, + "annotations": { + "type": "object", + "properties": { + "prime-rl.nvidia.com/config-sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "prime-rl.nvidia.com/manifest-sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "prime-rl.nvidia.com/topology-sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + } + }, + "required": [ + "prime-rl.nvidia.com/config-sha256", + "prime-rl.nvidia.com/manifest-sha256", + "prime-rl.nvidia.com/topology-sha256" + ], + "additionalProperties": {"type": "string"} + }, + "data": { + "type": "object", + "minProperties": 2, + "additionalProperties": {"type": "string"} + } + }, + "anyOf": [ + {"maxProperties": 0}, + {"required": ["name", "sha256", "canonicalData", "annotations", "data"]} + ], + "additionalProperties": false + }, + "topologyBinding": { + "type": "object", + "properties": { + "sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "canonical": {"type": "string", "minLength": 1} + }, + "required": ["sha256", "canonical"], + "additionalProperties": false + }, + "manifestCanonical": {"type": "string", "minLength": 1}, "resource": {"type": "object"} - } + }, + "additionalProperties": false } } }, @@ -159,7 +219,17 @@ "inference": { "properties": { "dynamoGraph": { - "required": ["clientTopology", "engineConfig", "resource"] + "properties": { + "engineConfig": {"minProperties": 5}, + "resource": {"minProperties": 1} + }, + "required": [ + "clientTopology", + "engineConfig", + "topologyBinding", + "manifestCanonical", + "resource" + ] } }, "required": ["dynamoGraph"] diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index a37357ed2a..cb44434940 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -22,8 +22,10 @@ from prime_rl.utils.config import cli ENGINE_MOUNT_PATH = "/etc/prime-rl/dynamo" +ENGINE_CONFIG_HASH_ANNOTATION = "prime-rl.nvidia.com/config-sha256" MANIFEST_HASH_ANNOTATION = "prime-rl.nvidia.com/manifest-sha256" MANIFEST_HASH_SCOPE_ANNOTATION = "prime-rl.nvidia.com/manifest-sha256-scope" +TOPOLOGY_HASH_ANNOTATION = "prime-rl.nvidia.com/topology-sha256" MANIFEST_HASH_SCOPE = ( "resource; json.dumps(sort_keys=true,indent=2)+newline; " "exclude=/metadata/annotations/prime-rl.nvidia.com~1manifest-sha256" @@ -225,7 +227,7 @@ def _canonical_json(value: Any) -> bytes: return (json.dumps(value, indent=2, sort_keys=True) + "\n").encode() -def _resource_manifest_hash(resource: dict[str, Any]) -> str: +def _resource_manifest_canonical(resource: dict[str, Any]) -> bytes: annotations = resource["metadata"]["annotations"] scoped_annotations = {key: value for key, value in annotations.items() if key != MANIFEST_HASH_ANNOTATION} scoped_resource = { @@ -235,7 +237,20 @@ def _resource_manifest_hash(resource: dict[str, Any]) -> str: "annotations": scoped_annotations, }, } - return _sha256_bytes(_canonical_json(scoped_resource)) + return _canonical_json(scoped_resource) + + +def _worker_topology_binding(services: dict[str, Any]) -> dict[str, Any]: + return { + service_name: { + "role": service["subComponentType"], + "replicas": service["replicas"], + "requestsGpu": service["resources"]["requests"]["gpu"], + "limitsGpu": service["resources"]["limits"]["gpu"], + } + for service_name, service in sorted(services.items()) + if service["componentType"] == "worker" + } def _worker_env(process: DynamoProcessSpec) -> list[dict[str, Any]]: @@ -410,17 +425,9 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) chat_template_content = resolve_chat_template_content(config) if chat_template_content is not None: engine_data[CHAT_TEMPLATE_ASSET] = chat_template_content - engine_hash = _sha256_bytes(_canonical_json(engine_data)) + engine_canonical = _canonical_json(engine_data) + engine_hash = _sha256_bytes(engine_canonical) config_map_name = f"{options.release_name}-dynamo-engine-{engine_hash[:12]}" - - annotations = { - "prime-rl.nvidia.com/config-sha256": engine_hash, - "prime-rl.nvidia.com/dynamo-sha": options.dynamo_sha, - "prime-rl.nvidia.com/image-digest": options.image_digest, - "prime-rl.nvidia.com/prime-sha": options.prime_sha, - "prime-rl.nvidia.com/run-name": options.run_name, - MANIFEST_HASH_SCOPE_ANNOTATION: MANIFEST_HASH_SCOPE, - } runtime_chat_template_path = ( Path(ENGINE_MOUNT_PATH) / CHAT_TEMPLATE_ASSET if chat_template_content is not None else None ) @@ -485,6 +492,33 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) config_map_name=config_map_name, engine_file="decode-engine.json", ) + client_topology = { + "schema_version": 1, + "admin_api": "dynamo", + "base_url": [f"http://{options.release_name}-frontend.{options.namespace}.svc.cluster.local:8000/v1"], + "rl_base_url": [f"http://{options.release_name}-frontend-rl.{options.namespace}.svc.cluster.local:8001"], + "dynamo_worker_roles": list(config.dynamo_worker_roles), + "dynamo_gpus_per_worker": config.dynamo_gpus_per_worker, + } + worker_services = { + "VllmDecodeWorker": decode, + "VllmPrefillWorker": prefill, + } + topology_binding = { + "clientTopology": client_topology, + "workerServices": _worker_topology_binding(worker_services), + } + topology_canonical = _canonical_json(topology_binding) + topology_hash = _sha256_bytes(topology_canonical) + annotations = { + ENGINE_CONFIG_HASH_ANNOTATION: engine_hash, + "prime-rl.nvidia.com/dynamo-sha": options.dynamo_sha, + "prime-rl.nvidia.com/image-digest": options.image_digest, + "prime-rl.nvidia.com/prime-sha": options.prime_sha, + "prime-rl.nvidia.com/run-name": options.run_name, + MANIFEST_HASH_SCOPE_ANNOTATION: MANIFEST_HASH_SCOPE, + TOPOLOGY_HASH_ANNOTATION: topology_hash, + } resource: dict[str, Any] = { "apiVersion": "nvidia.com/v1alpha1", "kind": "DynamoGraphDeployment", @@ -508,7 +542,8 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) for service in (prefill, decode): _add_pvc(resource, service, options.shared_pvc, "/data") - manifest_hash = _resource_manifest_hash(resource) + manifest_canonical = _resource_manifest_canonical(resource) + manifest_hash = _sha256_bytes(manifest_canonical) annotations = {**annotations, MANIFEST_HASH_ANNOTATION: manifest_hash} resource = { **resource, @@ -527,24 +562,19 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "inference": { "mode": "dynamoGraph", "dynamoGraph": { - "clientTopology": { - "schema_version": 1, - "admin_api": "dynamo", - "base_url": [ - f"http://{options.release_name}-frontend.{options.namespace}.svc.cluster.local:8000/v1" - ], - "rl_base_url": [ - f"http://{options.release_name}-frontend-rl.{options.namespace}.svc.cluster.local:8001" - ], - "dynamo_worker_roles": list(config.dynamo_worker_roles), - "dynamo_gpus_per_worker": config.dynamo_gpus_per_worker, - }, + "clientTopology": client_topology, "engineConfig": { "name": config_map_name, "sha256": engine_hash, + "canonicalData": engine_canonical.decode(), "annotations": annotations, "data": engine_data, }, + "topologyBinding": { + "sha256": topology_hash, + "canonical": topology_canonical.decode(), + }, + "manifestCanonical": manifest_canonical.decode(), "resource": resource, }, }, diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index 9bdb2b9e37..59c1dbf3b8 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -118,6 +118,24 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): "dynamo_worker_roles": ["prefill", "prefill", "decode", "decode"], "dynamo_gpus_per_worker": 1, } + topology_binding = graph["topologyBinding"] + assert hashlib.sha256(topology_binding["canonical"].encode()).hexdigest() == topology_binding["sha256"] + assert json.loads(topology_binding["canonical"])["clientTopology"] == graph["clientTopology"] + assert json.loads(topology_binding["canonical"])["workerServices"] == { + "VllmDecodeWorker": { + "limitsGpu": "1", + "replicas": 2, + "requestsGpu": "1", + "role": "decode", + }, + "VllmPrefillWorker": { + "limitsGpu": "1", + "replicas": 2, + "requestsGpu": "1", + "role": "prefill", + }, + } + assert resource["metadata"]["annotations"]["prime-rl.nvidia.com/topology-sha256"] == topology_binding["sha256"] for service in services.values(): pod_spec = service["extraPodSpec"] assert pod_spec["mainContainer"]["image"] == options.image @@ -235,6 +253,8 @@ def test_dgd_artifacts_are_deterministic_and_manifest_verifies(tmp_path: Path): unhashed_resource = deepcopy(resource) del unhashed_resource["metadata"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"] canonical_resource = (json.dumps(unhashed_resource, indent=2, sort_keys=True) + "\n").encode() + values = json.loads(paths["values"].read_text()) + assert values["inference"]["dynamoGraph"]["manifestCanonical"].encode() == canonical_resource assert hashlib.sha256(canonical_resource).hexdigest() == expected_manifest_hash @@ -252,6 +272,7 @@ def test_dgd_embeds_and_mounts_content_addressed_chat_template(tmp_path: Path): (json.dumps(engine_config["data"], indent=2, sort_keys=True) + "\n").encode() ).hexdigest() assert engine_config["sha256"] == expected_hash + assert engine_config["canonicalData"] == json.dumps(engine_config["data"], indent=2, sort_keys=True) + "\n" assert engine_config["name"].endswith(expected_hash[:12]) assert frontend["mainContainer"]["args"][-1] == "/etc/prime-rl/dynamo/chat-template.jinja" assert frontend["mainContainer"]["volumeMounts"] == [ diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index 948afb45b0..583404e7a0 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -87,6 +87,20 @@ def rendered_resource(rendered: str, kind: str, name: str) -> dict: ) +def write_values_mutation( + source: Path, + target: Path, + path: tuple[str, ...], + replacement: object, +) -> None: + values = json.loads(source.read_text()) + parent = values + for key in path[:-1]: + parent = parent[key] + parent[path[-1]] = replacement + target.write_text(json.dumps(values)) + + def test_native_chart_still_renders_inference_statefulset(): rendered = helm_template() assert "name: p4-math-inference\n" in rendered @@ -171,6 +185,131 @@ def test_dgd_chart_preserves_exact_content_addressed_configmap_bytes(tmp_path: P assert hashlib.sha256(canonical_data).hexdigest() == expected_hash +@pytest.mark.parametrize( + ("case", "path", "replacement", "error_fragment"), + [ + ( + "engine-data", + ("inference", "dynamoGraph", "engineConfig", "data", "prefill-engine.json"), + "tampered", + "engineConfig.data must match its canonical payload", + ), + ( + "engine-sha", + ("inference", "dynamoGraph", "engineConfig", "sha256"), + "0" * 64, + "engineConfig.sha256 must match its canonical payload", + ), + ( + "engine-name", + ("inference", "dynamoGraph", "engineConfig", "name"), + "p4-math-dynamo-engine-000000000000", + "engineConfig.name must be content-addressed", + ), + ( + "dgd-config-sha", + ( + "inference", + "dynamoGraph", + "resource", + "metadata", + "annotations", + "prime-rl.nvidia.com/config-sha256", + ), + "0" * 64, + "DynamoGraphDeployment config-sha256 must match engineConfig.sha256", + ), + ( + "manifest-sha", + ( + "inference", + "dynamoGraph", + "resource", + "metadata", + "annotations", + "prime-rl.nvidia.com/manifest-sha256", + ), + "0" * 64, + "manifest-sha256 must match its canonical payload", + ), + ], +) +def test_dgd_chart_rejects_content_identity_mutations( + case: str, + path: tuple[str, ...], + replacement: object, + error_fragment: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + mutation = tmp_path / f"{case}.json" + write_values_mutation(paths["values"], mutation, path, replacement) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert error_fragment in error.value.stderr + + +@pytest.mark.parametrize( + ("case", "path", "replacement"), + [ + ( + "client-roles", + ("inference", "dynamoGraph", "clientTopology", "dynamo_worker_roles"), + ["prefill", "decode", "decode", "decode"], + ), + ( + "client-gpus", + ("inference", "dynamoGraph", "clientTopology", "dynamo_gpus_per_worker"), + 2, + ), + ( + "worker-replicas", + ( + "inference", + "dynamoGraph", + "resource", + "spec", + "services", + "VllmDecodeWorker", + "replicas", + ), + 3, + ), + ( + "worker-gpus", + ( + "inference", + "dynamoGraph", + "resource", + "spec", + "services", + "VllmPrefillWorker", + "resources", + "limits", + "gpu", + ), + "2", + ), + ], +) +def test_dgd_chart_rejects_topology_mutations( + case: str, + path: tuple[str, ...], + replacement: object, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + mutation = tmp_path / f"{case}.json" + write_values_mutation(paths["values"], mutation, path, replacement) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert "topology binding" in error.value.stderr + + def test_dgd_chart_rejects_release_name_mismatch(tmp_path: Path): paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) From 5608d98afd108c7cd415c94f3f740973f05d0245 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 18:47:25 -0700 Subject: [PATCH 26/30] fix(k8s): bind generated workload contracts --- k8s/prime-rl/templates/deployment.yaml | 32 +- .../templates/dynamo-graph-deployment.yaml | 124 +++- k8s/prime-rl/templates/dynamo-rl-service.yaml | 1 + k8s/prime-rl/templates/service.yaml | 6 + k8s/prime-rl/values.schema.json | 104 +++- src/prime_rl/inference/dgd.py | 88 ++- tests/unit/inference/helm_dgd_test_utils.py | 150 +++++ tests/unit/inference/test_dgd.py | 114 +++- tests/unit/inference/test_helm_dgd.py | 553 ++++++++++-------- .../unit/inference/test_helm_dgd_integrity.py | 339 +++++++++++ 10 files changed, 1231 insertions(+), 280 deletions(-) create mode 100644 tests/unit/inference/helm_dgd_test_utils.py create mode 100644 tests/unit/inference/test_helm_dgd_integrity.py diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index 89b4470a95..a1993900e2 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -1,4 +1,15 @@ +{{- $dgdWorkload := dict }} +{{- if eq .Values.inference.mode "dynamoGraph" }} +{{- $workloadBinding := required "inference.dynamoGraph.workloadBinding is required" .Values.inference.dynamoGraph.workloadBinding }} +{{- $workloadCanonical := required "inference.dynamoGraph.workloadBinding.canonical is required" $workloadBinding.canonical }} +{{- $dgdWorkload = mustFromJson $workloadCanonical }} +{{- end }} {{- if .Values.orchestrator.enabled }} +{{- $orchestratorPlacement := .Values.orchestrator }} +{{- if eq .Values.inference.mode "dynamoGraph" }} +{{- $orchestratorWorkload := required "canonical orchestrator workload is required" (index $dgdWorkload "orchestrator") }} +{{- $orchestratorPlacement = required "canonical orchestrator placement is required" (index $orchestratorWorkload "placement") }} +{{- end }} apiVersion: apps/v1 kind: StatefulSet metadata: @@ -28,14 +39,14 @@ spec: - name: {{ . | quote }} {{- end }} {{- end }} - {{- with .Values.orchestrator.nodeSelector }} + {{- with $orchestratorPlacement.nodeSelector }} nodeSelector: {{- toYaml . | nindent 8 }} {{- end }} - {{- if .Values.orchestrator.runtimeClassName }} - runtimeClassName: {{ .Values.orchestrator.runtimeClassName }} + {{- if $orchestratorPlacement.runtimeClassName }} + runtimeClassName: {{ $orchestratorPlacement.runtimeClassName }} {{- end }} - {{- with .Values.orchestrator.tolerations }} + {{- with $orchestratorPlacement.tolerations }} tolerations: {{- toYaml . | nindent 8 }} {{- end }} @@ -302,6 +313,11 @@ spec: {{- end }} --- {{- if .Values.trainer.enabled }} +{{- $trainerPlacement := .Values.trainer }} +{{- if eq .Values.inference.mode "dynamoGraph" }} +{{- $trainerWorkload := required "canonical trainer workload is required" (index $dgdWorkload "trainer") }} +{{- $trainerPlacement = required "canonical trainer placement is required" (index $trainerWorkload "placement") }} +{{- end }} apiVersion: apps/v1 kind: StatefulSet metadata: @@ -332,14 +348,14 @@ spec: - name: {{ . | quote }} {{- end }} {{- end }} - {{- if .Values.trainer.runtimeClassName }} - runtimeClassName: {{ .Values.trainer.runtimeClassName }} + {{- if $trainerPlacement.runtimeClassName }} + runtimeClassName: {{ $trainerPlacement.runtimeClassName }} {{- end }} - {{- with .Values.trainer.nodeSelector }} + {{- with $trainerPlacement.nodeSelector }} nodeSelector: {{- toYaml . | nindent 8 }} {{- end }} - {{- with .Values.trainer.tolerations }} + {{- with $trainerPlacement.tolerations }} tolerations: {{- toYaml . | nindent 8 }} {{- end }} diff --git a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml index 321b470343..55a4cce9cd 100644 --- a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml +++ b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml @@ -1,4 +1,14 @@ -{{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} +{{- $candidateGraph := default (dict) .Values.inference.dynamoGraph }} +{{- if and (hasKey $candidateGraph "workloadBinding") (ne .Values.inference.mode "dynamoGraph") }} +{{- fail "generated DynamoGraph contract requires dynamoGraph mode" }} +{{- end }} +{{- if eq .Values.inference.mode "dynamoGraph" }} +{{- if not (kindIs "bool" .Values.inference.enabled) }} +{{- fail "inference.enabled must be boolean true in dynamoGraph mode" }} +{{- end }} +{{- if not .Values.inference.enabled }} +{{- fail "inference.enabled must be boolean true in dynamoGraph mode" }} +{{- end }} {{- $graph := required "inference.dynamoGraph is required" .Values.inference.dynamoGraph }} {{- $resource := required "inference.dynamoGraph.resource is required" $graph.resource }} {{- $resourceName := required "inference.dynamoGraph.resource.metadata.name is required" $resource.metadata.name }} @@ -91,6 +101,115 @@ {{- fail (printf "DynamoGraphDeployment service %s GPU resources must match the client topology binding" $serviceName) }} {{- end }} {{- end }} +{{- $workloadBinding := required "inference.dynamoGraph.workloadBinding is required" $graph.workloadBinding }} +{{- $workloadCanonical := required "inference.dynamoGraph.workloadBinding.canonical is required" $workloadBinding.canonical }} +{{- $workloadHash := required "inference.dynamoGraph.workloadBinding.sha256 is required" $workloadBinding.sha256 }} +{{- if ne (sha256sum $workloadCanonical) $workloadHash }} +{{- fail "workload binding sha256 must match its canonical payload" }} +{{- end }} +{{- if ne (required "DynamoGraphDeployment workload-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/workload-sha256")) $workloadHash }} +{{- fail "DynamoGraphDeployment workload binding annotation must match workloadBinding.sha256" }} +{{- end }} +{{- $workload := mustFromJson $workloadCanonical }} +{{- if ne (len $workload) 4 }} +{{- fail "workload binding must contain only controllerMode, orchestrator, storage, and trainer" }} +{{- end }} +{{- $controllerMode := required "canonical controllerMode is required" (index $workload "controllerMode") }} +{{- if and (ne $controllerMode "chartManaged") (ne $controllerMode "external") }} +{{- fail "canonical controllerMode must be chartManaged or external" }} +{{- end }} +{{- if ne (required "inference.dynamoGraph.controllerMode is required" $graph.controllerMode) $controllerMode }} +{{- fail "inference.dynamoGraph.controllerMode must match the workload binding" }} +{{- end }} +{{- $frontend := required "DynamoGraphDeployment Frontend is required" (index $services "Frontend") }} +{{- $frontendPod := required "DynamoGraphDeployment Frontend extraPodSpec is required" $frontend.extraPodSpec }} +{{- $orchestratorWorkload := required "canonical orchestrator workload is required" (index $workload "orchestrator") }} +{{- $trainerWorkload := required "canonical trainer workload is required" (index $workload "trainer") }} +{{- $storageWorkload := required "canonical storage workload is required" (index $workload "storage") }} +{{- if or (not (hasKey $orchestratorWorkload "enabled")) (not (hasKey $trainerWorkload "enabled")) }} +{{- fail "canonical controller workloads must define enabled" }} +{{- end }} +{{- $orchestratorEnabled := index $orchestratorWorkload "enabled" }} +{{- $trainerEnabled := index $trainerWorkload "enabled" }} +{{- if ne $orchestratorEnabled .Values.orchestrator.enabled }} +{{- fail "orchestrator.enabled must match the workload binding" }} +{{- end }} +{{- if ne $trainerEnabled .Values.trainer.enabled }} +{{- fail "trainer.enabled must match the workload binding" }} +{{- end }} +{{- $orchestratorGPU := required "canonical orchestrator gpu contract is required" (index $orchestratorWorkload "gpu") }} +{{- if not (hasKey $orchestratorGPU "enabled") }} +{{- fail "canonical orchestrator gpu contract must define enabled" }} +{{- end }} +{{- if index $orchestratorGPU "enabled" }} +{{- fail "canonical orchestrator GPU capability must be disabled" }} +{{- end }} +{{- range $componentName, $resources := dict "orchestrator" .Values.orchestrator.resources "trainer" .Values.trainer.resources }} +{{- range $resourceScope, $entries := dict "limits" (default (dict) $resources.limits) "requests" (default (dict) $resources.requests) }} +{{- range $resourceName, $_ := $entries }} +{{- if hasPrefix "nvidia.com/" $resourceName }} +{{- fail (printf "%s resources cannot set NVIDIA extended resource %s in %s" $componentName $resourceName $resourceScope) }} +{{- end }} +{{- end }} +{{- end }} +{{- end }} +{{- $trainerGPU := required "canonical trainer gpu contract is required" (index $trainerWorkload "gpu") }} +{{- if or (not (hasKey $trainerGPU "enabled")) (not (hasKey $trainerGPU "count")) }} +{{- fail "canonical trainer gpu contract must define enabled and count" }} +{{- end }} +{{- $trainerGPUEnabled := index $trainerGPU "enabled" }} +{{- $trainerGPUCount := int (index $trainerGPU "count") }} +{{- if eq $controllerMode "chartManaged" }} +{{- if or (not $orchestratorEnabled) (not $trainerEnabled) (not $trainerGPUEnabled) (lt $trainerGPUCount 1) }} +{{- fail "chartManaged mode requires orchestrator, trainer, and trainer GPU capability" }} +{{- end }} +{{- else }} +{{- if or $orchestratorEnabled $trainerEnabled $trainerGPUEnabled (ne $trainerGPUCount 0) }} +{{- fail "external mode forbids chart-managed controller workloads and trainer GPU capability" }} +{{- end }} +{{- end }} +{{- $actualTrainerGPU := dict "enabled" .Values.trainer.gpu.enabled "count" (int .Values.trainer.gpu.count) }} +{{- if ne (toJson $trainerGPU) (toJson $actualTrainerGPU) }} +{{- fail "trainer GPU configuration must match the workload binding" }} +{{- end }} +{{- if or (not (hasKey $storageWorkload "enabled")) (not (hasKey $storageWorkload "existingClaim")) (not (hasKey $storageWorkload "mountPath")) }} +{{- fail "canonical storage workload must define enabled, existingClaim, and mountPath" }} +{{- end }} +{{- $storageEnabled := index $storageWorkload "enabled" }} +{{- $storageClaim := index $storageWorkload "existingClaim" }} +{{- if and (eq $controllerMode "chartManaged") (not $storageEnabled) }} +{{- fail "chartManaged mode requires chart storage" }} +{{- end }} +{{- if and (eq $controllerMode "external") $storageEnabled (empty $storageClaim) }} +{{- fail "external mode requires an existing claim when chart storage is enabled" }} +{{- end }} +{{- $actualStorage := dict "enabled" .Values.storage.enabled "existingClaim" .Values.storage.existingClaim "mountPath" .Values.storage.mountPath }} +{{- if ne (toJson $storageWorkload) (toJson $actualStorage) }} +{{- fail "storage configuration must match the workload binding" }} +{{- end }} +{{- $orchestratorPlacement := required "canonical orchestrator placement is required" (index $orchestratorWorkload "placement") }} +{{- $expectedOrchestratorPlacement := dict + "nodeSelector" (required "Frontend nodeSelector is required" $frontendPod.nodeSelector) + "tolerations" (required "Frontend tolerations are required" $frontendPod.tolerations) }} +{{- if ne (toJson $orchestratorPlacement) (toJson $expectedOrchestratorPlacement) }} +{{- fail "orchestrator placement must match the manifest-bound DGD frontend placement" }} +{{- end }} +{{- $trainerPlacement := required "canonical trainer placement is required" (index $trainerWorkload "placement") }} +{{- range $serviceName, $service := dict "VllmDecodeWorker" $decode "VllmPrefillWorker" $prefill }} +{{- $workerPod := required (printf "%s extraPodSpec is required" $serviceName) $service.extraPodSpec }} +{{- $expectedTrainerPlacement := dict + "runtimeClassName" (required (printf "%s runtimeClassName is required" $serviceName) $workerPod.runtimeClassName) + "nodeSelector" (required (printf "%s nodeSelector is required" $serviceName) $workerPod.nodeSelector) + "tolerations" (required (printf "%s tolerations are required" $serviceName) $workerPod.tolerations) }} +{{- if ne (toJson $trainerPlacement) (toJson $expectedTrainerPlacement) }} +{{- fail "trainer placement must match every manifest-bound DGD worker placement" }} +{{- end }} +{{- end }} +{{- $frontendLabels := required "DynamoGraphDeployment Frontend extraPodMetadata.labels is required" $frontend.extraPodMetadata.labels }} +{{- $expectedFrontendLabels := dict "app.kubernetes.io/name" (include "prime-rl.name" .) "app.kubernetes.io/instance" .Release.Name }} +{{- if ne (toJson $frontendLabels) (toJson $expectedFrontendLabels) }} +{{- fail "DynamoGraphDeployment Frontend labels must match the chart selector labels" }} +{{- end }} {{- $manifestCanonical := required "inference.dynamoGraph.manifestCanonical is required" $graph.manifestCanonical }} {{- $manifestHash := required "DynamoGraphDeployment manifest-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/manifest-sha256") }} {{- if ne (sha256sum $manifestCanonical) $manifestHash }} @@ -102,6 +221,9 @@ {{- if ne (required "engineConfig topology-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/topology-sha256")) $topologyHash }} {{- fail "engineConfig topology-sha256 annotation must match topologyBinding.sha256" }} {{- end }} +{{- if ne (required "engineConfig workload-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/workload-sha256")) $workloadHash }} +{{- fail "engineConfig workload-sha256 annotation must match workloadBinding.sha256" }} +{{- end }} {{- if ne (toJson $engineAnnotations) (toJson $resourceAnnotations) }} {{- fail "engineConfig annotations must match DynamoGraphDeployment identity annotations" }} {{- end }} diff --git a/k8s/prime-rl/templates/dynamo-rl-service.yaml b/k8s/prime-rl/templates/dynamo-rl-service.yaml index 371f5a3d08..fb2597eb40 100644 --- a/k8s/prime-rl/templates/dynamo-rl-service.yaml +++ b/k8s/prime-rl/templates/dynamo-rl-service.yaml @@ -12,6 +12,7 @@ metadata: spec: type: ClusterIP selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} nvidia.com/dynamo-graph-deployment-name: {{ .Release.Name }} nvidia.com/dynamo-component: Frontend nvidia.com/dynamo-component-type: frontend diff --git a/k8s/prime-rl/templates/service.yaml b/k8s/prime-rl/templates/service.yaml index 6f7bf66be9..cc620ff187 100644 --- a/k8s/prime-rl/templates/service.yaml +++ b/k8s/prime-rl/templates/service.yaml @@ -12,6 +12,7 @@ metadata: spec: type: {{ .Values.orchestrator.service.type }} selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: orchestrator ports: @@ -37,6 +38,7 @@ metadata: spec: clusterIP: None # Headless service for StatefulSet selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: orchestrator ports: @@ -60,6 +62,7 @@ metadata: spec: type: {{ .Values.inference.service.type }} selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: inference ports: @@ -84,6 +87,7 @@ metadata: spec: type: {{ .Values.trainer.service.type }} selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: trainer ports: @@ -111,6 +115,7 @@ metadata: spec: clusterIP: None # Headless service for StatefulSet selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: inference ports: @@ -131,6 +136,7 @@ metadata: spec: clusterIP: None # Headless service for StatefulSet selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: trainer ports: diff --git a/k8s/prime-rl/values.schema.json b/k8s/prime-rl/values.schema.json index 5afebd8fdc..d6de07f9c7 100644 --- a/k8s/prime-rl/values.schema.json +++ b/k8s/prime-rl/values.schema.json @@ -81,12 +81,17 @@ "prime-rl.nvidia.com/topology-sha256": { "type": "string", "pattern": "^[0-9a-f]{64}$" + }, + "prime-rl.nvidia.com/workload-sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" } }, "required": [ "prime-rl.nvidia.com/config-sha256", "prime-rl.nvidia.com/manifest-sha256", - "prime-rl.nvidia.com/topology-sha256" + "prime-rl.nvidia.com/topology-sha256", + "prime-rl.nvidia.com/workload-sha256" ], "additionalProperties": {"type": "string"} }, @@ -114,6 +119,22 @@ "required": ["sha256", "canonical"], "additionalProperties": false }, + "controllerMode": { + "type": "string", + "enum": ["chartManaged", "external"] + }, + "workloadBinding": { + "type": "object", + "properties": { + "sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "canonical": {"type": "string", "minLength": 1} + }, + "required": ["sha256", "canonical"], + "additionalProperties": false + }, "manifestCanonical": {"type": "string", "minLength": 1}, "resource": {"type": "object"} }, @@ -145,6 +166,7 @@ "orchestrator": { "type": "object", "properties": { + "enabled": {"type": "boolean"}, "runtimeClassName": {"type": "string"}, "nodeSelector": { "type": "object", @@ -172,6 +194,15 @@ "trainer": { "type": "object", "properties": { + "enabled": {"type": "boolean"}, + "gpu": { + "type": "object", + "properties": { + "enabled": {"type": "boolean"}, + "count": {"type": "integer", "minimum": 0} + }, + "required": ["enabled", "count"] + }, "runtimeClassName": {"type": "string"}, "nodeSelector": { "type": "object", @@ -218,6 +249,7 @@ }, "inference": { "properties": { + "enabled": {"const": true}, "dynamoGraph": { "properties": { "engineConfig": {"minProperties": 5}, @@ -227,28 +259,72 @@ "clientTopology", "engineConfig", "topologyBinding", + "controllerMode", + "workloadBinding", "manifestCanonical", "resource" ] } }, - "required": ["dynamoGraph"] - }, - "orchestrator": { + "required": ["enabled", "dynamoGraph"] + } + }, + "oneOf": [ + { "properties": { - "nodeSelector": {"minProperties": 2}, - "tolerations": {"minItems": 2} - }, - "required": ["nodeSelector", "tolerations"] + "inference": { + "properties": { + "dynamoGraph": { + "properties": {"controllerMode": {"const": "chartManaged"}} + } + } + }, + "orchestrator": { + "properties": {"enabled": {"const": true}}, + "required": ["enabled"] + }, + "trainer": { + "properties": { + "enabled": {"const": true}, + "gpu": { + "properties": { + "enabled": {"const": true}, + "count": {"minimum": 1} + } + } + }, + "required": ["enabled", "gpu"] + } + } }, - "trainer": { + { "properties": { - "nodeSelector": {"minProperties": 3}, - "tolerations": {"minItems": 3} - }, - "required": ["runtimeClassName", "nodeSelector", "tolerations"] + "inference": { + "properties": { + "dynamoGraph": { + "properties": {"controllerMode": {"const": "external"}} + } + } + }, + "orchestrator": { + "properties": {"enabled": {"const": false}}, + "required": ["enabled"] + }, + "trainer": { + "properties": { + "enabled": {"const": false}, + "gpu": { + "properties": { + "enabled": {"const": false}, + "count": {"const": 0} + } + } + }, + "required": ["enabled", "gpu"] + } + } } - }, + ], "required": ["image", "inference", "orchestrator", "trainer"] } }, diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index cb44434940..35fc81fa3c 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -26,6 +26,7 @@ MANIFEST_HASH_ANNOTATION = "prime-rl.nvidia.com/manifest-sha256" MANIFEST_HASH_SCOPE_ANNOTATION = "prime-rl.nvidia.com/manifest-sha256-scope" TOPOLOGY_HASH_ANNOTATION = "prime-rl.nvidia.com/topology-sha256" +WORKLOAD_HASH_ANNOTATION = "prime-rl.nvidia.com/workload-sha256" MANIFEST_HASH_SCOPE = ( "resource; json.dumps(sort_keys=true,indent=2)+newline; " "exclude=/metadata/annotations/prime-rl.nvidia.com~1manifest-sha256" @@ -148,8 +149,9 @@ def __post_init__(self) -> None: raise ValueError(f"GPU scheduling fields must not be empty: {empty}") if not self.tolerations: raise ValueError("GPU scheduling requires at least one toleration") - if not any(toleration.key == "nvidia.com/gpu" for toleration in self.tolerations): - raise ValueError("GPU scheduling must tolerate the nvidia.com/gpu taint") + required_gpu_toleration = KubernetesToleration(key="nvidia.com/gpu") + if required_gpu_toleration not in self.tolerations: + raise ValueError("GPU scheduling requires nvidia.com/gpu Exists NoSchedule") @property def image_node_selector(self) -> dict[str, str]: @@ -173,6 +175,7 @@ def image_tolerations(self) -> tuple[KubernetesToleration, ...]: operator="Equal", value=self.architecture, ), + KubernetesToleration(key="nvidia.com/gpu"), KubernetesToleration(key="prime-rl", operator="Equal", value="true"), ) return _unique_tolerations((*required, *self.additional_image_tolerations)) @@ -201,6 +204,8 @@ class DynamoGraphRenderOptions: image_digest: str run_name: str gpu_scheduling: GPUSchedulingProfile + external_controller: bool = False + trainer_gpu_count: int = 1 model_cache_pvc: str | None = None shared_pvc: str | None = None image_pull_secrets: tuple[str, ...] = () @@ -217,6 +222,8 @@ def __post_init__(self) -> None: image_tag = self.image.rsplit("@", 1)[0] if self.prime_sha[:12] not in image_tag or self.dynamo_sha[:12] not in image_tag: raise ValueError("DGD image tag must include the Prime and Dynamo commit suffixes") + if self.trainer_gpu_count < 1: + raise ValueError("trainer_gpu_count must be at least one") def _sha256_bytes(value: bytes) -> str: @@ -253,6 +260,13 @@ def _worker_topology_binding(services: dict[str, Any]) -> dict[str, Any]: } +def _release_pod_labels(options: DynamoGraphRenderOptions) -> dict[str, str]: + return { + "app.kubernetes.io/name": "prime-rl", + "app.kubernetes.io/instance": options.release_name, + } + + def _worker_env(process: DynamoProcessSpec) -> list[dict[str, Any]]: values: list[dict[str, Any]] = [ {"name": name, "value": value} for name, value in sorted(process.environment().items()) @@ -338,6 +352,7 @@ def _worker_service( "componentType": "worker", "subComponentType": role, "replicas": replicas, + "extraPodMetadata": {"labels": _release_pod_labels(options)}, "sharedMemory": {"size": "64Gi"}, "resources": { "requests": {"gpu": str(config.deployment.gpus_per_node)}, @@ -474,6 +489,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) frontend = { "componentType": "frontend", "replicas": 1, + "extraPodMetadata": {"labels": _release_pod_labels(options)}, "extraPodSpec": frontend_pod_spec, } prefill = _worker_service( @@ -510,6 +526,39 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) } topology_canonical = _canonical_json(topology_binding) topology_hash = _sha256_bytes(topology_canonical) + chart_controller_enabled = not options.external_controller + controller_mode = "external" if options.external_controller else "chartManaged" + storage_enabled = chart_controller_enabled or options.shared_pvc is not None + workload = { + "controllerMode": controller_mode, + "orchestrator": { + "enabled": chart_controller_enabled, + "gpu": {"enabled": False}, + "placement": { + "nodeSelector": options.gpu_scheduling.image_node_selector, + "tolerations": options.gpu_scheduling.image_toleration_manifests, + }, + }, + "trainer": { + "enabled": chart_controller_enabled, + "gpu": { + "enabled": chart_controller_enabled, + "count": options.trainer_gpu_count if chart_controller_enabled else 0, + }, + "placement": { + "runtimeClassName": options.gpu_scheduling.runtime_class_name, + "nodeSelector": options.gpu_scheduling.node_selector, + "tolerations": options.gpu_scheduling.toleration_manifests, + }, + }, + "storage": { + "enabled": storage_enabled, + "existingClaim": options.shared_pvc or "", + "mountPath": "/data", + }, + } + workload_canonical = _canonical_json(workload) + workload_hash = _sha256_bytes(workload_canonical) annotations = { ENGINE_CONFIG_HASH_ANNOTATION: engine_hash, "prime-rl.nvidia.com/dynamo-sha": options.dynamo_sha, @@ -518,6 +567,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "prime-rl.nvidia.com/run-name": options.run_name, MANIFEST_HASH_SCOPE_ANNOTATION: MANIFEST_HASH_SCOPE, TOPOLOGY_HASH_ANNOTATION: topology_hash, + WORKLOAD_HASH_ANNOTATION: workload_hash, } resource: dict[str, Any] = { "apiVersion": "nvidia.com/v1alpha1", @@ -560,8 +610,10 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "pullSecrets": list(options.image_pull_secrets), }, "inference": { + "enabled": True, "mode": "dynamoGraph", "dynamoGraph": { + "controllerMode": controller_mode, "clientTopology": client_topology, "engineConfig": { "name": config_map_name, @@ -574,20 +626,25 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "sha256": topology_hash, "canonical": topology_canonical.decode(), }, + "workloadBinding": { + "sha256": workload_hash, + "canonical": workload_canonical.decode(), + }, "manifestCanonical": manifest_canonical.decode(), "resource": resource, }, }, + "orchestrator": {"enabled": chart_controller_enabled}, "trainer": { - "runtimeClassName": options.gpu_scheduling.runtime_class_name, - "nodeSelector": options.gpu_scheduling.node_selector, - "tolerations": options.gpu_scheduling.toleration_manifests, - }, - "orchestrator": { - "nodeSelector": options.gpu_scheduling.image_node_selector, - "tolerations": options.gpu_scheduling.image_toleration_manifests, + "enabled": chart_controller_enabled, + "gpu": { + "enabled": chart_controller_enabled, + "count": options.trainer_gpu_count if chart_controller_enabled else 0, + }, }, } + if options.external_controller: + values["storage"] = {"enabled": False} if options.model_cache_pvc: values["modelCache"] = { "enabled": True, @@ -643,6 +700,17 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--gpu-product", required=True) parser.add_argument("--gpu-node-pool", required=True) parser.add_argument("--gpu-node-pool-label", default="cloud.google.com/gke-nodepool") + parser.add_argument( + "--external-controller", + action="store_true", + help="Render only DGD inference workloads; an external controller owns orchestration and training", + ) + parser.add_argument( + "--trainer-gpus", + type=int, + default=1, + help="Exact GPU request and limit for the chart-managed trainer", + ) parser.add_argument( "--image-toleration", action="append", @@ -685,6 +753,8 @@ def main() -> None: additional_image_tolerations=tuple(args.image_toleration), additional_gpu_tolerations=tuple(args.gpu_toleration), ), + external_controller=args.external_controller, + trainer_gpu_count=args.trainer_gpus, model_cache_pvc=args.model_cache_pvc, shared_pvc=args.shared_pvc, image_pull_secrets=tuple(args.image_pull_secret), diff --git a/tests/unit/inference/helm_dgd_test_utils.py b/tests/unit/inference/helm_dgd_test_utils.py new file mode 100644 index 0000000000..83fad508b0 --- /dev/null +++ b/tests/unit/inference/helm_dgd_test_utils.py @@ -0,0 +1,150 @@ +import hashlib +import json +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference.dgd import DynamoGraphRenderOptions, GPUSchedulingProfile + +HELM = shutil.which("helm") +CHART = Path(__file__).parents[3] / "k8s" / "prime-rl" +PRIME_SHA = "1" * 40 +DYNAMO_SHA = "2" * 40 +IMAGE_DIGEST = f"sha256:{'3' * 64}" +GPU_SCHEDULING = GPUSchedulingProfile( + runtime_class_name="nvidia", + architecture="arm64", + product="NVIDIA-GB200", + node_pool="customer-gpu-o7v", +) + + +def inference_config( + weight_broadcast: str = "nccl", + *, + chat_template: str | None = None, +) -> InferenceConfig: + model = {"chat_template": chat_template} if chat_template is not None else {} + return InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "model": model, + "weight_broadcast": {"type": weight_broadcast}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + ) + + +def render_options( + tmp_path: Path, + *, + external_controller: bool = False, + release_name: str = "p4-math", + shared_pvc: str | None = None, + trainer_gpu_count: int = 1, +) -> DynamoGraphRenderOptions: + return DynamoGraphRenderOptions( + release_name=release_name, + namespace="bis-vllm", + image=f"nvcr.io/example/prime:prime-{PRIME_SHA[:12]}-dynamo-{DYNAMO_SHA[:12]}@{IMAGE_DIGEST}", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, + external_controller=external_controller, + trainer_gpu_count=trainer_gpu_count, + model_cache_pvc="model-cache", + hf_token_secret="hf-token-secret", + shared_pvc=shared_pvc, + image_pull_secrets=("nvcrimagepullsecret",), + ) + + +def helm_template(*args: str, release_name: str = "p4-math") -> str: + if HELM is None: + pytest.skip("helm is not installed") + return subprocess.run( + [HELM, "template", release_name, str(CHART), *args], + check=True, + capture_output=True, + text=True, + ).stdout + + +def rendered_documents(rendered: str) -> list[dict]: + return [document for document in yaml.safe_load_all(rendered) if document] + + +def rendered_resource(rendered: str, kind: str, name: str) -> dict: + return next( + document + for document in rendered_documents(rendered) + if document.get("kind") == kind and document.get("metadata", {}).get("name") == name + ) + + +def labels_match(selector: dict[str, str], labels: dict[str, str]) -> bool: + return selector.items() <= labels.items() + + +def toleration_identity(toleration: dict[str, str]) -> tuple[str, str, str | None, str | None]: + return ( + toleration["key"], + toleration["operator"], + toleration.get("value"), + toleration.get("effect"), + ) + + +def write_values_mutation( + source: Path, + target: Path, + path: tuple[str, ...], + replacement: object, +) -> None: + values = json.loads(source.read_text()) + parent = values + for key in path[:-1]: + parent = parent[key] + parent[path[-1]] = replacement + target.write_text(json.dumps(values)) + + +def canonical_json(value: object) -> str: + return json.dumps(value, indent=2, sort_keys=True) + "\n" + + +def rewrite_valid_integrity(values: dict, workload: dict) -> None: + graph = values["inference"]["dynamoGraph"] + workload_canonical = canonical_json(workload) + workload_hash = hashlib.sha256(workload_canonical.encode()).hexdigest() + graph["workloadBinding"] = { + "canonical": workload_canonical, + "sha256": workload_hash, + } + + resource = graph["resource"] + resource_annotations = resource["metadata"]["annotations"] + resource_annotations["prime-rl.nvidia.com/workload-sha256"] = workload_hash + graph["engineConfig"]["annotations"]["prime-rl.nvidia.com/workload-sha256"] = workload_hash + + scoped_resource = json.loads(json.dumps(resource)) + scoped_resource["metadata"]["annotations"].pop("prime-rl.nvidia.com/manifest-sha256") + manifest_canonical = canonical_json(scoped_resource) + manifest_hash = hashlib.sha256(manifest_canonical.encode()).hexdigest() + graph["manifestCanonical"] = manifest_canonical + resource_annotations["prime-rl.nvidia.com/manifest-sha256"] = manifest_hash + graph["engineConfig"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"] = manifest_hash diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index 59c1dbf3b8..27d417b1c1 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -57,7 +57,13 @@ def inference_config( ) -def render_options(tmp_path: Path) -> DynamoGraphRenderOptions: +def render_options( + tmp_path: Path, + *, + external_controller: bool = False, + shared_pvc: str | None = "p4-shared-data", + trainer_gpu_count: int = 1, +) -> DynamoGraphRenderOptions: return DynamoGraphRenderOptions( release_name="p4-math", namespace="bis-vllm", @@ -68,8 +74,10 @@ def render_options(tmp_path: Path) -> DynamoGraphRenderOptions: image_digest=IMAGE_DIGEST, run_name="p4-run", gpu_scheduling=GPU_SCHEDULING, + external_controller=external_controller, + trainer_gpu_count=trainer_gpu_count, model_cache_pvc="model-cache", - shared_pvc="p4-shared-data", + shared_pvc=shared_pvc, image_pull_secrets=("nvcrimagepullsecret",), hf_token_secret="hf-token-secret", ) @@ -136,6 +144,9 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): }, } assert resource["metadata"]["annotations"]["prime-rl.nvidia.com/topology-sha256"] == topology_binding["sha256"] + workload_binding = graph["workloadBinding"] + assert hashlib.sha256(workload_binding["canonical"].encode()).hexdigest() == workload_binding["sha256"] + assert resource["metadata"]["annotations"]["prime-rl.nvidia.com/workload-sha256"] == workload_binding["sha256"] for service in services.values(): pod_spec = service["extraPodSpec"] assert pod_spec["mainContainer"]["image"] == options.image @@ -152,6 +163,11 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): "value": "arm64", "effect": "NoSchedule", }, + { + "key": "nvidia.com/gpu", + "operator": "Exists", + "effect": "NoSchedule", + }, { "key": "prime-rl", "operator": "Equal", @@ -159,16 +175,17 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): "effect": "NoSchedule", }, ] - gpu_tolerations = [ - *image_tolerations, - {"key": "nvidia.com/gpu", "operator": "Exists", "effect": "NoSchedule"}, - ] + gpu_tolerations = image_tolerations image_selector = { "kubernetes.io/arch": "arm64", "cloud.google.com/gke-nodepool": "customer-gpu-o7v", } gpu_selector = {**image_selector, "nvidia.com/gpu.product": "NVIDIA-GB200"} frontend_pod = services["Frontend"]["extraPodSpec"] + assert services["Frontend"]["extraPodMetadata"]["labels"] == { + "app.kubernetes.io/name": "prime-rl", + "app.kubernetes.io/instance": "p4-math", + } assert frontend_pod["nodeSelector"] == image_selector assert frontend_pod["tolerations"] == image_tolerations assert "runtimeClassName" not in frontend_pod @@ -181,13 +198,32 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): assert worker_pod["runtimeClassName"] == "nvidia" assert worker_pod["nodeSelector"] == gpu_selector assert worker_pod["tolerations"] == gpu_tolerations - assert values["orchestrator"] == { - "nodeSelector": image_selector, - "tolerations": image_tolerations, + workload = json.loads(workload_binding["canonical"]) + assert workload["controllerMode"] == "chartManaged" + assert workload["storage"] == { + "enabled": True, + "existingClaim": "p4-shared-data", + "mountPath": "/data", + } + assert workload["orchestrator"] == { + "enabled": True, + "gpu": {"enabled": False}, + "placement": { + "nodeSelector": image_selector, + "tolerations": image_tolerations, + }, + } + assert workload["trainer"] == { + "enabled": True, + "gpu": {"enabled": True, "count": 1}, + "placement": { + "runtimeClassName": "nvidia", + "nodeSelector": gpu_selector, + "tolerations": gpu_tolerations, + }, } - assert values["trainer"]["runtimeClassName"] == "nvidia" - assert values["trainer"]["nodeSelector"] == gpu_selector - assert values["trainer"]["tolerations"] == gpu_tolerations + assert values["orchestrator"] == {"enabled": True} + assert values["trainer"] == {"enabled": True, "gpu": {"enabled": True, "count": 1}} prefill_args = services["VllmPrefillWorker"]["extraPodSpec"]["mainContainer"]["args"] decode_args = services["VllmDecodeWorker"]["extraPodSpec"]["mainContainer"]["args"] @@ -422,6 +458,26 @@ def test_cli_toleration_parser_is_typed_and_rejects_unknown_fields(): _parse_kubernetes_toleration('{"key":"dedicated","operator":"NotARealOperator"}') +@pytest.mark.parametrize( + "toleration", + [ + KubernetesToleration(key="nvidia.com/gpu", effect="NoExecute"), + KubernetesToleration(key="nvidia.com/gpu", operator="Equal", value="true"), + ], +) +def test_gpu_scheduling_requires_exact_nodepool_access_toleration( + toleration: KubernetesToleration, +): + with pytest.raises(ValueError, match="nvidia.com/gpu Exists NoSchedule"): + GPUSchedulingProfile( + runtime_class_name="nvidia", + architecture="arm64", + product="NVIDIA-GB200", + node_pool="customer-gpu-o7v", + tolerations=(toleration,), + ) + + def test_cli_accepts_typed_additional_image_and_gpu_tolerations(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( sys, @@ -449,6 +505,9 @@ def test_cli_accepts_typed_additional_image_and_gpu_tolerations(monkeypatch: pyt "NVIDIA-GB200", "--gpu-node-pool", "customer-gpu-o7v", + "--external-controller", + "--trainer-gpus", + "4", "--image-toleration", '{"key":"image-extra","operator":"Exists"}', "--gpu-toleration", @@ -460,6 +519,8 @@ def test_cli_accepts_typed_additional_image_and_gpu_tolerations(monkeypatch: pyt assert args.image_toleration == [KubernetesToleration(key="image-extra")] assert args.gpu_toleration == [KubernetesToleration(key="gpu-extra", operator="Equal", value="true")] + assert args.external_controller is True + assert args.trainer_gpus == 4 def test_gpu_scheduling_changes_manifest_identity(tmp_path: Path): @@ -476,6 +537,35 @@ def test_gpu_scheduling_changes_manifest_identity(tmp_path: Path): ) +def test_external_controller_binding_disables_chart_workloads(tmp_path: Path): + values = build_dgd_values( + inference_config(), + render_options(tmp_path, external_controller=True, shared_pvc=None), + ) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + + assert graph["controllerMode"] == "external" + assert workload["controllerMode"] == "external" + assert workload["orchestrator"]["enabled"] is False + assert workload["orchestrator"]["gpu"] == {"enabled": False} + assert workload["trainer"]["enabled"] is False + assert workload["trainer"]["gpu"] == {"enabled": False, "count": 0} + assert workload["storage"] == { + "enabled": False, + "existingClaim": "", + "mountPath": "/data", + } + assert values["orchestrator"] == {"enabled": False} + assert values["trainer"] == {"enabled": False, "gpu": {"enabled": False, "count": 0}} + assert values["storage"] == {"enabled": False} + + +def test_trainer_gpu_count_must_be_positive(): + with pytest.raises(ValueError, match="trainer_gpu_count"): + replace(render_options(Path("/tmp/not-written")), trainer_gpu_count=0) + + def test_dgd_rejects_native_backend(tmp_path: Path): config = InferenceConfig.model_validate({}) with pytest.raises(ValueError, match="Dynamo disaggregated"): diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index 583404e7a0..a9347ea0bb 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -1,104 +1,25 @@ import hashlib import json -import shutil import subprocess from pathlib import Path import pytest -import yaml - -from prime_rl.configs.inference import InferenceConfig -from prime_rl.inference.dgd import DynamoGraphRenderOptions, GPUSchedulingProfile, write_dgd_artifacts - -HELM = shutil.which("helm") -CHART = Path(__file__).parents[3] / "k8s" / "prime-rl" -PRIME_SHA = "1" * 40 -DYNAMO_SHA = "2" * 40 -IMAGE_DIGEST = f"sha256:{'3' * 64}" -GPU_SCHEDULING = GPUSchedulingProfile( - runtime_class_name="nvidia", - architecture="arm64", - product="NVIDIA-GB200", - node_pool="customer-gpu-o7v", -) - - -def inference_config( - weight_broadcast: str = "nccl", - *, - chat_template: str | None = None, -) -> InferenceConfig: - model = {"chat_template": chat_template} if chat_template is not None else {} - return InferenceConfig.model_validate( - { - "backend": {"type": "dynamo"}, - "model": model, - "weight_broadcast": {"type": weight_broadcast}, - "deployment": { - "type": "disaggregated", - "gpus_per_node": 1, - "prefill_nodes_per_replica": 1, - "decode_nodes_per_replica": 1, - "num_prefill_replicas": 2, - "num_decode_replicas": 2, - }, - } - ) - - -def render_options(tmp_path: Path, *, shared_pvc: str | None = None) -> DynamoGraphRenderOptions: - return DynamoGraphRenderOptions( - release_name="p4-math", - namespace="bis-vllm", - image=f"nvcr.io/example/prime:prime-{PRIME_SHA[:12]}-dynamo-{DYNAMO_SHA[:12]}@{IMAGE_DIGEST}", - output_dir=tmp_path, - prime_sha=PRIME_SHA, - dynamo_sha=DYNAMO_SHA, - image_digest=IMAGE_DIGEST, - run_name="p4-run", - gpu_scheduling=GPU_SCHEDULING, - model_cache_pvc="model-cache", - hf_token_secret="hf-token-secret", - shared_pvc=shared_pvc, - image_pull_secrets=("nvcrimagepullsecret",), - ) - - -def helm_template(*args: str, release_name: str = "p4-math") -> str: - if HELM is None: - pytest.skip("helm is not installed") - return subprocess.run( - [HELM, "template", release_name, str(CHART), *args], - check=True, - capture_output=True, - text=True, - ).stdout - -def rendered_documents(rendered: str) -> list[dict]: - return [document for document in yaml.safe_load_all(rendered) if document] - - -def rendered_resource(rendered: str, kind: str, name: str) -> dict: - return next( - document - for document in rendered_documents(rendered) - if document.get("kind") == kind and document.get("metadata", {}).get("name") == name - ) - - -def write_values_mutation( - source: Path, - target: Path, - path: tuple[str, ...], - replacement: object, -) -> None: - values = json.loads(source.read_text()) - parent = values - for key in path[:-1]: - parent = parent[key] - parent[path[-1]] = replacement - target.write_text(json.dumps(values)) +from prime_rl.inference.dgd import DynamoGraphRenderOptions, write_dgd_artifacts +from tests.unit.inference.helm_dgd_test_utils import ( + DYNAMO_SHA, + GPU_SCHEDULING, + IMAGE_DIGEST, + PRIME_SHA, + helm_template, + inference_config, + labels_match, + render_options, + rendered_documents, + rendered_resource, + rewrite_valid_integrity, + toleration_identity, +) def test_native_chart_still_renders_inference_statefulset(): @@ -139,16 +60,139 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert rendered.count("claimName: model-cache") == 2 assert rendered.count("name: HF_TOKEN") == 5 assert rendered.count("name: HF_HOME") == 5 - assert rendered.count("cloud.google.com/gke-nodepool: customer-gpu-o7v") == 5 - assert rendered.count("kubernetes.io/arch: arm64") == 5 - assert rendered.count("nvidia.com/gpu.product: NVIDIA-GB200") == 3 - assert rendered.count("runtimeClassName: nvidia") == 3 - assert rendered.count("key: kubernetes.io/arch") == 5 - assert rendered.count("key: nvidia.com/gpu") == 3 - assert rendered.count("key: prime-rl") == 5 + chart_pods = { + component: rendered_resource(rendered, "StatefulSet", f"p4-math-{component}")["spec"]["template"]["spec"] + for component in ("orchestrator", "trainer") + } + dgd_pods = { + component: rendered_graph["spec"]["services"][service]["extraPodSpec"] + for component, service in { + "frontend": "Frontend", + "prefill": "VllmPrefillWorker", + "decode": "VllmDecodeWorker", + }.items() + } + pods = {**chart_pods, **dgd_pods} + image_selector = { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + } + gpu_selector = {**image_selector, "nvidia.com/gpu.product": "NVIDIA-GB200"} + required_tolerations = { + ("kubernetes.io/arch", "Equal", "arm64", "NoSchedule"), + ("nvidia.com/gpu", "Exists", None, "NoSchedule"), + ("prime-rl", "Equal", "true", "NoSchedule"), + } + + assert set(pods) == {"orchestrator", "trainer", "frontend", "prefill", "decode"} + for component, pod in pods.items(): + assert {toleration_identity(item) for item in pod["tolerations"]} == required_tolerations + if component in {"trainer", "prefill", "decode"}: + assert pod["nodeSelector"] == gpu_selector + assert pod["runtimeClassName"] == "nvidia" + else: + assert pod["nodeSelector"] == image_selector + assert "runtimeClassName" not in pod + + assert "nvidia.com/gpu" not in chart_pods["orchestrator"]["containers"][0]["resources"].get("requests", {}) + assert chart_pods["trainer"]["containers"][0]["resources"]["requests"]["nvidia.com/gpu"] == 1 assert not any(kind in rendered for kind in ("kind: ClusterRole", "kind: CustomResourceDefinition")) +def test_external_controller_mode_renders_only_five_dgd_inference_pods(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config(), + render_options(tmp_path, external_controller=True), + ) + rendered = helm_template("-f", str(paths["values"])) + graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + + assert "kind: StatefulSet" not in rendered + assert "kind: PersistentVolumeClaim" not in rendered + assert "p4-math-shared-data" not in rendered + assert sum(service["replicas"] for service in graph["spec"]["services"].values()) == 5 + assert rendered.count("kind: DynamoGraphDeployment") == 1 + assert rendered.count("name: p4-math-frontend-rl") == 1 + + +def test_chart_managed_trainer_uses_exact_bound_gpu_resources(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config(), + render_options(tmp_path, trainer_gpu_count=4), + ) + values = json.loads(paths["values"].read_text()) + rendered = helm_template("-f", str(paths["values"])) + trainer = rendered_resource(rendered, "StatefulSet", "p4-math-trainer") + resources = trainer["spec"]["template"]["spec"]["containers"][0]["resources"] + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + + assert workload["trainer"]["gpu"] == {"enabled": True, "count": 4} + assert resources["requests"]["nvidia.com/gpu"] == 4 + assert resources["limits"]["nvidia.com/gpu"] == 4 + + +def test_legacy_statefulset_selectors_remain_upgrade_compatible(): + documents = rendered_documents(helm_template()) + controllers = [document for document in documents if document["kind"] == "StatefulSet"] + + assert len(controllers) == 3 + for controller in controllers: + role = controller["metadata"]["labels"]["role"] + assert controller["spec"]["selector"]["matchLabels"] == { + "app": "prime-rl", + "example": "reverse-text", + "role": role, + } + assert controller["spec"]["template"]["metadata"]["labels"]["app.kubernetes.io/instance"] == "p4-math" + + +def test_chart_service_selectors_are_release_disjoint(): + releases = {release: rendered_documents(helm_template(release_name=release)) for release in ("alpha", "beta")} + pod_labels = { + release: [ + document["spec"]["template"]["metadata"]["labels"] + for document in documents + if document["kind"] == "StatefulSet" + ] + for release, documents in releases.items() + } + + for release, documents in releases.items(): + other_release = "beta" if release == "alpha" else "alpha" + services = [document for document in documents if document["kind"] == "Service"] + assert len(services) == 6 + + for service in services: + selector = service["spec"]["selector"] + assert selector["app.kubernetes.io/instance"] == release + assert any(labels_match(selector, labels) for labels in pod_labels[release]) + assert not any(labels_match(selector, labels) for labels in pod_labels[other_release]) + + +def test_dgd_rl_service_selector_is_release_disjoint(tmp_path: Path): + release_pods: dict[str, dict[str, str]] = {} + release_services: dict[str, dict[str, str]] = {} + for release in ("alpha", "beta"): + output_dir = tmp_path / release + paths = write_dgd_artifacts(inference_config(), render_options(output_dir, release_name=release)) + rendered = helm_template("-f", str(paths["values"]), release_name=release) + graph = rendered_resource(rendered, "DynamoGraphDeployment", release) + release_pods[release] = { + **graph["spec"]["services"]["Frontend"]["extraPodMetadata"]["labels"], + "nvidia.com/dynamo-graph-deployment-name": release, + "nvidia.com/dynamo-component": "Frontend", + "nvidia.com/dynamo-component-type": "frontend", + } + release_services[release] = rendered_resource(rendered, "Service", f"{release}-frontend-rl")["spec"]["selector"] + + for release in ("alpha", "beta"): + other_release = "beta" if release == "alpha" else "alpha" + selector = release_services[release] + assert selector["app.kubernetes.io/instance"] == release + assert labels_match(selector, release_pods[release]) + assert not labels_match(selector, release_pods[other_release]) + + def test_dgd_chart_renders_chat_template_configmap_and_frontend_mount(tmp_path: Path): paths = write_dgd_artifacts( inference_config(chat_template="template-marker: {{ messages }}"), @@ -185,170 +229,178 @@ def test_dgd_chart_preserves_exact_content_addressed_configmap_bytes(tmp_path: P assert hashlib.sha256(canonical_data).hexdigest() == expected_hash -@pytest.mark.parametrize( - ("case", "path", "replacement", "error_fragment"), - [ - ( - "engine-data", - ("inference", "dynamoGraph", "engineConfig", "data", "prefill-engine.json"), - "tampered", - "engineConfig.data must match its canonical payload", - ), - ( - "engine-sha", - ("inference", "dynamoGraph", "engineConfig", "sha256"), - "0" * 64, - "engineConfig.sha256 must match its canonical payload", - ), - ( - "engine-name", - ("inference", "dynamoGraph", "engineConfig", "name"), - "p4-math-dynamo-engine-000000000000", - "engineConfig.name must be content-addressed", - ), - ( - "dgd-config-sha", - ( - "inference", - "dynamoGraph", - "resource", - "metadata", - "annotations", - "prime-rl.nvidia.com/config-sha256", - ), - "0" * 64, - "DynamoGraphDeployment config-sha256 must match engineConfig.sha256", - ), - ( - "manifest-sha", - ( - "inference", - "dynamoGraph", - "resource", - "metadata", - "annotations", - "prime-rl.nvidia.com/manifest-sha256", - ), - "0" * 64, - "manifest-sha256 must match its canonical payload", - ), - ], -) -def test_dgd_chart_rejects_content_identity_mutations( - case: str, - path: tuple[str, ...], - replacement: object, - error_fragment: str, - tmp_path: Path, -): +def test_dgd_chart_uses_canonical_workload_contract_as_sole_authority(tmp_path: Path): paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) - mutation = tmp_path / f"{case}.json" - write_values_mutation(paths["values"], mutation, path, replacement) - - with pytest.raises(subprocess.CalledProcessError) as error: - helm_template("-f", str(mutation)) - - assert error_fragment in error.value.stderr + values = json.loads(paths["values"].read_text()) + workload_binding = values["inference"]["dynamoGraph"]["workloadBinding"] + workload = json.loads(workload_binding["canonical"]) + + assert hashlib.sha256(workload_binding["canonical"].encode()).hexdigest() == workload_binding["sha256"] + assert values["orchestrator"] == {"enabled": True} + assert values["trainer"] == {"enabled": True, "gpu": {"enabled": True, "count": 1}} + assert workload == { + "controllerMode": "chartManaged", + "orchestrator": { + "enabled": True, + "gpu": {"enabled": False}, + "placement": { + "nodeSelector": { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + }, + "tolerations": [ + { + "effect": "NoSchedule", + "key": "kubernetes.io/arch", + "operator": "Equal", + "value": "arm64", + }, + {"effect": "NoSchedule", "key": "nvidia.com/gpu", "operator": "Exists"}, + { + "effect": "NoSchedule", + "key": "prime-rl", + "operator": "Equal", + "value": "true", + }, + ], + }, + }, + "trainer": { + "enabled": True, + "gpu": {"count": 1, "enabled": True}, + "placement": { + "nodeSelector": { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + "nvidia.com/gpu.product": "NVIDIA-GB200", + }, + "runtimeClassName": "nvidia", + "tolerations": [ + { + "effect": "NoSchedule", + "key": "kubernetes.io/arch", + "operator": "Equal", + "value": "arm64", + }, + {"effect": "NoSchedule", "key": "nvidia.com/gpu", "operator": "Exists"}, + { + "effect": "NoSchedule", + "key": "prime-rl", + "operator": "Equal", + "value": "true", + }, + ], + }, + }, + "storage": { + "enabled": True, + "existingClaim": "", + "mountPath": "/data", + }, + } @pytest.mark.parametrize( - ("case", "path", "replacement"), + ("component", "selector_key"), [ - ( - "client-roles", - ("inference", "dynamoGraph", "clientTopology", "dynamo_worker_roles"), - ["prefill", "decode", "decode", "decode"], - ), - ( - "client-gpus", - ("inference", "dynamoGraph", "clientTopology", "dynamo_gpus_per_worker"), - 2, - ), - ( - "worker-replicas", - ( - "inference", - "dynamoGraph", - "resource", - "spec", - "services", - "VllmDecodeWorker", - "replicas", - ), - 3, - ), - ( - "worker-gpus", - ( - "inference", - "dynamoGraph", - "resource", - "spec", - "services", - "VllmPrefillWorker", - "resources", - "limits", - "gpu", - ), - "2", - ), + ("orchestrator", "kubernetes.io/arch"), + ("orchestrator", "cloud.google.com/gke-nodepool"), + ("trainer", "kubernetes.io/arch"), + ("trainer", "cloud.google.com/gke-nodepool"), + ("trainer", "nvidia.com/gpu.product"), ], ) -def test_dgd_chart_rejects_topology_mutations( - case: str, - path: tuple[str, ...], - replacement: object, +def test_dgd_chart_rejects_rehashed_chart_selector_mutations( + component: str, + selector_key: str, tmp_path: Path, ): paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) - mutation = tmp_path / f"{case}.json" - write_values_mutation(paths["values"], mutation, path, replacement) + values = json.loads(paths["values"].read_text()) + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + selector = workload[component]["placement"]["nodeSelector"] + selector[f"tampered.example/{selector_key.rsplit('/', 1)[-1]}"] = selector.pop(selector_key) + rewrite_valid_integrity(values, workload) + mutation = tmp_path / f"{component}-selector.json" + mutation.write_text(json.dumps(values)) with pytest.raises(subprocess.CalledProcessError) as error: helm_template("-f", str(mutation)) - assert "topology binding" in error.value.stderr + assert f"{component} placement must match" in error.value.stderr -def test_dgd_chart_rejects_release_name_mismatch(tmp_path: Path): - paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) - - with pytest.raises(subprocess.CalledProcessError) as error: - helm_template("-f", str(paths["values"]), release_name="other-release") - - assert "must match embedded DynamoGraphDeployment metadata.name" in error.value.stderr - - -def test_dgd_chart_rejects_namespace_mismatch(tmp_path: Path): +def test_dgd_chart_rejects_rehashed_runtime_class_mutation(tmp_path: Path): paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + workload["trainer"]["placement"]["runtimeClassName"] = "tampered" + rewrite_valid_integrity(values, workload) + mutation = tmp_path / "runtime-class.json" + mutation.write_text(json.dumps(values)) with pytest.raises(subprocess.CalledProcessError) as error: - helm_template("-f", str(paths["values"]), "--set", "namespace=other-namespace") + helm_template("-f", str(mutation)) - assert "must match embedded DynamoGraphDeployment metadata.namespace" in error.value.stderr + assert "trainer placement must match" in error.value.stderr -@pytest.mark.parametrize( - ("component", "name"), - [ - ("orchestrator", "DYN_RL_TOPOLOGY"), - ("orchestrator", "HF_TOKEN"), - ("trainer", "HF_HOME"), - ], -) -def test_dgd_chart_rejects_raw_env_that_overrides_typed_contract( +@pytest.mark.parametrize("component", ["orchestrator", "trainer"]) +@pytest.mark.parametrize("toleration_key", ["kubernetes.io/arch", "nvidia.com/gpu", "prime-rl"]) +def test_dgd_chart_rejects_every_rehashed_required_toleration_mutation( component: str, - name: str, + toleration_key: str, tmp_path: Path, ): paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) - overlay = tmp_path / f"{component}-{name}.json" - overlay.write_text(json.dumps({component: {"env": [{"name": name, "value": "override"}]}})) + values = json.loads(paths["values"].read_text()) + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + toleration = next(item for item in workload[component]["placement"]["tolerations"] if item["key"] == toleration_key) + toleration["effect"] = "NoExecute" + rewrite_valid_integrity(values, workload) + mutation = tmp_path / f"{component}-{toleration_key.replace('/', '-')}.json" + mutation.write_text(json.dumps(values)) with pytest.raises(subprocess.CalledProcessError) as error: - helm_template("-f", str(paths["values"]), "-f", str(overlay)) + helm_template("-f", str(mutation)) - assert f"{component}.env cannot override generated {name}" in error.value.stderr + assert f"{component} placement must match" in error.value.stderr + + +def test_dgd_chart_ignores_legacy_component_placement_overlays(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + overlay = tmp_path / "legacy-placement.json" + overlay.write_text( + json.dumps( + { + "orchestrator": { + "nodeSelector": {"tampered": "true"}, + "runtimeClassName": "tampered", + "tolerations": [{"key": "tampered", "operator": "Exists"}], + }, + "trainer": { + "nodeSelector": {"tampered": "true"}, + "runtimeClassName": "tampered", + "tolerations": [{"key": "tampered", "operator": "Exists"}], + }, + } + ) + ) + rendered = helm_template("-f", str(paths["values"]), "-f", str(overlay)) + + orchestrator = rendered_resource(rendered, "StatefulSet", "p4-math-orchestrator")["spec"]["template"]["spec"] + trainer = rendered_resource(rendered, "StatefulSet", "p4-math-trainer")["spec"]["template"]["spec"] + assert orchestrator["nodeSelector"] == { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + } + assert "runtimeClassName" not in orchestrator + assert trainer["nodeSelector"] == { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + "nvidia.com/gpu.product": "NVIDIA-GB200", + } + assert trainer["runtimeClassName"] == "nvidia" def test_filesystem_broadcast_reuses_existing_claim_without_rendering_pvc(tmp_path: Path): @@ -362,6 +414,35 @@ def test_filesystem_broadcast_reuses_existing_claim_without_rendering_pvc(tmp_pa assert rendered.count("claimName: p4-shared-data") == 2 +def test_external_filesystem_broadcast_binds_existing_claim_without_rendering_pvc(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config("filesystem"), + render_options( + tmp_path, + external_controller=True, + shared_pvc="p4-shared-data", + ), + ) + rendered = helm_template("-f", str(paths["values"])) + values = json.loads(paths["values"].read_text()) + graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + services = graph["spec"]["services"] + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + + assert "kind: PersistentVolumeClaim" not in rendered + assert workload["storage"] == { + "enabled": True, + "existingClaim": "p4-shared-data", + "mountPath": "/data", + } + assert graph["spec"]["pvcs"] == [ + {"create": False, "name": "model-cache"}, + {"create": False, "name": "p4-shared-data"}, + ] + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + assert {"name": "p4-shared-data", "mountPoint": "/data"} in services[role]["volumeMounts"] + + def test_dgd_chart_rejects_mutable_prime_runtime_image(tmp_path: Path): paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) diff --git a/tests/unit/inference/test_helm_dgd_integrity.py b/tests/unit/inference/test_helm_dgd_integrity.py new file mode 100644 index 0000000000..0925f2bbed --- /dev/null +++ b/tests/unit/inference/test_helm_dgd_integrity.py @@ -0,0 +1,339 @@ +import json +import subprocess +from pathlib import Path + +import pytest + +from prime_rl.inference.dgd import write_dgd_artifacts +from tests.unit.inference.helm_dgd_test_utils import ( + helm_template, + inference_config, + render_options, + rewrite_valid_integrity, + write_values_mutation, +) + + +@pytest.mark.parametrize( + ("case", "path", "replacement", "error_fragment"), + [ + ( + "engine-data", + ("inference", "dynamoGraph", "engineConfig", "data", "prefill-engine.json"), + "tampered", + "engineConfig.data must match its canonical payload", + ), + ( + "engine-sha", + ("inference", "dynamoGraph", "engineConfig", "sha256"), + "0" * 64, + "engineConfig.sha256 must match its canonical payload", + ), + ( + "engine-name", + ("inference", "dynamoGraph", "engineConfig", "name"), + "p4-math-dynamo-engine-000000000000", + "engineConfig.name must be content-addressed", + ), + ( + "dgd-config-sha", + ( + "inference", + "dynamoGraph", + "resource", + "metadata", + "annotations", + "prime-rl.nvidia.com/config-sha256", + ), + "0" * 64, + "DynamoGraphDeployment config-sha256 must match engineConfig.sha256", + ), + ( + "manifest-sha", + ( + "inference", + "dynamoGraph", + "resource", + "metadata", + "annotations", + "prime-rl.nvidia.com/manifest-sha256", + ), + "0" * 64, + "manifest-sha256 must match its canonical payload", + ), + ( + "workload-sha", + ("inference", "dynamoGraph", "workloadBinding", "sha256"), + "0" * 64, + "workload binding sha256 must match its canonical payload", + ), + ( + "dgd-workload-sha", + ( + "inference", + "dynamoGraph", + "resource", + "metadata", + "annotations", + "prime-rl.nvidia.com/workload-sha256", + ), + "0" * 64, + "workload binding annotation must match workloadBinding.sha256", + ), + ], +) +def test_dgd_chart_rejects_content_identity_mutations( + case: str, + path: tuple[str, ...], + replacement: object, + error_fragment: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + mutation = tmp_path / f"{case}.json" + write_values_mutation(paths["values"], mutation, path, replacement) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert error_fragment in error.value.stderr + + +@pytest.mark.parametrize( + ("case", "path", "replacement"), + [ + ( + "client-roles", + ("inference", "dynamoGraph", "clientTopology", "dynamo_worker_roles"), + ["prefill", "decode", "decode", "decode"], + ), + ( + "client-gpus", + ("inference", "dynamoGraph", "clientTopology", "dynamo_gpus_per_worker"), + 2, + ), + ( + "worker-replicas", + ( + "inference", + "dynamoGraph", + "resource", + "spec", + "services", + "VllmDecodeWorker", + "replicas", + ), + 3, + ), + ( + "worker-gpus", + ( + "inference", + "dynamoGraph", + "resource", + "spec", + "services", + "VllmPrefillWorker", + "resources", + "limits", + "gpu", + ), + "2", + ), + ], +) +def test_dgd_chart_rejects_topology_mutations( + case: str, + path: tuple[str, ...], + replacement: object, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + mutation = tmp_path / f"{case}.json" + write_values_mutation(paths["values"], mutation, path, replacement) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert "topology binding" in error.value.stderr + + +def test_dgd_chart_rejects_release_name_mismatch(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), release_name="other-release") + + assert "must match embedded DynamoGraphDeployment metadata.name" in error.value.stderr + + +def test_dgd_chart_rejects_namespace_mismatch(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), "--set", "namespace=other-namespace") + + assert "must match embedded DynamoGraphDeployment metadata.namespace" in error.value.stderr + + +@pytest.mark.parametrize( + ("override_flag", "enabled_value"), + [ + ("--set", "false"), + ("--set-string", "false"), + ("--set-string", "true"), + ], +) +def test_dgd_chart_requires_boolean_true_inference_with_schema_skipped( + override_flag: str, + enabled_value: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "--skip-schema-validation", + "-f", + str(paths["values"]), + override_flag, + f"inference.enabled={enabled_value}", + ) + + assert "inference.enabled must be boolean true in dynamoGraph mode" in error.value.stderr + + +def test_dgd_chart_cannot_switch_mode_and_skip_integrity_validation(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "--skip-schema-validation", + "-f", + str(paths["values"]), + "--set", + "inference.mode=statefulset", + ) + + assert "generated DynamoGraph contract requires dynamoGraph mode" in error.value.stderr + + +@pytest.mark.parametrize( + ("external_controller", "overrides", "error_fragment"), + [ + (False, ("--set", "orchestrator.enabled=false"), "orchestrator.enabled must match"), + (False, ("--set", "trainer.enabled=false"), "trainer.enabled must match"), + (False, ("--set", "trainer.gpu.enabled=false"), "trainer GPU configuration must match"), + (False, ("--set", "trainer.gpu.count=2"), "trainer GPU configuration must match"), + ( + False, + ("--set", r"orchestrator.resources.requests.nvidia\.com/gpu=1"), + "orchestrator resources cannot set NVIDIA extended resource", + ), + ( + False, + ("--set", r"orchestrator.resources.limits.nvidia\.com/mig-1g\.23gb=1"), + "orchestrator resources cannot set NVIDIA extended resource", + ), + ( + False, + ("--set", r"orchestrator.resources.requests.nvidia\.com/gpu\.shared=1"), + "orchestrator resources cannot set NVIDIA extended resource", + ), + ( + False, + ("--set", r"trainer.resources.requests.nvidia\.com/gpu=2"), + "trainer resources cannot set NVIDIA extended resource", + ), + ( + False, + ("--set", r"trainer.resources.limits.nvidia\.com/mig-1g\.23gb=1"), + "trainer resources cannot set NVIDIA extended resource", + ), + (True, ("--set", "orchestrator.enabled=true"), "orchestrator.enabled must match"), + (True, ("--set", "trainer.enabled=true"), "trainer.enabled must match"), + (True, ("--set", "trainer.gpu.enabled=true"), "trainer GPU configuration must match"), + (True, ("--set", "storage.enabled=true"), "storage configuration must match"), + ( + True, + ("--set", "inference.dynamoGraph.controllerMode=chartManaged"), + "controllerMode must match", + ), + ], +) +def test_dgd_chart_rejects_workload_contract_overlays_with_schema_skipped( + external_controller: bool, + overrides: tuple[str, ...], + error_fragment: str, + tmp_path: Path, +): + paths = write_dgd_artifacts( + inference_config(), + render_options(tmp_path, external_controller=external_controller), + ) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "--skip-schema-validation", + "-f", + str(paths["values"]), + *overrides, + ) + + assert error_fragment in error.value.stderr + + +@pytest.mark.parametrize("external_controller", [False, True]) +def test_dgd_chart_rejects_rehashed_workload_mode_contradictions( + external_controller: bool, + tmp_path: Path, +): + paths = write_dgd_artifacts( + inference_config(), + render_options(tmp_path, external_controller=external_controller), + ) + values = json.loads(paths["values"].read_text()) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + + if external_controller: + workload["trainer"]["enabled"] = True + workload["trainer"]["gpu"] = {"enabled": True, "count": 1} + values["trainer"] = {"enabled": True, "gpu": {"enabled": True, "count": 1}} + error_fragment = "external mode forbids chart-managed controller workloads" + else: + workload["orchestrator"]["enabled"] = False + values["orchestrator"]["enabled"] = False + error_fragment = "chartManaged mode requires orchestrator, trainer" + + rewrite_valid_integrity(values, workload) + mutation = tmp_path / "contradictory-workload.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("--skip-schema-validation", "-f", str(mutation)) + + assert error_fragment in error.value.stderr + + +@pytest.mark.parametrize( + ("component", "name"), + [ + ("orchestrator", "DYN_RL_TOPOLOGY"), + ("orchestrator", "HF_TOKEN"), + ("trainer", "HF_HOME"), + ], +) +def test_dgd_chart_rejects_raw_env_that_overrides_typed_contract( + component: str, + name: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + overlay = tmp_path / f"{component}-{name}.json" + overlay.write_text(json.dumps({component: {"env": [{"name": name, "value": "override"}]}})) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), "-f", str(overlay)) + + assert f"{component}.env cannot override generated {name}" in error.value.stderr From 657c8fbaada2c83ec45b20953f95d60569f3b108 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Thu, 9 Jul 2026 20:30:43 -0700 Subject: [PATCH 27/30] fix(k8s): bind generated runtime contract --- k8s/README.md | 28 +++ k8s/prime-rl/templates/deployment.yaml | 10 +- .../templates/dynamo-graph-deployment.yaml | 93 +++++++- .../templates/release-validation.yaml | 3 + k8s/prime-rl/values.schema.json | 51 ++++- k8s/prime-rl/values.yaml | 6 +- src/prime_rl/inference/dgd.py | 209 ++++++------------ src/prime_rl/inference/dgd_cli.py | 107 +++++++++ .../inference/dgd_controller_contract.py | 148 +++++++++++++ tests/unit/inference/helm_dgd_test_utils.py | 12 +- tests/unit/inference/test_dgd.py | 164 ++++++++++++-- tests/unit/inference/test_helm_dgd.py | 112 +++++----- .../unit/inference/test_helm_dgd_integrity.py | 171 +++++++++++++- 13 files changed, 854 insertions(+), 260 deletions(-) create mode 100644 k8s/prime-rl/templates/release-validation.yaml create mode 100644 src/prime_rl/inference/dgd_cli.py create mode 100644 src/prime_rl/inference/dgd_controller_contract.py diff --git a/k8s/README.md b/k8s/README.md index 01e251a85b..87f9586e58 100644 --- a/k8s/README.md +++ b/k8s/README.md @@ -59,6 +59,34 @@ helm install my-exp ./prime-rl \ --set config.secrets.name=prime-rl-secrets ``` +### Generated DynamoGraph deployments + +`dynamo-dgd` produces internally consistent Helm values for Dynamo disaggregated inference. The source SHAs, image tag, and image digest are caller-supplied evidence: the renderer checks that they agree, but it does not authenticate source or image provenance. Verify those inputs through the release process before rendering and protect the generated values from modification. + +Chart-managed mode requires explicit runnable controller commands; the renderer supplies safe execution defaults of one replica and `autoStart: true` and refuses to emit a sleeper deployment: + +```bash +uv run dynamo-dgd inference.toml \ + --release-name my-exp \ + --namespace my-namespace \ + --image "$IMAGE@$IMAGE_DIGEST" \ + --image-digest "$IMAGE_DIGEST" \ + --prime-sha "$PRIME_SHA" \ + --dynamo-sha "$DYNAMO_SHA" \ + --output-dir ./artifacts \ + --gpu-architecture arm64 \ + --gpu-product NVIDIA-GB200 \ + --gpu-node-pool prime-gpu \ + --orchestrator-command 'uv run orchestrator @ /app/configs/debug/orch.toml --output-dir /data/outputs' \ + --trainer-command 'uv run trainer @ /app/configs/debug/rl/train.toml --output-dir /data/outputs' + +helm install my-exp ./prime-rl \ + --namespace my-namespace \ + -f ./artifacts/dynamo-helm-values.json +``` + +Use `--external-controller` when another system runs orchestration and training; generated values then render no controller StatefulSets. Generated releases are limited to 41 characters so every derived Kubernetes Service name remains valid. The image tag must contain the first 12 characters of both caller-supplied source SHAs, and the digest must match `--image-digest`. + ## Uninstalling ```bash diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index a1993900e2..3f433a41c5 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -57,7 +57,7 @@ spec: {{- if .Values.orchestrator.autoStart }} command: ["/bin/bash", "-c"] args: - - {{ .Values.orchestrator.command }} + - {{ .Values.orchestrator.command | quote }} {{- else }} command: ["sleep", "infinity"] {{- end }} @@ -114,7 +114,7 @@ spec: secretKeyRef: name: {{ .Values.huggingFace.tokenSecretName }} key: {{ .Values.huggingFace.tokenSecretKey }} - optional: true + optional: false {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: @@ -239,7 +239,7 @@ spec: secretKeyRef: name: {{ .Values.huggingFace.tokenSecretName }} key: {{ .Values.huggingFace.tokenSecretKey }} - optional: true + optional: false {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: @@ -366,7 +366,7 @@ spec: {{- if .Values.trainer.autoStart }} command: ["/bin/bash", "-c"] args: - - {{ .Values.trainer.command }} + - {{ .Values.trainer.command | quote }} {{- else }} command: ["sleep", "infinity"] {{- end }} @@ -427,7 +427,7 @@ spec: secretKeyRef: name: {{ .Values.huggingFace.tokenSecretName }} key: {{ .Values.huggingFace.tokenSecretKey }} - optional: true + optional: false {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: diff --git a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml index 55a4cce9cd..4cee4d4eb3 100644 --- a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml +++ b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml @@ -1,5 +1,5 @@ {{- $candidateGraph := default (dict) .Values.inference.dynamoGraph }} -{{- if and (hasKey $candidateGraph "workloadBinding") (ne .Values.inference.mode "dynamoGraph") }} +{{- if and (ne .Values.inference.mode "dynamoGraph") (gt (len $candidateGraph) 0) }} {{- fail "generated DynamoGraph contract requires dynamoGraph mode" }} {{- end }} {{- if eq .Values.inference.mode "dynamoGraph" }} @@ -19,6 +19,9 @@ {{- if ne .Values.namespace $resourceNamespace }} {{- fail (printf "values namespace %q must match embedded DynamoGraphDeployment metadata.namespace %q" .Values.namespace $resourceNamespace) }} {{- end }} +{{- if ne .Release.Namespace $resourceNamespace }} +{{- fail (printf "Helm Release.Namespace %q must match values and embedded DynamoGraphDeployment namespace %q" .Release.Namespace $resourceNamespace) }} +{{- end }} {{- $topology := required "inference.dynamoGraph.clientTopology is required" $graph.clientTopology }} {{- $baseURLs := required "inference.dynamoGraph.clientTopology.base_url is required" $topology.base_url }} {{- $rlBaseURLs := required "inference.dynamoGraph.clientTopology.rl_base_url is required" $topology.rl_base_url }} @@ -111,8 +114,31 @@ {{- fail "DynamoGraphDeployment workload binding annotation must match workloadBinding.sha256" }} {{- end }} {{- $workload := mustFromJson $workloadCanonical }} -{{- if ne (len $workload) 4 }} -{{- fail "workload binding must contain only controllerMode, orchestrator, storage, and trainer" }} +{{- $workloadKeys := list "config" "controllerMode" "huggingFace" "image" "modelCache" "orchestrator" "storage" "trainer" }} +{{- if ne (len $workload) (len $workloadKeys) }} +{{- fail "workload binding must contain the complete chart runtime contract" }} +{{- end }} +{{- range $key := $workloadKeys }} +{{- if not (hasKey $workload $key) }} +{{- fail (printf "workload binding is missing %s" $key) }} +{{- end }} +{{- end }} +{{- $imageWorkload := index $workload "image" }} +{{- $actualImage := dict "reference" .Values.image.reference "pullPolicy" .Values.image.pullPolicy "pullSecrets" .Values.image.pullSecrets }} +{{- if ne (toJson $imageWorkload) (toJson $actualImage) }} +{{- fail "image configuration must match the workload binding" }} +{{- end }} +{{- $configWorkload := index $workload "config" }} +{{- if ne (toJson $configWorkload) (toJson .Values.config) }} +{{- fail "config configuration must match the workload binding" }} +{{- end }} +{{- $huggingFaceWorkload := index $workload "huggingFace" }} +{{- if ne (toJson $huggingFaceWorkload) (toJson .Values.huggingFace) }} +{{- fail "huggingFace configuration must match the workload binding" }} +{{- end }} +{{- $modelCacheWorkload := index $workload "modelCache" }} +{{- if ne (toJson $modelCacheWorkload) (toJson .Values.modelCache) }} +{{- fail "modelCache configuration must match the workload binding" }} {{- end }} {{- $controllerMode := required "canonical controllerMode is required" (index $workload "controllerMode") }} {{- if and (ne $controllerMode "chartManaged") (ne $controllerMode "external") }} @@ -126,6 +152,34 @@ {{- $orchestratorWorkload := required "canonical orchestrator workload is required" (index $workload "orchestrator") }} {{- $trainerWorkload := required "canonical trainer workload is required" (index $workload "trainer") }} {{- $storageWorkload := required "canonical storage workload is required" (index $workload "storage") }} +{{- $actualOrchestrator := dict + "enabled" .Values.orchestrator.enabled + "replicas" .Values.orchestrator.replicas + "autoStart" .Values.orchestrator.autoStart + "command" .Values.orchestrator.command + "resources" .Values.orchestrator.resources + "service" .Values.orchestrator.service + "env" .Values.orchestrator.env }} +{{- if ne (toJson (omit $orchestratorWorkload "gpu" "placement")) (toJson $actualOrchestrator) }} +{{- fail "orchestrator configuration must match the workload binding" }} +{{- end }} +{{- $actualTrainer := dict + "enabled" .Values.trainer.enabled + "replicas" .Values.trainer.replicas + "autoStart" .Values.trainer.autoStart + "command" .Values.trainer.command + "gpu" .Values.trainer.gpu + "pytorchCudaAllocConf" .Values.trainer.pytorchCudaAllocConf + "resources" .Values.trainer.resources + "service" .Values.trainer.service + "env" .Values.trainer.env + "probes" .Values.trainer.probes }} +{{- if ne (toJson (omit $trainerWorkload "placement")) (toJson $actualTrainer) }} +{{- fail "trainer configuration must match the workload binding" }} +{{- end }} +{{- if ne (toJson $storageWorkload) (toJson .Values.storage) }} +{{- fail "storage configuration must match the workload binding" }} +{{- end }} {{- if or (not (hasKey $orchestratorWorkload "enabled")) (not (hasKey $trainerWorkload "enabled")) }} {{- fail "canonical controller workloads must define enabled" }} {{- end }} @@ -163,14 +217,21 @@ {{- if or (not $orchestratorEnabled) (not $trainerEnabled) (not $trainerGPUEnabled) (lt $trainerGPUCount 1) }} {{- fail "chartManaged mode requires orchestrator, trainer, and trainer GPU capability" }} {{- end }} +{{- range $componentName, $canonical := dict "orchestrator" $orchestratorWorkload "trainer" $trainerWorkload }} +{{- $command := index $canonical "command" }} +{{- if or (lt (int (index $canonical "replicas")) 1) (not (index $canonical "autoStart")) (not (regexMatch (printf "^uv[[:space:]]+run[[:space:]]+%s([[:space:]]|$)" $componentName) $command)) }} +{{- fail (printf "chartManaged %s execution requires positive replicas, autoStart=true, and an executable uv run %s command" $componentName $componentName) }} +{{- end }} +{{- end }} {{- else }} {{- if or $orchestratorEnabled $trainerEnabled $trainerGPUEnabled (ne $trainerGPUCount 0) }} {{- fail "external mode forbids chart-managed controller workloads and trainer GPU capability" }} {{- end }} +{{- range $componentName, $canonical := dict "orchestrator" $orchestratorWorkload "trainer" $trainerWorkload }} +{{- if or (ne (int (index $canonical "replicas")) 0) (index $canonical "autoStart") (not (empty (index $canonical "command"))) }} +{{- fail (printf "external mode requires zeroed %s execution" $componentName) }} +{{- end }} {{- end }} -{{- $actualTrainerGPU := dict "enabled" .Values.trainer.gpu.enabled "count" (int .Values.trainer.gpu.count) }} -{{- if ne (toJson $trainerGPU) (toJson $actualTrainerGPU) }} -{{- fail "trainer GPU configuration must match the workload binding" }} {{- end }} {{- if or (not (hasKey $storageWorkload "enabled")) (not (hasKey $storageWorkload "existingClaim")) (not (hasKey $storageWorkload "mountPath")) }} {{- fail "canonical storage workload must define enabled, existingClaim, and mountPath" }} @@ -183,10 +244,6 @@ {{- if and (eq $controllerMode "external") $storageEnabled (empty $storageClaim) }} {{- fail "external mode requires an existing claim when chart storage is enabled" }} {{- end }} -{{- $actualStorage := dict "enabled" .Values.storage.enabled "existingClaim" .Values.storage.existingClaim "mountPath" .Values.storage.mountPath }} -{{- if ne (toJson $storageWorkload) (toJson $actualStorage) }} -{{- fail "storage configuration must match the workload binding" }} -{{- end }} {{- $orchestratorPlacement := required "canonical orchestrator placement is required" (index $orchestratorWorkload "placement") }} {{- $expectedOrchestratorPlacement := dict "nodeSelector" (required "Frontend nodeSelector is required" $frontendPod.nodeSelector) @@ -245,6 +302,22 @@ {{- if not (regexMatch "^.+@sha256:[0-9a-f]{64}$" $image) }} {{- fail "image.reference must be pinned to a full sha256 digest in dynamoGraph mode" }} {{- end }} +{{- $referenceDigest := regexFind "sha256:[0-9a-f]{64}$" $image }} +{{- $annotatedDigest := required "DynamoGraphDeployment image-digest annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/image-digest") }} +{{- if ne $referenceDigest $annotatedDigest }} +{{- fail "image.reference digest must match the DynamoGraphDeployment image-digest annotation" }} +{{- end }} +{{- $imageWithoutDigest := trimSuffix (printf "@%s" $referenceDigest) $image }} +{{- $imageTag := regexFind "[^/:]+$" $imageWithoutDigest }} +{{- range $label, $annotation := dict "Prime" "prime-rl.nvidia.com/prime-sha" "Dynamo" "prime-rl.nvidia.com/dynamo-sha" }} +{{- $sourceSHA := required (printf "DynamoGraphDeployment %s SHA annotation is required" $label) (index $resourceAnnotations $annotation) }} +{{- if not (regexMatch "^[0-9a-f]{40}$" $sourceSHA) }} +{{- fail (printf "%s SHA annotation must be a full 40-character Git commit SHA" $label) }} +{{- end }} +{{- if not (contains (substr 0 12 $sourceSHA) $imageTag) }} +{{- fail (printf "%s SHA annotation must match a 12-character commit suffix in the image tag" $label) }} +{{- end }} +{{- end }} {{- range $serviceName, $service := $resource.spec.services }} {{- $serviceImage := required (printf "DynamoGraphDeployment service %s image is required" $serviceName) $service.extraPodSpec.mainContainer.image }} {{- if ne $image $serviceImage }} diff --git a/k8s/prime-rl/templates/release-validation.yaml b/k8s/prime-rl/templates/release-validation.yaml new file mode 100644 index 0000000000..3d6a4ba065 --- /dev/null +++ b/k8s/prime-rl/templates/release-validation.yaml @@ -0,0 +1,3 @@ +{{- if gt (len .Release.Name) 41 }} +{{- fail (printf "Helm release name %q must be at most 41 characters so generated Service names are valid" .Release.Name) }} +{{- end }} diff --git a/k8s/prime-rl/values.schema.json b/k8s/prime-rl/values.schema.json index d6de07f9c7..e90aec8ede 100644 --- a/k8s/prime-rl/values.schema.json +++ b/k8s/prime-rl/values.schema.json @@ -6,6 +6,10 @@ "type": "object", "properties": { "reference": {"type": "string"}, + "pullPolicy": { + "type": "string", + "enum": ["Always", "IfNotPresent", "Never"] + }, "pullSecrets": { "type": "array", "items": {"type": "string", "minLength": 1}, @@ -167,6 +171,9 @@ "type": "object", "properties": { "enabled": {"type": "boolean"}, + "replicas": {"type": "integer", "minimum": 0}, + "autoStart": {"type": "boolean"}, + "command": {"type": "string"}, "runtimeClassName": {"type": "string"}, "nodeSelector": { "type": "object", @@ -195,6 +202,9 @@ "type": "object", "properties": { "enabled": {"type": "boolean"}, + "replicas": {"type": "integer", "minimum": 0}, + "autoStart": {"type": "boolean"}, + "command": {"type": "string"}, "gpu": { "type": "object", "properties": { @@ -280,12 +290,24 @@ } }, "orchestrator": { - "properties": {"enabled": {"const": true}}, - "required": ["enabled"] + "properties": { + "enabled": {"const": true}, + "replicas": {"minimum": 1}, + "autoStart": {"const": true}, + "command": { + "pattern": "^uv[ \\t]+run[ \\t]+orchestrator([ \\t]|$)" + } + }, + "required": ["enabled", "replicas", "autoStart", "command"] }, "trainer": { "properties": { "enabled": {"const": true}, + "replicas": {"minimum": 1}, + "autoStart": {"const": true}, + "command": { + "pattern": "^uv[ \\t]+run[ \\t]+trainer([ \\t]|$)" + }, "gpu": { "properties": { "enabled": {"const": true}, @@ -293,7 +315,7 @@ } } }, - "required": ["enabled", "gpu"] + "required": ["enabled", "replicas", "autoStart", "command", "gpu"] } } }, @@ -307,12 +329,20 @@ } }, "orchestrator": { - "properties": {"enabled": {"const": false}}, - "required": ["enabled"] + "properties": { + "enabled": {"const": false}, + "replicas": {"const": 0}, + "autoStart": {"const": false}, + "command": {"const": ""} + }, + "required": ["enabled", "replicas", "autoStart", "command"] }, "trainer": { "properties": { "enabled": {"const": false}, + "replicas": {"const": 0}, + "autoStart": {"const": false}, + "command": {"const": ""}, "gpu": { "properties": { "enabled": {"const": false}, @@ -320,12 +350,21 @@ } } }, - "required": ["enabled", "gpu"] + "required": ["enabled", "replicas", "autoStart", "command", "gpu"] } } } ], "required": ["image", "inference", "orchestrator", "trainer"] + }, + "else": { + "properties": { + "inference": { + "properties": { + "dynamoGraph": {"maxProperties": 0} + } + } + } } }, { diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml index 5ea9c848e4..22f0d0f23e 100644 --- a/k8s/prime-rl/values.yaml +++ b/k8s/prime-rl/values.yaml @@ -10,7 +10,7 @@ image: repository: primeintellect/prime-rl pullPolicy: IfNotPresent tag: "main" - # Generated DGD values set the full reviewed repository:tag@sha256 reference. + # Generated DGD values set the caller-supplied repository:tag@sha256 reference. reference: "" pullSecrets: [] @@ -75,9 +75,7 @@ inference: replicas: 1 # Generated by `dynamo-dgd` when mode is dynamoGraph. - dynamoGraph: - engineConfig: {} - resource: {} + dynamoGraph: {} # Auto-start configuration (set to false to use sleep infinity for debugging) autoStart: false diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index 35fc81fa3c..abe3474661 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -2,7 +2,6 @@ from __future__ import annotations -import argparse import hashlib import json import re @@ -11,6 +10,7 @@ from typing import Any, Literal from prime_rl.configs.inference import DisaggregatedInferenceDeploymentConfig, InferenceConfig +from prime_rl.inference.dgd_controller_contract import build_chart_runtime_contract from prime_rl.inference.dynamo import ( CHAT_TEMPLATE_ASSET, DynamoProcessSpec, @@ -19,7 +19,6 @@ resolve_chat_template_content, write_role_engine_configs, ) -from prime_rl.utils.config import cli ENGINE_MOUNT_PATH = "/etc/prime-rl/dynamo" ENGINE_CONFIG_HASH_ANNOTATION = "prime-rl.nvidia.com/config-sha256" @@ -61,8 +60,12 @@ } ) _DGD_RESERVED_ENV_PREFIXES = ("DYN_HEALTH_CHECK_", "DYN_SYSTEM_") -_RAW_HUGGING_FACE_TOKEN_KEYS = frozenset({"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"}) _MODEL_CACHE_MOUNT_PATH = "/model-cache" +_CREDENTIAL_ENV_KEY_PATTERNS = ( + re.compile(r"(?:^|_)(?:TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIALS?)(?:_|$)"), + re.compile(r"(?:^|_)(?:API|ACCESS|PRIVATE|SECRET)_KEY(?:_|$)"), +) +MAX_RELEASE_NAME_LENGTH = 41 @dataclass(frozen=True, slots=True) @@ -206,12 +209,22 @@ class DynamoGraphRenderOptions: gpu_scheduling: GPUSchedulingProfile external_controller: bool = False trainer_gpu_count: int = 1 + orchestrator_replicas: int = 1 + trainer_replicas: int = 1 + orchestrator_command: str | None = None + trainer_command: str | None = None model_cache_pvc: str | None = None shared_pvc: str | None = None image_pull_secrets: tuple[str, ...] = () hf_token_secret: str | None = None def __post_init__(self) -> None: + if re.fullmatch(r"[a-z0-9](?:[-a-z0-9]*[a-z0-9])?", self.release_name) is None: + raise ValueError("release_name must be a lowercase Kubernetes DNS label") + if len(self.release_name) > MAX_RELEASE_NAME_LENGTH: + raise ValueError( + f"release_name must be at most {MAX_RELEASE_NAME_LENGTH} characters so generated Service names are valid" + ) for name, value in (("prime_sha", self.prime_sha), ("dynamo_sha", self.dynamo_sha)): if re.fullmatch(r"[0-9a-f]{40}", value) is None: raise ValueError(f"{name} must be a full 40-character Git commit SHA") @@ -219,11 +232,24 @@ def __post_init__(self) -> None: raise ValueError("image_digest must be a full sha256 digest") if not self.image.endswith(f"@{self.image_digest}"): raise ValueError("DGD image must be pinned to image_digest") - image_tag = self.image.rsplit("@", 1)[0] + image_name = self.image.rsplit("@", 1)[0] + image_tag = image_name.rsplit("/", 1)[-1].partition(":")[2] if self.prime_sha[:12] not in image_tag or self.dynamo_sha[:12] not in image_tag: raise ValueError("DGD image tag must include the Prime and Dynamo commit suffixes") if self.trainer_gpu_count < 1: raise ValueError("trainer_gpu_count must be at least one") + if not self.external_controller: + for component, replicas, command in ( + ("orchestrator", self.orchestrator_replicas, self.orchestrator_command), + ("trainer", self.trainer_replicas, self.trainer_command), + ): + if replicas < 1: + raise ValueError(f"{component}_replicas must be at least one in chart-managed mode") + expected_prefix = f"uv run {component}" + if command is None or re.match(rf"^uv\s+run\s+{component}(?:\s|$)", command) is None: + raise ValueError(f"{component}_command must start with {expected_prefix!r} in chart-managed mode") + elif self.orchestrator_command is not None or self.trainer_command is not None: + raise ValueError("external_controller cannot define chart-managed controller commands") def _sha256_bytes(value: bytes) -> str: @@ -297,7 +323,7 @@ def _apply_pod_credentials( "secretKeyRef": { "name": options.hf_token_secret, "key": "HF_TOKEN", - "optional": True, + "optional": False, } }, } @@ -403,11 +429,13 @@ def _validate_typed_credentials( ] ) for source, environment in environment_sources: - token_conflicts = sorted(_RAW_HUGGING_FACE_TOKEN_KEYS & environment.keys()) - if token_conflicts: + credential_conflicts = sorted( + key for key in environment if any(pattern.search(key.upper()) for pattern in _CREDENTIAL_ENV_KEY_PATTERNS) + ) + if credential_conflicts: raise ValueError( - f"{source} env_vars contains raw Hugging Face credentials {token_conflicts}; " - "use DynamoGraphRenderOptions.hf_token_secret" + f"{source} env_vars contains raw credentials {credential_conflicts}; " + "use a typed Kubernetes SecretKeyRef instead" ) hf_home = environment.get("HF_HOME") if options.model_cache_pvc and hf_home is not None and hf_home != _MODEL_CACHE_MOUNT_PATH: @@ -526,37 +554,29 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) } topology_canonical = _canonical_json(topology_binding) topology_hash = _sha256_bytes(topology_canonical) - chart_controller_enabled = not options.external_controller controller_mode = "external" if options.external_controller else "chartManaged" - storage_enabled = chart_controller_enabled or options.shared_pvc is not None - workload = { - "controllerMode": controller_mode, - "orchestrator": { - "enabled": chart_controller_enabled, - "gpu": {"enabled": False}, - "placement": { - "nodeSelector": options.gpu_scheduling.image_node_selector, - "tolerations": options.gpu_scheduling.image_toleration_manifests, - }, + workload, chart_values = build_chart_runtime_contract( + controller_mode=controller_mode, + image_reference=options.image, + image_pull_secrets=options.image_pull_secrets, + orchestrator_replicas=options.orchestrator_replicas, + trainer_replicas=options.trainer_replicas, + orchestrator_command=options.orchestrator_command, + trainer_command=options.trainer_command, + trainer_gpu_count=options.trainer_gpu_count, + orchestrator_placement={ + "nodeSelector": options.gpu_scheduling.image_node_selector, + "tolerations": options.gpu_scheduling.image_toleration_manifests, }, - "trainer": { - "enabled": chart_controller_enabled, - "gpu": { - "enabled": chart_controller_enabled, - "count": options.trainer_gpu_count if chart_controller_enabled else 0, - }, - "placement": { - "runtimeClassName": options.gpu_scheduling.runtime_class_name, - "nodeSelector": options.gpu_scheduling.node_selector, - "tolerations": options.gpu_scheduling.toleration_manifests, - }, + trainer_placement={ + "runtimeClassName": options.gpu_scheduling.runtime_class_name, + "nodeSelector": options.gpu_scheduling.node_selector, + "tolerations": options.gpu_scheduling.toleration_manifests, }, - "storage": { - "enabled": storage_enabled, - "existingClaim": options.shared_pvc or "", - "mountPath": "/data", - }, - } + shared_pvc=options.shared_pvc, + model_cache_pvc=options.model_cache_pvc, + hf_token_secret=options.hf_token_secret, + ) workload_canonical = _canonical_json(workload) workload_hash = _sha256_bytes(workload_canonical) annotations = { @@ -604,11 +624,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) } values: dict[str, Any] = { "namespace": options.namespace, - "image": { - "reference": options.image, - "pullPolicy": "IfNotPresent", - "pullSecrets": list(options.image_pull_secrets), - }, + **chart_values, "inference": { "enabled": True, "mode": "dynamoGraph", @@ -634,34 +650,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) "resource": resource, }, }, - "orchestrator": {"enabled": chart_controller_enabled}, - "trainer": { - "enabled": chart_controller_enabled, - "gpu": { - "enabled": chart_controller_enabled, - "count": options.trainer_gpu_count if chart_controller_enabled else 0, - }, - }, } - if options.external_controller: - values["storage"] = {"enabled": False} - if options.model_cache_pvc: - values["modelCache"] = { - "enabled": True, - "existingClaim": options.model_cache_pvc, - "mountPath": "/model-cache", - } - if options.hf_token_secret: - values["huggingFace"] = { - "tokenSecretName": options.hf_token_secret, - "tokenSecretKey": "HF_TOKEN", - } - if options.shared_pvc: - values["storage"] = { - "enabled": True, - "existingClaim": options.shared_pvc, - "mountPath": "/data", - } return values @@ -684,84 +673,16 @@ def write_dgd_artifacts(config: InferenceConfig, options: DynamoGraphRenderOptio return paths -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("inference_config", type=Path) - parser.add_argument("--release-name", required=True) - parser.add_argument("--namespace", required=True) - parser.add_argument("--image", required=True) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--prime-sha", required=True) - parser.add_argument("--dynamo-sha", required=True) - parser.add_argument("--image-digest", required=True) - parser.add_argument("--run-name") - parser.add_argument("--gpu-runtime-class", default="nvidia") - parser.add_argument("--gpu-architecture", required=True) - parser.add_argument("--gpu-product", required=True) - parser.add_argument("--gpu-node-pool", required=True) - parser.add_argument("--gpu-node-pool-label", default="cloud.google.com/gke-nodepool") - parser.add_argument( - "--external-controller", - action="store_true", - help="Render only DGD inference workloads; an external controller owns orchestration and training", - ) - parser.add_argument( - "--trainer-gpus", - type=int, - default=1, - help="Exact GPU request and limit for the chart-managed trainer", - ) - parser.add_argument( - "--image-toleration", - action="append", - default=[], - type=_parse_kubernetes_toleration, - help='Additional image-pod toleration as JSON, e.g. \'{"key":"dedicated","operator":"Exists"}\'', - ) - parser.add_argument( - "--gpu-toleration", - action="append", - default=[], - type=_parse_kubernetes_toleration, - help='Additional GPU-pod toleration as JSON, e.g. \'{"key":"capacity","operator":"Exists"}\'', - ) - parser.add_argument("--model-cache-pvc") - parser.add_argument("--shared-pvc") - parser.add_argument("--image-pull-secret", action="append", default=[]) - parser.add_argument("--hf-token-secret") - return parser.parse_args() +def _parse_args(): + from prime_rl.inference.dgd_cli import parse_args + + return parse_args() def main() -> None: - args = _parse_args() - config = cli(InferenceConfig, args=["@", str(args.inference_config)]) - options = DynamoGraphRenderOptions( - release_name=args.release_name, - namespace=args.namespace, - image=args.image, - output_dir=args.output_dir, - prime_sha=args.prime_sha, - dynamo_sha=args.dynamo_sha, - image_digest=args.image_digest, - run_name=args.run_name or args.release_name, - gpu_scheduling=GPUSchedulingProfile( - runtime_class_name=args.gpu_runtime_class, - architecture=args.gpu_architecture, - product=args.gpu_product, - node_pool=args.gpu_node_pool, - node_pool_label=args.gpu_node_pool_label, - additional_image_tolerations=tuple(args.image_toleration), - additional_gpu_tolerations=tuple(args.gpu_toleration), - ), - external_controller=args.external_controller, - trainer_gpu_count=args.trainer_gpus, - model_cache_pvc=args.model_cache_pvc, - shared_pvc=args.shared_pvc, - image_pull_secrets=tuple(args.image_pull_secret), - hf_token_secret=args.hf_token_secret, - ) - for path in write_dgd_artifacts(config, options).values(): - print(path) + from prime_rl.inference.dgd_cli import main as cli_main + + cli_main() if __name__ == "__main__": diff --git a/src/prime_rl/inference/dgd_cli.py b/src/prime_rl/inference/dgd_cli.py new file mode 100644 index 0000000000..16e371f5d3 --- /dev/null +++ b/src/prime_rl/inference/dgd_cli.py @@ -0,0 +1,107 @@ +"""Command-line interface for rendering Prime Dynamo graph deployments.""" + +import argparse +from pathlib import Path + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference.dgd import ( + DynamoGraphRenderOptions, + GPUSchedulingProfile, + _parse_kubernetes_toleration, + write_dgd_artifacts, +) +from prime_rl.utils.config import cli + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("inference_config", type=Path) + parser.add_argument("--release-name", required=True) + parser.add_argument("--namespace", required=True) + parser.add_argument("--image", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--prime-sha", required=True) + parser.add_argument("--dynamo-sha", required=True) + parser.add_argument("--image-digest", required=True) + parser.add_argument("--run-name") + parser.add_argument("--gpu-runtime-class", default="nvidia") + parser.add_argument("--gpu-architecture", required=True) + parser.add_argument("--gpu-product", required=True) + parser.add_argument("--gpu-node-pool", required=True) + parser.add_argument("--gpu-node-pool-label", default="cloud.google.com/gke-nodepool") + parser.add_argument( + "--external-controller", + action="store_true", + help="Render only DGD inference workloads; an external controller owns orchestration and training", + ) + parser.add_argument( + "--trainer-gpus", + type=int, + default=1, + help="Exact GPU request and limit for the chart-managed trainer", + ) + parser.add_argument("--orchestrator-replicas", type=int, default=1) + parser.add_argument("--trainer-replicas", type=int, default=1) + parser.add_argument( + "--orchestrator-command", + help="Chart-managed command; must start with 'uv run orchestrator'", + ) + parser.add_argument( + "--trainer-command", + help="Chart-managed command; must start with 'uv run trainer'", + ) + parser.add_argument( + "--image-toleration", + action="append", + default=[], + type=_parse_kubernetes_toleration, + help='Additional image-pod toleration as JSON, e.g. \'{"key":"dedicated","operator":"Exists"}\'', + ) + parser.add_argument( + "--gpu-toleration", + action="append", + default=[], + type=_parse_kubernetes_toleration, + help='Additional GPU-pod toleration as JSON, e.g. \'{"key":"capacity","operator":"Exists"}\'', + ) + parser.add_argument("--model-cache-pvc") + parser.add_argument("--shared-pvc") + parser.add_argument("--image-pull-secret", action="append", default=[]) + parser.add_argument("--hf-token-secret") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + config = cli(InferenceConfig, args=["@", str(args.inference_config)]) + options = DynamoGraphRenderOptions( + release_name=args.release_name, + namespace=args.namespace, + image=args.image, + output_dir=args.output_dir, + prime_sha=args.prime_sha, + dynamo_sha=args.dynamo_sha, + image_digest=args.image_digest, + run_name=args.run_name or args.release_name, + gpu_scheduling=GPUSchedulingProfile( + runtime_class_name=args.gpu_runtime_class, + architecture=args.gpu_architecture, + product=args.gpu_product, + node_pool=args.gpu_node_pool, + node_pool_label=args.gpu_node_pool_label, + additional_image_tolerations=tuple(args.image_toleration), + additional_gpu_tolerations=tuple(args.gpu_toleration), + ), + external_controller=args.external_controller, + trainer_gpu_count=args.trainer_gpus, + orchestrator_replicas=args.orchestrator_replicas, + trainer_replicas=args.trainer_replicas, + orchestrator_command=args.orchestrator_command, + trainer_command=args.trainer_command, + model_cache_pvc=args.model_cache_pvc, + shared_pvc=args.shared_pvc, + image_pull_secrets=tuple(args.image_pull_secret), + hf_token_secret=args.hf_token_secret, + ) + for path in write_dgd_artifacts(config, options).values(): + print(path) diff --git a/src/prime_rl/inference/dgd_controller_contract.py b/src/prime_rl/inference/dgd_controller_contract.py new file mode 100644 index 0000000000..d7be20674a --- /dev/null +++ b/src/prime_rl/inference/dgd_controller_contract.py @@ -0,0 +1,148 @@ +"""Canonical Helm controller and PVC values for DynamoGraph deployments.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from typing import Literal + +ControllerMode = Literal["chartManaged", "external"] + + +def _orchestrator_service() -> dict[str, object]: + return { + "enabled": True, + "type": "ClusterIP", + "port": 8000, + "ncclPort": 29501, + } + + +def _trainer_service() -> dict[str, object]: + return { + "enabled": True, + "type": "ClusterIP", + "port": 8000, + "ncclPort": 29501, + } + + +def _trainer_probes() -> dict[str, object]: + return { + "enabled": False, + "startup": { + "periodSeconds": 10, + "failureThreshold": 60, + "timeoutSeconds": 30, + }, + "liveness": { + "periodSeconds": 30, + "failureThreshold": 6, + "timeoutSeconds": 30, + }, + "readiness": { + "periodSeconds": 10, + "failureThreshold": 3, + "timeoutSeconds": 30, + }, + } + + +def _controller_values(component: Mapping[str, object]) -> dict[str, object]: + return {key: deepcopy(value) for key, value in component.items() if key not in {"gpu", "placement"}} + + +def build_chart_runtime_contract( + *, + controller_mode: ControllerMode, + image_reference: str, + image_pull_secrets: Sequence[str], + orchestrator_replicas: int, + trainer_replicas: int, + orchestrator_command: str | None, + trainer_command: str | None, + trainer_gpu_count: int, + orchestrator_placement: Mapping[str, object], + trainer_placement: Mapping[str, object], + shared_pvc: str | None, + model_cache_pvc: str | None, + hf_token_secret: str | None, +) -> tuple[dict[str, object], dict[str, dict[str, object]]]: + """Return one immutable-by-construction contract and its exact chart values.""" + enabled = controller_mode == "chartManaged" + image = { + "reference": image_reference, + "pullPolicy": "IfNotPresent", + "pullSecrets": list(image_pull_secrets), + } + config = { + "example": "reverse-text", + "secrets": { + "enabled": False, + "name": "prime-rl-secrets", + }, + } + hugging_face = { + "tokenSecretName": hf_token_secret or "", + "tokenSecretKey": "HF_TOKEN", + } + model_cache = { + "enabled": model_cache_pvc is not None, + "existingClaim": model_cache_pvc or "", + "mountPath": "/model-cache", + } + storage = { + "enabled": enabled or shared_pvc is not None, + "existingClaim": shared_pvc or "", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", + "mountPath": "/data", + } + orchestrator = { + "enabled": enabled, + "replicas": orchestrator_replicas if enabled else 0, + "autoStart": enabled, + "command": orchestrator_command if enabled else "", + "gpu": {"enabled": False}, + "placement": deepcopy(orchestrator_placement), + "resources": {"requests": {"memory": "2Gi", "cpu": "1"}}, + "service": _orchestrator_service(), + "env": [], + } + trainer = { + "enabled": enabled, + "replicas": trainer_replicas if enabled else 0, + "autoStart": enabled, + "command": trainer_command if enabled else "", + "gpu": { + "enabled": enabled, + "count": trainer_gpu_count if enabled else 0, + }, + "placement": deepcopy(trainer_placement), + "pytorchCudaAllocConf": "expandable_segments:True", + "resources": {"requests": {"memory": "4Gi", "cpu": "1"}}, + "service": _trainer_service(), + "env": [], + "probes": _trainer_probes(), + } + workload = { + "controllerMode": controller_mode, + "image": image, + "config": config, + "huggingFace": hugging_face, + "modelCache": model_cache, + "orchestrator": orchestrator, + "storage": storage, + "trainer": trainer, + } + values = { + "image": deepcopy(image), + "config": deepcopy(config), + "huggingFace": deepcopy(hugging_face), + "modelCache": deepcopy(model_cache), + "orchestrator": _controller_values(orchestrator), + "storage": deepcopy(storage), + "trainer": _controller_values(trainer) | {"gpu": deepcopy(trainer["gpu"])}, + } + return workload, values diff --git a/tests/unit/inference/helm_dgd_test_utils.py b/tests/unit/inference/helm_dgd_test_utils.py index 83fad508b0..2d24e77c43 100644 --- a/tests/unit/inference/helm_dgd_test_utils.py +++ b/tests/unit/inference/helm_dgd_test_utils.py @@ -15,6 +15,8 @@ PRIME_SHA = "1" * 40 DYNAMO_SHA = "2" * 40 IMAGE_DIGEST = f"sha256:{'3' * 64}" +ORCHESTRATOR_COMMAND = "uv run orchestrator @ /app/configs/debug/orch.toml --output-dir /data/outputs" +TRAINER_COMMAND = "uv run trainer @ /app/configs/debug/rl/train.toml --output-dir /data/outputs" GPU_SCHEDULING = GPUSchedulingProfile( runtime_class_name="nvidia", architecture="arm64", @@ -66,6 +68,8 @@ def render_options( gpu_scheduling=GPU_SCHEDULING, external_controller=external_controller, trainer_gpu_count=trainer_gpu_count, + orchestrator_command=None if external_controller else ORCHESTRATOR_COMMAND, + trainer_command=None if external_controller else TRAINER_COMMAND, model_cache_pvc="model-cache", hf_token_secret="hf-token-secret", shared_pvc=shared_pvc, @@ -73,11 +77,15 @@ def render_options( ) -def helm_template(*args: str, release_name: str = "p4-math") -> str: +def helm_template( + *args: str, + release_name: str = "p4-math", + release_namespace: str = "bis-vllm", +) -> str: if HELM is None: pytest.skip("helm is not installed") return subprocess.run( - [HELM, "template", release_name, str(CHART), *args], + [HELM, "template", release_name, str(CHART), "--namespace", release_namespace, *args], check=True, capture_output=True, text=True, diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index 27d417b1c1..d3ddd830ce 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -22,6 +22,8 @@ PRIME_SHA = "1" * 40 DYNAMO_SHA = "2" * 40 IMAGE_DIGEST = f"sha256:{'3' * 64}" +ORCHESTRATOR_COMMAND = "uv run orchestrator @ /app/configs/debug/orch.toml --output-dir /data/outputs" +TRAINER_COMMAND = "uv run trainer @ /app/configs/debug/rl/train.toml --output-dir /data/outputs" GPU_SCHEDULING = GPUSchedulingProfile( runtime_class_name="nvidia", architecture="arm64", @@ -61,11 +63,12 @@ def render_options( tmp_path: Path, *, external_controller: bool = False, + release_name: str = "p4-math", shared_pvc: str | None = "p4-shared-data", trainer_gpu_count: int = 1, ) -> DynamoGraphRenderOptions: return DynamoGraphRenderOptions( - release_name="p4-math", + release_name=release_name, namespace="bis-vllm", image=f"nvcr.io/example/prime:prime-{PRIME_SHA[:12]}-dynamo-{DYNAMO_SHA[:12]}@{IMAGE_DIGEST}", output_dir=tmp_path, @@ -76,6 +79,8 @@ def render_options( gpu_scheduling=GPU_SCHEDULING, external_controller=external_controller, trainer_gpu_count=trainer_gpu_count, + orchestrator_command=None if external_controller else ORCHESTRATOR_COMMAND, + trainer_command=None if external_controller else TRAINER_COMMAND, model_cache_pvc="model-cache", shared_pvc=shared_pvc, image_pull_secrets=("nvcrimagepullsecret",), @@ -107,6 +112,9 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): assert values["storage"] == { "enabled": True, "existingClaim": "p4-shared-data", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", "mountPath": "/data", } assert values["modelCache"] == { @@ -203,27 +211,25 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): assert workload["storage"] == { "enabled": True, "existingClaim": "p4-shared-data", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", "mountPath": "/data", } - assert workload["orchestrator"] == { - "enabled": True, - "gpu": {"enabled": False}, - "placement": { - "nodeSelector": image_selector, - "tolerations": image_tolerations, - }, + assert workload["orchestrator"]["gpu"] == {"enabled": False} + assert workload["orchestrator"]["placement"] == { + "nodeSelector": image_selector, + "tolerations": image_tolerations, } - assert workload["trainer"] == { - "enabled": True, - "gpu": {"enabled": True, "count": 1}, - "placement": { - "runtimeClassName": "nvidia", - "nodeSelector": gpu_selector, - "tolerations": gpu_tolerations, - }, + assert workload["trainer"]["placement"] == { + "runtimeClassName": "nvidia", + "nodeSelector": gpu_selector, + "tolerations": gpu_tolerations, + } + assert values["orchestrator"] == { + key: value for key, value in workload["orchestrator"].items() if key not in {"gpu", "placement"} } - assert values["orchestrator"] == {"enabled": True} - assert values["trainer"] == {"enabled": True, "gpu": {"enabled": True, "count": 1}} + assert values["trainer"] == {key: value for key, value in workload["trainer"].items() if key != "placement"} prefill_args = services["VllmPrefillWorker"]["extraPodSpec"]["mainContainer"]["args"] decode_args = services["VllmDecodeWorker"]["extraPodSpec"]["mainContainer"]["args"] @@ -347,6 +353,9 @@ def test_filesystem_broadcast_requires_one_shared_existing_claim(tmp_path: Path) assert values["storage"] == { "enabled": True, "existingClaim": "p4-shared-data", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", "mountPath": "/data", } assert resource["spec"]["pvcs"] == [ @@ -405,8 +414,19 @@ def test_dgd_rejects_operator_owned_environment(scope: str, key: str, tmp_path: @pytest.mark.parametrize("scope", ["global", "prefill", "decode"]) -@pytest.mark.parametrize("key", ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"]) -def test_dgd_rejects_raw_hugging_face_credentials(scope: str, key: str, tmp_path: Path): +@pytest.mark.parametrize( + "key", + [ + "HF_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "AWS_SECRET_ACCESS_KEY", + "NVIDIA_API_KEY", + "WANDB_API_KEY", + "MODEL_REGISTRY_PASSWORD", + "OIDC_CLIENT_SECRET", + ], +) +def test_dgd_rejects_raw_credentials(scope: str, key: str, tmp_path: Path): config_data = inference_config().model_dump(mode="python") if scope == "global": config_data["env_vars"][key] = "plaintext-secret" @@ -414,7 +434,7 @@ def test_dgd_rejects_raw_hugging_face_credentials(scope: str, key: str, tmp_path config_data["deployment"][f"{scope}_env_vars"] = {key: "plaintext-secret"} config = InferenceConfig.model_validate(config_data) - with pytest.raises(ValueError, match=rf"{scope}.*{key}.*hf_token_secret"): + with pytest.raises(ValueError, match=rf"{scope}.*{key}.*SecretKeyRef"): build_dgd_values(config, render_options(tmp_path)) @@ -548,17 +568,28 @@ def test_external_controller_binding_disables_chart_workloads(tmp_path: Path): assert graph["controllerMode"] == "external" assert workload["controllerMode"] == "external" assert workload["orchestrator"]["enabled"] is False + assert workload["orchestrator"]["replicas"] == 0 + assert workload["orchestrator"]["autoStart"] is False + assert workload["orchestrator"]["command"] == "" assert workload["orchestrator"]["gpu"] == {"enabled": False} assert workload["trainer"]["enabled"] is False + assert workload["trainer"]["replicas"] == 0 + assert workload["trainer"]["autoStart"] is False + assert workload["trainer"]["command"] == "" assert workload["trainer"]["gpu"] == {"enabled": False, "count": 0} assert workload["storage"] == { "enabled": False, "existingClaim": "", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", "mountPath": "/data", } - assert values["orchestrator"] == {"enabled": False} - assert values["trainer"] == {"enabled": False, "gpu": {"enabled": False, "count": 0}} - assert values["storage"] == {"enabled": False} + assert values["orchestrator"] == { + key: value for key, value in workload["orchestrator"].items() if key not in {"gpu", "placement"} + } + assert values["trainer"] == {key: value for key, value in workload["trainer"].items() if key != "placement"} + assert values["storage"] == workload["storage"] def test_trainer_gpu_count_must_be_positive(): @@ -566,6 +597,91 @@ def test_trainer_gpu_count_must_be_positive(): replace(render_options(Path("/tmp/not-written")), trainer_gpu_count=0) +@pytest.mark.parametrize( + ("change", "error"), + [ + ({"orchestrator_replicas": 0}, "orchestrator_replicas"), + ({"trainer_replicas": 0}, "trainer_replicas"), + ({"orchestrator_command": None}, "orchestrator_command"), + ({"trainer_command": "sleep infinity"}, "trainer_command"), + ({"trainer_command": "uv run trainer-impersonator"}, "trainer_command"), + ], +) +def test_chart_managed_controller_execution_must_be_runnable(change: dict[str, object], error: str): + with pytest.raises(ValueError, match=error): + replace(render_options(Path("/tmp/not-written")), **change) + + +def test_external_controller_rejects_ignored_chart_commands(): + with pytest.raises(ValueError, match="external_controller.*commands"): + replace(render_options(Path("/tmp/not-written")), external_controller=True) + + +def test_release_name_boundary_preserves_generated_service_names(): + boundary = "a" * 41 + + assert replace(render_options(Path("/tmp/not-written")), release_name=boundary).release_name == boundary + + with pytest.raises(ValueError, match="at most 41 characters"): + replace(render_options(Path("/tmp/not-written")), release_name="a" * 42) + + +def test_image_commit_suffixes_must_be_in_the_image_tag(tmp_path: Path): + with pytest.raises(ValueError, match="commit suffixes"): + replace( + render_options(tmp_path, external_controller=True), + image=(f"nvcr.io/prime-{PRIME_SHA[:12]}/dynamo-{DYNAMO_SHA[:12]}/runtime:reviewed@{IMAGE_DIGEST}"), + ) + + +def test_chart_runtime_binding_covers_every_rendered_controller_input(tmp_path: Path): + values = build_dgd_values(inference_config(), render_options(tmp_path)) + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + + assert set(workload) == { + "config", + "controllerMode", + "huggingFace", + "image", + "modelCache", + "orchestrator", + "storage", + "trainer", + } + assert workload["image"] == { + "reference": values["image"]["reference"], + "pullPolicy": values["image"]["pullPolicy"], + "pullSecrets": values["image"]["pullSecrets"], + } + assert workload["storage"] == values["storage"] + assert workload["modelCache"] == values["modelCache"] + assert workload["huggingFace"] == values["huggingFace"] + assert workload["config"] == values["config"] + assert workload["orchestrator"] | {"placement": None, "gpu": None} == ( + values["orchestrator"] | {"placement": None, "gpu": None} + ) + assert workload["trainer"] | {"placement": None} == values["trainer"] | {"placement": None} + + +def test_explicit_hf_secret_references_are_required(tmp_path: Path): + values = build_dgd_values(inference_config(), render_options(tmp_path)) + services = values["inference"]["dynamoGraph"]["resource"]["spec"]["services"] + + for service in services.values(): + hf_token = next(item for item in service["extraPodSpec"]["mainContainer"]["env"] if item["name"] == "HF_TOKEN") + assert hf_token["valueFrom"]["secretKeyRef"]["optional"] is False + + +def test_dgd_readme_uses_runtime_image_config_paths_and_states_trust_boundary(): + repository = Path(__file__).parents[3] + readme = (repository / "k8s" / "README.md").read_text() + + for runtime_path in ("/app/configs/debug/orch.toml", "/app/configs/debug/rl/train.toml"): + assert runtime_path in readme + assert (repository / runtime_path.removeprefix("/app/")).is_file() + assert "does not authenticate source or image provenance" in readme + + def test_dgd_rejects_native_backend(tmp_path: Path): config = InferenceConfig.model_validate({}) with pytest.raises(ValueError, match="Dynamo disaggregated"): diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index a9347ea0bb..c641f6af64 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -10,7 +10,9 @@ DYNAMO_SHA, GPU_SCHEDULING, IMAGE_DIGEST, + ORCHESTRATOR_COMMAND, PRIME_SHA, + TRAINER_COMMAND, helm_template, inference_config, labels_match, @@ -34,6 +36,28 @@ def test_chart_rejects_unknown_inference_mode(): helm_template("--set", "inference.mode=typo") +def test_native_chart_rejects_invalid_image_pull_policy(): + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("--set", "image.pullPolicy=Sometimes") + + assert "image/pullPolicy" in error.value.stderr + assert "'Always', 'IfNotPresent', 'Never'" in error.value.stderr + + +def test_chart_release_name_boundary_keeps_every_resource_name_valid(): + release_name = "a" * 41 + rendered = helm_template(release_name=release_name, release_namespace="default") + + assert all(len(document["metadata"]["name"]) <= 63 for document in rendered_documents(rendered)) + + +def test_chart_rejects_release_name_that_would_overflow_service_names(): + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template(release_name="a" * 42, release_namespace="default") + + assert "at most 41 characters" in error.value.stderr + + def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_path: Path): options = render_options(tmp_path) paths = write_dgd_artifacts(inference_config(), options) @@ -87,6 +111,9 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert set(pods) == {"orchestrator", "trainer", "frontend", "prefill", "decode"} for component, pod in pods.items(): assert {toleration_identity(item) for item in pod["tolerations"]} == required_tolerations + container = pod["containers"][0] if component in chart_pods else pod["mainContainer"] + hf_token = next(item for item in container["env"] if item["name"] == "HF_TOKEN") + assert hf_token["valueFrom"]["secretKeyRef"]["optional"] is False if component in {"trainer", "prefill", "decode"}: assert pod["nodeSelector"] == gpu_selector assert pod["runtimeClassName"] == "nvidia" @@ -96,6 +123,10 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert "nvidia.com/gpu" not in chart_pods["orchestrator"]["containers"][0]["resources"].get("requests", {}) assert chart_pods["trainer"]["containers"][0]["resources"]["requests"]["nvidia.com/gpu"] == 1 + assert chart_pods["orchestrator"]["containers"][0]["args"] == [ORCHESTRATOR_COMMAND] + assert chart_pods["trainer"]["containers"][0]["args"] == [TRAINER_COMMAND] + assert rendered_resource(rendered, "StatefulSet", "p4-math-orchestrator")["spec"]["replicas"] == 1 + assert rendered_resource(rendered, "StatefulSet", "p4-math-trainer")["spec"]["replicas"] == 1 assert not any(kind in rendered for kind in ("kind: ClusterRole", "kind: CustomResourceDefinition")) @@ -236,67 +267,21 @@ def test_dgd_chart_uses_canonical_workload_contract_as_sole_authority(tmp_path: workload = json.loads(workload_binding["canonical"]) assert hashlib.sha256(workload_binding["canonical"].encode()).hexdigest() == workload_binding["sha256"] - assert values["orchestrator"] == {"enabled": True} - assert values["trainer"] == {"enabled": True, "gpu": {"enabled": True, "count": 1}} - assert workload == { - "controllerMode": "chartManaged", - "orchestrator": { - "enabled": True, - "gpu": {"enabled": False}, - "placement": { - "nodeSelector": { - "cloud.google.com/gke-nodepool": "customer-gpu-o7v", - "kubernetes.io/arch": "arm64", - }, - "tolerations": [ - { - "effect": "NoSchedule", - "key": "kubernetes.io/arch", - "operator": "Equal", - "value": "arm64", - }, - {"effect": "NoSchedule", "key": "nvidia.com/gpu", "operator": "Exists"}, - { - "effect": "NoSchedule", - "key": "prime-rl", - "operator": "Equal", - "value": "true", - }, - ], - }, - }, - "trainer": { - "enabled": True, - "gpu": {"count": 1, "enabled": True}, - "placement": { - "nodeSelector": { - "cloud.google.com/gke-nodepool": "customer-gpu-o7v", - "kubernetes.io/arch": "arm64", - "nvidia.com/gpu.product": "NVIDIA-GB200", - }, - "runtimeClassName": "nvidia", - "tolerations": [ - { - "effect": "NoSchedule", - "key": "kubernetes.io/arch", - "operator": "Equal", - "value": "arm64", - }, - {"effect": "NoSchedule", "key": "nvidia.com/gpu", "operator": "Exists"}, - { - "effect": "NoSchedule", - "key": "prime-rl", - "operator": "Equal", - "value": "true", - }, - ], - }, - }, - "storage": { - "enabled": True, - "existingClaim": "", - "mountPath": "/data", - }, + assert workload["controllerMode"] == "chartManaged" + for key in ("config", "huggingFace", "image", "modelCache", "storage"): + assert workload[key] == values[key] + assert values["orchestrator"] == { + key: value for key, value in workload["orchestrator"].items() if key not in {"gpu", "placement"} + } + assert values["trainer"] == {key: value for key, value in workload["trainer"].items() if key != "placement"} + assert workload["orchestrator"]["placement"]["nodeSelector"] == { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + } + assert workload["trainer"]["placement"]["nodeSelector"] == { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + "nvidia.com/gpu.product": "NVIDIA-GB200", } @@ -433,6 +418,9 @@ def test_external_filesystem_broadcast_binds_existing_claim_without_rendering_pv assert workload["storage"] == { "enabled": True, "existingClaim": "p4-shared-data", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", "mountPath": "/data", } assert graph["spec"]["pvcs"] == [ @@ -467,7 +455,7 @@ def test_dgd_chart_rejects_runtime_image_that_differs_from_workers(tmp_path: Pat f"image.reference=nvcr.io/example/prime:reviewed@{other_digest}", ) - assert "must use the same image.reference" in error.value.stderr + assert "image configuration must match the workload binding" in error.value.stderr def test_dgd_rejects_image_without_matching_digest(tmp_path: Path): diff --git a/tests/unit/inference/test_helm_dgd_integrity.py b/tests/unit/inference/test_helm_dgd_integrity.py index 0925f2bbed..3f9940ff68 100644 --- a/tests/unit/inference/test_helm_dgd_integrity.py +++ b/tests/unit/inference/test_helm_dgd_integrity.py @@ -176,6 +176,67 @@ def test_dgd_chart_rejects_namespace_mismatch(tmp_path: Path): assert "must match embedded DynamoGraphDeployment metadata.namespace" in error.value.stderr +def test_dgd_chart_rejects_consistently_rehashed_image_digest_drift(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + drifted_image = values["image"]["reference"].rsplit("@", 1)[0] + f"@sha256:{'4' * 64}" + values["image"]["reference"] = drifted_image + workload["image"]["reference"] = drifted_image + for service in graph["resource"]["spec"]["services"].values(): + service["extraPodSpec"]["mainContainer"]["image"] = drifted_image + rewrite_valid_integrity(values, workload) + mutation = tmp_path / "rehashed-image-drift.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert "image-digest annotation" in error.value.stderr + + +@pytest.mark.parametrize( + ("annotation", "label"), + [ + ("prime-rl.nvidia.com/prime-sha", "Prime"), + ("prime-rl.nvidia.com/dynamo-sha", "Dynamo"), + ], +) +def test_dgd_chart_rejects_rehashed_source_sha_not_present_in_image_tag( + annotation: str, + label: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + graph["resource"]["metadata"]["annotations"][annotation] = "4" * 40 + graph["engineConfig"]["annotations"][annotation] = "4" * 40 + rewrite_valid_integrity(values, workload) + mutation = tmp_path / f"rehashed-{label.lower()}-sha.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert f"{label} SHA annotation" in error.value.stderr + + +def test_dgd_chart_rejects_release_namespace_mismatch(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "-f", + str(paths["values"]), + release_namespace="other-namespace", + ) + + assert "Helm Release.Namespace" in error.value.stderr + + @pytest.mark.parametrize( ("override_flag", "enabled_value"), [ @@ -218,6 +279,28 @@ def test_dgd_chart_cannot_switch_mode_and_skip_integrity_validation(tmp_path: Pa assert "generated DynamoGraph contract requires dynamoGraph mode" in error.value.stderr +@pytest.mark.parametrize("skip_schema", [False, True]) +def test_dgd_chart_cannot_delete_workload_sentinel_and_switch_mode( + skip_schema: bool, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + values["inference"]["mode"] = "statefulset" + del values["inference"]["dynamoGraph"]["workloadBinding"] + mutation = tmp_path / "deleted-sentinel-mode-switch.json" + mutation.write_text(json.dumps(values)) + args = ["-f", str(mutation)] + if skip_schema: + args.insert(0, "--skip-schema-validation") + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template(*args) + + expected = "DynamoGraph contract requires dynamoGraph mode" if skip_schema else "maxProperties" + assert expected in error.value.stderr + + @pytest.mark.parametrize( ("external_controller", "overrides", "error_fragment"), [ @@ -225,6 +308,12 @@ def test_dgd_chart_cannot_switch_mode_and_skip_integrity_validation(tmp_path: Pa (False, ("--set", "trainer.enabled=false"), "trainer.enabled must match"), (False, ("--set", "trainer.gpu.enabled=false"), "trainer GPU configuration must match"), (False, ("--set", "trainer.gpu.count=2"), "trainer GPU configuration must match"), + (False, ("--set", "orchestrator.replicas=0"), "orchestrator execution must match"), + (False, ("--set", "trainer.replicas=0"), "trainer execution must match"), + (False, ("--set", "orchestrator.autoStart=false"), "orchestrator execution must match"), + (False, ("--set", "trainer.autoStart=false"), "trainer execution must match"), + (False, ("--set-string", "orchestrator.command=sleep infinity"), "orchestrator execution must match"), + (False, ("--set-string", "trainer.command=sleep infinity"), "trainer execution must match"), ( False, ("--set", r"orchestrator.resources.requests.nvidia\.com/gpu=1"), @@ -280,7 +369,13 @@ def test_dgd_chart_rejects_workload_contract_overlays_with_schema_skipped( *overrides, ) - assert error_fragment in error.value.stderr + root = overrides[1].split(".", 1)[0] + complete_contract_error = ( + f"{root} configuration must match the workload binding" + if root in {"orchestrator", "storage", "trainer"} + else "" + ) + assert error_fragment in error.value.stderr or complete_contract_error in error.value.stderr @pytest.mark.parametrize("external_controller", [False, True]) @@ -299,7 +394,11 @@ def test_dgd_chart_rejects_rehashed_workload_mode_contradictions( if external_controller: workload["trainer"]["enabled"] = True workload["trainer"]["gpu"] = {"enabled": True, "count": 1} - values["trainer"] = {"enabled": True, "gpu": {"enabled": True, "count": 1}} + values["trainer"] = { + **values["trainer"], + "enabled": True, + "gpu": {"enabled": True, "count": 1}, + } error_fragment = "external mode forbids chart-managed controller workloads" else: workload["orchestrator"]["enabled"] = False @@ -316,6 +415,39 @@ def test_dgd_chart_rejects_rehashed_workload_mode_contradictions( assert error_fragment in error.value.stderr +@pytest.mark.parametrize( + ("component", "field", "replacement"), + [ + ("orchestrator", "replicas", 0), + ("trainer", "replicas", 0), + ("orchestrator", "autoStart", False), + ("trainer", "autoStart", False), + ("orchestrator", "command", "sleep infinity"), + ("trainer", "command", "sleep infinity"), + ], +) +def test_dgd_chart_rejects_rehashed_non_runnable_controller_execution( + component: str, + field: str, + replacement: object, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + workload[component][field] = replacement + values[component][field] = replacement + rewrite_valid_integrity(values, workload) + mutation = tmp_path / f"rehashed-{component}-{field}.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("--skip-schema-validation", "-f", str(mutation)) + + assert f"chartManaged {component} execution requires" in error.value.stderr + + @pytest.mark.parametrize( ("component", "name"), [ @@ -336,4 +468,37 @@ def test_dgd_chart_rejects_raw_env_that_overrides_typed_contract( with pytest.raises(subprocess.CalledProcessError) as error: helm_template("-f", str(paths["values"]), "-f", str(overlay)) - assert f"{component}.env cannot override generated {name}" in error.value.stderr + assert f"{component} configuration must match the workload binding" in error.value.stderr + + +@pytest.mark.parametrize( + ("overrides", "contract"), + [ + (("--set", "image.pullPolicy=Never"), "image"), + (("--set-json", "image.pullSecrets=[]"), "image"), + (("--set", "storage.storageClassName=other"), "storage"), + (("--set", "storage.size=1Gi"), "storage"), + (("--set", "modelCache.enabled=false"), "modelCache"), + (("--set", "huggingFace.tokenSecretName=other"), "huggingFace"), + (("--set", "config.example=other"), "config"), + (("--set", "config.secrets.enabled=true"), "config"), + (("--set", "orchestrator.resources.requests.memory=1Mi"), "orchestrator"), + (("--set", "orchestrator.service.port=9000"), "orchestrator"), + (("--set-json", 'orchestrator.env=[{"name":"PYTHONPATH","value":"/data/other"}]'), "orchestrator"), + (("--set", "trainer.resources.requests.memory=1Mi"), "trainer"), + (("--set", "trainer.service.ncclPort=9001"), "trainer"), + (("--set", "trainer.probes.enabled=true"), "trainer"), + (("--set", "trainer.pytorchCudaAllocConf=max_split_size_mb:64"), "trainer"), + ], +) +def test_dgd_chart_rejects_complete_controller_contract_drift( + overrides: tuple[str, str], + contract: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), *overrides) + + assert f"{contract} configuration must match the workload binding" in error.value.stderr From f271e9e8b692bc49ee8bdb8894b2b5a3e23f0f51 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 10 Jul 2026 06:58:14 -0700 Subject: [PATCH 28/30] fix(k8s): allow direct GPU resource scheduling --- .../templates/dynamo-graph-deployment.yaml | 4 +- src/prime_rl/inference/dgd.py | 40 +++++++++++-------- src/prime_rl/inference/dgd_cli.py | 10 ++++- tests/unit/inference/test_dgd.py | 14 +++++++ tests/unit/inference/test_helm_dgd.py | 29 ++++++++++++++ 5 files changed, 78 insertions(+), 19 deletions(-) diff --git a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml index 4cee4d4eb3..5e3fc16269 100644 --- a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml +++ b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml @@ -255,9 +255,11 @@ {{- range $serviceName, $service := dict "VllmDecodeWorker" $decode "VllmPrefillWorker" $prefill }} {{- $workerPod := required (printf "%s extraPodSpec is required" $serviceName) $service.extraPodSpec }} {{- $expectedTrainerPlacement := dict - "runtimeClassName" (required (printf "%s runtimeClassName is required" $serviceName) $workerPod.runtimeClassName) "nodeSelector" (required (printf "%s nodeSelector is required" $serviceName) $workerPod.nodeSelector) "tolerations" (required (printf "%s tolerations are required" $serviceName) $workerPod.tolerations) }} +{{- if hasKey $workerPod "runtimeClassName" }} +{{- $_ := set $expectedTrainerPlacement "runtimeClassName" (required (printf "%s runtimeClassName must not be empty" $serviceName) $workerPod.runtimeClassName) }} +{{- end }} {{- if ne (toJson $trainerPlacement) (toJson $expectedTrainerPlacement) }} {{- fail "trainer placement must match every manifest-bound DGD worker placement" }} {{- end }} diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index abe3474661..c11665d374 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -130,7 +130,7 @@ def _unique_tolerations( class GPUSchedulingProfile: """Image placement plus the stricter placement required by GPU consumers.""" - runtime_class_name: str + runtime_class_name: str | None architecture: str product: str node_pool: str @@ -141,7 +141,6 @@ class GPUSchedulingProfile: def __post_init__(self) -> None: required = { - "runtime_class_name": self.runtime_class_name, "architecture": self.architecture, "product": self.product, "node_pool": self.node_pool, @@ -150,6 +149,8 @@ def __post_init__(self) -> None: empty = [name for name, value in required.items() if not value] if empty: raise ValueError(f"GPU scheduling fields must not be empty: {empty}") + if self.runtime_class_name == "": + raise ValueError("runtime_class_name must be non-empty or None") if not self.tolerations: raise ValueError("GPU scheduling requires at least one toleration") required_gpu_toleration = KubernetesToleration(key="nvidia.com/gpu") @@ -187,6 +188,13 @@ def image_tolerations(self) -> tuple[KubernetesToleration, ...]: def image_toleration_manifests(self) -> list[dict[str, str]]: return [toleration.as_manifest() for toleration in self.image_tolerations] + @property + def image_placement(self) -> dict[str, Any]: + return { + "nodeSelector": self.image_node_selector, + "tolerations": self.image_toleration_manifests, + } + @property def gpu_tolerations(self) -> tuple[KubernetesToleration, ...]: return _unique_tolerations((*self.image_tolerations, *self.tolerations, *self.additional_gpu_tolerations)) @@ -195,6 +203,16 @@ def gpu_tolerations(self) -> tuple[KubernetesToleration, ...]: def toleration_manifests(self) -> list[dict[str, str]]: return [toleration.as_manifest() for toleration in self.gpu_tolerations] + @property + def gpu_placement(self) -> dict[str, Any]: + placement = { + "nodeSelector": self.node_selector, + "tolerations": self.toleration_manifests, + } + if self.runtime_class_name is not None: + placement["runtimeClassName"] = self.runtime_class_name + return placement + @dataclass(frozen=True) class DynamoGraphRenderOptions: @@ -362,9 +380,7 @@ def _worker_service( ], } pod_spec = { - "runtimeClassName": options.gpu_scheduling.runtime_class_name, - "nodeSelector": options.gpu_scheduling.node_selector, - "tolerations": options.gpu_scheduling.toleration_manifests, + **options.gpu_scheduling.gpu_placement, "volumes": [ { "name": "dynamo-engine-config", @@ -492,8 +508,7 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) ], } frontend_pod_spec = { - "nodeSelector": options.gpu_scheduling.image_node_selector, - "tolerations": options.gpu_scheduling.image_toleration_manifests, + **options.gpu_scheduling.image_placement, "mainContainer": frontend_container, } if chat_template_content is not None: @@ -564,15 +579,8 @@ def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) orchestrator_command=options.orchestrator_command, trainer_command=options.trainer_command, trainer_gpu_count=options.trainer_gpu_count, - orchestrator_placement={ - "nodeSelector": options.gpu_scheduling.image_node_selector, - "tolerations": options.gpu_scheduling.image_toleration_manifests, - }, - trainer_placement={ - "runtimeClassName": options.gpu_scheduling.runtime_class_name, - "nodeSelector": options.gpu_scheduling.node_selector, - "tolerations": options.gpu_scheduling.toleration_manifests, - }, + orchestrator_placement=options.gpu_scheduling.image_placement, + trainer_placement=options.gpu_scheduling.gpu_placement, shared_pvc=options.shared_pvc, model_cache_pvc=options.model_cache_pvc, hf_token_secret=options.hf_token_secret, diff --git a/src/prime_rl/inference/dgd_cli.py b/src/prime_rl/inference/dgd_cli.py index 16e371f5d3..4efbec20c1 100644 --- a/src/prime_rl/inference/dgd_cli.py +++ b/src/prime_rl/inference/dgd_cli.py @@ -24,7 +24,13 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--dynamo-sha", required=True) parser.add_argument("--image-digest", required=True) parser.add_argument("--run-name") - parser.add_argument("--gpu-runtime-class", default="nvidia") + runtime_class = parser.add_mutually_exclusive_group() + runtime_class.add_argument("--gpu-runtime-class", default="nvidia") + runtime_class.add_argument( + "--no-gpu-runtime-class", + action="store_true", + help="Request nvidia.com/gpu resources without setting a Kubernetes RuntimeClass", + ) parser.add_argument("--gpu-architecture", required=True) parser.add_argument("--gpu-product", required=True) parser.add_argument("--gpu-node-pool", required=True) @@ -84,7 +90,7 @@ def main() -> None: image_digest=args.image_digest, run_name=args.run_name or args.release_name, gpu_scheduling=GPUSchedulingProfile( - runtime_class_name=args.gpu_runtime_class, + runtime_class_name=None if args.no_gpu_runtime_class else args.gpu_runtime_class, architecture=args.gpu_architecture, product=args.gpu_product, node_pool=args.gpu_node_pool, diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index d3ddd830ce..08bddf097f 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -525,6 +525,7 @@ def test_cli_accepts_typed_additional_image_and_gpu_tolerations(monkeypatch: pyt "NVIDIA-GB200", "--gpu-node-pool", "customer-gpu-o7v", + "--no-gpu-runtime-class", "--external-controller", "--trainer-gpus", "4", @@ -540,6 +541,7 @@ def test_cli_accepts_typed_additional_image_and_gpu_tolerations(monkeypatch: pyt assert args.image_toleration == [KubernetesToleration(key="image-extra")] assert args.gpu_toleration == [KubernetesToleration(key="gpu-extra", operator="Equal", value="true")] assert args.external_controller is True + assert args.no_gpu_runtime_class is True assert args.trainer_gpus == 4 @@ -557,6 +559,18 @@ def test_gpu_scheduling_changes_manifest_identity(tmp_path: Path): ) +def test_gpu_runtime_class_can_be_explicitly_omitted(tmp_path: Path): + scheduling = replace(GPU_SCHEDULING, runtime_class_name=None) + options = replace(render_options(tmp_path), gpu_scheduling=scheduling) + + values = build_dgd_values(inference_config(), options) + services = values["inference"]["dynamoGraph"]["resource"]["spec"]["services"] + + assert "runtimeClassName" not in services["Frontend"]["extraPodSpec"] + assert "runtimeClassName" not in services["VllmPrefillWorker"]["extraPodSpec"] + assert "runtimeClassName" not in services["VllmDecodeWorker"]["extraPodSpec"] + + def test_external_controller_binding_disables_chart_workloads(tmp_path: Path): values = build_dgd_values( inference_config(), diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index c641f6af64..509639674a 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -1,6 +1,7 @@ import hashlib import json import subprocess +from dataclasses import replace from pathlib import Path import pytest @@ -146,6 +147,34 @@ def test_external_controller_mode_renders_only_five_dgd_inference_pods(tmp_path: assert rendered.count("name: p4-math-frontend-rl") == 1 +@pytest.mark.parametrize("external_controller", [False, True]) +def test_dgd_chart_renders_without_gpu_runtime_class( + tmp_path: Path, + external_controller: bool, +): + options = render_options(tmp_path, external_controller=external_controller) + options = replace( + options, + gpu_scheduling=replace(options.gpu_scheduling, runtime_class_name=None), + ) + paths = write_dgd_artifacts(inference_config(), options) + + rendered = helm_template("-f", str(paths["values"])) + graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + workload = json.loads( + json.loads(paths["values"].read_text())["inference"]["dynamoGraph"]["workloadBinding"]["canonical"] + ) + + assert "runtimeClassName" not in workload["trainer"]["placement"] + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + assert "runtimeClassName" not in graph["spec"]["services"][role]["extraPodSpec"] + if external_controller: + assert "kind: StatefulSet" not in rendered + else: + trainer = rendered_resource(rendered, "StatefulSet", "p4-math-trainer") + assert "runtimeClassName" not in trainer["spec"]["template"]["spec"] + + def test_chart_managed_trainer_uses_exact_bound_gpu_resources(tmp_path: Path): paths = write_dgd_artifacts( inference_config(), From 449cccaf1bd9da39a7fdcca4d281750a92e6b213 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 10 Jul 2026 07:39:49 -0700 Subject: [PATCH 29/30] fix(k8s): realize DGD PVC mounts in pod specs --- src/prime_rl/inference/dgd.py | 21 +++++- tests/unit/inference/test_dgd.py | 96 +++++++++++++++++++++++++-- tests/unit/inference/test_helm_dgd.py | 35 +++++++++- 3 files changed, 141 insertions(+), 11 deletions(-) diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py index c11665d374..e667aca7dc 100644 --- a/src/prime_rl/inference/dgd.py +++ b/src/prime_rl/inference/dgd.py @@ -410,7 +410,26 @@ def _add_pvc(resource: dict[str, Any], service: dict[str, Any], name: str | None pvcs = resource["spec"].setdefault("pvcs", []) if not any(pvc["name"] == name for pvc in pvcs): pvcs.append({"name": name, "create": False}) - service.setdefault("volumeMounts", []).append({"name": name, "mountPoint": mount_point}) + # Keep PodSpec projection as the single mount source. Older alpha operators + # do not realize service.volumeMounts, while current alpha-to-beta conversion + # appends those mounts to extraPodSpec and would otherwise create duplicates. + pod_spec = service["extraPodSpec"] + container_mount = {"name": name, "mountPath": mount_point} + container_mounts = pod_spec["mainContainer"].setdefault("volumeMounts", []) + if container_mount not in container_mounts: + if any(mount["mountPath"] == mount_point for mount in container_mounts): + raise ValueError(f"PVC {name!r} conflicts with an existing container mount at {mount_point!r}") + container_mounts.append(container_mount) + + pod_volume = { + "name": name, + "persistentVolumeClaim": {"claimName": name}, + } + pod_volumes = pod_spec.setdefault("volumes", []) + if pod_volume not in pod_volumes: + if any(volume["name"] == name for volume in pod_volumes): + raise ValueError(f"PVC {name!r} conflicts with an existing pod volume") + pod_volumes.append(pod_volume) def _validate_dgd_environment(config: InferenceConfig) -> None: diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py index 08bddf097f..67fbbdcb7b 100644 --- a/tests/unit/inference/test_dgd.py +++ b/tests/unit/inference/test_dgd.py @@ -12,6 +12,7 @@ DynamoGraphRenderOptions, GPUSchedulingProfile, KubernetesToleration, + _add_pvc, _parse_args, _parse_kubernetes_toleration, build_dgd_values, @@ -134,6 +135,16 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): "dynamo_worker_roles": ["prefill", "prefill", "decode", "decode"], "dynamo_gpus_per_worker": 1, } + expected_cache_mount = {"name": "model-cache", "mountPath": "/model-cache"} + expected_cache_volume = { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + } + for service in services.values(): + assert "volumeMounts" not in service + pod_spec = service["extraPodSpec"] + assert pod_spec["mainContainer"].get("volumeMounts", []).count(expected_cache_mount) == 1 + assert pod_spec.get("volumes", []).count(expected_cache_volume) == 1 topology_binding = graph["topologyBinding"] assert hashlib.sha256(topology_binding["canonical"].encode()).hexdigest() == topology_binding["sha256"] assert json.loads(topology_binding["canonical"])["clientTopology"] == graph["clientTopology"] @@ -262,9 +273,7 @@ def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): assert {key: prefill_env[key] for key in prefill_process.environment()} == prefill_process.environment() assert {key: frontend_env[key] for key in frontend_process.environment()} == frontend_process.environment() - assert {mount["name"] for mount in services["Frontend"]["volumeMounts"]} == {"model-cache"} - for role in ("VllmPrefillWorker", "VllmDecodeWorker"): - assert not any(mount["name"] == "p4-shared-data" for mount in services[role]["volumeMounts"]) + assert all("volumeMounts" not in service for service in services.values()) prefill = json.loads(graph["engineConfig"]["data"]["prefill-engine.json"]) decode = json.loads(graph["engineConfig"]["data"]["decode-engine.json"]) @@ -322,7 +331,8 @@ def test_dgd_embeds_and_mounts_content_addressed_chat_template(tmp_path: Path): "name": "dynamo-chat-template", "mountPath": "/etc/prime-rl/dynamo", "readOnly": True, - } + }, + {"name": "model-cache", "mountPath": "/model-cache"}, ] assert frontend["volumes"] == [ { @@ -331,7 +341,11 @@ def test_dgd_embeds_and_mounts_content_addressed_chat_template(tmp_path: Path): "name": engine_config["name"], "items": [{"key": "chat-template.jinja", "path": "chat-template.jinja"}], }, - } + }, + { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + }, ] assert not (tmp_path / "first" / "chat-template.jinja").exists() @@ -362,9 +376,19 @@ def test_filesystem_broadcast_requires_one_shared_existing_claim(tmp_path: Path) {"name": "model-cache", "create": False}, {"name": "p4-shared-data", "create": False}, ] - assert not any(mount["name"] == "p4-shared-data" for mount in services["Frontend"]["volumeMounts"]) + assert all("volumeMounts" not in service for service in services.values()) for role in ("VllmPrefillWorker", "VllmDecodeWorker"): - assert {"name": "p4-shared-data", "mountPoint": "/data"} in services[role]["volumeMounts"] + pod_spec = services[role]["extraPodSpec"] + assert pod_spec["mainContainer"]["volumeMounts"].count({"name": "p4-shared-data", "mountPath": "/data"}) == 1 + assert ( + pod_spec["volumes"].count( + { + "name": "p4-shared-data", + "persistentVolumeClaim": {"claimName": "p4-shared-data"}, + } + ) + == 1 + ) def test_filesystem_broadcast_rejects_missing_shared_claim(tmp_path: Path): @@ -374,6 +398,64 @@ def test_filesystem_broadcast_rejects_missing_shared_claim(tmp_path: Path): build_dgd_values(inference_config("filesystem"), options) +@pytest.mark.parametrize( + "existing_mount", + [ + ({"name": "other", "mountPath": "/model-cache"}, "existing container mount"), + ], +) +def test_add_pvc_rejects_container_mount_path_collision(existing_mount: tuple[dict[str, str], str]): + mount, message = existing_mount + resource = {"spec": {}} + service = { + "extraPodSpec": { + "mainContainer": {"volumeMounts": [mount]}, + "volumes": [], + } + } + + with pytest.raises(ValueError, match=message): + _add_pvc(resource, service, "model-cache", "/model-cache") + + +def test_filesystem_broadcast_can_reuse_model_cache_claim(tmp_path: Path): + options = replace(render_options(tmp_path), shared_pvc="model-cache") + values = build_dgd_values(inference_config("filesystem"), options) + resource = values["inference"]["dynamoGraph"]["resource"] + services = resource["spec"]["services"] + + assert resource["spec"]["pvcs"] == [{"name": "model-cache", "create": False}] + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + pod_spec = services[role]["extraPodSpec"] + mounts = pod_spec["mainContainer"]["volumeMounts"] + assert mounts.count({"name": "model-cache", "mountPath": "/model-cache"}) == 1 + assert mounts.count({"name": "model-cache", "mountPath": "/data"}) == 1 + assert ( + pod_spec["volumes"].count( + { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + } + ) + == 1 + ) + + +def test_add_pvc_rejects_pod_volume_collision(): + resource = {"spec": {}} + service = { + "extraPodSpec": { + "mainContainer": { + "volumeMounts": [{"name": "model-cache", "mountPath": "/model-cache"}], + }, + "volumes": [{"name": "model-cache", "emptyDir": {}}], + } + } + + with pytest.raises(ValueError, match="existing pod volume"): + _add_pvc(resource, service, "model-cache", "/model-cache") + + @pytest.mark.parametrize( ("scope", "key"), [ diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index 509639674a..191281c8cd 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -82,7 +82,7 @@ def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_pat assert rendered.count(f"image: {options.image}") == 3 assert rendered.count("nvcrimagepullsecret") == 5 assert rendered.count("name: DYN_RL_TOPOLOGY") == 2 - assert rendered.count("claimName: model-cache") == 2 + assert rendered.count("claimName: model-cache") == 5 assert rendered.count("name: HF_TOKEN") == 5 assert rendered.count("name: HF_HOME") == 5 chart_pods = { @@ -175,6 +175,24 @@ def test_dgd_chart_renders_without_gpu_runtime_class( assert "runtimeClassName" not in trainer["spec"]["template"]["spec"] +def test_dgd_chart_projects_model_cache_once_into_operator_pod_specs(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + rendered = helm_template("-f", str(paths["values"])) + graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + expected_mount = {"name": "model-cache", "mountPath": "/model-cache"} + expected_volume = { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + } + + for service in graph["spec"]["services"].values(): + assert "volumeMounts" not in service + pod_spec = service["extraPodSpec"] + assert pod_spec["mainContainer"].get("volumeMounts", []).count(expected_mount) == 1 + assert pod_spec.get("volumes", []).count(expected_volume) == 1 + + def test_chart_managed_trainer_uses_exact_bound_gpu_resources(tmp_path: Path): paths = write_dgd_artifacts( inference_config(), @@ -425,7 +443,7 @@ def test_filesystem_broadcast_reuses_existing_claim_without_rendering_pvc(tmp_pa rendered = helm_template("-f", str(paths["values"])) assert "kind: PersistentVolumeClaim" not in rendered - assert rendered.count("claimName: p4-shared-data") == 2 + assert rendered.count("claimName: p4-shared-data") == 4 def test_external_filesystem_broadcast_binds_existing_claim_without_rendering_pvc(tmp_path: Path): @@ -456,8 +474,19 @@ def test_external_filesystem_broadcast_binds_existing_claim_without_rendering_pv {"create": False, "name": "model-cache"}, {"create": False, "name": "p4-shared-data"}, ] + assert all("volumeMounts" not in service for service in services.values()) for role in ("VllmPrefillWorker", "VllmDecodeWorker"): - assert {"name": "p4-shared-data", "mountPoint": "/data"} in services[role]["volumeMounts"] + pod_spec = services[role]["extraPodSpec"] + assert pod_spec["mainContainer"]["volumeMounts"].count({"name": "p4-shared-data", "mountPath": "/data"}) == 1 + assert ( + pod_spec["volumes"].count( + { + "name": "p4-shared-data", + "persistentVolumeClaim": {"claimName": "p4-shared-data"}, + } + ) + == 1 + ) def test_dgd_chart_rejects_mutable_prime_runtime_image(tmp_path: Path): From be5b75cfd8fbbc3f86415603cd773629473a7282 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Fri, 10 Jul 2026 08:01:15 -0700 Subject: [PATCH 30/30] fix(k8s): select Grove-managed RL frontend --- k8s/prime-rl/templates/dynamo-rl-service.yaml | 1 - tests/unit/inference/test_helm_dgd.py | 11 ++++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/k8s/prime-rl/templates/dynamo-rl-service.yaml b/k8s/prime-rl/templates/dynamo-rl-service.yaml index fb2597eb40..371f5a3d08 100644 --- a/k8s/prime-rl/templates/dynamo-rl-service.yaml +++ b/k8s/prime-rl/templates/dynamo-rl-service.yaml @@ -12,7 +12,6 @@ metadata: spec: type: ClusterIP selector: - {{- include "prime-rl.selectorLabels" . | nindent 4 }} nvidia.com/dynamo-graph-deployment-name: {{ .Release.Name }} nvidia.com/dynamo-component: Frontend nvidia.com/dynamo-component-type: frontend diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py index 191281c8cd..ae528366ac 100644 --- a/tests/unit/inference/test_helm_dgd.py +++ b/tests/unit/inference/test_helm_dgd.py @@ -254,9 +254,10 @@ def test_dgd_rl_service_selector_is_release_disjoint(tmp_path: Path): output_dir = tmp_path / release paths = write_dgd_artifacts(inference_config(), render_options(output_dir, release_name=release)) rendered = helm_template("-f", str(paths["values"]), release_name=release) - graph = rendered_resource(rendered, "DynamoGraphDeployment", release) release_pods[release] = { - **graph["spec"]["services"]["Frontend"]["extraPodMetadata"]["labels"], + # Grove owns and rewrites the conventional app labels on realized + # pods, but Dynamo's identity labels remain stable. + "app.kubernetes.io/name": f"{release}-0-frontend", "nvidia.com/dynamo-graph-deployment-name": release, "nvidia.com/dynamo-component": "Frontend", "nvidia.com/dynamo-component-type": "frontend", @@ -266,7 +267,11 @@ def test_dgd_rl_service_selector_is_release_disjoint(tmp_path: Path): for release in ("alpha", "beta"): other_release = "beta" if release == "alpha" else "alpha" selector = release_services[release] - assert selector["app.kubernetes.io/instance"] == release + assert selector == { + "nvidia.com/dynamo-graph-deployment-name": release, + "nvidia.com/dynamo-component": "Frontend", + "nvidia.com/dynamo-component-type": "frontend", + } assert labels_match(selector, release_pods[release]) assert not labels_match(selector, release_pods[other_release])