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/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index 20ac441af0..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 @@ -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) @@ -426,12 +443,81 @@ 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 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_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": + 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.") + 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.") + 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).""" @@ -463,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/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 2c11a4ddb6..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.""" @@ -728,7 +730,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/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index c92e53889a..0a50e17e38 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: @@ -534,7 +540,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,15 +697,38 @@ 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_updates["admin_api"] = self.inference.backend.type if "dp_rank_count" not in client.model_fields_set: - if self.deployment.type == "multi_node": - client.dp_rank_count = 1 + if self.inference.backend.type == "dynamo" or self.deployment.type == "multi_node": + client_updates["dp_rank_count"] = 1 else: - client.dp_rank_count = self.inference.data_parallel_size_local or self.inference.parallel.dp + client_updates["dp_rank_count"] = self.inference.data_parallel_size_local or self.inference.parallel.dp + if self.inference.backend.type == "dynamo": + expected_topology = { + "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: + 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 c63bcc9c53..225d544013 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,18 @@ 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``.""" + + 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.""" @@ -255,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/entrypoints/inference.py b/src/prime_rl/entrypoints/inference.py index efc120f96e..3a6a62d637 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_dry_run_worker_specs, build_frontend_process + + 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())}") 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..ca9f9e0533 --- /dev/null +++ b/src/prime_rl/inference/dynamo.py @@ -0,0 +1,529 @@ +"""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" +CHAT_TEMPLATE_ASSET = "chat-template.jinja" + +_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( + { + "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", + } +) +_WORKER_EXTENSION_CLS = { + "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) +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 + role: Role + gpu_ids: tuple[str, ...] + system_port: int + 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) + 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}") + 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], ...]: + 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 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 = { + **config.env_vars, + "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=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, + engine_config: Path, + *, + 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", + } + 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=arguments, + environment_items=_environment_items(environment), + ) + + +def build_engine_config( + config: InferenceConfig, + 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) + 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.dynamo_local_dp + 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 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"] = { + "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, + 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 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 + 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(config.dynamo_worker_roles) + gpus_per_worker = config.dynamo_gpus_per_worker + else: + roles = list(config.dynamo_worker_roles) + gpus_per_worker = config.dynamo_gpus_per_worker + + 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]) + 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, + data_parallel_rpc_port=ports.data_parallel_rpc, + ), + ) + specs.append( + DynamoWorkerSpec( + name=name, + role=role, + gpu_ids=worker_gpus, + system_port=ports.system, + process=build_worker_process( + config, + role, + engine_path, + nixl_host="127.0.0.1", + nixl_port=ports.nixl, + 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": + gpu_count = len(config.dynamo_worker_roles) * config.dynamo_gpus_per_worker + else: + gpu_count = config.dynamo_gpus_per_worker + 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], +) -> dict[str, str]: + return spec.process.environment(base_environment) | { + "CUDA_VISIBLE_DEVICES": ",".join(spec.gpu_ids), + "DYN_SYSTEM_PORT": str(spec.system_port), + } + + +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.""" + environment = os.environ.copy() + environment.setdefault("DYN_DISCOVERY_BACKEND", "file") + environment.setdefault("DYN_EVENT_PLANE", "zmq") + environment.setdefault("DYN_FILE_KV_TTL_SECS", "1800") + 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 + + 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 = build_frontend_process(config) + frontend_env = frontend.environment(environment) | {"CUDA_VISIBLE_DEVICES": ""} + frontend_env.pop("DYN_SYSTEM_PORT", None) + + try: + 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.process.command(), env=worker_env, start_new_session=True)) + + exited_process = next((process for process in processes if process.poll() is not None), None) + while exited_process is None: + time.sleep(0.2) + 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: + for process in reversed(processes): + _terminate(process) diff --git a/src/prime_rl/inference/dynamo_admin.py b/src/prime_rl/inference/dynamo_admin.py new file mode 100644 index 0000000000..02fcaf8bec --- /dev/null +++ b/src/prime_rl/inference/dynamo_admin.py @@ -0,0 +1,489 @@ +"""Dynamo worker discovery and engine administration.""" + +from __future__ import annotations + +import asyncio +import os +from collections import Counter +from collections.abc import Awaitable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, TypeAlias +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.async_utils import gather_shielded +from prime_rl.utils.logger import get_logger + +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 +_RETRYABLE_DISCOVERY_HTTP_STATUS_CODES = frozenset({408, 409, 429}) +_REQUIRED_ROUTES = frozenset( + { + "init_weights_update_group", + "pause_generation", + "resume_generation", + "update_weights_from_disk", + "update_weights_from_distributed", + } +) + +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.""" + + 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") + + +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 + + +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") + 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: + 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 + + +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, + *, + model_name: str, + topology: DynamoTopology, +) -> tuple[DynamoWorker, ...]: + """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: + raise TimeoutError("Dynamo worker discovery deadline has already expired") + logger = get_logger() + last_error: Exception | None = None + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + + while (remaining := deadline - loop.time()) > 0: + try: + 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() + snapshots.append(_parse_snapshot(response.json(), model_name, topology)) + first = snapshots[0] + if any(snapshot != first for snapshot in snapshots[1:]): + raise _DiscoveryConvergenceError("Dynamo discovery frontends returned inconsistent worker snapshots") + workers = first[1] + 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: + await asyncio.sleep(min(DISCOVERY_POLL_INTERVAL_S, remaining)) + + raise TimeoutError(f"Dynamo workers were not ready after {timeout} seconds: {last_error!r}") + + +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( + "Dynamo worker membership changed after initialization: " + f"expected {expected_workers!r}, discovered {discovered_workers!r}" + ) + + +class DynamoAdminAPI: + """Typed adapter for Dynamo's per-worker ``/engine`` endpoints.""" + + 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: + 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, + 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), + wait=wait_exponential(multiplier=1, min=1, max=10), + reraise=True, + ): + with attempt: + 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], + *, + host: str, + port: int, + timeout: int, + inference_world_size: int | None, + 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: + 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})" + ) + 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() + try: + await self._settle_fanout( + ( + self._post( + client, + "pause_generation", + {"mode": "wait", "clear_cache": False}, + retry_transient=True, + ) + for client in clients + ), + "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) + 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" + 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: + # 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 + + 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/__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/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/orchestrator/algo/opsd.py b/src/prime_rl/orchestrator/algo/opsd.py index 737666bea9..a8c5791218 100644 --- a/src/prime_rl/orchestrator/algo/opsd.py +++ b/src/prime_rl/orchestrator/algo/opsd.py @@ -1,14 +1,15 @@ 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 + 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,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) :] - await asyncio.gather(*(score_sample(sample) for sample in rollout.samples)) + 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 settle_scores() + return + async with self.policy_gate.request(expected_version=rollout.policy_version): + await settle_scores() diff --git a/src/prime_rl/orchestrator/component_supervision.py b/src/prime_rl/orchestrator/component_supervision.py new file mode 100644 index 0000000000..5d1cff41ab --- /dev/null +++ b/src/prime_rl/orchestrator/component_supervision.py @@ -0,0 +1,107 @@ +"""Failure propagation for orchestrator background components.""" + +from __future__ import annotations + +import asyncio +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: + """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") + + +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 11a1ab4d34..c384e86686 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -2,22 +2,14 @@ - 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 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 +19,22 @@ 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.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, 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, @@ -44,7 +43,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 +55,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 +74,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 +86,9 @@ 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._policy_update: tuple[int, PolicyUpdateToken] | None = None self.max_inflight = max_inflight_rollouts self.inflight_permits = 0 @@ -151,19 +98,17 @@ 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)) 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 +212,178 @@ 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 + + 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(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(step) + + @property + def policy_update_pending(self) -> bool: + return self.policy_gate.pending + + 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.""" 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: + 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: 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]], + 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 + } + 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)) + 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()} + for task in tasks: + task.cancel() + + results: list[object] = [] + cancellation: asyncio.CancelledError | None = None + 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( + emit_policy_cancellation_markers( + groups, + metadata_by_group, + out_q=self.out_q, + stopped=self.stopped, + metrics=self.metrics, + ) + ) + + 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)) + 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 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 fill_inflight(self) -> None: """Schedule new rollouts up to ``max_inflight``, honoring @@ -309,26 +392,24 @@ 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: + 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 is only entered when the orchestrator triggers - # eval, which requires ``eval_source`` to be configured + # PREFER_EVAL implies a configured eval source. 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 + # 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") + 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") + scheduled = await self.try_schedule("train", epoch=epoch) if not scheduled: return @@ -339,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 @@ -356,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 @@ -390,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 @@ -409,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() @@ -420,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 @@ -428,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 @@ -496,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: @@ -536,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 @@ -549,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 @@ -599,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): @@ -619,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_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/dispatcher_transactions.py b/src/prime_rl/orchestrator/dispatcher_transactions.py new file mode 100644 index 0000000000..fd92956c6c --- /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.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, + 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/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 c17ed9e2c7..f883b1bcd7 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, 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 from prime_rl.orchestrator.eval_sink import EvalSink from prime_rl.orchestrator.eval_source import EvalSource @@ -50,6 +52,8 @@ 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_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 ( @@ -71,7 +75,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 @@ -90,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``. @@ -123,6 +121,7 @@ class Orchestrator: train_source: TrainSource train_sink: TrainSink dispatcher: RolloutDispatcher + policy_gate: MutablePolicyGate watcher: WeightWatcher lag_monitor: EventLoopLagMonitor periodic_logger: PeriodicLogger @@ -210,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 ) @@ -242,7 +247,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)})" @@ -293,8 +301,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, @@ -351,6 +358,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=self.policy_gate.enabled, + policy_gate=self.policy_gate, ) self.train_sink = TrainSink( config, @@ -455,14 +464,14 @@ 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() 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 @@ -491,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 @@ -705,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).""" @@ -849,6 +707,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..8c10b2de5b --- /dev/null +++ b/src/prime_rl/orchestrator/policy_gate.py @@ -0,0 +1,133 @@ +"""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 dataclasses import dataclass + +from prime_rl.orchestrator.types import Policy + + +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 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._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_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_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 self._pending_token is None and epoch.value == self._epoch + + @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_token is not None: + 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) -> PolicyUpdateToken: + """Close admission after every in-progress scheduling commit.""" + if not self.enabled: + return PolicyUpdateToken(step=step, epoch=self._epoch) + async with self._admission_lock: + if self._pending_token is not None: + raise RuntimeError(f"A policy update is already pending while preparing step {step}") + self._epoch += 1 + token = PolicyUpdateToken(step=step, epoch=self._epoch) + self._pending_token = token + return token + + 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: + 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.""" + 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..6429bb6732 --- /dev/null +++ b/src/prime_rl/orchestrator/pool_identity.py @@ -0,0 +1,54 @@ +"""Serving-resource identity used by the mutable-policy barrier.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import verifiers.v1 as vf + + 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 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. + + 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) + # 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: + return True + if left_requests & right_requests: + return True + if not left_admin or not right_admin: + return True + 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 e3d0e93198..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]): @@ -172,8 +174,14 @@ 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. 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: ... 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/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 c01d349f40..966d69dfdb 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: 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}" - ) + # 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 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,34 @@ 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: - 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 self._notify_update_succeeded(entered_observers, next_step) + + @staticmethod + async def _notify_update_succeeded(observers: list[VersionObserver], step: int) -> 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): + await observer.on_new_version(step) + + @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/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index ae8c1cde74..19b0e6443a 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -5,17 +5,25 @@ 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 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 from prime_rl.configs.shared import ClientConfig +from prime_rl.inference.dynamo_admin import ( + DynamoAdminAPI, + DynamoTopology, + DynamoWorker, + discover_workers, + 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`` @@ -29,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).""" @@ -38,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.""" @@ -73,6 +99,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.""" ... @@ -115,8 +152,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 FixedInferencePool: + """Base capability for pools whose endpoint set is fixed for their lifetime.""" def __init__( self, @@ -136,7 +173,7 @@ 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._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) @@ -147,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 @@ -166,22 +207,183 @@ 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(FixedInferencePool): + """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 + deadline = _readiness_deadline(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=_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"), ) - await maybe_check_has_model(self._admin_clients, model_name, skip_model_check=self._skip_model_check) async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step) - async def 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 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, + ) 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(FixedInferencePool): + """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, + ): + 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, + 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._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 + deadline = _readiness_deadline(ready_timeout) + 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, + model_name, + skip_model_check=self._skip_model_check, + timeout=_remaining_readiness_timeout(deadline, "model registration"), + ) + 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_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: + 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, + gpus_per_worker=self._topology.gpus_per_worker, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) + + 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( @@ -193,6 +395,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 @@ -205,7 +410,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, @@ -250,14 +456,15 @@ 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 + 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 = { @@ -281,45 +488,144 @@ 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, skip_model_check: bool = False + admin_clients: list[AsyncClient], + 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") - results = await asyncio.gather(*[admin_client.get("/v1/models") for admin_client in admin_clients]) - for admin_client, result in zip(admin_clients, results): - 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}") + + 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() + except Exception as error: + if not _is_retryable_readiness_error(error): + raise + last_error = error + 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} " + f"before the {timeout}-second readiness deadline; last error: {last_error!r}" + ) + 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") async def check_health( - admin_clients: list[AsyncClient], interval: int = 1, log_interval: int = 10, timeout: int = 1800 + admin_clients: list[AsyncClient], + interval: float = 1, + log_interval: float = 10, + 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: - 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: - await admin_client.get("/health") - logger.debug(f"Inference pool is ready after {wait_time} seconds") - return - except NotFoundError: - logger.warning("The route /health does not exist. Skipping health check.") + response = await admin_client.get( + "/health", + timeout=httpx.Timeout(min(remaining, 10.0)), + ) + 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 Exception as e: - if wait_time % log_interval == 0 and wait_time > 0: + if not _is_retryable_readiness_error(e): + raise + 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/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index ef47dee774..208f9ebcd6 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, @@ -224,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() @@ -511,3 +516,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/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 new file mode 100644 index 0000000000..b08ead2a6a --- /dev/null +++ b/tests/unit/inference/test_dynamo.py @@ -0,0 +1,385 @@ +import json +from pathlib import Path + +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, + build_frontend_process, + build_local_worker_specs, + build_worker_environment, + build_worker_process, + write_role_engine_configs, +) + + +def disaggregated_config(**overrides) -> InferenceConfig: + data = { + "backend": {"type": "dynamo"}, + "weight_broadcast": {"type": "nccl"}, + "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, + }, + } + 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, + "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}, + "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", + [ + "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) + + +@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"]) + + 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 + prefill_configs = [ + json.loads(Path(spec.process.arguments[1]).read_text()) for spec in specs if spec.role == "prefill" + ] + 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 + 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, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 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_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={ + "type": "disaggregated", + "gpus_per_node": 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"}, + "decode_env_vars": {"ROLE_SETTING": "decode"}, + } + ) + 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 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" + + +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 == ["", "", "", ""] + + +@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 = [] + + 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(child_code 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 == 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 new file mode 100644 index 0000000000..408e08cc17 --- /dev/null +++ b/tests/unit/inference/test_dynamo_admin.py @@ -0,0 +1,586 @@ +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, + 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) + + +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(): + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "namespace": "test", + "workers": [ + worker(2, system_url="http://worker-b:8081"), + worker(1, system_url="http://worker-a:8081"), + ], + }, + ) + + client = async_client(handler, "http://frontend:8001") + try: + 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(): + 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(ValueError, match="missing RL routes"): + 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 +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),) + restarted = [ + DynamoWorker(2, "backend", "agg", "http://worker:8081", "test-model", ROUTES), + ] + with pytest.raises(RuntimeError, match="membership changed"): + validate_worker_membership(expected, restarted) + + +@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, + gpus_per_worker=2, + 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 + 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 +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", + ] + + +@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"] + + +@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')"] 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..d3486af182 --- /dev/null +++ b/tests/unit/inference/test_dynamo_admin_barriers.py @@ -0,0 +1,269 @@ +"""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_without_resuming_indeterminate_workers( + 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 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 +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 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 +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 7e7a427112..62b6de4f87 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,13 @@ 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, 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 FROZEN = {"name": "org/ref-model", "base_url": ["http://ref:8001/v1"]} @@ -78,6 +81,91 @@ 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() + + +@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() + + +@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 73768a7698..811480a293 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -1,12 +1,252 @@ 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.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 +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_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() @@ -96,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 new file mode 100644 index 0000000000..e17c4e02df --- /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://other-frontend/v1", "http://worker:8081"), + True, + ), + ( + _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_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 new file mode 100644 index 0000000000..637f3b988a --- /dev/null +++ b/tests/unit/orchestrator/test_weight_update_barrier.py @@ -0,0 +1,750 @@ +import asyncio +import uuid +from pathlib import Path +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, Rollout +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_invalidates_slow_scheduling_before_short_commit(): + dispatcher = _dispatcher() + dispatcher.mode = DispatcherMode.PREFER_EVAL + scheduling_started = asyncio.Event() + allow_schedule_to_commit = asyncio.Event() + request_started = asyncio.Event() + + async def active_request() -> None: + request_started.set() + await asyncio.Future() + + async def schedule_one(_kind: str, *, epoch) -> bool: + scheduling_started.set() + await allow_schedule_to_commit.wait() + 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] + + fill = asyncio.create_task(dispatcher.fill_inflight()) + await scheduling_started.wait() + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await asyncio.sleep(0) + + await barrier + assert dispatcher.policy_update_pending + + allow_schedule_to_commit.set() + await fill + assert not request_started.is_set() + assert not dispatcher.inflight + + # 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() + 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 + + 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.TraceTask(type="Task", data=vf.TaskData(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 + + +@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" + ] + + 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): + 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_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) + 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 dispatcher.policy_update_pending + + +@pytest.mark.asyncio +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) + (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, + ) + + 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 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_configs.py b/tests/unit/test_configs.py index 49243a9f98..54c46d2075 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 @@ -211,6 +217,218 @@ 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, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + ) + + assert config.enable_prefix_caching is True + 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( + { + "backend": {"type": "dynamo"}, + "enable_prefix_caching": False, + "deployment": {"type": "disaggregated"}, + } + ) + + +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"}}) + + +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, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "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.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/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 40de4cfee6..7053bf69aa 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -1,12 +1,24 @@ 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.utils.client import _is_retryable_lora_error, load_lora_adapter, setup_clients +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, + check_health, + load_lora_adapter, + maybe_check_has_model, + 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(): @@ -108,3 +120,258 @@ 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()) + + +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", + dynamo_worker_roles=("agg",), + dynamo_gpus_per_worker=1, + ), + model_name="test-model", + ) + ) + + assert isinstance(pool, DynamoInferencePool) + 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( + 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") + 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 = [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 == 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 +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( + 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, strict=False, **_kwargs): + assert strict is True + 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:]))