From e63b21d66454a21e1c46003a30525a166b4f036e Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 2 Sep 2026 16:40:26 -0600 Subject: [PATCH 1/2] feat(slurm): harden distributed allocation runtime Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/runtime/__init__.py | 7 +- .../slurm/runtime/backpressure.py | 234 ++++++++++++++ .../data_designer/slurm/runtime/controller.py | 148 +++++---- .../slurm/runtime/endpoint_steps.py | 113 +++++++ .../data_designer/slurm/runtime/entrypoint.py | 4 +- .../src/data_designer/slurm/runtime/models.py | 3 +- .../data_designer/slurm/runtime/network.py | 27 ++ .../data_designer/slurm/runtime/node_spec.py | 207 +++++++++++++ .../slurm/runtime/node_worker.py | 259 ++++++++++++++++ .../data_designer/slurm/runtime/preflight.py | 185 +++++++++-- .../src/data_designer/slurm/runtime/probes.py | 15 +- .../src/data_designer/slurm/runtime/proxy.py | 15 +- .../slurm/runtime/server_steps.py | 208 +++++++++++++ .../slurm/runtime/step_factory.py | 177 +++++++++++ .../src/data_designer/slurm/runtime/steps.py | 289 ++---------------- .../tests/runtime/conftest.py | 18 +- .../tests/runtime/test_backpressure.py | 82 +++++ .../tests/runtime/test_controller.py | 90 +++++- .../tests/runtime/test_node_worker.py | 168 ++++++++++ .../tests/runtime/test_preflight.py | 43 ++- .../tests/runtime/test_steps.py | 84 ++++- 21 files changed, 1985 insertions(+), 391 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/endpoint_steps.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/server_steps.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/step_factory.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_backpressure.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_node_worker.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/__init__.py index 02864bfce..0c5253a3b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/__init__.py @@ -10,10 +10,14 @@ if TYPE_CHECKING: from data_designer.slurm.runtime.bundle import stage_runtime_bundle # noqa: F401 - from data_designer.slurm.runtime.controller import OneNodeAllocationController # noqa: F401 + from data_designer.slurm.runtime.controller import AllocationController, OneNodeAllocationController # noqa: F401 from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode # noqa: F401 _LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "AllocationController": ( + "data_designer.slurm.runtime.controller", + "AllocationController", + ), "OneNodeAllocationController": ( "data_designer.slurm.runtime.controller", "OneNodeAllocationController", @@ -24,6 +28,7 @@ } __all__ = [ + "AllocationController", "OneNodeAllocationController", "SlurmRuntimeError", "SlurmRuntimeErrorCode", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py new file mode 100644 index 000000000..68490d34c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM middleware that exposes bounded queue admission to client AIMD.""" + +from __future__ import annotations + +import importlib +import json +import math +import os +import threading +import time +from collections.abc import Awaitable, Callable, Iterable, Mapping +from dataclasses import dataclass +from http import HTTPStatus +from typing import Any, Protocol + +MAX_WAITING_REQUESTS_ENVIRONMENT = "DD_VLLM_MAX_WAITING_REQUESTS" +RETRY_AFTER_SECONDS_ENVIRONMENT = "DD_VLLM_RETRY_AFTER_SECONDS" +_METRIC_NAME = "vllm:num_requests_waiting" +_EXEMPT_PATHS = ("/health", "/ready", "/metrics", "/version", "/v1/models", "/ping") + +AsgiMessage = dict[str, Any] +AsgiScope = dict[str, Any] +AsgiReceive = Callable[[], Awaitable[AsgiMessage]] +AsgiSend = Callable[[AsgiMessage], Awaitable[None]] +AsgiApp = Callable[[AsgiScope, AsgiReceive, AsgiSend], Awaitable[None]] + + +class QueueDepthReader(Protocol): + """Read the latest aggregate waiting-request count.""" + + def __call__(self) -> int | None: + """Return queue depth, or ``None`` when metrics are unavailable.""" + ... + + +@dataclass(frozen=True, slots=True) +class QueueBackpressureSettings: + """Validated worker-owned queue admission settings.""" + + max_waiting_requests: int + retry_after_seconds: int | None + poll_interval_seconds: float = 0.1 + stale_after_seconds: float = 1.0 + + def __post_init__(self) -> None: + if type(self.max_waiting_requests) is not int or self.max_waiting_requests < 0: + raise ValueError("maximum waiting requests must be non-negative") + if self.retry_after_seconds is not None and ( + type(self.retry_after_seconds) is not int or self.retry_after_seconds <= 0 + ): + raise ValueError("retry-after seconds must be positive or absent") + if self.poll_interval_seconds <= 0 or self.stale_after_seconds <= 0: + raise ValueError("queue sampler intervals must be positive") + + @classmethod + def from_environment(cls, environment: Mapping[str, str] | None = None) -> QueueBackpressureSettings: + """Load the policy transported by the structured runtime step.""" + source = os.environ if environment is None else environment + maximum = _parse_non_negative_integer(source.get(MAX_WAITING_REQUESTS_ENVIRONMENT), default=128) + retry_value = source.get(RETRY_AFTER_SECONDS_ENVIRONMENT) + retry_after = None if retry_value == "" else _parse_positive_integer(retry_value, default=1) + return cls(maximum, retry_after) + + +@dataclass(frozen=True, slots=True) +class QueueSnapshot: + """One sampled queue depth and its monotonic observation time.""" + + depth: int | None + observed_at: float + + +class QueueBackpressureController: + """Cache queue metrics away from the request path and decide admission.""" + + def __init__( + self, + settings: QueueBackpressureSettings | None = None, + reader: QueueDepthReader | None = None, + *, + start_sampler: bool = True, + ) -> None: + self.settings = settings or QueueBackpressureSettings.from_environment() + self._reader = reader or read_vllm_queue_depth + self._start_sampler = start_sampler + self._snapshot = QueueSnapshot(None, 0.0) + self._lock = threading.Lock() + self._thread: threading.Thread | None = None + + def sample_once(self) -> QueueSnapshot: + """Refresh the cached queue depth once.""" + depth = self._reader() + if depth is not None: + depth = max(0, depth) + snapshot = QueueSnapshot(depth, time.monotonic()) + with self._lock: + self._snapshot = snapshot + return snapshot + + def should_reject(self) -> tuple[bool, QueueSnapshot]: + """Return the fail-open admission decision and supporting snapshot.""" + self._ensure_sampler() + with self._lock: + snapshot = self._snapshot + stale = time.monotonic() - snapshot.observed_at > self.settings.stale_after_seconds + reject = snapshot.depth is not None and not stale and snapshot.depth > self.settings.max_waiting_requests + return reject, snapshot + + def _ensure_sampler(self) -> None: + if not self._start_sampler or self._thread is not None: + return + with self._lock: + if self._thread is None: + self._thread = threading.Thread(target=self._sample_forever, name="dd-vllm-queue-depth", daemon=True) + self._thread.start() + + def _sample_forever(self) -> None: + while True: + self.sample_once() + time.sleep(self.settings.poll_interval_seconds) + + +class QueueDepthBackpressureMiddleware: + """Reject non-health HTTP requests with 429 above the resolved queue threshold.""" + + def __init__(self, app: AsgiApp, controller: QueueBackpressureController | None = None) -> None: + self.app = app + self.controller = controller or QueueBackpressureController() + + async def __call__(self, scope: AsgiScope, receive: AsgiReceive, send: AsgiSend) -> None: + """Apply queue admission without changing exempt or accepted requests.""" + path = str(scope.get("path", "")) + if scope.get("type") != "http" or path in _EXEMPT_PATHS: + await self.app(scope, receive, send) + return + reject, snapshot = self.controller.should_reject() + if not reject: + await self.app(scope, receive, send) + return + await _send_overload(send, self.controller.settings, snapshot) + + +def read_vllm_queue_depth() -> int | None: + """Return aggregate vLLM queue depth from supported metrics registries.""" + for metrics in (_read_vllm_metrics(), _read_prometheus_metrics()): + values = _collect_metric_values(metrics) + if values: + return max(0, int(sum(values))) + return None + + +async def _send_overload( + send: AsgiSend, + settings: QueueBackpressureSettings, + snapshot: QueueSnapshot, +) -> None: + body = json.dumps( + { + "error": { + "message": "serving queue admission threshold exceeded", + "type": "rate_limit_exceeded", + "code": HTTPStatus.TOO_MANY_REQUESTS.value, + "queue_depth": snapshot.depth, + "max_waiting_requests": settings.max_waiting_requests, + } + }, + separators=(",", ":"), + ).encode() + headers = [(b"content-type", b"application/json")] + if settings.retry_after_seconds is not None: + headers.append((b"retry-after", str(settings.retry_after_seconds).encode())) + await send({"type": "http.response.start", "status": HTTPStatus.TOO_MANY_REQUESTS.value, "headers": headers}) + await send({"type": "http.response.body", "body": body}) + + +def _read_vllm_metrics() -> Iterable[object]: + try: + module = importlib.import_module("vllm.v1.metrics.reader") + return module.get_metrics_snapshot() + except (ImportError, AttributeError, RuntimeError): + return () + + +def _read_prometheus_metrics() -> Iterable[object]: + try: + module = importlib.import_module("prometheus_client") + return tuple(sample for family in module.REGISTRY.collect() for sample in getattr(family, "samples", ())) + except (ImportError, AttributeError, RuntimeError): + return () + + +def _collect_metric_values(metrics: Iterable[object]) -> tuple[float, ...]: + names = {_METRIC_NAME, _METRIC_NAME.replace(":", "_")} + values: list[float] = [] + for metric in metrics: + if str(getattr(metric, "name", "")) not in names: + continue + try: + value = float(getattr(metric, "value")) + if math.isfinite(value) and value >= 0: + values.append(value) + except (AttributeError, TypeError, ValueError): + continue + return tuple(values) + + +def _parse_non_negative_integer(value: str | None, *, default: int) -> int: + if value is None: + return default + if not value.isascii() or not value.isdigit(): + raise ValueError("queue threshold is invalid") + return int(value) + + +def _parse_positive_integer(value: str | None, *, default: int) -> int: + parsed = _parse_non_negative_integer(value, default=default) + if parsed <= 0: + raise ValueError("retry-after seconds must be positive") + return parsed + + +__all__ = [ + "MAX_WAITING_REQUESTS_ENVIRONMENT", + "QueueBackpressureController", + "QueueBackpressureSettings", + "QueueDepthBackpressureMiddleware", + "QueueDepthReader", + "QueueSnapshot", + "RETRY_AFTER_SECONDS_ENVIRONMENT", + "read_vllm_queue_depth", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index 2341d617b..21f7131d1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""One-node allocation orchestration with one process and cleanup owner.""" +"""Distributed allocation orchestration with one process and cleanup owner.""" from __future__ import annotations @@ -14,17 +14,16 @@ from data_designer.slurm.client import ClientResult from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.runtime.endpoint_steps import build_endpoint_steps from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.logs import bind_execution_logs, execution_log_directory from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStep -from data_designer.slurm.runtime.preflight import AllocationPreflight +from data_designer.slurm.runtime.preflight import AllocationLayout, AllocationPreflight from data_designer.slurm.runtime.probes import ReadinessProber from data_designer.slurm.runtime.records import load_complete_client_candidate -from data_designer.slurm.runtime.steps import ( - ClientStepBuilder, - build_endpoint_steps, - build_vllm_steps, -) +from data_designer.slurm.runtime.server_steps import build_vllm_steps +from data_designer.slurm.runtime.step_factory import place_step_on_host +from data_designer.slurm.runtime.steps import ClientStepBuilder from data_designer.slurm.runtime.supervisor import ManagedStep, RuntimeClock, StepSupervisor from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment from data_designer.slurm.serving.resolver import resolve_vllm_server @@ -94,11 +93,13 @@ class _RunOutcome: @dataclass(frozen=True, slots=True) class _RuntimeTopology: deployments: tuple[ResolvedVllmServerDeployment, ...] + layout: AllocationLayout + server_steps: tuple[RuntimeStep, ...] endpoint_steps: tuple[tuple[RuntimeStep, RuntimeEndpoint], ...] endpoints: tuple[RuntimeEndpoint, ...] -class OneNodeAllocationController: +class AllocationController: """Execute one planned shard attempt and finalize only complete candidates.""" def __init__( @@ -210,7 +211,7 @@ def _persist_terminal_outcome(self, outcome: _RunOutcome) -> AttemptManifest: def _execute(self) -> tuple[ArtifactReference, datetime]: topology = self._prepare_runtime() required_processes = self._start_runtime(topology) - return self._run_client(topology.endpoints, required_processes) + return self._run_client(topology.endpoints, topology.layout, required_processes) def _prepare_runtime(self) -> _RuntimeTopology: self._validate_attempt_state() @@ -221,7 +222,7 @@ def _prepare_runtime(self) -> _RuntimeTopology: restarting = self._readiness is not None if restarting: self._begin_execution(deployments, ReadinessState.RESTARTING) - self._preflight.verify(self._context, self._environment) + layout = self._preflight.verify(self._context, self._environment) self._mark_attempt_running() if not restarting: self._begin_execution(deployments, ReadinessState.PENDING) @@ -233,21 +234,62 @@ def _prepare_runtime(self) -> _RuntimeTopology: self._context.attempt_directory, self._environment, self._runtime_proxy_path, + layout, ) ) endpoints = tuple(endpoint for _, endpoint in endpoint_steps) + server_preflight_steps, server_steps = self._build_server_steps(deployments, layout) client_preflight = self._bind_logs( - self._client_steps.build_preflight_step( - self._context.plan, - self._context.shard, - self._attempt, - self._context.attempt_directory, - endpoints, - self._environment, + place_step_on_host( + self._client_steps.build_preflight_step( + self._context.plan, + self._context.shard, + self._attempt, + self._context.attempt_directory, + endpoints, + self._environment, + ), + layout.get_host(self._context.plan.client.host_node_index), ) ) self._supervisor.wait(self._supervisor.start(client_preflight)) - return _RuntimeTopology(deployments, endpoint_steps, endpoints) + for step in server_preflight_steps: + self._supervisor.wait(self._supervisor.start(step)) + return _RuntimeTopology(deployments, layout, server_steps, endpoint_steps, endpoints) + + def _build_server_steps( + self, + deployments: tuple[ResolvedVllmServerDeployment, ...], + layout: AllocationLayout, + ) -> tuple[tuple[RuntimeStep, ...], tuple[RuntimeStep, ...]]: + runtime_node_worker_path = self._runtime_proxy_path.with_name("node_worker.py") + preflight_steps: list[RuntimeStep] = [] + server_steps: list[RuntimeStep] = [] + for deployment in deployments: + preflight_steps.extend( + self._bind_logs(step) + for step in build_vllm_steps( + deployment, + self._context.plan, + self._context.attempt_directory, + self._environment, + runtime_node_worker_path, + layout, + preflight=True, + ) + ) + server_steps.extend( + self._bind_logs(step) + for step in build_vllm_steps( + deployment, + self._context.plan, + self._context.attempt_directory, + self._environment, + runtime_node_worker_path, + layout, + ) + ) + return tuple(preflight_steps), tuple(server_steps) def _begin_execution( self, @@ -268,12 +310,12 @@ def _start_runtime(self, topology: _RuntimeTopology) -> tuple[ManagedStep, ...]: if status.state in {ReadinessState.PENDING, ReadinessState.RESTARTING}: status.state = ReadinessState.STARTING self._publish_readiness(ReadinessState.STARTING) - server_processes = self._start_servers(topology.deployments) - for status, required in zip(self._statuses, server_processes, strict=True): - self._wait_for_backends(status, required) + server_processes = self._start_server_steps(topology.server_steps) + for status in self._statuses: + self._wait_for_backends(status, topology.layout, server_processes) self._publish_readiness(ReadinessState.STARTING) endpoint_processes = tuple(self._supervisor.start(step) for step, _ in topology.endpoint_steps) - required_processes = tuple(process for group in server_processes for process in group) + endpoint_processes + required_processes = server_processes + endpoint_processes self._wait_for_endpoints(topology.endpoints, required_processes) for status in self._statuses: status.state = ReadinessState.READY @@ -282,19 +324,30 @@ def _start_runtime(self, topology: _RuntimeTopology) -> tuple[ManagedStep, ...]: self._publish_readiness(ReadinessState.READY) return required_processes + def _start_server_steps(self, steps: tuple[RuntimeStep, ...]) -> tuple[ManagedStep, ...]: + managed: list[ManagedStep] = [] + for step in steps: + self._supervisor.require_running(tuple(managed)) + managed.append(self._supervisor.start(step)) + return tuple(managed) + def _run_client( self, endpoints: tuple[RuntimeEndpoint, ...], + layout: AllocationLayout, required_processes: tuple[ManagedStep, ...], ) -> tuple[ArtifactReference, datetime]: generation = self._bind_logs( - self._client_steps.build_generation_step( - self._context.plan, - self._context.shard, - self._attempt, - self._context.attempt_directory, - endpoints, - self._environment, + place_step_on_host( + self._client_steps.build_generation_step( + self._context.plan, + self._context.shard, + self._attempt, + self._context.attempt_directory, + endpoints, + self._environment, + ), + layout.get_host(self._context.plan.client.host_node_index), ) ) generation_started_at = self._now() @@ -364,33 +417,12 @@ def _mark_attempt_running(self) -> None: ) self._attempt = self._state.update_attempt(self._attempt) - def _start_servers( + def _wait_for_backends( self, - deployments: tuple[ResolvedVllmServerDeployment, ...], - ) -> tuple[tuple[ManagedStep, ...], ...]: - all_managed: list[tuple[ManagedStep, ...]] = [] - for deployment in deployments: - started_at = self._clock.monotonic() - managed: list[ManagedStep] = [] - steps = tuple( - self._bind_logs(step) - for step in build_vllm_steps( - deployment, - self._context.plan, - self._context.attempt_directory, - self._environment, - ) - ) - for process, step in zip(deployment.processes, steps, strict=True): - self._supervisor.require_running(tuple(managed)) - remaining_delay = process.launch_delay_seconds - (self._clock.monotonic() - started_at) - if remaining_delay > 0: - self._clock.sleep(remaining_delay) - managed.append(self._supervisor.start(step)) - all_managed.append(tuple(managed)) - return tuple(all_managed) - - def _wait_for_backends(self, status: _DeploymentStatus, required: tuple[ManagedStep, ...]) -> None: + status: _DeploymentStatus, + layout: AllocationLayout, + required: tuple[ManagedStep, ...], + ) -> None: probes = status.deployment.readiness_probes deadline = self._clock.monotonic() + max(probe.deadline_seconds for probe in probes) previous_ready = -1 @@ -398,7 +430,7 @@ def _wait_for_backends(self, status: _DeploymentStatus, required: tuple[ManagedS self._supervisor.require_running(required) ready = sum( self._prober.is_ready( - "127.0.0.1", + layout.get_host(probe.node_index), probe.port, probe.path, timeout_seconds=self._probe_timeout_seconds, @@ -616,4 +648,6 @@ def _normalize_failure(error: BaseException) -> SlurmRuntimeError: ) -__all__ = ["OneNodeAllocationController", "RuntimeStateStore"] +OneNodeAllocationController = AllocationController + +__all__ = ["AllocationController", "OneNodeAllocationController", "RuntimeStateStore"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/endpoint_steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/endpoint_steps.py new file mode 100644 index 000000000..3afe8ea51 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/endpoint_steps.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compose logical endpoint proxies from resolved remote lane heads.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path + +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.runtime.step_factory import SrunStepOptions, base_environment, build_srun_step +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment + + +def build_endpoint_steps( + deployments: tuple[ResolvedVllmServerDeployment, ...], + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], + runtime_proxy_path: Path, + layout: AllocationLayout | None = None, +) -> tuple[tuple[RuntimeStep, RuntimeEndpoint], ...]: + """Build one package-owned logical-endpoint process per deployment.""" + selected_layout = layout or _one_node_layout(plan) + return tuple( + _build_endpoint_step( + deployment, + plan, + attempt_directory, + source_environment, + runtime_proxy_path, + selected_layout, + ) + for deployment in deployments + ) + + +def _build_endpoint_step( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], + runtime_proxy_path: Path, + layout: AllocationLayout, +) -> tuple[RuntimeStep, RuntimeEndpoint]: + endpoint = RuntimeEndpoint( + model_alias=deployment.model_alias, + served_model_name=deployment.served_model_name, + host="127.0.0.1", + port=deployment.logical_endpoint.port, + ) + environment = base_environment(source_environment) + environment["PYTHONPATH"] = get_container_path(plan, runtime_proxy_path.parents[3].as_posix()) + step = build_srun_step( + step_id=f"{deployment.deployment_id}-endpoint", + role=RuntimeStepRole.ENDPOINT, + command=_build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port, layout), + environment=environment, + plan=plan, + attempt_directory=attempt_directory, + options=SrunStepOptions( + image_path=plan.client.image.path, + container_environment=("PYTHONPATH",), + node_hosts=(layout.get_host(plan.client.host_node_index),), + ), + ) + return step, endpoint + + +def _build_endpoint_command( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + runtime_proxy_path: Path, + port: int, + layout: AllocationLayout, +) -> tuple[str, ...]: + backend_hosts = tuple(layout.get_host(backend.node_index) for backend in deployment.backend_endpoints) + backends = tuple( + f"http://{host}:{backend.port}" + for host, backend in zip(backend_hosts, deployment.backend_endpoints, strict=True) + ) + retry_after_seconds = deployment.launch_policy.queue_backpressure.retry_after_seconds + retry_arguments = ("--retry-after-seconds", str(retry_after_seconds)) if retry_after_seconds is not None else () + return ( + "python3", + get_container_path(plan, runtime_proxy_path.as_posix()), + "--listen-port", + str(port), + *(argument for host in sorted(set(backend_hosts)) for argument in ("--allowed-host", host)), + *retry_arguments, + *(argument for backend in backends for argument in ("--backend", backend)), + ) + + +def _one_node_layout(plan: ResolvedSlurmRunPlan) -> AllocationLayout: + node_indices = { + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + } + if node_indices != {0}: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "multi-node step construction requires a verified allocation layout", + ) + return AllocationLayout(("127.0.0.1",)) + + +__all__ = ["build_endpoint_steps"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index 30e2fd077..05a904a85 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -12,7 +12,7 @@ from pathlib import Path from data_designer.slurm.planning import PlannedShard -from data_designer.slurm.runtime.controller import OneNodeAllocationController +from data_designer.slurm.runtime.controller import AllocationController from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.preflight import SystemAllocationPreflight @@ -50,7 +50,7 @@ def _run(plan_path: Path, attempt_directory: Path, environment: Mapping[str, str clock = SystemRuntimeClock() signals = TerminationSignalCoordinator() supervisor = StepSupervisor(SubprocessStepRunner(), signals=signals, clock=clock) - controller = OneNodeAllocationController( + controller = AllocationController( context, runtime_proxy_path=Path(__file__).with_name("proxy.py"), state=writer, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py index 50d2557fe..38c06ac49 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py @@ -25,6 +25,7 @@ class RuntimeStepRole(str, Enum): """Lifecycle role of one allocation-local Slurm step.""" CLIENT_PREFLIGHT = "client_preflight" + SERVER_PREFLIGHT = "server_preflight" SERVER = "server" ENDPOINT = "endpoint" CLIENT = "client" @@ -104,7 +105,7 @@ def __post_init__(self) -> None: if type(self.served_model_name) is not str or not self.served_model_name: raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "endpoint model name is invalid") if self.host != "127.0.0.1": - raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "one-node endpoints must use loopback") + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "logical endpoints must use loopback") if type(self.port) is not int or not 1 <= self.port <= 65535: raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "endpoint port is invalid") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py new file mode 100644 index 000000000..d17e4311a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared validation for scheduler-derived runtime network identities.""" + +from __future__ import annotations + +import re + +_HOST_NAME_PATTERN = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._-]{0,251}[A-Za-z0-9])?$") + + +def validate_host_name(host: str) -> str: + """Return a safe scheduler host name or raise ``ValueError``.""" + if type(host) is not str or _HOST_NAME_PATTERN.fullmatch(host) is None: + raise ValueError("allocation host identity is invalid") + return host + + +def validate_network_port(port: int) -> int: + """Return a valid TCP port or raise ``ValueError``.""" + if type(port) is not int or not 1 <= port <= 65535: + raise ValueError("allocation network port is invalid") + return port + + +__all__ = ["validate_host_name", "validate_network_port"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py new file mode 100644 index 000000000..e708a5489 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict transport records for one coordinated deployment's node workers.""" + +from __future__ import annotations + +import base64 +import binascii +import json +from dataclasses import dataclass + +from data_designer.slurm.runtime.network import validate_host_name, validate_network_port + +_MAXIMUM_SPEC_BYTES = 64 * 1024 + + +@dataclass(frozen=True, slots=True) +class NodeProcessSpec: + """One resolved lane process owned by a node worker.""" + + process_id: str + command: tuple[str, ...] + gpu_indices: tuple[int, ...] + launch_delay_seconds: int + + def __post_init__(self) -> None: + if ( + type(self.process_id) is not str + or not self.process_id + or any(ord(value) < 32 or ord(value) == 127 for value in self.process_id) + ): + raise ValueError("node process identity is invalid") + if ( + type(self.command) is not tuple + or not self.command + or any(type(value) is not str or not value or "\0" in value for value in self.command) + ): + raise ValueError("node process command is invalid") + if ( + type(self.gpu_indices) is not tuple + or not self.gpu_indices + or self.gpu_indices != tuple(sorted(set(self.gpu_indices))) + or any(type(value) is not int or value < 0 for value in self.gpu_indices) + ): + raise ValueError("node process GPU indices are invalid") + if type(self.launch_delay_seconds) is not int or self.launch_delay_seconds < 0: + raise ValueError("node process launch delay is invalid") + + +@dataclass(frozen=True, slots=True) +class NodeSpec: + """Resolved work and ports for one node in a coordinated step.""" + + node_index: int + host: str + ports: tuple[int, ...] + processes: tuple[NodeProcessSpec, ...] + + def __post_init__(self) -> None: + if type(self.node_index) is not int or self.node_index < 0: + raise ValueError("node index is invalid") + validate_host_name(self.host) + if type(self.ports) is not tuple or type(self.processes) is not tuple: + raise ValueError("node ports and processes must be tuples") + for port in self.ports: + validate_network_port(port) + if self.ports != tuple(sorted(set(self.ports))): + raise ValueError("node ports must be sorted and unique") + if not self.processes or len({process.process_id for process in self.processes}) != len(self.processes): + raise ValueError("node process identities must be present and unique") + + +@dataclass(frozen=True, slots=True) +class NodeWorkerSpec: + """Strict, bounded input shared by every task in one deployment step.""" + + schema_version: int + resolved_gpus_per_node: int + nodes: tuple[NodeSpec, ...] + + def __post_init__(self) -> None: + if type(self.schema_version) is not int or self.schema_version != 1: + raise ValueError("node worker schema version is unsupported") + if type(self.resolved_gpus_per_node) is not int or self.resolved_gpus_per_node <= 0: + raise ValueError("resolved GPU count is invalid") + if ( + type(self.nodes) is not tuple + or not self.nodes + or tuple(node.node_index for node in self.nodes) != tuple(sorted({node.node_index for node in self.nodes})) + ): + raise ValueError("node worker nodes must use sorted unique indices") + process_ids = tuple(process.process_id for node in self.nodes for process in node.processes) + if len({node.host for node in self.nodes}) != len(self.nodes) or len(set(process_ids)) != len(process_ids): + raise ValueError("node worker hosts and process identities must be unique") + for node in self.nodes: + if any(index >= self.resolved_gpus_per_node for process in node.processes for index in process.gpu_indices): + raise ValueError("node process GPU index is outside the allocation") + assigned_gpus = tuple(index for process in node.processes for index in process.gpu_indices) + if len(set(assigned_gpus)) != len(assigned_gpus): + raise ValueError("node processes must not share GPUs") + + +def encode_node_worker_spec(spec: NodeWorkerSpec) -> str: + """Serialize one validated node-worker specification for an argv boundary.""" + payload = { + "schema_version": spec.schema_version, + "resolved_gpus_per_node": spec.resolved_gpus_per_node, + "nodes": [ + { + "node_index": node.node_index, + "host": node.host, + "ports": list(node.ports), + "processes": [ + { + "process_id": process.process_id, + "command": list(process.command), + "gpu_indices": list(process.gpu_indices), + "launch_delay_seconds": process.launch_delay_seconds, + } + for process in node.processes + ], + } + for node in spec.nodes + ], + } + serialized = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + if len(serialized) > _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification exceeds its size limit") + return base64.urlsafe_b64encode(serialized).decode() + + +def decode_node_worker_spec(encoded: str) -> NodeWorkerSpec: + """Decode and validate an untrusted node-worker specification.""" + if type(encoded) is not str or not encoded or len(encoded) > 2 * _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification is invalid") + try: + serialized = base64.b64decode(encoded, altchars=b"-_", validate=True) + if len(serialized) > _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification exceeds its size limit") + payload = json.loads(serialized) + except (binascii.Error, json.JSONDecodeError, UnicodeEncodeError) as error: + raise ValueError("node worker specification is invalid") from error + return _parse_worker_spec(payload) + + +def _parse_worker_spec(payload: object) -> NodeWorkerSpec: + root = _require_mapping(payload, {"schema_version", "resolved_gpus_per_node", "nodes"}) + nodes = tuple(_parse_node(value) for value in _require_list(root["nodes"])) + return NodeWorkerSpec( + schema_version=_require_integer(root["schema_version"]), + resolved_gpus_per_node=_require_integer(root["resolved_gpus_per_node"]), + nodes=nodes, + ) + + +def _parse_node(payload: object) -> NodeSpec: + value = _require_mapping(payload, {"node_index", "host", "ports", "processes"}) + return NodeSpec( + node_index=_require_integer(value["node_index"]), + host=_require_string(value["host"]), + ports=tuple(_require_integer(port) for port in _require_list(value["ports"])), + processes=tuple(_parse_process(process) for process in _require_list(value["processes"])), + ) + + +def _parse_process(payload: object) -> NodeProcessSpec: + value = _require_mapping(payload, {"process_id", "command", "gpu_indices", "launch_delay_seconds"}) + return NodeProcessSpec( + process_id=_require_string(value["process_id"]), + command=tuple(_require_string(argument) for argument in _require_list(value["command"])), + gpu_indices=tuple(_require_integer(index) for index in _require_list(value["gpu_indices"])), + launch_delay_seconds=_require_integer(value["launch_delay_seconds"]), + ) + + +def _require_mapping(payload: object, keys: set[str]) -> dict[str, object]: + if not isinstance(payload, dict) or set(payload) != keys or not all(type(key) is str for key in payload): + raise ValueError("node worker object shape is invalid") + return payload + + +def _require_list(payload: object) -> list[object]: + if not isinstance(payload, list): + raise ValueError("node worker list value is invalid") + return payload + + +def _require_string(payload: object) -> str: + if type(payload) is not str: + raise ValueError("node worker string value is invalid") + return payload + + +def _require_integer(payload: object) -> int: + if type(payload) is not int: + raise ValueError("node worker integer value is invalid") + return payload + + +__all__ = [ + "NodeProcessSpec", + "NodeSpec", + "NodeWorkerSpec", + "decode_node_worker_spec", + "encode_node_worker_spec", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py new file mode 100644 index 000000000..649489727 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the resolved vLLM lanes assigned to one physical allocation node.""" + +from __future__ import annotations + +import argparse +import os +import signal +import socket +import subprocess +import sys +import time +from collections.abc import Mapping, Sequence + +from data_designer.slurm.runtime.node_spec import ( + NodeProcessSpec as NodeProcessSpec, +) +from data_designer.slurm.runtime.node_spec import ( + NodeSpec as NodeSpec, +) +from data_designer.slurm.runtime.node_spec import ( + NodeWorkerSpec as NodeWorkerSpec, +) +from data_designer.slurm.runtime.node_spec import ( + decode_node_worker_spec as decode_node_worker_spec, +) +from data_designer.slurm.runtime.node_spec import ( + encode_node_worker_spec as encode_node_worker_spec, +) + +_POLL_INTERVAL_SECONDS = 0.1 +_TERMINATION_GRACE_SECONDS = 10.0 + + +class NodeProcessSupervisor: + """Own every lane child on one node and stop them as a unit.""" + + def __init__(self, *, environment: Mapping[str, str]) -> None: + self._environment = dict(environment) + self._children: list[subprocess.Popen[bytes]] = [] + self._cleanup_started = False + self._cleanup_complete = False + self._stop_requested = False + + @property + def cleanup_complete(self) -> bool: + """Return whether every registered lane child has exited.""" + return self._cleanup_complete + + def run(self, node: NodeSpec, visible_gpus: tuple[str, ...]) -> int: + """Launch local lanes and fail the node task when any lane exits.""" + started_at = time.monotonic() + try: + for process in node.processes: + self._wait_for_launch(process.launch_delay_seconds, started_at) + if self._stop_requested: + return 1 + self._children.append(self._start_process(process, visible_gpus)) + return self._wait_for_first_exit() + finally: + self.cleanup() + + def request_stop(self) -> None: + """Ask the poll loop to enter cleanup at its next boundary.""" + self._stop_requested = True + + def cleanup(self) -> None: + """Idempotently terminate every registered lane child.""" + if self._cleanup_complete: + return + self._cleanup_started = True + live = tuple(child for child in reversed(self._children) if child.poll() is None) + _signal_children(live, signal.SIGTERM) + deadline = time.monotonic() + _TERMINATION_GRACE_SECONDS + for child in live: + _wait_until_deadline(child, deadline) + _signal_children(tuple(child for child in live if child.poll() is None), signal.SIGKILL) + for child in live: + try: + child.wait(timeout=_TERMINATION_GRACE_SECONDS) + except (OSError, subprocess.SubprocessError): + pass + self._cleanup_complete = all(child.poll() is not None for child in self._children) + if not self._cleanup_complete: + raise RuntimeError("node process cleanup is incomplete") + + def _start_process(self, process: NodeProcessSpec, visible_gpus: tuple[str, ...]) -> subprocess.Popen[bytes]: + if self._cleanup_started: + raise RuntimeError("cannot launch a node process after cleanup") + environment = dict(self._environment) + environment["CUDA_VISIBLE_DEVICES"] = ",".join(visible_gpus[index] for index in process.gpu_indices) + return subprocess.Popen( + process.command, + stdin=subprocess.DEVNULL, + env=environment, + shell=False, + start_new_session=True, + close_fds=True, + ) + + def _wait_for_launch(self, delay_seconds: int, started_at: float) -> None: + while not self._stop_requested: + remaining = delay_seconds - (time.monotonic() - started_at) + if remaining <= 0: + return + time.sleep(min(remaining, _POLL_INTERVAL_SECONDS)) + + def _wait_for_first_exit(self) -> int: + while True: + if self._stop_requested: + return 1 + for child in self._children: + returncode = child.poll() + if returncode is not None: + return returncode if returncode != 0 else 1 + time.sleep(_POLL_INTERVAL_SECONDS) + + +def main(arguments: Sequence[str] | None = None) -> int: + """Preflight or serve the node-local slice of a coordinated deployment.""" + parser = argparse.ArgumentParser(prog="data-designer-slurm-node-worker") + parser.add_argument("operation", choices=("preflight", "serve")) + parser.add_argument("--spec", required=True) + parsed = parser.parse_args(arguments) + try: + spec = decode_node_worker_spec(parsed.spec) + node = _select_node(spec, os.environ) + visible_gpus = _parse_visible_gpus(os.environ.get("CUDA_VISIBLE_DEVICES")) + _verify_node(spec, node, visible_gpus, os.environ) + if parsed.operation == "preflight": + _verify_ports(node.ports) + return 0 + return _run_node(node, visible_gpus) + except (OSError, RuntimeError, ValueError, subprocess.SubprocessError): + print("node worker failed at a validated runtime boundary", file=sys.stderr) + return 70 + + +def _run_node(node: NodeSpec, visible_gpus: tuple[str, ...]) -> int: + supervisor = NodeProcessSupervisor(environment=os.environ) + interrupted: list[int] = [] + + def handle_termination(signum: int, frame: object) -> None: + del frame + if not interrupted: + interrupted.append(signum) + supervisor.request_stop() + + previous = {selected: signal.signal(selected, handle_termination) for selected in (signal.SIGINT, signal.SIGTERM)} + try: + return _run_until_exit_or_signal(supervisor, node, visible_gpus, interrupted) + finally: + try: + supervisor.cleanup() + finally: + for selected, handler in previous.items(): + signal.signal(selected, handler) + + +def _run_until_exit_or_signal( + supervisor: NodeProcessSupervisor, + node: NodeSpec, + visible_gpus: tuple[str, ...], + interrupted: list[int], +) -> int: + if interrupted: + return 128 + interrupted[0] + result = supervisor.run(node, visible_gpus) + if interrupted: + return 128 + interrupted[0] + return result + + +def _select_node(spec: NodeWorkerSpec, environment: Mapping[str, str]) -> NodeSpec: + process_id = _parse_non_negative_integer(environment.get("SLURM_PROCID"), "SLURM_PROCID") + if process_id >= len(spec.nodes): + raise ValueError("Slurm process identity is outside the deployment") + return spec.nodes[process_id] + + +def _verify_node( + spec: NodeWorkerSpec, + node: NodeSpec, + visible_gpus: tuple[str, ...], + environment: Mapping[str, str], +) -> None: + if len(visible_gpus) != spec.resolved_gpus_per_node: + raise ValueError("node GPU visibility does not match the resolved plan") + task_count = _parse_non_negative_integer(environment.get("SLURM_NTASKS"), "SLURM_NTASKS") + if task_count != len(spec.nodes): + raise ValueError("Slurm task count does not match the deployment") + scheduler_host = environment.get("SLURMD_NODENAME") + if scheduler_host is not None and scheduler_host != node.host: + raise ValueError("Slurm node identity does not match the deployment") + + +def _verify_ports(ports: tuple[int, ...]) -> None: + reservations: list[socket.socket] = [] + try: + for port in ports: + reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + reservations.append(reservation) + reservation.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0) + reservation.bind(("0.0.0.0", port)) + finally: + for reservation in reservations: + reservation.close() + + +def _parse_visible_gpus(value: str | None) -> tuple[str, ...]: + if value is None or not value.strip(): + return () + values = tuple(item.strip() for item in value.split(",")) + if not all(values) or len(values) != len(set(values)) or any("\0" in item for item in values): + raise ValueError("node GPU visibility is invalid") + return values + + +def _parse_non_negative_integer(value: str | None, name: str) -> int: + if value is None or not value.isascii() or not value.isdigit(): + raise ValueError(f"{name} is unavailable or invalid") + return int(value) + + +def _signal_process_group(process: subprocess.Popen[bytes], selected: signal.Signals) -> None: + try: + os.killpg(process.pid, selected) + except ProcessLookupError: + pass + + +def _signal_children(children: tuple[subprocess.Popen[bytes], ...], selected: signal.Signals) -> None: + for child in children: + try: + _signal_process_group(child, selected) + except OSError: + continue + + +def _wait_until_deadline(process: subprocess.Popen[bytes], deadline: float) -> None: + while process.poll() is None and time.monotonic() < deadline: + time.sleep(_POLL_INTERVAL_SECONDS) + + +if __name__ == "__main__": # pragma: no cover - exercised as a managed step + raise SystemExit(main()) + + +__all__ = [ + "NodeProcessSpec", + "NodeProcessSupervisor", + "NodeSpec", + "NodeWorkerSpec", + "decode_node_worker_spec", + "encode_node_worker_spec", + "main", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py index 476d62a07..2889780f6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -10,38 +10,125 @@ import re import socket import stat +import subprocess from collections.abc import Mapping +from dataclasses import dataclass from pathlib import Path from typing import Protocol from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.network import validate_host_name from data_designer.slurm.runtime.paths import get_container_path _DIGEST_CHUNK_SIZE = 1024 * 1024 _GPU_COUNT_PATTERN = re.compile(r"^(?:gpu(?::[^:]+)?):([0-9]+)$") +_NODE_LIST_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.\-,\[\]]{0,4095}$") + + +@dataclass(frozen=True, slots=True) +class AllocationLayout: + """Verified allocation node identities in planner index order.""" + + node_hosts: tuple[str, ...] + + def __post_init__(self) -> None: + if ( + type(self.node_hosts) is not tuple + or not self.node_hosts + or len(self.node_hosts) != len(set(self.node_hosts)) + ): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "allocation node identities are invalid") + try: + for host in self.node_hosts: + validate_host_name(host) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.PREFLIGHT_FAILED, str(error)) from error + + def get_host(self, node_index: int) -> str: + """Return the verified host assigned to one planner node index.""" + if type(node_index) is not int or node_index < 0: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "resolved node index is outside the allocation", + ) + try: + return self.node_hosts[node_index] + except (IndexError, TypeError): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "resolved node index is outside the allocation", + ) from None + + +class AllocationHostResolver(Protocol): + """Resolve the scheduler node expression to ordered host identities.""" + + def resolve(self, node_list: str, environment: Mapping[str, str]) -> tuple[str, ...]: + """Return one host name per allocation node in scheduler order.""" + ... + + +class ScontrolAllocationHostResolver: + """Resolve Slurm host expressions through the scheduler-owned CLI.""" + + def resolve(self, node_list: str, environment: Mapping[str, str]) -> tuple[str, ...]: + """Expand a validated allocation node expression without a shell.""" + if type(node_list) is not str or _NODE_LIST_PATTERN.fullmatch(node_list) is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "scheduler allocation node list is unavailable or invalid", + ) + command_environment = { + "PATH": environment.get("PATH") or os.defpath, + "LC_ALL": "C", + } + try: + completed = subprocess.run( + ("scontrol", "show", "hostnames", node_list), + capture_output=True, + text=True, + check=False, + timeout=30, + env=command_environment, + ) + except (OSError, subprocess.SubprocessError) as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocation host identities could not be resolved", + ) from error + if completed.returncode != 0: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocation host identities could not be resolved", + ) + return tuple(line.strip() for line in completed.stdout.splitlines() if line.strip()) class AllocationPreflight(Protocol): - """Verify one allocation without starting package-managed processes.""" + """Verify one allocation and resolve its node identities before launch.""" - def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> None: - """Raise a normalized error when allocation facts disagree with the plan.""" + def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> AllocationLayout: + """Return verified allocation layout or raise a normalized error.""" ... class SystemAllocationPreflight: """Production verification of scheduler, filesystem, GPU, and port facts.""" - def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> None: + def __init__(self, host_resolver: AllocationHostResolver | None = None) -> None: + self._host_resolver = host_resolver or ScontrolAllocationHostResolver() + + def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> AllocationLayout: """Verify every launch-critical fact before model services start.""" try: - self._verify_scheduler(context, environment) + node_count = self._verify_scheduler(context, environment) + layout = self._resolve_layout(node_count, environment) self._verify_attempt_directory(context.attempt_directory) get_container_path(context.plan, context.attempt_directory.as_posix(), require_writable=True) self._verify_artifacts(context) - self._verify_ports(context) + verify_local_ports(context) except SlurmRuntimeError: raise except (OSError, ValueError) as error: @@ -49,22 +136,24 @@ def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "allocation preflight could not verify launch inputs", ) from error + return layout @staticmethod - def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str]) -> None: + def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str]) -> int: node_indices = { context.plan.client.host_node_index, *(index for deployment in context.plan.deployments for index in deployment.node_indices), } - if node_indices != {0}: + node_count = max(node_indices) + 1 + if node_indices != set(range(node_count)): raise SlurmRuntimeError( SlurmRuntimeErrorCode.PREFLIGHT_FAILED, - "one-node runtime received a multi-node execution plan", + "resolved allocation node indices are not complete", ) expected = { "SLURM_ARRAY_JOB_ID": context.attempt.scheduler.array_job_id, "SLURM_ARRAY_TASK_ID": context.shard.array_task_index, - "SLURM_JOB_NUM_NODES": 1, + "SLURM_JOB_NUM_NODES": node_count, "SLURM_NODEID": 0, } for name, value in expected.items(): @@ -79,6 +168,22 @@ def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str] SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "allocation GPU visibility does not match the resolved plan", ) + return node_count + + def _resolve_layout(self, node_count: int, environment: Mapping[str, str]) -> AllocationLayout: + node_list = environment.get("SLURM_JOB_NODELIST") + if node_list is None or _NODE_LIST_PATTERN.fullmatch(node_list) is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "scheduler allocation node list is unavailable or invalid", + ) + layout = AllocationLayout(self._host_resolver.resolve(node_list, environment)) + if len(layout.node_hosts) != node_count: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "scheduler allocation host count does not match the resolved plan", + ) + return layout @staticmethod def _verify_attempt_directory(attempt_directory: Path) -> None: @@ -112,31 +217,37 @@ def _verify_artifacts(context: AllocationContext) -> None: references.extend(reference for reference in optional_references if reference is not None) unique_references = {(reference.path, reference.sha256): reference for reference in references} for reference in unique_references.values(): - _verify_artifact(reference) + verify_artifact(reference) - @staticmethod - def _verify_ports(context: AllocationContext) -> None: - ports = tuple(port.port for port in context.plan.client.ports) + tuple( - port.port for deployment in context.plan.deployments for port in deployment.ports - ) - reservations: list[socket.socket] = [] - try: - for port in ports: - reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - reservations.append(reservation) - reservation.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0) - reservation.bind(("127.0.0.1", port)) - except OSError as error: - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.PREFLIGHT_FAILED, - "one or more resolved allocation ports are unavailable", - ) from error - finally: - for reservation in reservations: - reservation.close() + +def verify_local_ports(context: AllocationContext) -> None: + """Verify ports owned by the local controller node before nested steps start.""" + local_node_index = context.plan.client.host_node_index + ports = tuple(port.port for port in context.plan.client.ports if port.node_index == local_node_index) + tuple( + port.port + for deployment in context.plan.deployments + for port in deployment.ports + if port.node_index == local_node_index + ) + reservations: list[socket.socket] = [] + try: + for port in ports: + reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + reservations.append(reservation) + reservation.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0) + reservation.bind(("0.0.0.0", port)) + except OSError as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "one or more resolved allocation ports are unavailable", + ) from error + finally: + for reservation in reservations: + reservation.close() -def _verify_artifact(reference: ArtifactReference) -> None: +def verify_artifact(reference: ArtifactReference) -> None: + """Verify one regular artifact without following or racing a replacement.""" path = Path(reference.path) before = path.lstat() if not stat.S_ISREG(before.st_mode): @@ -187,4 +298,12 @@ def _parse_gpu_count(value: str | None) -> int: return len(values) -__all__ = ["AllocationPreflight", "SystemAllocationPreflight"] +__all__ = [ + "AllocationHostResolver", + "AllocationLayout", + "AllocationPreflight", + "ScontrolAllocationHostResolver", + "SystemAllocationPreflight", + "verify_artifact", + "verify_local_ports", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/probes.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/probes.py index 6d56ea4e3..0a42d3a3a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/probes.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/probes.py @@ -1,16 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Bounded loopback readiness probes for allocation-local endpoints.""" +"""Bounded readiness probes for allocation-local and remote-node endpoints.""" from __future__ import annotations import http.client from typing import Protocol +from data_designer.slurm.runtime.network import validate_host_name, validate_network_port + class ReadinessProber(Protocol): - """Probe one loopback HTTP endpoint.""" + """Probe one validated HTTP endpoint.""" def is_ready(self, host: str, port: int, path: str, *, timeout_seconds: float) -> bool: """Return whether the endpoint produced HTTP 200 within the timeout.""" @@ -21,8 +23,13 @@ class HttpReadinessProber: """Production stdlib HTTP prober with no redirect or proxy behavior.""" def is_ready(self, host: str, port: int, path: str, *, timeout_seconds: float) -> bool: - """Probe a reviewed loopback address and fully close the connection.""" - if host != "127.0.0.1" or not path.startswith("/") or timeout_seconds <= 0: + """Probe a validated allocation address and fully close the connection.""" + try: + validate_host_name(host) + validate_network_port(port) + except ValueError: + return False + if not path.startswith("/") or path.startswith("//") or timeout_seconds <= 0: return False connection = http.client.HTTPConnection(host, port, timeout=timeout_seconds) try: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py index 64dd51fe2..a9f6198d3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py @@ -13,6 +13,8 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from urllib.parse import SplitResult, urlsplit +from data_designer.slurm.runtime.network import validate_host_name + _MAXIMUM_REQUEST_BYTES = 64 * 1024 * 1024 _MAXIMUM_RESPONSE_BYTES = 64 * 1024 * 1024 _HOP_HEADERS = frozenset( @@ -191,9 +193,14 @@ def main(arguments: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(prog="data-designer-slurm-proxy") parser.add_argument("--listen-port", required=True, type=int) parser.add_argument("--backend", action="append", required=True) + parser.add_argument("--allowed-host", action="append", required=True) parser.add_argument("--retry-after-seconds", type=int) parsed = parser.parse_args(arguments) - backends = tuple(_parse_backend(value) for value in parsed.backend) + try: + allowed_hosts = frozenset(validate_host_name(value) for value in parsed.allowed_host) + except ValueError as error: + parser.error(str(error)) + backends = tuple(_parse_backend(value, allowed_hosts) for value in parsed.backend) if not 1 <= parsed.listen_port <= 65535: parser.error("listen port must be between 1 and 65535") if parsed.retry_after_seconds is not None and parsed.retry_after_seconds <= 0: @@ -204,7 +211,7 @@ def main(arguments: Sequence[str] | None = None) -> int: return 0 -def _parse_backend(value: str) -> _Backend: +def _parse_backend(value: str, allowed_hosts: frozenset[str] = frozenset({"127.0.0.1"})) -> _Backend: parsed: SplitResult = urlsplit(value) try: port = parsed.port @@ -212,7 +219,7 @@ def _parse_backend(value: str) -> _Backend: raise argparse.ArgumentTypeError("backend port is invalid") from error if ( parsed.scheme != "http" - or parsed.hostname != "127.0.0.1" + or parsed.hostname not in allowed_hosts or parsed.username is not None or parsed.password is not None or parsed.path not in {"", "/"} @@ -220,7 +227,7 @@ def _parse_backend(value: str) -> _Backend: or parsed.fragment or port is None ): - raise argparse.ArgumentTypeError("backends must be loopback HTTP origins") + raise argparse.ArgumentTypeError("backends must be allowed HTTP origins") return _Backend(parsed.hostname, port) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_steps.py new file mode 100644 index 000000000..12cb54f18 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_steps.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compose resolved vLLM deployments into coordinated multi-node steps.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path + +from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.node_spec import NodeProcessSpec, NodeSpec, NodeWorkerSpec, encode_node_worker_spec +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.runtime.step_factory import SrunStepOptions, base_environment, build_srun_step +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.vllm import ResolvedVllmProcess + + +def build_vllm_steps( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], + runtime_node_worker_path: Path | None = None, + layout: AllocationLayout | None = None, + *, + preflight: bool = False, +) -> tuple[RuntimeStep, ...]: + """Build one coordinated multi-node step for a resolved vLLM deployment.""" + selected_layout = layout or _one_node_layout(plan) + worker_path = runtime_node_worker_path or attempt_directory / "runtime/data_designer/slurm/runtime/node_worker.py" + worker_spec = _build_node_worker_spec(deployment, selected_layout) + operation = "preflight" if preflight else "serve" + role = RuntimeStepRole.SERVER_PREFLIGHT if preflight else RuntimeStepRole.SERVER + environment, container_environment = _build_vllm_environment( + deployment, + source_environment, + plan, + worker_path, + ) + step = build_srun_step( + step_id=f"{deployment.deployment_id}-{operation}", + role=role, + command=( + "python3", + get_container_path(plan, worker_path.as_posix()), + operation, + "--spec", + encode_node_worker_spec(worker_spec), + ), + environment=environment, + plan=plan, + attempt_directory=attempt_directory, + options=SrunStepOptions( + image_path=deployment.image.path, + container_environment=container_environment, + node_hosts=tuple(selected_layout.get_host(index) for index in deployment.node_indices), + gpu_count=deployment.gpus_per_node, + kill_on_bad_exit=True, + ), + ) + return (step,) + + +def _build_node_worker_spec( + deployment: ResolvedVllmServerDeployment, + layout: AllocationLayout, +) -> NodeWorkerSpec: + nodes = tuple( + NodeSpec( + node_index=node_index, + host=layout.get_host(node_index), + ports=_get_node_ports(deployment, node_index), + processes=tuple( + NodeProcessSpec( + process_id=process.process_id, + command=_build_vllm_command(deployment, process, layout), + gpu_indices=tuple(process.gpu_indices), + launch_delay_seconds=process.launch_delay_seconds, + ) + for process in deployment.processes + if process.node_index == node_index + ), + ) + for node_index in deployment.node_indices + ) + return NodeWorkerSpec(schema_version=1, resolved_gpus_per_node=deployment.gpus_per_node, nodes=nodes) + + +def _get_node_ports(deployment: ResolvedVllmServerDeployment, node_index: int) -> tuple[int, ...]: + http_ports = tuple(endpoint.port for endpoint in deployment.backend_endpoints if endpoint.node_index == node_index) + rendezvous_ports = tuple( + process.rendezvous.port + for process in deployment.processes + if process.node_index == node_index and process.pipeline_rank == 0 and process.rendezvous is not None + ) + return tuple(sorted((*http_ports, *rendezvous_ports))) + + +def _build_vllm_command( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + layout: AllocationLayout, +) -> tuple[str, ...]: + backend = deployment.backend_endpoints[process.deployment_replica_index] + command: tuple[str, ...] = ( + deployment.executable_path, + "serve", + deployment.model, + "--served-model-name", + deployment.served_model_name, + "--host", + "0.0.0.0", + "--port", + str(backend.port), + "--tensor-parallel-size", + str(process.tensor_parallel), + "--distributed-executor-backend", + "uni" if process.tensor_parallel == 1 else "mp", + "--data-parallel-backend", + "mp", + "--middleware", + "data_designer.slurm.runtime.backpressure.QueueDepthBackpressureMiddleware", + ) + command += _distributed_arguments(process, layout) + if deployment.launch_policy.enable_expert_parallel: + command += ("--enable-expert-parallel",) + return command + deployment.launch_policy.extra_args + + +def _distributed_arguments(process: ResolvedVllmProcess, layout: AllocationLayout) -> tuple[str, ...]: + if process.pipeline_parallel == 1: + return () + rendezvous = process.rendezvous + if rendezvous is None: # pragma: no cover - resolved contracts enforce this + raise AssertionError("distributed process has no rendezvous") + arguments = ( + "--pipeline-parallel-size", + str(process.pipeline_parallel), + "--nnodes", + str(process.pipeline_parallel), + "--node-rank", + str(process.pipeline_rank), + "--master-addr", + layout.get_host(rendezvous.master_node_index), + "--master-port", + str(rendezvous.port), + "--distributed-timeout-seconds", + str(rendezvous.timeout_seconds), + ) + return arguments + (("--headless",) if process.pipeline_rank > 0 else ()) + + +def _build_vllm_environment( + deployment: ResolvedVllmServerDeployment, + source_environment: Mapping[str, str], + plan: ResolvedSlurmRunPlan, + runtime_node_worker_path: Path, +) -> tuple[dict[str, str], tuple[str, ...]]: + environment = base_environment(source_environment) + container_environment: list[str] = [] + for name, binding in deployment.launch_policy.environment.items(): + if isinstance(binding, LiteralEnvironmentBinding): + environment[name] = binding.value + elif isinstance(binding, SecretRef): + try: + environment[name] = source_environment[binding.environment] + except KeyError: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + f"required server secret environment {binding.environment!r} is unavailable", + ) from None + else: # pragma: no cover - persisted contracts reject unknown bindings + raise AssertionError(f"unhandled environment binding: {type(binding)!r}") + container_environment.append(name) + runtime_root = runtime_node_worker_path.parents[3] + environment["PYTHONPATH"] = get_container_path(plan, runtime_root.as_posix()) + queue_policy = deployment.launch_policy.queue_backpressure + environment[MAX_WAITING_REQUESTS_ENVIRONMENT] = str(queue_policy.max_waiting_requests) + environment[RETRY_AFTER_SECONDS_ENVIRONMENT] = ( + "" if queue_policy.retry_after_seconds is None else str(queue_policy.retry_after_seconds) + ) + container_environment.extend(("PYTHONPATH", MAX_WAITING_REQUESTS_ENVIRONMENT, RETRY_AFTER_SECONDS_ENVIRONMENT)) + return environment, tuple(container_environment) + + +def _one_node_layout(plan: ResolvedSlurmRunPlan) -> AllocationLayout: + node_indices = { + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + } + if node_indices != {0}: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "multi-node step construction requires a verified allocation layout", + ) + return AllocationLayout(("127.0.0.1",)) + + +__all__ = ["build_vllm_steps"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_factory.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_factory.py new file mode 100644 index 000000000..83b578083 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_factory.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compose validated allocation-local commands into structured ``srun`` steps.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path + +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.network import validate_host_name + +_SLURM_ENVIRONMENT_NAMES = ( + "SLURM_ARRAY_JOB_ID", + "SLURM_ARRAY_TASK_ID", + "SLURM_CLUSTER_NAME", + "SLURM_CONF", + "SLURM_CPUS_ON_NODE", + "SLURM_JOB_CPUS_PER_NODE", + "SLURM_JOB_GPUS", + "SLURM_JOB_ID", + "SLURM_JOB_NODELIST", + "SLURM_JOB_NUM_NODES", + "SLURM_NODEID", + "SLURM_PROCID", + "SLURM_SUBMIT_DIR", +) + + +@dataclass(frozen=True, slots=True) +class SrunStepOptions: + """Scheduler and container options for one structured runtime step.""" + + image_path: str + container_environment: tuple[str, ...] = () + node_hosts: tuple[str, ...] = () + gpu_count: int | None = None + kill_on_bad_exit: bool = False + + def __post_init__(self) -> None: + if type(self.image_path) is not str or not self.image_path: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime image path is invalid") + if type(self.container_environment) is not tuple or any( + type(name) is not str or not name or any(character in name for character in (",", "=", "\0")) + for name in self.container_environment + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime container environment names are invalid", + ) + if type(self.node_hosts) is not tuple: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime node hosts are invalid") + try: + for host in self.node_hosts: + validate_host_name(host) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, str(error)) from error + if self.node_hosts and len(self.node_hosts) != len(set(self.node_hosts)): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime node hosts must be unique") + if self.gpu_count is not None and (type(self.gpu_count) is not int or self.gpu_count <= 0): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime GPU count is invalid") + if type(self.kill_on_bad_exit) is not bool: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime failure policy is invalid") + + +def build_srun_step( + *, + step_id: str, + role: RuntimeStepRole, + command: tuple[str, ...], + environment: Mapping[str, str], + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + options: SrunStepOptions, +) -> RuntimeStep: + """Build one shell-free, resource-scoped ``srun`` step.""" + srun_command = _build_srun_prefix(plan, options) + _add_srun_container_options(srun_command, plan, options.container_environment) + srun_command.extend(("--", *command)) + log_root = attempt_directory / "logs" + return RuntimeStep( + step_id=step_id, + role=role, + command=tuple(srun_command), + environment=environment, + stdout_path=log_root / f"{step_id}.out", + stderr_path=log_root / f"{step_id}.err", + ) + + +def base_environment(source_environment: Mapping[str, str]) -> dict[str, str]: + """Return the bounded environment shared by all package-owned steps.""" + environment = {"PATH": source_environment.get("PATH") or os.defpath, "LC_ALL": "C"} + environment.update( + (name, source_environment[name]) for name in _SLURM_ENVIRONMENT_NAMES if name in source_environment + ) + return environment + + +def place_step_on_host(step: RuntimeStep, host: str) -> RuntimeStep: + """Constrain a package-built single-task ``srun`` step to one verified host.""" + try: + validate_host_name(host) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, str(error)) from error + if step.command[0] != "srun": + return step + command = list(step.command) + command.insert(1, f"--nodelist={host}") + return RuntimeStep( + step_id=step.step_id, + role=step.role, + command=tuple(command), + environment=step.environment, + stdout_path=step.stdout_path, + stderr_path=step.stderr_path, + ) + + +def _build_srun_prefix(plan: ResolvedSlurmRunPlan, options: SrunStepOptions) -> list[str]: + task_count = len(options.node_hosts) or 1 + command = [ + "srun", + f"--nodes={task_count}", + f"--ntasks={task_count}", + "--ntasks-per-node=1", + "--exact", + "--overlap", + "--unbuffered", + "--export=ALL", + f"--cpus-per-task={plan.client.authored.cpus}", + f"--container-image={options.image_path}", + ] + if options.node_hosts: + command.append(f"--nodelist={','.join(options.node_hosts)}") + if options.kill_on_bad_exit: + command.append("--kill-on-bad-exit=1") + if options.gpu_count is None: + command.append("--gres=none") + else: + command.extend((f"--gpus-per-task={options.gpu_count}", "--gpu-bind=none")) + return command + + +def _add_srun_container_options( + srun_command: list[str], + plan: ResolvedSlurmRunPlan, + container_environment: tuple[str, ...], +) -> None: + mounts = _render_mounts(plan) + if mounts: + srun_command.append(f"--container-mounts={mounts}") + if container_environment: + srun_command.append(f"--container-env={','.join(sorted(set(container_environment)))}") + + +def _render_mounts(plan: ResolvedSlurmRunPlan) -> str: + rendered: list[str] = [] + for mount in plan.container_mounts: + if any(character in mount.source or character in mount.target for character in (",", ":")): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "container mount paths cannot contain comma or colon delimiters", + ) + value = f"{mount.source}:{mount.target}" + if mount.read_only: + value += ":ro" + rendered.append(value) + return ",".join(rendered) + + +__all__ = ["SrunStepOptions", "base_environment", "build_srun_step", "place_step_on_host"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py index ed998331a..0cdb2d7a8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -1,41 +1,32 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Build every allocation-local process as one structured ``srun`` step.""" +"""Build coordinated allocation-local processes as structured ``srun`` steps.""" from __future__ import annotations -import os from collections.abc import Mapping from pathlib import Path from typing import Protocol from pydantic import BaseModel -from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.config.environment import SecretRef from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.endpoint_steps import build_endpoint_steps as build_endpoint_steps from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStep, RuntimeStepRole from data_designer.slurm.runtime.paths import get_container_path -from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment -from data_designer.slurm.serving.vllm import ResolvedVllmProcess -from data_designer.slurm.state import AttemptManifest - -_SLURM_ENVIRONMENT_NAMES = ( - "SLURM_ARRAY_JOB_ID", - "SLURM_ARRAY_TASK_ID", - "SLURM_CLUSTER_NAME", - "SLURM_CONF", - "SLURM_CPUS_ON_NODE", - "SLURM_JOB_CPUS_PER_NODE", - "SLURM_JOB_GPUS", - "SLURM_JOB_ID", - "SLURM_JOB_NODELIST", - "SLURM_JOB_NUM_NODES", - "SLURM_NODEID", - "SLURM_PROCID", - "SLURM_SUBMIT_DIR", +from data_designer.slurm.runtime.server_steps import build_vllm_steps as build_vllm_steps +from data_designer.slurm.runtime.step_factory import ( + SrunStepOptions, + base_environment, + build_srun_step, +) +from data_designer.slurm.runtime.step_factory import ( + place_step_on_host as place_step_on_host, ) +from data_designer.slurm.state import AttemptManifest class ClientStepBuilder(Protocol): @@ -127,15 +118,17 @@ def _build_step( ) -> RuntimeStep: command = _build_client_command(operation, plan, shard, attempt, attempt_directory, endpoints) secret_names, environment = _build_client_environment(plan, source_environment) - return _build_srun_step( + return build_srun_step( step_id=step_id, role=role, - image_path=plan.client.image.path, command=command, environment=environment, - container_environment=secret_names, plan=plan, attempt_directory=attempt_directory, + options=SrunStepOptions( + image_path=plan.client.image.path, + container_environment=secret_names, + ), ) @@ -174,7 +167,7 @@ def _build_client_environment( source_environment: Mapping[str, str], ) -> tuple[tuple[str, ...], dict[str, str]]: secret_names = _collect_client_secret_environment_names(plan) - environment = _base_environment(source_environment) + environment = base_environment(source_environment) for name in secret_names: try: environment[name] = source_environment[name] @@ -186,257 +179,11 @@ def _build_client_environment( return secret_names, environment -def build_vllm_steps( - deployment: ResolvedVllmServerDeployment, - plan: ResolvedSlurmRunPlan, - attempt_directory: Path, - source_environment: Mapping[str, str], -) -> tuple[RuntimeStep, ...]: - """Build one structured server step per resolved vLLM process.""" - return tuple( - _build_vllm_step(deployment, process, plan, attempt_directory, source_environment) - for process in deployment.processes - ) - - -def build_endpoint_steps( - deployments: tuple[ResolvedVllmServerDeployment, ...], - plan: ResolvedSlurmRunPlan, - attempt_directory: Path, - source_environment: Mapping[str, str], - runtime_proxy_path: Path, -) -> tuple[tuple[RuntimeStep, RuntimeEndpoint], ...]: - """Build one package-owned logical-endpoint process per deployment.""" - steps: list[tuple[RuntimeStep, RuntimeEndpoint]] = [] - for deployment in deployments: - steps.append(_build_endpoint_step(deployment, plan, attempt_directory, source_environment, runtime_proxy_path)) - return tuple(steps) - - -def _build_endpoint_step( - deployment: ResolvedVllmServerDeployment, - plan: ResolvedSlurmRunPlan, - attempt_directory: Path, - source_environment: Mapping[str, str], - runtime_proxy_path: Path, -) -> tuple[RuntimeStep, RuntimeEndpoint]: - endpoint = RuntimeEndpoint( - model_alias=deployment.model_alias, - served_model_name=deployment.served_model_name, - host="127.0.0.1", - port=deployment.logical_endpoint.port, - ) - command = _build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port) - step = _build_srun_step( - step_id=f"{deployment.deployment_id}-endpoint", - role=RuntimeStepRole.ENDPOINT, - image_path=plan.client.image.path, - command=command, - environment=_base_environment(source_environment), - container_environment=(), - plan=plan, - attempt_directory=attempt_directory, - ) - return step, endpoint - - -def _build_endpoint_command( - deployment: ResolvedVllmServerDeployment, - plan: ResolvedSlurmRunPlan, - runtime_proxy_path: Path, - port: int, -) -> tuple[str, ...]: - backends = tuple(f"http://127.0.0.1:{backend.port}" for backend in deployment.backend_endpoints) - retry_after_seconds = deployment.launch_policy.queue_backpressure.retry_after_seconds - retry_arguments = ("--retry-after-seconds", str(retry_after_seconds)) if retry_after_seconds is not None else () - return ( - "python3", - get_container_path(plan, runtime_proxy_path.as_posix()), - "--listen-port", - str(port), - *retry_arguments, - *(argument for backend in backends for argument in ("--backend", backend)), - ) - - def plan_path(plan: ResolvedSlurmRunPlan) -> str: """Return the canonical persisted resolved-plan path.""" return str(Path(plan.authored_config.path).with_name("resolved-plan.json")) -def _build_vllm_step( - deployment: ResolvedVllmServerDeployment, - process: ResolvedVllmProcess, - plan: ResolvedSlurmRunPlan, - attempt_directory: Path, - source_environment: Mapping[str, str], -) -> RuntimeStep: - _validate_local_vllm_process(process) - command = _build_vllm_command(deployment, process) - environment, container_environment = _build_vllm_environment(deployment, source_environment) - return _build_srun_step( - step_id=process.process_id, - role=RuntimeStepRole.SERVER, - image_path=deployment.image.path, - command=command, - environment=environment, - container_environment=container_environment, - plan=plan, - attempt_directory=attempt_directory, - gpu_indices=tuple(process.gpu_indices), - ) - - -def _validate_local_vllm_process(process: ResolvedVllmProcess) -> None: - if process.pipeline_parallel != 1 or process.node_index != 0 or process.http_port is None: - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.INVALID_CONTEXT, - "one-node runtime received a distributed vLLM process", - ) - - -def _build_vllm_command( - deployment: ResolvedVllmServerDeployment, - process: ResolvedVllmProcess, -) -> tuple[str, ...]: - if process.http_port is None: # pragma: no cover - validated before command construction - raise AssertionError("vLLM HTTP port is unavailable") - command: tuple[str, ...] = ( - deployment.executable_path, - "serve", - deployment.model, - "--served-model-name", - deployment.served_model_name, - "--host", - "127.0.0.1", - "--port", - str(process.http_port), - "--tensor-parallel-size", - str(process.tensor_parallel), - ) - if deployment.launch_policy.enable_expert_parallel: - command += ("--enable-expert-parallel",) - return command + deployment.launch_policy.extra_args - - -def _build_vllm_environment( - deployment: ResolvedVllmServerDeployment, - source_environment: Mapping[str, str], -) -> tuple[dict[str, str], tuple[str, ...]]: - environment = _base_environment(source_environment) - container_environment: list[str] = [] - for name, binding in deployment.launch_policy.environment.items(): - if isinstance(binding, LiteralEnvironmentBinding): - environment[name] = binding.value - elif isinstance(binding, SecretRef): - try: - environment[name] = source_environment[binding.environment] - except KeyError: - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.PREFLIGHT_FAILED, - f"required server secret environment {binding.environment!r} is unavailable", - ) from None - else: # pragma: no cover - persisted contracts reject unknown bindings - raise AssertionError(f"unhandled environment binding: {type(binding)!r}") - container_environment.append(name) - return environment, tuple(container_environment) - - -def _build_srun_step( - *, - step_id: str, - role: RuntimeStepRole, - image_path: str, - command: tuple[str, ...], - environment: Mapping[str, str], - container_environment: tuple[str, ...], - plan: ResolvedSlurmRunPlan, - attempt_directory: Path, - gpu_indices: tuple[int, ...] = (), -) -> RuntimeStep: - srun_command = _build_srun_prefix(plan, image_path) - _add_srun_resources(srun_command, gpu_indices) - _add_srun_container_options(srun_command, plan, container_environment) - srun_command.extend(("--", *command)) - log_root = attempt_directory / "logs" - return RuntimeStep( - step_id=step_id, - role=role, - command=tuple(srun_command), - environment=environment, - stdout_path=log_root / f"{step_id}.out", - stderr_path=log_root / f"{step_id}.err", - ) - - -def _build_srun_prefix(plan: ResolvedSlurmRunPlan, image_path: str) -> list[str]: - return [ - "srun", - "--nodes=1", - "--ntasks=1", - "--exact", - "--overlap", - "--unbuffered", - "--export=ALL", - f"--cpus-per-task={plan.client.authored.cpus}", - f"--container-image={image_path}", - ] - - -def _add_srun_resources(srun_command: list[str], gpu_indices: tuple[int, ...]) -> None: - if gpu_indices: - srun_command.extend( - ( - f"--gpus-per-task={len(gpu_indices)}", - f"--gpu-bind=mask_gpu:{_get_gpu_mask(gpu_indices):#x}", - ) - ) - else: - srun_command.append("--gres=none") - - -def _add_srun_container_options( - srun_command: list[str], - plan: ResolvedSlurmRunPlan, - container_environment: tuple[str, ...], -) -> None: - mounts = _render_mounts(plan) - if mounts: - srun_command.append(f"--container-mounts={mounts}") - if container_environment: - srun_command.append(f"--container-env={','.join(sorted(set(container_environment)))}") - - -def _get_gpu_mask(gpu_indices: tuple[int, ...]) -> int: - return sum(1 << index for index in gpu_indices) - - -def _render_mounts(plan: ResolvedSlurmRunPlan) -> str: - rendered: list[str] = [] - for mount in plan.container_mounts: - if any(character in mount.source or character in mount.target for character in (",", ":")): - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.INVALID_CONTEXT, - "container mount paths cannot contain comma or colon delimiters", - ) - value = f"{mount.source}:{mount.target}" - if mount.read_only: - value += ":ro" - rendered.append(value) - return ",".join(rendered) - - -def _base_environment(source_environment: Mapping[str, str]) -> dict[str, str]: - environment = { - "PATH": source_environment.get("PATH") or os.defpath, - "LC_ALL": "C", - } - environment.update( - (name, source_environment[name]) for name in _SLURM_ENVIRONMENT_NAMES if name in source_environment - ) - return environment - - def _collect_client_secret_environment_names(plan: ResolvedSlurmRunPlan) -> tuple[str, ...]: names: set[str] = set() _collect_secret_environment_names(plan.client.authored.dependencies.index_credentials, names) diff --git a/packages/data-designer-slurm/tests/runtime/conftest.py b/packages/data-designer-slurm/tests/runtime/conftest.py index 9dc83da6d..a882a2557 100644 --- a/packages/data-designer-slurm/tests/runtime/conftest.py +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -16,6 +16,7 @@ from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.preflight import AllocationLayout from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, @@ -81,11 +82,13 @@ def __init__(self, failure: BaseException | None = None) -> None: self.failure = failure self.calls = 0 - def verify(self, context: AllocationContext, environment: object) -> None: - del context, environment + def verify(self, context: AllocationContext, environment: object) -> AllocationLayout: + del environment self.calls += 1 if self.failure is not None: raise self.failure + node_count = max(index for deployment in context.plan.deployments for index in deployment.node_indices) + 1 + return AllocationLayout(tuple(f"node-{index}" for index in range(node_count))) class FakeClientStepBuilder: @@ -116,9 +119,18 @@ def build_generation_step( @pytest.fixture def runtime_case(tmp_path: Path, single_node_plan: ResolvedSlurmRunPlan) -> RuntimeCase: + return _build_runtime_case(tmp_path, single_node_plan) + + +@pytest.fixture +def multi_node_runtime_case(tmp_path: Path, multi_node_plan: ResolvedSlurmRunPlan) -> RuntimeCase: + return _build_runtime_case(tmp_path, multi_node_plan) + + +def _build_runtime_case(tmp_path: Path, source_plan: ResolvedSlurmRunPlan) -> RuntimeCase: workspace = tmp_path / "workspace" workspace.mkdir(mode=0o700) - plan = relocate_plan(single_node_plan, workspace) + plan = relocate_plan(source_plan, workspace) created_at = datetime(2026, 9, 1, 12, tzinfo=timezone.utc) plan_reference = ArtifactReference( path=Path(plan.authored_config.path).with_name("resolved-plan.json").as_posix(), diff --git a/packages/data-designer-slurm/tests/runtime/test_backpressure.py b/packages/data-designer-slurm/tests/runtime/test_backpressure.py new file mode 100644 index 000000000..9f2a80c72 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_backpressure.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, + AsgiMessage, + QueueBackpressureController, + QueueBackpressureSettings, + QueueDepthBackpressureMiddleware, +) + + +def test_backpressure_settings_are_strict_and_support_absent_retry_header() -> None: + settings = QueueBackpressureSettings.from_environment( + { + MAX_WAITING_REQUESTS_ENVIRONMENT: "7", + RETRY_AFTER_SECONDS_ENVIRONMENT: "", + } + ) + assert settings.max_waiting_requests == 7 + assert settings.retry_after_seconds is None + with pytest.raises(ValueError, match="threshold"): + QueueBackpressureSettings.from_environment({MAX_WAITING_REQUESTS_ENVIRONMENT: "-1"}) + + +@pytest.mark.asyncio +async def test_middleware_rejects_fresh_overload_and_preserves_health() -> None: + downstream_calls: list[str] = [] + messages: list[AsgiMessage] = [] + + async def app(scope: dict[str, object], receive: object, send: object) -> None: + del receive, send + downstream_calls.append(str(scope["path"])) + + async def receive() -> AsgiMessage: + return {"type": "http.request"} + + async def send(message: AsgiMessage) -> None: + messages.append(message) + + controller = QueueBackpressureController( + QueueBackpressureSettings(2, 3), + reader=lambda: 4, + start_sampler=False, + ) + controller.sample_once() + middleware = QueueDepthBackpressureMiddleware(app, controller) + + await middleware({"type": "http", "path": "/v1/chat/completions"}, receive, send) + await middleware({"type": "http", "path": "/health"}, receive, send) + + assert messages[0]["status"] == 429 + assert (b"retry-after", b"3") in messages[0]["headers"] + assert downstream_calls == ["/health"] + + +@pytest.mark.asyncio +async def test_middleware_fails_open_when_metrics_are_unavailable() -> None: + called = False + + async def app(scope: dict[str, object], receive: object, send: object) -> None: + nonlocal called + del scope, receive, send + called = True + + async def receive() -> AsgiMessage: + return {"type": "http.request"} + + async def send(message: AsgiMessage) -> None: + raise AssertionError(f"unexpected middleware response: {message}") + + controller = QueueBackpressureController(reader=lambda: None, start_sampler=False) + controller.sample_once() + middleware = QueueDepthBackpressureMiddleware(app, controller) + await middleware({"type": "http", "path": "/v1/completions"}, receive, send) + + assert called diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index de4b363c5..7e602b394 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -17,7 +17,8 @@ from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.runtime.controller import OneNodeAllocationController from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode -from data_designer.slurm.runtime.models import RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.models import AllocationContext, RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.preflight import AllocationLayout from data_designer.slurm.runtime.signals import TerminationSignalCoordinator from data_designer.slurm.runtime.supervisor import StepSupervisor from data_designer.slurm.state import ( @@ -76,12 +77,14 @@ def __init__( failed_role: RuntimeStepRole | None = None, fail_cleanup: bool = False, incomplete_cleanup: bool = False, + failed_step_id: str | None = None, ) -> None: self.generation_hook = generation_hook self.server_returncode = server_returncode self.failed_role = failed_role self.fail_cleanup = fail_cleanup self.incomplete_cleanup = incomplete_cleanup + self.failed_step_id = failed_step_id self.steps: list[RuntimeStep] = [] self.processes: list[_FakeProcess] = [] @@ -90,7 +93,7 @@ def start(self, step: RuntimeStep) -> _FakeProcess: if step.role is RuntimeStepRole.CLIENT and self.generation_hook is not None: self.generation_hook() is_managed_service = step.role in {RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT} - if step.role is self.failed_role: + if step.role is self.failed_role or step.step_id == self.failed_step_id: returncode = 23 elif step.role is RuntimeStepRole.SERVER: returncode = self.server_returncode @@ -112,10 +115,12 @@ def __init__(self, *, ready: bool, clock: FakeClock, advance_on_failure: float = self.clock = clock self.advance_on_failure = advance_on_failure self.calls = 0 + self.hosts: list[str] = [] def is_ready(self, host: str, port: int, path: str, *, timeout_seconds: float) -> bool: - del host, port, path, timeout_seconds + del port, path, timeout_seconds self.calls += 1 + self.hosts.append(host) if not self.ready and self.advance_on_failure: self.clock.advance(self.advance_on_failure) return self.ready @@ -125,9 +130,11 @@ def is_ready(self, host: str, port: int, path: str, *, timeout_seconds: float) - class _RestartAwarePreflight: state: FakeStateStore - def verify(self, context: object, environment: object) -> None: - del context, environment + def verify(self, context: AllocationContext, environment: object) -> AllocationLayout: + del environment assert self.state.readiness[-1].state is ReadinessState.RESTARTING + node_count = max(index for deployment in context.plan.deployments for index in deployment.node_indices) + 1 + return AllocationLayout(tuple(f"node-{index}" for index in range(node_count))) def _supervisor( @@ -170,6 +177,7 @@ def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_c assert result.candidate_output is not None assert [step.role for step in runner.steps] == [ RuntimeStepRole.CLIENT_PREFLIGHT, + RuntimeStepRole.SERVER_PREFLIGHT, RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT, RuntimeStepRole.CLIENT, @@ -182,6 +190,37 @@ def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_c assert {step.stdout_path.parent.name for step in runner.steps} == {"execution-00000001"} +def test_controller_runs_multi_node_deployments_and_probes_remote_heads( + multi_node_runtime_case: RuntimeCase, +) -> None: + runtime_case = multi_node_runtime_case + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + prober = _FakeProber(ready=True, clock=clock) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=prober, + clock=clock, + environment={"HF_TOKEN": "server-secret"}, + ) + + result = controller.run() + + assert result.state is AttemptLifecycleState.SUCCEEDED + server_steps = tuple(step for step in runner.steps if step.role is RuntimeStepRole.SERVER) + assert len(server_steps) == 2 + assert "--nodelist=node-0,node-1" in server_steps[0].command + assert "--nodelist=node-2" in server_steps[1].command + assert {"node-0", "node-2", "127.0.0.1"}.issubset(prober.hosts) + assert all(process.poll() is not None for process in runner.processes) + + def test_controller_publishes_result_before_success_with_real_state_writer( runtime_case: RuntimeCase, authored_run_single: DataDesignerSlurmConfig, @@ -302,6 +341,7 @@ def test_requeued_running_attempt_publishes_restart_epoch_and_uses_fresh_logs( assert state.readiness[-1].state is ReadinessState.STOPPED assert [step.role for step in runner.steps] == [ RuntimeStepRole.CLIENT_PREFLIGHT, + RuntimeStepRole.SERVER_PREFLIGHT, RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT, RuntimeStepRole.CLIENT, @@ -333,6 +373,39 @@ def test_required_server_exit_fails_and_cleans_partial_start(runtime_case: Runti assert state.readiness[-1].state is ReadinessState.STOPPED +def test_one_distributed_deployment_failure_stops_all_deployments( + multi_node_runtime_case: RuntimeCase, +) -> None: + runtime_case = multi_node_runtime_case + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner(failed_step_id="deployment-00001-serve") + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment={"HF_TOKEN": "server-secret"}, + ) + + with pytest.raises(SlurmRuntimeError, match="status 23"): + controller.run() + + server_processes = tuple( + process + for step, process in zip(runner.steps, runner.processes, strict=True) + if step.role is RuntimeStepRole.SERVER + ) + assert len(server_processes) == 2 + assert all(process.poll() is not None for process in server_processes) + assert all(step.role is not RuntimeStepRole.ENDPOINT for step in runner.steps) + assert state.readiness[-1].state is ReadinessState.STOPPED + + def test_interrupted_failed_readiness_write_cannot_bypass_cleanup(runtime_case: RuntimeCase) -> None: clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) @@ -395,7 +468,12 @@ def test_readiness_timeout_fails_and_terminates_server(runtime_case: RuntimeCase @pytest.mark.parametrize( "failed_role", - (RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.ENDPOINT, RuntimeStepRole.CLIENT), + ( + RuntimeStepRole.SERVER_PREFLIGHT, + RuntimeStepRole.CLIENT_PREFLIGHT, + RuntimeStepRole.ENDPOINT, + RuntimeStepRole.CLIENT, + ), ) def test_managed_step_failure_fails_attempt_and_cleans_started_services( runtime_case: RuntimeCase, diff --git a/packages/data-designer-slurm/tests/runtime/test_node_worker.py b/packages/data-designer-slurm/tests/runtime/test_node_worker.py new file mode 100644 index 000000000..7ed48c58c --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_node_worker.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from data_designer.slurm.runtime import node_worker as runtime_node_worker +from data_designer.slurm.runtime.node_spec import ( + NodeProcessSpec, + NodeSpec, + NodeWorkerSpec, + decode_node_worker_spec, + encode_node_worker_spec, +) +from data_designer.slurm.runtime.node_worker import NodeProcessSupervisor + + +def test_node_worker_spec_round_trips_and_rejects_untrusted_payload() -> None: + spec = _worker_spec((NodeProcessSpec("lane-0", ("true",), (0,), 0),)) + assert decode_node_worker_spec(encode_node_worker_spec(spec)) == spec + with pytest.raises(ValueError, match="invalid"): + decode_node_worker_spec("not-base64") + + +def test_partial_startup_failure_cleans_already_started_lane(monkeypatch: pytest.MonkeyPatch) -> None: + process = _FakeProcess(pid=41) + calls = 0 + signals: list[int] = [] + + def start(*arguments: object, **keywords: object) -> _FakeProcess: + nonlocal calls + del arguments, keywords + calls += 1 + if calls == 2: + raise OSError("injected partial startup") + return process + + monkeypatch.setattr(runtime_node_worker.subprocess, "Popen", start) + + def terminate(child: _FakeProcess, selected: int) -> None: + signals.append(selected) + child.returncode = -15 + + monkeypatch.setattr(runtime_node_worker, "_signal_process_group", terminate) + supervisor = NodeProcessSupervisor(environment={}) + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=( + NodeProcessSpec("lane-0", ("first",), (0,), 0), + NodeProcessSpec("lane-1", ("second",), (1,), 0), + ), + ) + + with pytest.raises(OSError, match="partial startup"): + supervisor.run(node, ("0", "1")) + + assert signals + assert supervisor.cleanup_complete + + +def test_follower_failure_terminates_sibling_without_orphan(tmp_path: Path) -> None: + ready = tmp_path / "ready" + stopped = tmp_path / "stopped" + long_running = ( + "import signal,time; from pathlib import Path; " + f"ready=Path({ready.as_posix()!r}); stopped=Path({stopped.as_posix()!r}); " + "signal.signal(signal.SIGTERM, lambda *_: (stopped.write_text('stopped'), exit(0))); " + "ready.write_text('ready'); time.sleep(30)" + ) + failing = "import time; time.sleep(0.2); raise SystemExit(23)" + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=( + NodeProcessSpec("head", (sys.executable, "-c", long_running), (0,), 0), + NodeProcessSpec("follower", (sys.executable, "-c", failing), (1,), 0), + ), + ) + supervisor = NodeProcessSupervisor(environment={}) + + assert supervisor.run(node, ("0", "1")) == 23 + supervisor.cleanup() + + assert ready.exists() + assert stopped.read_text() == "stopped" + assert supervisor.cleanup_complete + + +def test_cancellation_request_stops_all_node_lanes(tmp_path: Path) -> None: + ready = tmp_path / "ready" + stopped = tmp_path / "stopped" + child = ( + "import signal,time; from pathlib import Path; " + f"ready=Path({ready.as_posix()!r}); stopped=Path({stopped.as_posix()!r}); " + "signal.signal(signal.SIGTERM, lambda *_: (stopped.write_text('stopped'), exit(0))); " + "ready.write_text('ready'); time.sleep(30)" + ) + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=(NodeProcessSpec("lane", (sys.executable, "-c", child), (0,), 0),), + ) + supervisor = NodeProcessSupervisor(environment={}) + with ThreadPoolExecutor(max_workers=1) as executor: + result = executor.submit(supervisor.run, node, ("0",)) + _wait_for_file(ready) + supervisor.request_stop() + assert result.result(timeout=5) == 1 + + assert stopped.read_text() == "stopped" + assert supervisor.cleanup_complete + + +def test_cancellation_interrupts_a_pending_launch_delay() -> None: + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=(NodeProcessSpec("delayed", ("must-not-start",), (0,), 30),), + ) + supervisor = NodeProcessSupervisor(environment={}) + with ThreadPoolExecutor(max_workers=1) as executor: + result = executor.submit(supervisor.run, node, ("0",)) + time.sleep(0.05) + supervisor.request_stop() + assert result.result(timeout=2) == 1 + + assert supervisor.cleanup_complete + + +@dataclass(slots=True) +class _FakeProcess: + pid: int + returncode: int | None = None + + def poll(self) -> int | None: + return self.returncode + + def wait(self, timeout: float | None = None) -> int: + del timeout + self.returncode = -15 + return self.returncode + + +def _worker_spec(processes: tuple[NodeProcessSpec, ...]) -> NodeWorkerSpec: + return NodeWorkerSpec( + schema_version=1, + resolved_gpus_per_node=8, + nodes=(NodeSpec(node_index=0, host="compute-001", ports=(18000,), processes=processes),), + ) + + +def _wait_for_file(path: Path) -> None: + deadline = time.monotonic() + 3 + while not path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + assert path.exists() diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index e4be62e92..3d80c8eb8 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -5,6 +5,7 @@ import hashlib import os +import subprocess from pathlib import Path import pytest @@ -13,7 +14,37 @@ from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.runtime import preflight as runtime_preflight from data_designer.slurm.runtime.errors import SlurmRuntimeError -from data_designer.slurm.runtime.preflight import SystemAllocationPreflight, _verify_artifact +from data_designer.slurm.runtime.preflight import ( + AllocationLayout, + ScontrolAllocationHostResolver, + SystemAllocationPreflight, + verify_artifact, + verify_local_ports, +) + + +def test_allocation_layout_rejects_negative_indices() -> None: + layout = AllocationLayout(("compute-001", "compute-002")) + with pytest.raises(SlurmRuntimeError, match="outside"): + layout.get_host(-1) + + +def test_scontrol_resolver_is_shell_free_and_rejects_option_injection(monkeypatch: pytest.MonkeyPatch) -> None: + commands: list[tuple[str, ...]] = [] + + def run(command: tuple[str, ...], **options: object) -> subprocess.CompletedProcess[str]: + del options + commands.append(command) + return subprocess.CompletedProcess(command, 0, stdout="compute-001\ncompute-002\n", stderr="") + + monkeypatch.setattr(runtime_preflight.subprocess, "run", run) + resolver = ScontrolAllocationHostResolver() + + assert resolver.resolve("compute-[001-002]", {"PATH": "/usr/bin"}) == ("compute-001", "compute-002") + assert commands == [("scontrol", "show", "hostnames", "compute-[001-002]")] + with pytest.raises(SlurmRuntimeError, match="node list"): + resolver.resolve("--help", {}) + assert len(commands) == 1 def test_scheduler_preflight_accepts_exact_one_node_gpu_shape(runtime_case: RuntimeCase) -> None: @@ -62,17 +93,17 @@ def test_artifact_verification_checks_exact_bytes_and_rejects_symlink(tmp_path: artifact.write_bytes(content) reference = ArtifactReference(path=artifact.as_posix(), sha256=hashlib.sha256(content).hexdigest()) - _verify_artifact(reference) + verify_artifact(reference) artifact.write_bytes(b"changed") with pytest.raises(OSError, match="digest"): - _verify_artifact(reference) + verify_artifact(reference) target = tmp_path / "target.bin" target.write_bytes(content) artifact.unlink() artifact.symlink_to(target) with pytest.raises(OSError, match="regular file"): - _verify_artifact(reference) + verify_artifact(reference) def test_artifact_verification_rejects_path_replacement_during_read( @@ -100,7 +131,7 @@ def replace_after_descriptor_validation(descriptor: int) -> os.stat_result: monkeypatch.setattr(runtime_preflight.os, "fstat", replace_after_descriptor_validation) with pytest.raises(OSError, match="replaced while it was verified"): - _verify_artifact(reference) + verify_artifact(reference) assert replaced @@ -127,7 +158,7 @@ def close(self) -> None: monkeypatch.setattr("data_designer.slurm.runtime.preflight.socket.socket", _UnavailableSocket) with pytest.raises(SlurmRuntimeError, match="ports are unavailable"): - SystemAllocationPreflight._verify_ports(runtime_case.context) + verify_local_ports(runtime_case.context) assert _UnavailableSocket.closed diff --git a/packages/data-designer-slurm/tests/runtime/test_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py index b01ed41aa..1cc2a5075 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -12,8 +12,14 @@ from data_designer.slurm.config import QueueBackpressureConfig from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStepRole +from data_designer.slurm.runtime.node_spec import decode_node_worker_spec +from data_designer.slurm.runtime.preflight import AllocationLayout from data_designer.slurm.runtime.steps import ( DefaultClientStepBuilder, build_endpoint_steps, @@ -77,9 +83,9 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti server = server_steps[0] assert server.role is RuntimeStepRole.SERVER - assert f"--gpus-per-task={deployment.topology.tensor_parallel}" in server.command - expected_mask = sum(1 << index for index in deployment.processes[0].gpu_indices) - assert f"--gpu-bind=mask_gpu:{expected_mask:#x}" in server.command + assert f"--gpus-per-task={deployment.gpus_per_node}" in server.command + assert "--gpu-bind=none" in server.command + assert "--kill-on-bad-exit=1" in server.command assert "CUDA_VISIBLE_DEVICES" not in server.environment assert all("CUDA_VISIBLE_DEVICES" not in argument for argument in server.command) assert all("--gpus-per-task=" not in argument for step in client_steps for argument in step.command) @@ -160,6 +166,78 @@ def test_endpoint_step_uses_resolved_retry_policy_and_backends(runtime_case: Run assert "--max-waiting-requests" not in disabled_step.command +def test_multi_node_deployment_is_one_coordinated_step_with_ranked_lane_commands( + tmp_path: Path, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + plan = relocate_plan(multi_node_plan, workspace) + deployment = resolve_vllm_server(plan, plan.deployments[0].deployment_id) + layout = AllocationLayout(("compute-001", "compute-002", "compute-003")) + + (step,) = build_vllm_steps( + deployment, + plan, + workspace / "attempt", + {"HF_TOKEN": "server-secret"}, + workspace / "runtime/data_designer/slurm/runtime/node_worker.py", + layout, + ) + + assert "--nodes=2" in step.command + assert "--ntasks=2" in step.command + assert "--ntasks-per-node=1" in step.command + assert "--nodelist=compute-001,compute-002" in step.command + assert "--kill-on-bad-exit=1" in step.command + encoded = step.command[step.command.index("--spec") + 1] + spec = decode_node_worker_spec(encoded) + assert tuple(node.host for node in spec.nodes) == ("compute-001", "compute-002") + head, follower = (node.processes[0] for node in spec.nodes) + assert "--node-rank" in head.command + assert head.command[head.command.index("--node-rank") + 1] == "0" + assert head.command[head.command.index("--master-addr") + 1] == "compute-001" + assert head.command[head.command.index("--distributed-executor-backend") + 1] == "mp" + assert head.command[head.command.index("--data-parallel-backend") + 1] == "mp" + assert "--headless" not in head.command + assert follower.command[follower.command.index("--node-rank") + 1] == "1" + assert "--headless" in follower.command + assert head.launch_delay_seconds == follower.launch_delay_seconds + assert step.environment[MAX_WAITING_REQUESTS_ENVIRONMENT] == str( + deployment.launch_policy.queue_backpressure.max_waiting_requests + ) + assert step.environment[RETRY_AFTER_SECONDS_ENVIRONMENT] == str( + deployment.launch_policy.queue_backpressure.retry_after_seconds + ) + + +def test_multi_node_endpoint_targets_resolved_remote_heads( + tmp_path: Path, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + plan = relocate_plan(multi_node_plan, workspace) + deployments = tuple(resolve_vllm_server(plan, item.deployment_id) for item in plan.deployments) + layout = AllocationLayout(("compute-001", "compute-002", "compute-003")) + + endpoint_steps = build_endpoint_steps( + deployments, + plan, + workspace / "attempt", + {}, + workspace / "runtime/data_designer/slurm/runtime/proxy.py", + layout, + ) + + first_command = endpoint_steps[0][0].command + second_command = endpoint_steps[1][0].command + assert "http://compute-001:18000" in first_command + assert "--allowed-host" in first_command + assert "http://compute-003:18007" in second_command + assert all("--nodelist=compute-001" in step.command for step, _ in endpoint_steps) + + def test_client_receives_client_secrets_without_server_only_secrets( tmp_path: Path, multi_node_plan: ResolvedSlurmRunPlan, From b8eb0dcaa631b8a0205dc689a7881a6715118a79 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 3 Sep 2026 15:03:49 -0600 Subject: [PATCH 2/2] fix Slurm queue sampler recovery Signed-off-by: Nabin Mulepati --- .../slurm/runtime/backpressure.py | 7 ++-- .../tests/runtime/test_backpressure.py | 36 +++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py index 68490d34c..2d7a3acb7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py @@ -91,8 +91,11 @@ def __init__( self._thread: threading.Thread | None = None def sample_once(self) -> QueueSnapshot: - """Refresh the cached queue depth once.""" - depth = self._reader() + """Refresh the cached queue depth, failing open on reader errors.""" + try: + depth = self._reader() + except Exception: + depth = None if depth is not None: depth = max(0, depth) snapshot = QueueSnapshot(depth, time.monotonic()) diff --git a/packages/data-designer-slurm/tests/runtime/test_backpressure.py b/packages/data-designer-slurm/tests/runtime/test_backpressure.py index 9f2a80c72..4478a3690 100644 --- a/packages/data-designer-slurm/tests/runtime/test_backpressure.py +++ b/packages/data-designer-slurm/tests/runtime/test_backpressure.py @@ -80,3 +80,39 @@ async def send(message: AsgiMessage) -> None: await middleware({"type": "http", "path": "/v1/completions"}, receive, send) assert called + + +def test_sampler_fails_open_after_reader_error_and_recovers() -> None: + readings: list[Exception | int] = [RuntimeError("metrics unavailable"), 4] + + def read_queue_depth() -> int: + reading = readings.pop(0) + if isinstance(reading, Exception): + raise reading + return reading + + controller = QueueBackpressureController( + QueueBackpressureSettings(2, 3), + reader=read_queue_depth, + start_sampler=False, + ) + + failed_snapshot = controller.sample_once() + failed_reject, _ = controller.should_reject() + recovered_snapshot = controller.sample_once() + recovered_reject, _ = controller.should_reject() + + assert failed_snapshot.depth is None + assert not failed_reject + assert recovered_snapshot.depth == 4 + assert recovered_reject + + +def test_sampler_preserves_process_control_exceptions() -> None: + def interrupt() -> int: + raise KeyboardInterrupt + + controller = QueueBackpressureController(reader=interrupt, start_sampler=False) + + with pytest.raises(KeyboardInterrupt): + controller.sample_once()