diff --git a/fern/versions/latest/pages/slurm/overview.mdx b/fern/versions/latest/pages/slurm/overview.mdx index eac729fac..89d797cef 100644 --- a/fern/versions/latest/pages/slurm/overview.mdx +++ b/fern/versions/latest/pages/slurm/overview.mdx @@ -12,7 +12,7 @@ The submit environment needs: - `sbatch`, `squeue`, `sacct`, `scancel`, `scontrol`, and `sinfo` - Access to a shared workspace that is visible from the submit host and every eligible compute node -Compute nodes need access to the same shared Python environment at the exact path used by the submit command. Local serving requires Enroot 3.5 or later for both digest-pinned OCI imports and existing SQSH inspection, and serving images must contain vLLM. Hosted-only runs require access to their external endpoints and do not need serving images or Enroot. +Compute nodes need `srun`, `scancel`, and `scontrol`, plus access to the same shared Python environment at the exact path used by the submit command. Local serving requires Enroot 3.5 or later for both digest-pinned OCI imports and existing SQSH inspection, and serving images must contain vLLM. Hosted-only runs require access to their external endpoints and do not need serving images or Enroot. Your Slurm administrator can confirm the account, GPU partition, CPU image-build partition, GPU request mode, and shared filesystem paths for your cluster. diff --git a/fern/versions/latest/pages/slurm/troubleshooting.mdx b/fern/versions/latest/pages/slurm/troubleshooting.mdx index e1a7476ca..14570ad99 100644 --- a/fern/versions/latest/pages/slurm/troubleshooting.mdx +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -81,6 +81,20 @@ Inspect the attempt readiness record and Slurm stderr under the managed run dire ## A shard failed +### Backend health is lost during generation + +After startup, the allocation checks each logical endpoint's aggregate backend health every second. Health probes use separate connection capacity so they can run while inference connections are saturated. If a required backend becomes unhealthy, including while its process remains alive, the monitor records `runtime_health_lost` in the attempt readiness evidence and fails the attempt. Inspect the backend logs before retrying; a healthy spare replica does not make the deployment ready when another required replica has failed. + +### Inspect routing and overload behavior + +Endpoint stderr contains structured `slurm_proxy_request` events. Each event identifies attempted upstream origins, outcomes, connection reuse, connection duration, time to response headers, total attempt duration, final HTTP status, and whether the response completed. Total attempt duration includes connection establishment and response streaming. A reused connection has a connection duration of zero; unavailable timings are `null`. Logs exclude request paths, authorization headers, prompts, and generated content. + +The allocation-local endpoint's `/metrics` response includes request and attempt counts, per-backend activity, connection opens/reuse/failures/evictions, retries by reason, final HTTP status counts, and timeouts. Evictions include expired, stale, and discarded connections; normal session shutdown is excluded. Timeouts include terminal attempts and interrupted response streams, even when no retry follows. These counters describe this endpoint process and reset when it restarts. + +During cleanup, health monitoring stops first. The controller sends `TERM` directly to each endpoint's Slurm task step using its recorded job and step IDs. Logical endpoints then stop accepting requests and drain active responses for up to 10 seconds while their `srun` launchers and serving backends remain available. The controller allows 12 seconds for draining and step exit before forcing remaining endpoints to stop and terminating other steps. Responses exceeding the drain deadline can be interrupted. + +### Retry a failed shard + Use `status` to find the failed task and attempt. Batch stderr identifies the allocation runtime phase and exit status; synchronous step failures also report their durable stdout and stderr paths. Preview a sparse retry before submission: ```bash diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh index c5e43336a..e58072a29 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh @@ -3,6 +3,9 @@ DD_MANAGED_PIDS=() DD_REQUIRED_PIDS=() +DD_ENDPOINT_PIDS=() +declare -A DD_ENDPOINT_STEP_FILES=() +DD_HEALTH_PID= DD_CLEANUP_COMPLETE=0 dd_start_runtime_timer() { @@ -44,25 +47,66 @@ dd_wait_for_client() { dd_cleanup_steps() { ((DD_CLEANUP_COMPLETE == 0)) || return 0 - local index pid position - local -a indices=() - ((${#DD_MANAGED_PIDS[@]} == 0)) || indices=("${!DD_MANAGED_PIDS[@]}") - for ((position = ${#indices[@]} - 1; position >= 0; position--)); do - index=${indices[position]} - pid=${DD_MANAGED_PIDS[index]} + if [[ -n ${DD_HEALTH_PID} ]]; then + dd_stop_steps 1 "${DD_HEALTH_PID}" + fi + dd_drain_endpoints + dd_stop_steps 1 "${DD_MANAGED_PIDS[@]+"${DD_MANAGED_PIDS[@]}"}" + DD_CLEANUP_COMPLETE=1 +} + +dd_drain_endpoints() { + local pid identity_file step_identity + for pid in "${DD_ENDPOINT_PIDS[@]+"${DD_ENDPOINT_PIDS[@]}"}"; do + kill -0 "${pid}" 2>/dev/null || continue + identity_file=${DD_ENDPOINT_STEP_FILES[${pid}]:-} + step_identity= + if [[ -f ${identity_file} && ! -L ${identity_file} ]]; then + IFS= read -r step_identity < "${identity_file}" || true + fi + if [[ ${step_identity} =~ ^([1-9][0-9]*)\.([0-9]+)$ && ${BASH_REMATCH[1]} == "${SLURM_JOB_ID:-}" ]]; then + # TERM to srun becomes KILL; signal only this endpoint's task step. + scancel --quiet --signal=TERM "${step_identity}" || true + fi + done + # Keep launchers and backends alive through the proxy's 10-second drain. + dd_wait_steps 12 "${DD_ENDPOINT_PIDS[@]+"${DD_ENDPOINT_PIDS[@]}"}" + dd_stop_steps 1 "${DD_ENDPOINT_PIDS[@]+"${DD_ENDPOINT_PIDS[@]}"}" +} + +dd_stop_steps() { + local grace_seconds=$1 + shift + (($# > 0)) || return 0 + local pid + for pid in "$@"; do kill -0 "${pid}" 2>/dev/null && kill -TERM "${pid}" 2>/dev/null || true done - dd_sleep 0.2 - for ((position = ${#indices[@]} - 1; position >= 0; position--)); do - index=${indices[position]} - pid=${DD_MANAGED_PIDS[index]} + dd_wait_steps "${grace_seconds}" "$@" + for pid in "$@"; do kill -0 "${pid}" 2>/dev/null && kill -KILL "${pid}" 2>/dev/null || true done - for ((position = ${#indices[@]} - 1; position >= 0; position--)); do - index=${indices[position]} - wait "${DD_MANAGED_PIDS[index]}" 2>/dev/null || true + for pid in "$@"; do + wait "${pid}" 2>/dev/null || true + done +} + +dd_wait_steps() { + local grace_seconds=$1 + shift + (($# > 0)) || return 0 + local pid pending deadline=$((SECONDS + grace_seconds)) + while ((SECONDS < deadline)); do + pending=0 + for pid in "$@"; do + if kill -0 "${pid}" 2>/dev/null; then + pending=1 + break + fi + done + ((pending != 0)) || break + dd_sleep 0.1 done - DD_CLEANUP_COMPLETE=1 } dd_stop_runtime_timer() { 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 70df7feaf..e050641aa 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 @@ -17,6 +17,7 @@ from data_designer.slurm.runtime.bootstrap import build_runtime_manifest from data_designer.slurm.runtime.context import load_allocation_context from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.health import monitor_runtime_health from data_designer.slurm.runtime.logs import execution_log_directory from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.ports import resolve_allocation_deployments @@ -52,6 +53,8 @@ def main(arguments: Sequence[str] | None = None) -> int: _ready(parsed, os.environ) elif parsed.operation == "client": _client(parsed, os.environ) + elif parsed.operation == "monitor": + _monitor(parsed, os.environ) elif parsed.operation == "succeed": _succeed(parsed, os.environ) else: @@ -73,7 +76,7 @@ def main(arguments: Sequence[str] | None = None) -> int: def _parse_arguments(arguments: Sequence[str] | None) -> argparse.Namespace: parser = argparse.ArgumentParser(prog="data-designer-slurm-runtime") subparsers = parser.add_subparsers(dest="operation", required=True) - for operation in ("prepare", "ready", "succeed", "fail"): + for operation in ("prepare", "ready", "succeed", "fail", "monitor"): subparser = subparsers.add_parser(operation) _add_context_arguments(subparser) prepare = subparsers.choices["prepare"] @@ -150,6 +153,17 @@ def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> Non ) +def _monitor(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = _load_context(arguments, environment) + readiness = writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) + deployments = _resolve_deployments(context, environment) + monitor_runtime_health( + writer, + readiness, + {deployment.deployment_id: deployment.logical_endpoint.port for deployment in deployments}, + ) + + def _client( arguments: argparse.Namespace, environment: Mapping[str, str], @@ -215,7 +229,18 @@ def _client( def _succeed(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: context, writer = _load_context(arguments, environment) - stopped_at = _write_stopped_readiness(context, writer) + # Cleanup has stopped the monitor, so its last persisted health evidence is final. + readiness = writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) + if readiness.state is ReadinessState.FAILED or any( + deployment.state is ReadinessState.FAILED + or (deployment.last_probe is not None and deployment.last_probe.outcome is ProbeOutcome.FAILURE) + for deployment in readiness.deployments + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "cannot finalize a successful allocation after required inference service failure", + ) + stopped_at = _write_stopped_readiness(context, writer, previous=readiness) attempt = writer.load_attempt(context.shard.shard_id, context.attempt.attempt_id) if attempt.candidate_output is None: raise SlurmRuntimeError( @@ -316,6 +341,10 @@ def _write_failed_and_stopped_readiness(context: AllocationContext, writer: Slur previous = _load_optional_readiness(context, writer) if previous is None: return max(datetime.now(UTC), context.attempt.updated_at) + if previous.state is ReadinessState.STOPPED: + return _now(context.attempt, previous) + if previous.state is ReadinessState.FAILED: + return _write_stopped_readiness(context, writer, previous=previous) failed_at = _now(context.attempt, previous) failed = writer.write_readiness( AttemptReadiness( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh index f1598795b..31e25c701 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -96,6 +96,9 @@ dd_slurm_run_allocation() { dd_require_running DD_RUNTIME_PHASE=ready dd_run_bound_control_phase ready + if ((DD_EXPECTED_GPUS > 0)); then + dd_start_health_monitor + fi DD_RUNTIME_PHASE=client dd_runtime_log 'Starting Data Designer client...' @@ -133,7 +136,7 @@ dd_run_bound_control_phase() { dd_verify_host_context() { local tool - for tool in bash sha256sum tar jq srun scontrol getent curl find; do + for tool in bash sha256sum tar jq srun scancel scontrol getent curl find; do command -v "${tool}" >/dev/null || { printf 'required allocation tool %q is unavailable\n' "${tool}" >&2 return 69 @@ -234,9 +237,11 @@ dd_start_endpoints() { for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do DD_STEP_STARTED_AT[${step_id}]=${SECONDS} dd_read_step "${DD_RUNTIME_MANIFEST}" "${step_id}" - rm -f -- "${DD_STEP_STDOUT}.listening" - dd_start_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + rm -f -- "${DD_STEP_STDOUT}.listening" "${DD_STEP_STDOUT}.step-id" + dd_start_step "${DD_RUNTIME_MANIFEST}" "${step_id}" "${DD_STEP_STDOUT}.step-id" dd_register_required_pid "${DD_LAST_PID}" + DD_ENDPOINT_PIDS+=("${DD_LAST_PID}") + DD_ENDPOINT_STEP_FILES[${DD_LAST_PID}]=${DD_STEP_STDOUT}.step-id dd_require_running done } diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/health.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/health.py new file mode 100644 index 000000000..b73264c6f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/health.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Continuous aggregate health checks for a running allocation.""" + +from __future__ import annotations + +import http.client +import json +import time +from collections.abc import Mapping +from datetime import UTC, datetime + +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.state import ( + AttemptReadiness, + DeploymentReadiness, + ProbeEvidence, + ProbeOutcome, + ReadinessState, + SlurmStateWriter, +) + +_HEALTH_TIMEOUT_SECONDS = 6.0 +_MAXIMUM_HEALTH_BYTES = 4096 + + +def monitor_runtime_health( + writer: SlurmStateWriter, + readiness: AttemptReadiness, + endpoint_ports: Mapping[str, int], + *, + poll_interval_seconds: float = 1.0, +) -> None: + """Publish health evidence until readiness stops or a required endpoint fails. + + The allocation controller owns process termination and attempt finalization. + This monitor fails on service loss even if all serving processes remain alive. + """ + if poll_interval_seconds <= 0: + raise ValueError("health poll interval must be positive") + if set(endpoint_ports) != {deployment.deployment_id for deployment in readiness.deployments}: + raise ValueError("health endpoints must match readiness deployments") + while readiness.state is ReadinessState.READY: + readiness = writer.write_readiness(probe_runtime_health(readiness, endpoint_ports)) + if readiness.state is ReadinessState.FAILED: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.STEP_FAILED, "required inference service lost health") + time.sleep(poll_interval_seconds) + readiness = writer.load_readiness(readiness.shard_id, readiness.attempt_id) + + +def probe_runtime_health(readiness: AttemptReadiness, endpoint_ports: Mapping[str, int]) -> AttemptReadiness: + """Check every logical endpoint and preserve per-deployment backend counts.""" + deployments: list[DeploymentReadiness] = [] + for deployment in readiness.deployments: + count = _ready_backend_count(endpoint_ports[deployment.deployment_id], deployment.expected_backends) + healthy = count == deployment.expected_backends + observed_at = datetime.now(UTC) + deployments.append( + deployment.model_copy( + update={ + "state": ReadinessState.READY if healthy else ReadinessState.FAILED, + "ready_backends": count, + "last_probe": ProbeEvidence( + observed_at=observed_at, + outcome=ProbeOutcome.SUCCESS if healthy else ProbeOutcome.FAILURE, + reason_code="runtime_healthy" if healthy else "runtime_health_lost", + redacted_message=f"{count} of {deployment.expected_backends} backends ready", + ), + } + ) + ) + return readiness.model_copy( + update={ + "revision": readiness.revision + 1, + "updated_at": datetime.now(UTC), + "state": ( + ReadinessState.READY + if all(deployment.state is ReadinessState.READY for deployment in deployments) + else ReadinessState.FAILED + ), + "deployments": tuple(deployments), + } + ) + + +def _ready_backend_count(port: int, expected: int) -> int: + connection = http.client.HTTPConnection("127.0.0.1", port, timeout=_HEALTH_TIMEOUT_SECONDS) + try: + connection.request("GET", "/health", headers={"Connection": "close"}) + response = connection.getresponse() + body = response.read(_MAXIMUM_HEALTH_BYTES + 1) + if response.status not in {200, 503} or len(body) > _MAXIMUM_HEALTH_BYTES: + return 0 + payload = json.loads(body) + if not isinstance(payload, dict) or type(payload.get("backends")) is not int or payload["backends"] != expected: + return 0 + count = payload.get("backends_ready") + if type(count) is not int or not 0 <= count <= expected: + return 0 + return count if response.status == 200 or count < expected else 0 + except (OSError, http.client.HTTPException, ValueError): + return 0 + finally: + connection.close() 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 ffdee6a64..677a4120a 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 @@ -15,10 +15,18 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path +from types import SimpleNamespace +from typing import TYPE_CHECKING from urllib.parse import SplitResult, urlsplit from aiohttp import ClientError, ClientResponse, ClientSession, ClientTimeout, TCPConnector, TraceConfig, web +if TYPE_CHECKING: + from aiohttp.client_proto import ResponseHandler + from aiohttp.client_reqrep import ConnectionKey + from aiohttp.connector import Connection + from aiohttp.tracing import Trace + from data_designer.slurm.runtime.network import validate_host_name _LOGGER = logging.getLogger("data_designer.slurm.runtime.proxy") @@ -81,9 +89,6 @@ def acquire(self, excluded: frozenset[int]) -> int: self._active[index] += 1 return index - def acquire_index(self, index: int) -> None: - self._active[index] += 1 - def release(self, index: int) -> None: if self._active[index] <= 0: raise RuntimeError("backend activity underflow") @@ -102,6 +107,8 @@ def __init__(self) -> None: self.connection_opened = 0 self.connection_reused = 0 self.connection_failed = 0 + self.connection_evicted: int = 0 + self.timeouts: int = 0 self.retries: Counter[str] = Counter() self.final_statuses: Counter[int] = Counter() self.backend_requests: Counter[int] = Counter() @@ -116,9 +123,11 @@ def snapshot(self, pool: _BackendPool, connector: TCPConnector) -> dict[str, obj "opened": self.connection_opened, "reused": self.connection_reused, "failed": self.connection_failed, + "evicted": self.connection_evicted, "idle": idle_connections, }, "retries": dict(sorted(self.retries.items())), + "timeouts": self.timeouts, "final_statuses": {str(key): value for key, value in sorted(self.final_statuses.items())}, "backends": [ { @@ -131,6 +140,73 @@ def snapshot(self, pool: _BackendPool, connector: TCPConnector) -> dict[str, obj } +class _ProxyConnector(TCPConnector): + """Observe pool evictions, for which aiohttp has no public trace signal. + + Keep these hooks limited to the supported aiohttp 3.x connector lifecycle. + Shutdown closes are excluded; stale, expired, and discarded connections count. + """ + + def __init__(self, metrics: _ProxyMetrics, backend_count: int) -> None: + self._metrics: _ProxyMetrics = metrics + super().__init__( + limit=_MAXIMUM_CONNECTIONS, + limit_per_host=max(1, _MAXIMUM_CONNECTIONS // backend_count), + keepalive_timeout=_KEEPALIVE_TIMEOUT_SECONDS, + force_close=False, + socket_factory=_create_socket, + ) + + def _cleanup(self) -> None: + before = sum(len(connections) for connections in self._conns.values()) + super()._cleanup() + self._metrics.connection_evicted += before - sum(len(connections) for connections in self._conns.values()) + + async def _get(self, key: ConnectionKey, traces: list[Trace]) -> Connection | None: + before = len(self._conns.get(key, ())) + connection = await super()._get(key, traces) + self._metrics.connection_evicted += max(0, before - len(self._conns.get(key, ())) - int(connection is not None)) + return connection + + def _release(self, key: ConnectionKey, protocol: ResponseHandler, *, should_close: bool = False) -> None: + if not self._closed and (should_close or protocol.should_close): + self._metrics.connection_evicted += 1 + super()._release(key, protocol, should_close=should_close) + + +@dataclass(slots=True) +class _AttemptTrace: + backend_index: int + upstream: str + started_at: float + outcome: str = "cancelled" + connection_started_at: float | None = None + connection_seconds: float | None = None + response_headers_seconds: float | None = None + response_seconds: float = 0.0 + reused: bool = False + + def finish(self, outcome: str) -> None: + now = time.monotonic() + self.outcome = outcome + self.response_seconds = now - self.started_at + if self.connection_seconds is None and self.connection_started_at is not None: + self.connection_seconds = now - self.connection_started_at + + def snapshot(self) -> dict[str, object]: + return { + "backend_index": self.backend_index, + "upstream": self.upstream, + "outcome": self.outcome, + "connection_reused": self.reused, + "connection_seconds": None if self.connection_seconds is None else round(self.connection_seconds, 6), + "response_headers_seconds": ( + None if self.response_headers_seconds is None else round(self.response_headers_seconds, 6) + ), + "response_seconds": round(self.response_seconds, 6), + } + + class _ProxyApplication: """One logical endpoint with a shared upstream connection pool.""" @@ -146,6 +222,7 @@ def __init__( self.metrics = _ProxyMetrics() self.connector: TCPConnector | None = None self.session: ClientSession | None = None + self.health_session: ClientSession | None = None def create(self) -> web.Application: application = web.Application(client_max_size=_MAXIMUM_REQUEST_BYTES) @@ -159,14 +236,7 @@ def create(self) -> web.Application: async def _start(self, application: web.Application) -> None: del application trace = self._trace_config() - per_backend_limit = max(1, _MAXIMUM_CONNECTIONS // len(self.pool.backends)) - self.connector = TCPConnector( - limit=_MAXIMUM_CONNECTIONS, - limit_per_host=per_backend_limit, - keepalive_timeout=_KEEPALIVE_TIMEOUT_SECONDS, - force_close=False, - socket_factory=_create_socket, - ) + self.connector = _ProxyConnector(self.metrics, len(self.pool.backends)) self.session = ClientSession( connector=self.connector, timeout=ClientTimeout( @@ -178,21 +248,43 @@ async def _start(self, application: web.Application) -> None: raise_for_status=False, trace_configs=[trace], ) + # Health probes must not wait behind saturated inference connection slots. + self.health_session = ClientSession( + connector=TCPConnector(limit=len(self.pool.backends), limit_per_host=1, socket_factory=_create_socket), + timeout=ClientTimeout(total=_CONNECT_TIMEOUT_SECONDS), + auto_decompress=False, + ) async def _stop(self, application: web.Application) -> None: del application if self.session is not None: await self.session.close() + if self.health_session is not None: + await self.health_session.close() def _trace_config(self) -> TraceConfig: trace = TraceConfig() - async def connection_opened(*_: object) -> None: + async def connection_started(session: ClientSession, context: SimpleNamespace, params: object) -> None: + del session, params + if isinstance(context.trace_request_ctx, _AttemptTrace): + context.trace_request_ctx.connection_started_at = time.monotonic() + + async def connection_opened(session: ClientSession, context: SimpleNamespace, params: object) -> None: + del session, params self.metrics.connection_opened += 1 + if isinstance(attempt := context.trace_request_ctx, _AttemptTrace): + if attempt.connection_started_at is not None: + attempt.connection_seconds = time.monotonic() - attempt.connection_started_at - async def connection_reused(*_: object) -> None: + async def connection_reused(session: ClientSession, context: SimpleNamespace, params: object) -> None: + del session, params self.metrics.connection_reused += 1 + if isinstance(attempt := context.trace_request_ctx, _AttemptTrace): + attempt.reused = True + attempt.connection_seconds = 0.0 + trace.on_connection_create_start.append(connection_started) trace.on_connection_create_end.append(connection_opened) trace.on_connection_reuseconn.append(connection_reused) return trace @@ -213,20 +305,17 @@ async def _health(self, request: web.Request) -> web.Response: ) async def _probe_backend(self, index: int, backend: _Backend) -> bool: - session = self._require_session() - self.pool.acquire_index(index) + del index + if self.health_session is None: + raise RuntimeError("proxy health session is not started") try: - async with session.get( + async with self.health_session.get( f"{backend.origin}{self.health_path}", - timeout=ClientTimeout(total=_CONNECT_TIMEOUT_SECONDS), - trace_request_ctx={"backend_index": index}, ) as response: await response.read() return response.status == 200 except (TimeoutError, ClientError): return False - finally: - self.pool.release(index) async def _metrics(self, request: web.Request) -> web.Response: del request @@ -238,7 +327,7 @@ async def _forward(self, request: web.Request) -> web.StreamResponse: self.metrics.requests += 1 body = await request.read() excluded: set[int] = set() - outcomes: list[str] = [] + outcomes: list[_AttemptTrace] = [] attempt_limit = min(_MAXIMUM_UPSTREAM_ATTEMPTS, len(self.pool.backends)) final_response: ClientResponse | None = None final_index: int | None = None @@ -250,12 +339,17 @@ async def _forward(self, request: web.Request) -> web.StreamResponse: backend = self.pool.backends[index] self.metrics.attempts += 1 self.metrics.backend_requests[index] += 1 + attempt_trace = _AttemptTrace(index, backend.origin, time.monotonic()) + outcomes.append(attempt_trace) try: - response = await self._request_backend(request, backend, body, index) + response = await self._request_backend(request, backend, body, attempt_trace) + attempt_trace.response_headers_seconds = time.monotonic() - attempt_trace.started_at except (TimeoutError, ClientError) as error: self.metrics.connection_failed += 1 reason = "timeout" if isinstance(error, asyncio.TimeoutError) else "connection_error" - outcomes.append(reason) + attempt_trace.finish(reason) + if reason == "timeout": + self.metrics.timeouts += 1 self.pool.release(index) excluded.add(index) if attempt + 1 < attempt_limit: @@ -275,10 +369,12 @@ async def _forward(self, request: web.Request) -> web.StreamResponse: self._record_final(502, outcomes, started_at) return result except BaseException: + attempt_trace.finish("cancelled") + self._record_final(499, outcomes, started_at, response_completed=False) self.pool.release(index) raise - outcomes.append(str(response.status)) + attempt_trace.finish(str(response.status)) should_retry = response.status in _RETRYABLE_STATUS_CODES and attempt + 1 < attempt_limit if should_retry: reason = f"http_{response.status}" @@ -300,8 +396,19 @@ async def _forward(self, request: web.Request) -> web.StreamResponse: raise RuntimeError("proxy attempt loop produced no final response") try: result = await self._stream_response(request, final_response) + attempt_trace.finish(str(final_response.status)) self._record_final(final_response.status, outcomes, started_at) return result + except (TimeoutError, ClientError, ConnectionError, asyncio.CancelledError) as error: + reason = "stream_error" + if isinstance(error, TimeoutError): + reason = "stream_timeout" + self.metrics.timeouts += 1 + elif isinstance(error, asyncio.CancelledError): + reason = "cancelled" + attempt_trace.finish(reason) + self._record_final(final_response.status, outcomes, started_at, response_completed=False) + raise finally: final_response.release() self.pool.release(final_index) @@ -311,7 +418,7 @@ async def _request_backend( request: web.Request, backend: _Backend, body: bytes, - index: int, + attempt_trace: _AttemptTrace, ) -> ClientResponse: session = self._require_session() headers = _filter_request_headers(request.headers) @@ -321,7 +428,7 @@ async def _request_backend( data=body, headers=headers, allow_redirects=False, - trace_request_ctx={"backend_index": index}, + trace_request_ctx=attempt_trace, ) async def _stream_response(self, request: web.Request, response: ClientResponse) -> web.StreamResponse: @@ -337,7 +444,14 @@ async def _stream_response(self, request: web.Request, response: ClientResponse) await downstream.write_eof() return downstream - def _record_final(self, status: int, outcomes: list[str], started_at: float) -> None: + def _record_final( + self, + status: int, + outcomes: list[_AttemptTrace], + started_at: float, + *, + response_completed: bool = True, + ) -> None: self.metrics.final_statuses[status] += 1 _LOGGER.info( "slurm_proxy_request %s", @@ -345,7 +459,9 @@ def _record_final(self, status: int, outcomes: list[str], started_at: float) -> { "attempts": len(outcomes), "final_status": status, - "upstream_outcomes": outcomes, + "upstream_outcomes": [attempt.outcome for attempt in outcomes], + "upstream_attempts": [attempt.snapshot() for attempt in outcomes], + "response_completed": response_completed, "elapsed_seconds": round(time.monotonic() - started_at, 6), "failover": len(outcomes) > 1, }, @@ -395,6 +511,7 @@ def _create_socket( def main(arguments: Sequence[str] | None = None) -> int: """Serve one loopback logical endpoint until the step is terminated.""" + logging.basicConfig(level=logging.INFO, format="%(message)s") 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) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh index c4ecbaaec..e8a73e65c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh @@ -127,8 +127,22 @@ dd_build_srun_command() { dd_start_step() { local manifest=$1 local step_id=$2 + local task_identity_file=${3:-} dd_read_step "${manifest}" "${step_id}" dd_build_srun_command + if [[ -n ${task_identity_file} ]]; then + [[ ${DD_STEP_EXECUTION} == native ]] || return 64 + DD_STEP_COMMAND=( + bash -c ' + set -Eeuo pipefail + [[ ${SLURM_JOB_ID} =~ ^[1-9][0-9]*$ && ${SLURM_STEP_ID} =~ ^[0-9]+$ ]] + umask 077 + printf "%s.%s\n" "${SLURM_JOB_ID}" "${SLURM_STEP_ID}" > "$1" + shift + exec "$@" + ' data-designer-endpoint "${task_identity_file}" "${DD_STEP_COMMAND[@]}" + ) + fi ( dd_materialize_step_environment exec "${DD_SRUN_COMMAND[@]}" -- "${DD_STEP_COMMAND[@]}" @@ -155,9 +169,27 @@ dd_run_step() { } dd_run_control_phase() { + dd_build_control_command "$@" + ( + dd_scope_control_environment + exec "${DD_CONTROL_COMMAND[@]}" + ) +} + +dd_start_health_monitor() { + dd_build_control_command monitor "${DD_RUNTIME_RETRY_ARGUMENTS[@]+"${DD_RUNTIME_RETRY_ARGUMENTS[@]}"}" + ( + dd_scope_control_environment + exec "${DD_CONTROL_COMMAND[@]}" + ) & + DD_HEALTH_PID=$! + dd_register_required_pid "${DD_HEALTH_PID}" +} + +dd_build_control_command() { local operation=$1 shift - local -a command=( + DD_CONTROL_COMMAND=( srun --nodes=1 --ntasks=1 @@ -179,8 +211,4 @@ dd_run_control_phase() { "${DD_ATTEMPT_PATH}" "$@" ) - ( - dd_scope_control_environment - exec "${command[@]}" - ) } diff --git a/packages/data-designer-slurm/tests/runtime/test_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py index ad8361a77..1fa601003 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bundle.py +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -41,6 +41,7 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p assert names[6] == "data_designer/slurm/__init__.py" assert "data_designer/slurm/runtime/controller.py" in names assert "data_designer/slurm/runtime/entrypoint.py" in names + assert "data_designer/slurm/runtime/health.py" in names assert "data_designer/slurm/state/store.py" in names assert "data_designer/slurm/contracts.py" in names assert archive.getmember("entrypoint.sh").mode == 0o500 diff --git a/packages/data-designer-slurm/tests/runtime/test_cleanup.py b/packages/data-designer-slurm/tests/runtime/test_cleanup.py new file mode 100644 index 000000000..ac224819a --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_cleanup.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import shlex +import subprocess +import sys +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize("exit_status", (0, 1, 130), ids=("completion", "failure", "cancellation")) +def test_allocation_keeps_backend_alive_until_endpoint_drains(tmp_path: Path, exit_status: int) -> None: + runtime = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + events = tmp_path / "events" + worker = tmp_path / "worker.py" + scheduler_bin = _write_fake_scheduler(tmp_path) + worker.write_text( + "import os, signal, sys, time\n" + "from pathlib import Path\n" + "role, events, ready = sys.argv[1:4]\n" + "def stop(*args):\n" + " if role == 'endpoint':\n" + " time.sleep(0.4)\n" + " os.kill(int(sys.argv[4]), 0)\n" + " with open(events, 'a') as stream: stream.write(role + '\\n')\n" + " sys.exit(0)\n" + "signal.signal(signal.SIGTERM, stop)\n" + "Path(ready).touch()\n" + "while True: time.sleep(0.05)\n" + ) + command = f""" +set -Eeuo pipefail +source {shlex.quote(str(runtime / "entrypoint.sh"))} +dd_cleanup_allocation_scratch() {{ :; }} +trap dd_runtime_exit EXIT +{shlex.quote(sys.executable)} {shlex.quote(str(worker))} backend {shlex.quote(str(events))} {tmp_path}/backend.ready & +backend=$! +DD_MANAGED_PIDS+=("$backend") +while [[ ! -f {tmp_path}/backend.ready ]]; do sleep 0.01; done +SLURM_JOB_ID=4101 +DD_RUNTIME_MANIFEST=unused +DD_STEP_STDOUT={shlex.quote(str(tmp_path / "endpoint.out"))} +DD_STEP_STDERR={shlex.quote(str(tmp_path / "endpoint.err"))} +dd_read_step_ids() {{ DD_STEP_IDS=(endpoint); }} +dd_read_step() {{ + DD_STEP_EXECUTION=native + DD_STEP_COMMAND=({shlex.quote(sys.executable)} {shlex.quote(str(worker))} endpoint {shlex.quote(str(events))} {tmp_path}/endpoint.ready "$backend") +}} +dd_build_srun_command() {{ DD_SRUN_COMMAND=(srun); }} +dd_materialize_step_environment() {{ :; }} +dd_start_endpoints +while [[ ! -f {tmp_path}/endpoint.ready ]]; do sleep 0.01; done +exit {exit_status} +""" + result = subprocess.run( + ["bash", "-c", command], + capture_output=True, + text=True, + timeout=5, + env={**os.environ, "PATH": f"{scheduler_bin}:{os.environ['PATH']}"}, + ) + assert result.returncode == exit_status, result.stderr + assert events.read_text().splitlines() == ["task_term", "endpoint", "backend"] + assert (tmp_path / "endpoint.out.step-id").read_text() == "4101.7\n" + + +@pytest.mark.parametrize("identity", ("4101.7", "4102.7", "4101", "4101.batch", "4101.7 --all", None)) +def test_endpoint_signal_is_bound_to_a_numeric_step_in_this_job(tmp_path: Path, identity: str | None) -> None: + runtime = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + identity_file = tmp_path / "endpoint.step-id" + if identity is not None: + identity_file.write_text(identity + "\n") + command = f""" +set -Eeuo pipefail +source {shlex.quote(str(runtime / "cleanup.sh"))} +SLURM_JOB_ID=4101 +DD_ENDPOINT_PIDS=("$$") +DD_ENDPOINT_STEP_FILES[$$]={shlex.quote(str(identity_file))} +scancel() {{ printf '%s\\n' "$*" > {shlex.quote(str(tmp_path / "signal"))}; }} +dd_wait_steps() {{ :; }} +dd_stop_steps() {{ :; }} +dd_drain_endpoints +""" + result = subprocess.run(["bash", "-c", command], capture_output=True, text=True, timeout=3) + assert result.returncode == 0, result.stderr + signal = tmp_path / "signal" + assert signal.exists() is (identity == "4101.7") + if signal.exists(): + assert signal.read_text() == "--quiet --signal=TERM 4101.7\n" + + +def test_endpoint_cleanup_forces_launcher_only_after_drain_deadline(tmp_path: Path) -> None: + runtime = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote(str(runtime / "cleanup.sh"))} +(trap '' TERM; touch {tmp_path}/ready; while :; do sleep 0.1; done) & +pid=$! +while [[ ! -f {tmp_path}/ready ]]; do sleep 0.01; done +SLURM_JOB_ID=4101 +printf '4101.7\\n' > {tmp_path}/endpoint.step-id +DD_ENDPOINT_PIDS=("$pid") +DD_ENDPOINT_STEP_FILES[$pid]={tmp_path}/endpoint.step-id +scancel() {{ [[ $* == '--quiet --signal=TERM 4101.7' ]]; kill -TERM "$pid"; }} +dd_sleep() {{ kill -0 "$pid"; SECONDS=$((SECONDS + 1)); }} +started=$SECONDS +dd_drain_endpoints +((SECONDS - started >= 12)) +! kill -0 "$pid" 2>/dev/null +""" + result = subprocess.run(["bash", "-c", command], capture_output=True, text=True, timeout=3) + assert result.returncode == 0, result.stderr + + +def _write_fake_scheduler(directory: Path) -> Path: + """Model srun's TERM-to-KILL forwarding separately from task signaling.""" + scheduler_bin = directory / "bin" + scheduler_bin.mkdir() + launcher = scheduler_bin / "srun" + launcher.write_text( + f"#!{sys.executable}\n" + "import os, signal, subprocess, sys\n" + "from pathlib import Path\n" + "root = Path(__file__).resolve().parents[1]\n" + "child = subprocess.Popen(sys.argv[sys.argv.index('--') + 1:], " + "env={**os.environ, 'SLURM_JOB_ID': '4101', 'SLURM_STEP_ID': '7'})\n" + "def stop(*args):\n" + " with (root / 'events').open('a') as stream: stream.write('launcher_term\\n')\n" + " child.kill()\n" + "signal.signal(signal.SIGTERM, stop)\n" + "(root / 'task.pid').write_text(str(child.pid))\n" + "status = child.wait()\n" + "sys.exit(status if status >= 0 else 128 - status)\n" + ) + cancel = scheduler_bin / "scancel" + cancel.write_text( + f"#!{sys.executable}\n" + "import os, signal, sys\n" + "from pathlib import Path\n" + "root = Path(__file__).resolve().parents[1]\n" + "assert sys.argv[1:] == ['--quiet', '--signal=TERM', '4101.7']\n" + "with (root / 'events').open('a') as stream: stream.write('task_term\\n')\n" + "os.kill(int((root / 'task.pid').read_text()), signal.SIGTERM)\n" + ) + for path in (launcher, cancel): + path.chmod(0o700) + return scheduler_bin + + +def test_cleanup_kills_unresponsive_step_after_deadline(tmp_path: Path) -> None: + runtime = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote(str(runtime / "cleanup.sh"))} +(trap '' TERM; touch {tmp_path}/ready; while :; do sleep 0.1; done) & +pid=$! +while [[ ! -f {tmp_path}/ready ]]; do sleep 0.01; done +dd_stop_steps 1 "$pid" +! kill -0 "$pid" 2>/dev/null +""" + result = subprocess.run(["bash", "-c", command], capture_output=True, text=True, timeout=3) + assert result.returncode == 0, result.stderr diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 3fc08f990..adea56570 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -293,6 +293,10 @@ def finalize_winner(self, *args: object, **kwargs: object) -> None: ) assert entrypoint.main(_phase_arguments("succeed", runtime_case)) == 70 assert state.attempt.state is AttemptLifecycleState.RUNNING + assert state.readiness[-1].state is ReadinessState.STOPPED + assert entrypoint.main(_phase_arguments("fail", runtime_case)) == 0 + assert state.attempt.state is AttemptLifecycleState.FAILED + assert state.readiness[-1].state is ReadinessState.STOPPED def _patch_runtime_context( diff --git a/packages/data-designer-slurm/tests/runtime/test_health.py b/packages/data-designer-slurm/tests/runtime/test_health.py new file mode 100644 index 000000000..99f15edb6 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_health.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import threading +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import UTC, datetime +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest +from conftest import FakeStateStore, RuntimeCase + +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.runtime import entrypoint +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.health import monitor_runtime_health, probe_runtime_health +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptReadiness, + DeploymentReadiness, + EndpointPublicationState, + ReadinessState, +) + + +@pytest.fixture +def healthy_readiness(runtime_case: RuntimeCase) -> AttemptReadiness: + attempt = runtime_case.context.attempt + now = datetime.now(UTC) + return AttemptReadiness( + schema_version=1, + run_id=attempt.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + revision=1, + updated_at=now, + state=ReadinessState.READY, + deployments=( + DeploymentReadiness( + deployment_id="deployment-00000", + model_alias="generator", + state=ReadinessState.READY, + expected_backends=2, + ready_backends=2, + endpoint_publication=EndpointPublicationState.PUBLISHED, + ), + ), + ) + + +@contextmanager +def _health_endpoint(responses: list[tuple[int, bytes]]) -> Iterator[int]: + class Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + assert self.path == "/health" + status, body = responses[0] if len(responses) == 1 else responses.pop(0) + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server.server_port + finally: + server.shutdown() + server.server_close() + thread.join(timeout=2) + + +def test_monitor_detects_health_loss_without_process_exit_and_preserves_failure_evidence( + runtime_case: RuntimeCase, + healthy_readiness: AttemptReadiness, + monkeypatch: pytest.MonkeyPatch, +) -> None: + state = FakeStateStore(runtime_case.context.attempt) + state.write_readiness(healthy_readiness) + responses = [ + (200, b'{"backends":2,"backends_ready":2}'), + (503, b'{"backends":2,"backends_ready":1}'), + ] + with _health_endpoint(responses) as port: + with pytest.raises(SlurmRuntimeError, match="lost health"): + monitor_runtime_health(state, healthy_readiness, {"deployment-00000": port}, poll_interval_seconds=0.01) + assert state.readiness[-2].state is ReadinessState.READY + assert state.readiness[-1].state is ReadinessState.FAILED + assert state.readiness[-1].deployments[0].ready_backends == 1 + # The endpoint is still alive and responsive after the health transition. + assert probe_runtime_health(healthy_readiness, {"deployment-00000": port}).state is ReadinessState.FAILED + + monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args, **kwargs: (runtime_case.context, state)) + assert ( + entrypoint.main( + [ + "fail", + "--plan", + str(runtime_case.workspace / "plan.json"), + "--attempt-dir", + str(runtime_case.context.attempt_directory), + ] + ) + == 0 + ) + assert state.attempt.state is AttemptLifecycleState.FAILED + assert state.readiness[-1].state is ReadinessState.STOPPED + assert state.readiness[-1].deployments[0].last_probe.reason_code == "runtime_health_lost" + + +@pytest.mark.parametrize( + ("status", "body"), + [ + (500, b"internal"), + (200, b"not json"), + (200, b"[]"), + (200, b'{"backends":2,"backends_ready":true}'), + (200, b'{"backends":2,"backends_ready":3}'), + (200, b'{"backends":3,"backends_ready":2}'), + (503, b'{"backends":2,"backends_ready":2}'), + (200, b" " * 4097), + ], +) +def test_invalid_health_response_fails_closed(healthy_readiness: AttemptReadiness, status: int, body: bytes) -> None: + with _health_endpoint([(status, body)]) as port: + result = probe_runtime_health(healthy_readiness, {"deployment-00000": port}) + assert result.state is ReadinessState.FAILED + assert result.deployments[0].ready_backends == 0 + assert body.decode() not in result.model_dump_json() + + +def test_health_detects_endpoint_connection_loss(healthy_readiness: AttemptReadiness) -> None: + with _health_endpoint([(200, b'{"backends":2,"backends_ready":2}')]) as port: + assert probe_runtime_health(healthy_readiness, {"deployment-00000": port}).state is ReadinessState.READY + result = probe_runtime_health(healthy_readiness, {"deployment-00000": port}) + assert result.state is ReadinessState.FAILED + assert result.deployments[0].ready_backends == 0 + + +@pytest.mark.parametrize("already_stopped", (False, True)) +def test_success_rejects_health_failure_recorded_during_cleanup( + runtime_case: RuntimeCase, + healthy_readiness: AttemptReadiness, + monkeypatch: pytest.MonkeyPatch, + already_stopped: bool, +) -> None: + state = FakeStateStore( + runtime_case.context.attempt.model_copy( + update={ + "state": AttemptLifecycleState.RUNNING, + "candidate_output": ArtifactReference( + path=(runtime_case.context.attempt_directory / "output-manifest.json").as_posix(), + sha256="a" * 64, + ), + } + ) + ) + state.write_readiness(healthy_readiness) + with _health_endpoint([(503, b'{"backends":2,"backends_ready":1}')]) as port: + state.write_readiness(probe_runtime_health(healthy_readiness, {"deployment-00000": port})) + monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args, **kwargs: (runtime_case.context, state)) + if already_stopped: + entrypoint._write_stopped_readiness(runtime_case.context, state) + failure = state.readiness[-1] + arguments = [ + "--plan", + str(runtime_case.workspace / "plan.json"), + "--attempt-dir", + str(runtime_case.context.attempt_directory), + ] + assert entrypoint.main(["succeed", *arguments]) == 70 + assert not state.winners + assert state.attempt.state is AttemptLifecycleState.RUNNING + assert state.readiness[-1] is failure + assert entrypoint.main(["fail", *arguments]) == 0 + assert state.attempt.state is AttemptLifecycleState.FAILED + assert state.readiness[-1].deployments[0].last_probe.reason_code == "runtime_health_lost" + + +@pytest.mark.parametrize("invalid_interval", (False, True)) +def test_monitor_rejects_invalid_settings( + runtime_case: RuntimeCase, healthy_readiness: AttemptReadiness, invalid_interval: bool +) -> None: + state = FakeStateStore(runtime_case.context.attempt) + endpoints = {"deployment-00000": 1} if invalid_interval else {} + with pytest.raises(ValueError, match="interval|endpoints"): + monitor_runtime_health(state, healthy_readiness, endpoints, poll_interval_seconds=0 if invalid_interval else 1) diff --git a/packages/data-designer-slurm/tests/runtime/test_proxy.py b/packages/data-designer-slurm/tests/runtime/test_proxy.py index 50231eff8..a434dd58c 100644 --- a/packages/data-designer-slurm/tests/runtime/test_proxy.py +++ b/packages/data-designer-slurm/tests/runtime/test_proxy.py @@ -5,13 +5,15 @@ import argparse import asyncio +import json +import logging import socket from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager from pathlib import Path import pytest -from aiohttp import ClientSession, ClientTimeout, TCPConnector, web +from aiohttp import ClientError, ClientSession, ClientTimeout, TCPConnector, web from aiohttp.test_utils import TestServer from data_designer.slurm.runtime import proxy as runtime_proxy @@ -107,13 +109,15 @@ def run_app(application: web.Application, **options: object) -> None: options.pop(name) runner = web.AppRunner(application, **options) - def create_connector(**options: object) -> TCPConnector: + connector_type = runtime_proxy._ProxyConnector + + def create_connector(*args: object) -> TCPConnector: nonlocal connector - connector = TCPConnector(**options) + connector = connector_type(*args) return connector monkeypatch.setattr(runtime_proxy.web, "run_app", run_app) - monkeypatch.setattr(runtime_proxy, "TCPConnector", create_connector) + monkeypatch.setattr(runtime_proxy, "_ProxyConnector", create_connector) with socket.socket() as reservation: reservation.bind(("127.0.0.1", 0)) port = reservation.getsockname()[1] @@ -451,6 +455,203 @@ async def available(request: web.Request) -> web.Response: assert metrics["connections"]["idle"] == 1 +@pytest.mark.asyncio +async def test_proxy_logs_attempt_upstreams_and_timings_without_payloads(caplog: pytest.LogCaptureFixture) -> None: + async def overloaded(request: web.Request) -> web.Response: + return web.Response(status=429, text="private-upstream-body", headers={"Retry-After": "2"}) + + async def available(request: web.Request) -> web.Response: + return web.Response(text="private-generated-content") + + caplog.set_level(logging.INFO, logger=runtime_proxy.__name__) + async with _serve(_application(overloaded)) as first, _serve(_application(available)) as second: + upstreams = [_backend(first).origin, _backend(second).origin] + proxy = _ProxyApplication((_backend(first), _backend(second)), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post( + endpoint.make_url("/v1/chat/completions?private-query"), + data=b"private-prompt", + headers={"Authorization": "Bearer private-key"}, + ) as response: + assert await response.text() == "private-generated-content" + records = [record.getMessage() for record in caplog.records if record.name == runtime_proxy.__name__] + assert len(records) == 1 + assert "private-" not in records[0] + event = json.loads(records[0].split(" ", 1)[1]) + assert event["upstream_outcomes"] == ["429", "200"] + assert event["failover"] is True + assert event["response_completed"] is True + assert [attempt["upstream"] for attempt in event["upstream_attempts"]] == upstreams + for attempt in event["upstream_attempts"]: + assert attempt["connection_seconds"] >= 0 + assert attempt["response_seconds"] >= attempt["response_headers_seconds"] >= 0 + + +@pytest.mark.asyncio +async def test_proxy_counts_terminal_timeout_without_retry(total_timeout: None) -> None: + async def slow(request: web.Request) -> web.Response: + await asyncio.sleep(1) + return web.Response() + + async with _serve(_application(slow)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/complete")) as response: + assert response.status == 502 + async with client.get(endpoint.make_url("/metrics")) as response: + metrics = await response.json() + assert metrics["timeouts"] == 1 + assert metrics["retries"] == {} + assert metrics["final_statuses"] == {"502": 1} + + +@pytest.mark.asyncio +async def test_proxy_counts_expired_idle_connections(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(runtime_proxy, "_KEEPALIVE_TIMEOUT_SECONDS", 0.02) + + async def available(request: web.Request) -> web.Response: + return web.Response(text="complete") + + async with _serve(_application(available)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.get(endpoint.make_url("/complete")) as response: + await response.read() + await asyncio.sleep(0.08) + async with client.get(endpoint.make_url("/metrics")) as response: + metrics = await response.json() + assert metrics["connections"]["evicted"] == 1 + assert metrics["connections"]["idle"] == 0 + + +@pytest.mark.asyncio +async def test_proxy_counts_discarded_upstream_connections() -> None: + async def available(request: web.Request) -> web.Response: + return web.Response(text="complete", headers={"Connection": "close"}) + + async with _serve(_application(available)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.get(endpoint.make_url("/complete")) as response: + await response.read() + async with client.get(endpoint.make_url("/metrics")) as response: + metrics = await response.json() + assert metrics["connections"]["evicted"] == 1 + assert metrics["connections"]["idle"] == 0 + + +@pytest.mark.asyncio +async def test_proxy_discards_stale_connection_before_reuse() -> None: + transport: asyncio.Transport | None = None + + async def available(request: web.Request) -> web.Response: + nonlocal transport + transport = request.transport + return web.Response(text="complete") + + async with _serve(_application(available)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.get(endpoint.make_url("/complete")) as response: + assert await response.text() == "complete" + assert transport is not None + transport.close() + await asyncio.sleep(0.02) + async with client.get(endpoint.make_url("/complete")) as response: + assert await response.text() == "complete" + async with client.get(endpoint.make_url("/metrics")) as response: + metrics = await response.json() + assert metrics["connections"]["opened"] == 2 + assert metrics["connections"]["evicted"] == 1 + + +@pytest.mark.asyncio +async def test_proxy_records_stream_timeout_without_replaying( + total_timeout: None, caplog: pytest.LogCaptureFixture +) -> None: + calls = 0 + + async def streaming(request: web.Request) -> web.StreamResponse: + nonlocal calls + calls += 1 + response = web.StreamResponse() + await response.prepare(request) + await response.write(b"first") + await asyncio.sleep(0.2) + return response + + caplog.set_level(logging.INFO, logger=runtime_proxy.__name__) + async with _serve(_application(streaming)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + with pytest.raises(ClientError): + async with client.get(endpoint.make_url("/complete")) as response: + await response.read() + async with client.get(endpoint.make_url("/metrics")) as response: + assert (await response.json())["timeouts"] == 1 + assert calls == 1 + record = next(record for record in caplog.records if record.name == runtime_proxy.__name__) + event = json.loads(record.getMessage().split(" ", 1)[1]) + assert event["response_completed"] is False + assert event["upstream_outcomes"] == ["stream_timeout"] + + +@pytest.mark.asyncio +async def test_production_proxy_shutdown_drains_active_stream(monkeypatch: pytest.MonkeyPatch) -> None: + release = asyncio.Event() + + async def streaming(request: web.Request) -> web.StreamResponse: + response = web.StreamResponse() + await response.prepare(request) + await response.write(b"first") + await release.wait() + await response.write(b"last") + return response + + async with _serve(_application(streaming)) as backend, ClientSession() as client: + endpoint_context = _serve_production(_backend(backend), monkeypatch) + origin, _ = await endpoint_context.__aenter__() + try: + async with client.get(f"{origin}/complete") as response: + assert await response.content.readexactly(5) == b"first" + shutdown = asyncio.create_task(endpoint_context.__aexit__(None, None, None)) + await asyncio.sleep(0.05) + assert not shutdown.done() + release.set() + assert await response.read() == b"last" + await asyncio.wait_for(shutdown, 2) + finally: + release.set() + await endpoint_context.__aexit__(None, None, None) + + +@pytest.mark.asyncio +async def test_health_does_not_wait_for_saturated_inference_pool(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(runtime_proxy, "_MAXIMUM_CONNECTIONS", 1) + entered, release = asyncio.Event(), asyncio.Event() + + async def backend_handler(request: web.Request) -> web.Response: + if request.path == "/health": + return web.Response() + entered.set() + await release.wait() + return web.Response() + + async with _serve(_application(backend_handler)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + pending = asyncio.create_task(client.get(endpoint.make_url("/complete"))) + try: + await asyncio.wait_for(entered.wait(), 1) + async with await asyncio.wait_for(client.get(endpoint.make_url("/health")), 1) as response: + assert response.status == 200 + assert (await response.json())["backends_ready"] == 1 + finally: + release.set() + async with await pending as response: + await response.read() + + @pytest.mark.asyncio async def test_proxy_waits_for_a_pooled_connection_without_timing_out(monkeypatch: pytest.MonkeyPatch) -> None: first_started = asyncio.Event() diff --git a/packages/data-designer-slurm/tests/runtime/test_proxy_admission.py b/packages/data-designer-slurm/tests/runtime/test_proxy_admission.py new file mode 100644 index 000000000..3c5a196b1 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_proxy_admission.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from contextlib import AsyncExitStack + +import pytest +from aiohttp import web +from aiohttp.test_utils import TestServer + +from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode +from data_designer.engine.models.clients.adapters.openai_compatible import OpenAICompatibleClient +from data_designer.engine.models.clients.errors import ProviderError, ProviderErrorKind +from data_designer.engine.models.clients.model_request_executor import ModelRequestExecutor +from data_designer.engine.models.clients.retry import RetryConfig +from data_designer.engine.models.clients.types import ChatCompletionRequest +from data_designer.engine.models.request_admission.config import RequestAdmissionConfig +from data_designer.engine.models.request_admission.controller import AdaptiveRequestAdmissionController +from data_designer.engine.models.request_admission.resources import RequestDomain, RequestResourceKey +from data_designer.slurm.runtime.proxy import _Backend, _ProxyApplication + + +@pytest.mark.asyncio +@pytest.mark.parametrize("total_overload", (False, True), ids=("partial-overload", "total-overload")) +async def test_proxy_overload_reaches_real_client_admission(total_overload: bool) -> None: + calls: list[int] = [] + + def application(index: int) -> web.Application: + async def respond(request: web.Request) -> web.Response: + calls.append(index) + if total_overload or index == 0: + return web.json_response({"error": {"message": "overloaded"}}, status=429, headers={"Retry-After": "3"}) + return web.json_response({"choices": [{"message": {"role": "assistant", "content": "complete"}}]}) + + app = web.Application() + app.router.add_post("/v1/chat/completions", respond) + return app + + async with AsyncExitStack() as stack: + servers = [await stack.enter_async_context(TestServer(application(index))) for index in range(8)] + proxy = _ProxyApplication(tuple(_Backend(server.host, server.port) for server in servers), 1) + endpoint = await stack.enter_async_context(TestServer(proxy.create())) + resource = RequestResourceKey("slurm", "model", RequestDomain.CHAT) + admission = AdaptiveRequestAdmissionController( + RequestAdmissionConfig(multiplicative_decrease_factor=0.5, cooldown_seconds=0) + ) + admission.register(provider_name="slurm", model_id="model", alias="generator", max_parallel_requests=8) + client = OpenAICompatibleClient( + provider_name="slurm", + endpoint=str(endpoint.make_url("/v1")), + concurrency_mode=ClientConcurrencyMode.ASYNC, + retry_config=RetryConfig(max_retries=0), + ) + stack.push_async_callback(client.aclose) + executor = ModelRequestExecutor(client, admission, "slurm", "model", retry_config=RetryConfig(max_retries=0)) + request = ChatCompletionRequest(model="model", messages=[]) + if total_overload: + with pytest.raises(ProviderError) as failure: + await executor.acompletion(request) + assert failure.value.kind is ProviderErrorKind.RATE_LIMIT + assert failure.value.retry_after == 3 + assert calls == [0, 1, 2] + else: + assert (await executor.acompletion(request)).message.content == "complete" + assert calls == [0, 1] + pressure = admission.pressure.snapshot(resource) + assert pressure is not None + assert pressure.current_limit == (4 if total_overload else 8) + assert pressure.last_outcome == ("rate_limited" if total_overload else "success") + if total_overload: + assert pressure.cooldown_remaining_seconds > 2 + else: + assert pressure.cooldown_remaining_seconds == 0 + assert pressure.in_flight_count == 0 diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index 51d01cba7..ffb330797 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -66,6 +66,11 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( _write_executable(fake_bin / "getent", "#!/usr/bin/env bash\nexit 0\n") _write_executable(fake_bin / "scontrol", "#!/usr/bin/env bash\nprintf 'compute-001\\n'\n") _write_executable(fake_bin / "srun", _fake_srun()) + _write_executable( + fake_bin / "scancel", + '#!/usr/bin/env bash\n[[ $* == "--quiet --signal=TERM 4101.7" ]] || exit 64\n' + 'kill -TERM "$(cat "${FAKE_ENDPOINT_PID_FILE}")"\n', + ) command = f""" set -Eeuo pipefail DD_PLAN_SHA256={"a" * 64} @@ -85,6 +90,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( "FAKE_MANIFEST_SOURCE": manifest_path.as_posix(), "FAKE_GPU_MODE": gpu_request_mode, "FAKE_LISTENING_FILE": str(log_directory / "deployment-00000-endpoint.out.listening"), + "FAKE_ENDPOINT_PID_FILE": str(tmp_path / "endpoint.pid"), "SOURCE_TOKEN": "supersecret", "CUDA_VISIBLE_DEVICES": "0", "DD_SCRATCH_ROOT": scratch_root.as_posix(), @@ -583,13 +589,14 @@ def test_endpoint_probe_waits_for_listener_and_rejects_stale_marker(tmp_path: Pa runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" marker = tmp_path / "endpoint.out.listening" marker.touch() + marker.with_suffix(".step-id").write_text("stale") command = f""" source {shlex.quote(str(runtime_root / "entrypoint.sh"))} DD_RUNTIME_MANIFEST=unused DD_STEP_STDOUT={shlex.quote(str(tmp_path / "endpoint.out"))} dd_read_step_ids() {{ DD_STEP_IDS=(endpoint); }} dd_read_step() {{ DD_STEP_PROBE_FIELDS=(127.0.0.1 10000 /health 30); }} -dd_start_step() {{ [[ ! -f ${{DD_STEP_STDOUT}}.listening ]]; DD_LAST_PID=1; }} +dd_start_step() {{ [[ ! -f ${{DD_STEP_STDOUT}}.listening && ! -f ${{DD_STEP_STDOUT}}.step-id ]]; DD_LAST_PID=1; }} dd_register_required_pid() {{ :; }} dd_require_running() {{ :; }} dd_sleep() {{ @@ -752,6 +759,10 @@ def _fake_srun() -> str: [[ ${arguments[index]} == -- ]] && break done command=${arguments[index + 1]} +if [[ ${command} == bash ]]; then + printf '4101.7\\n' > "${arguments[index + 5]}" + command=${arguments[index + 6]} +fi if [[ ${command} == */python3 ]]; then operation=${arguments[index + 4]} [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] @@ -764,6 +775,10 @@ def _fake_srun() -> str: fi done fi + if [[ ${operation} == monitor ]]; then + trap 'exit 0' TERM INT + while :; do sleep 0.1; done + fi exit 0 fi case ${command} in @@ -785,6 +800,7 @@ def _fake_srun() -> str: ;; fake-endpoint) [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + printf '%s\\n' "$$" > "${FAKE_ENDPOINT_PID_FILE}" touch "${FAKE_LISTENING_FILE}" trap 'exit 0' TERM INT while :; do :; done diff --git a/plans/1103/live-drain-acceptance.json b/plans/1103/live-drain-acceptance.json new file mode 100644 index 000000000..9e72b4a18 --- /dev/null +++ b/plans/1103/live-drain-acceptance.json @@ -0,0 +1,49 @@ +{ + "passed": true, + "scheduler": "Slurm", + "gpus": 0, + "cases": { + "completion": { + "controller_exit_status": 0, + "active_response_completed": true, + "proxy_logged_response_completed": true, + "listener_released": true, + "task_identity_bound_to_allocation": true, + "drain_elapsed_seconds": 2.0000033378601074 + }, + "failure": { + "controller_exit_status": 1, + "active_response_completed": true, + "proxy_logged_response_completed": true, + "listener_released": true, + "task_identity_bound_to_allocation": true, + "drain_elapsed_seconds": 2.0000098149757832 + }, + "cancellation": { + "controller_exit_status": 130, + "active_response_completed": true, + "proxy_logged_response_completed": true, + "listener_released": true, + "task_identity_bound_to_allocation": true, + "drain_elapsed_seconds": 2.000048526097089 + } + }, + "runtime_source_sha256": { + "entrypoint.sh": "d3797d97b2017824be7a2f15d22e28d2a0412c516ace5f17b9459e56114c4544", + "scratch.sh": "65c9ca253f84d0f08bd2cf69914fe7c18959bea63fe0e30afb3763ccd3ff2963", + "step_runner.sh": "0e169ba03e8726d12e7519d7f32d545f0d64d44e1a0401b3ac277e05302d75d3", + "cleanup.sh": "10ab028c0de7159ad5527365c715ce0499f5af58d4809a2dde449ec8fae43f81", + "plan_reader.sh": "9c0cd73619c78919908f526033423363622b250991192fff14b830734856e514" + }, + "proxy_source_sha256": "c31d0c25bc692f6f32c32def6c5ac8c47f34448ab6501a386d869818e0c25a79", + "validation_date": "2026-10-29", + "allocation_state": "COMPLETED", + "allocation_exit_status": 0, + "endpoint_task_exit_statuses": [ + 0, + 0, + 0 + ], + "validation_scope": "Production Bash launch/endpoint/cleanup helpers with real srun and scancel, installed production proxy, and a synthetic streaming HTTP backend; not a full slurm execute run.", + "private_harness_sha256": "f9ca3d359ea77de946bb19a8f7e6a986348f4299d64d974d9cb246a07178e531" +} diff --git a/plans/1103/live-overload-acceptance.json b/plans/1103/live-overload-acceptance.json new file mode 100644 index 000000000..443c453c4 --- /dev/null +++ b/plans/1103/live-overload-acceptance.json @@ -0,0 +1,173 @@ +{ + "source_commit": "f499303aaf7867f2ee50b21544406ebaa6fa4084", + "replicas": 2, + "installed_dependency_inventory_sha256": "702d96d2ec5bc63136bd5b578fe20a9308dbd2d9c23806b02059febcc8e26846", + "installed_dependency_compatibility_check": "97 packages checked; all compatible", + "cases": { + "healthy_generation": { + "expected_records": 32, + "actual_records": 32, + "nonempty_answers": true, + "elapsed_seconds": 22.324767655692995, + "records_per_minute": 86.00313470722212, + "proxy_metrics": { + "requests": 33, + "attempts": 33, + "connections": { + "limit": 256, + "opened": 2, + "reused": 31, + "failed": 0, + "evicted": 0, + "idle": 2 + }, + "retries": {}, + "timeouts": 0, + "final_statuses": { + "200": 33 + }, + "backends": [ + { + "index": 0, + "active_requests": 0, + "requests": 17 + }, + { + "index": 1, + "active_requests": 0, + "requests": 16 + } + ] + }, + "proxy_usage": { + "cpu_seconds": 0.26, + "peak_rss_mib": 38.33984375 + } + }, + "partial_overload": { + "observed_waiting_requests": [ + 7.0 + ], + "admission_limit": 8, + "cooldown_remaining_seconds": 0.0, + "proxy_metrics": { + "requests": 1, + "attempts": 2, + "connections": { + "limit": 256, + "opened": 2, + "reused": 0, + "failed": 0, + "evicted": 0, + "idle": 2 + }, + "retries": { + "http_429": 1 + }, + "timeouts": 0, + "final_statuses": { + "200": 1 + }, + "backends": [ + { + "index": 0, + "active_requests": 0, + "requests": 1 + }, + { + "index": 1, + "active_requests": 0, + "requests": 1 + } + ] + }, + "recovery_succeeded": true, + "recovery_elapsed_seconds": 0.5283840745687485, + "cooldown_remaining_before_recovery_seconds": 0.0, + "proxy_usage": { + "cpu_seconds": 0.18, + "peak_rss_mib": 37.6484375 + } + }, + "total_overload": { + "observed_waiting_requests": [ + 7.0, + 7.0 + ], + "admission_limit": 4, + "cooldown_remaining_seconds": 2.999856797978282, + "proxy_metrics": { + "requests": 1, + "attempts": 2, + "connections": { + "limit": 256, + "opened": 2, + "reused": 0, + "failed": 0, + "evicted": 0, + "idle": 2 + }, + "retries": { + "http_429": 1 + }, + "timeouts": 0, + "final_statuses": { + "429": 1 + }, + "backends": [ + { + "index": 0, + "active_requests": 0, + "requests": 1 + }, + { + "index": 1, + "active_requests": 0, + "requests": 1 + } + ] + }, + "recovery_succeeded": true, + "recovery_elapsed_seconds": 3.0825562216341496, + "cooldown_remaining_before_recovery_seconds": 2.5813617296516895, + "proxy_usage": { + "cpu_seconds": 0.17, + "peak_rss_mib": 37.65234375 + } + } + }, + "wheel_sha256": { + "data_designer-0.9.8.dev84+f499303a-py3-none-any.whl": "119eb708291b02369b9ae5e661a357814fb321f729dad6c335fb80234f92484b", + "data_designer_config-0.9.8.dev84+f499303a-py3-none-any.whl": "14a62ad378a6353fd2e93d0520414a95b9deb7f493c72444a5f2bec02e95ee7b", + "data_designer_engine-0.9.8.dev84+f499303a-py3-none-any.whl": "c9f2e49df393ae14eb2dd7ffd4923410c7ef06dfdf32436ce2ca0cbd6880f651", + "data_designer_slurm-0.9.8.dev84+f499303a-py3-none-any.whl": "84900c24b022e2b0c9cc40db5561d17ccb9ea20d7f9d05895ca3e4f967a784d9" + }, + "installed_versions": { + "data-designer": "0.9.8.dev84+f499303a", + "data-designer-config": "0.9.8.dev84+f499303a", + "data-designer-engine": "0.9.8.dev84+f499303a", + "data-designer-slurm": "0.9.8.dev84+f499303a" + }, + "vllm_version": "0.22.0", + "passed": true, + "test_configuration": { + "gpus": 2, + "replicas": 2, + "max_num_seqs_per_replica": 1, + "max_waiting_requests": 0, + "retry_after_seconds": 3, + "load_requests_per_overloaded_replica": 8, + "load_completion_tokens": 1024, + "model": "openai/gpt-oss-20b" + }, + "source_archive_sha256": "55c2187a011a161386e1aba609758895758dd9974cac90fca1a0b2008b9ccd6f", + "dependency_lock_sha256": "6e6390270140396f81677df96039af4e0934c3dd0184cb995c8171c7940257d0", + "runtime_source_sha256": { + "proxy.py": "c31d0c25bc692f6f32c32def6c5ac8c47f34448ab6501a386d869818e0c25a79", + "health.py": "fc3d0df230585c3313f63df86a3157cb249ccb8f25ce53bd4f37cfe28354536d", + "cleanup.sh": "80d279273239fe5987012931ea99694f64ad425468c199d990a34934b7bb9997", + "entrypoint.sh": "8a1dc1dae28b3bead5ae13a4d11a980d791c5b953a2ea3443d75caa37187d1d8", + "entrypoint.py": "c6a1bb3586ff8b51bb81638c94dbbd12272c3eca9a847f29833416c2fd1df14d", + "step_runner.sh": "aea7694f77960bea454dcf7b469b7481d6c89ee3d153ca418b796cf5194cbef9" + } +} diff --git a/plans/1103/runtime-proxy-completion.md b/plans/1103/runtime-proxy-completion.md new file mode 100644 index 000000000..824929e45 --- /dev/null +++ b/plans/1103/runtime-proxy-completion.md @@ -0,0 +1,156 @@ +# Plan: Complete the Slurm inference proxy lifecycle + +## Scope and delivered behavior + +Issue #1103 carries forward the outstanding lifecycle, telemetry, and overload +acceptance requirements from #966 after the core proxy implementation in #967. + +- The coordinated allocation runtime starts a CPU-only health monitor after + startup readiness. It probes each logical endpoint's aggregate health, writes + revisioned backend counts and probe evidence, and fails the attempt when a + required service loses health while its process remains alive. Finalization + preserves that evidence. See `runtime/health.py:28` and `runtime/entrypoint.py:156` + under `packages/data-designer-slurm/src/data_designer/slurm/`. +- Health probes use their own bounded pool, independent of inference admission, + so saturated inference connections cannot starve health checks or distort + least-connections selection. See `runtime/proxy.py`. +- Cleanup stops monitoring, sends `TERM` to each exact endpoint task step with + `scancel`, and allows 12 seconds for their 10-second drain window and step exit. + The task publishes its job and step IDs before executing the proxy; cleanup + validates that the identity belongs to the current allocation. Launchers and + required serving steps stay available until endpoint termination. Remaining steps then stop with bounded + TERM/KILL escalation. See `runtime/cleanup.sh:48`. +- Request events contain per-attempt origins, connection/reuse timing, response + header and total durations, outcomes, and completion state. Counters include + expired/stale/discarded connections and terminal/streaming timeouts. Payloads, + request paths, and authorization headers are excluded. Eviction accounting + uses narrow aiohttp 3.x connector hooks because no public eviction signal exists. + +## Automated acceptance + +The final complete Slurm suite passed with 2,123 tests and one platform skip. The +focused health, proxy, and request-admission suite passed all 54 cases. + +Coverage includes real-process cleanup on completion, cancellation, and failure; +forced termination after a deadline; health loss without process exit and +retention of failure evidence; health checks under connection saturation; expired, +stale, and discarded connections; streaming timeout without replay; and draining +an active stream with production runner settings. + +An eight-replica loopback test drives the actual HTTP adapter and request-admission +controller. It verifies the three-attempt ceiling and end-to-end 429/cooldown +handling separately from the two-replica GPU acceptance below. + +## Live overload acceptance — 2026-10-28 + +The repeatable harness is `scripts/validate_slurm_proxy_overload.py`. The complete +sanitized report and immutable artifact identities are in +[`live-overload-acceptance.json`](live-overload-acceptance.json). + +The run installed four wheels built from commit +`f499303aaf7867f2ee50b21544406ebaa6fa4084` into a new Python 3.11 environment. +The third-party dependency inventory was pinned separately; `uv pip check` +verified compatibility of all 97 installed packages. The report records both +that inventory's digest and the source revision's `uv.lock` digest. +Two GPU-backed vLLM 0.22.0 processes served `openai/gpt-oss-20b` in one Slurm +allocation. The production queue middleware, proxy entrypoint, HTTP adapter, and +request-admission controller were used without synthesized overload responses. +The allocation completed with exit status zero and released its serving steps. + +For this bounded test, each replica used `max-num-seqs=1`, the supported queue +threshold override `DD_VLLM_MAX_WAITING_REQUESTS=0`, and `Retry-After: 3`. +Eight long inference requests per overloaded replica created real waiting queues. +The harness required positive vLLM waiting-request metrics before testing the +proxy. The production default threshold remains 128. + +| Case | Observed result | +| --- | --- | +| Healthy generation | 32 requested and 32 persisted records, all answers nonempty; 22.325 seconds, 86.003 records/minute | +| Healthy proxy traffic | 33 HTTP 200 responses including startup generation; zero failures, retries, or timeouts; 2 connections opened and 31 reused | +| Partial overload | Queue depth 7 on one replica; one local 429 retry, then success on the spare; admission remains 8 with no cooldown | +| Total overload | Queue depths 7 and 7; two upstream attempts, one local retry, final 429 and `Retry-After: 3` reach the real client | +| Admission and recovery | Admission decreases from 8 to 4; initial cooldown is 3.000 seconds; recovery succeeds after a 3.083-second request with 2.581 seconds of cooldown remaining at its start | +| Proxy resources | Healthy case: 0.26 CPU seconds and 38.340 MiB peak RSS; partial/total cases: 0.18/0.17 CPU seconds and approximately 37.65 MiB peak RSS | + +Overload-case proxy metric snapshots precede the recovery request. Background load +requests go directly to the serving backends and are excluded from proxy counts. +The throughput number describes this small serial validation workload, not a +capacity benchmark. Raw scheduler, serving, proxy, and generated-data artifacts +remain private; the checked-in report excludes infrastructure names and paths, +credentials, prompts, and completions. + +## Reproduction and acceptance boundaries + +Launch two idle vLLM replicas using the package's +`data_designer.slurm.runtime.backpressure.QueueDepthBackpressureMiddleware` with +the queue settings above and `--max-num-seqs 1`. Within that allocation, invoke +the harness with two `--backend` origins, their served `--model` name, the exact +`--source-commit`, the four-wheel `--wheels` directory, and a private `--output` +path. The script exits nonzero on any failed assertion and writes a sanitized +summary only after every scenario passes. + +The live run proves actual queue admission and overload propagation through +installed artifacts, including persisted Data Designer output. It does not run +the full managed `slurm execute` release matrix. Managed health/finalization and +shutdown paths have the automated process/state coverage described above. +The broader security, multi-node, retry, and final release-artifact matrix remains +owned by #870. The recorded live run applies to its exact source and wheel +identities. Subsequent review changes affect Python attempt finalization and shell +endpoint signal handling. The proxy, health monitor, backpressure middleware, and +overload acceptance harness remain byte-for-byte identical to the GPU-tested +source. Finalization changes have regression coverage; shell signal handling is +validated separately below. + +## Review follow-up + +Greptile identified a shutdown race: the monitor could persist a health failure +after the client's final process check, but success finalization would overwrite +it. Finalization now checks the monitor's saved evidence after cleanup has stopped +it and refuses to publish a winner after a recorded health failure, including +failure evidence retained in `STOPPED` readiness. + +Independent review also found that failure cleanup attempted an invalid +`STOPPED` → `FAILED` readiness transition after unsuccessful winner finalization. +Failure cleanup now retains terminal readiness while marking the attempt failed. +Regression cases exercise both the late health failure and recovery after winner +publication fails, and the bundle test explicitly verifies that the health module +is staged. + +Ten review passes covered success finalization, terminal-state recovery, health +protocol/state validation, shell monitoring and drain ordering, connector +accounting, streaming cancellation and telemetry, overload/admission behavior, +packaging/layering/secret scope, live-harness provenance, and the final diff and +validation record. Those passes preceded the additional Slurm signal-handling +finding described below. + +## Slurm signal-handling review + +Review identified that the cleanup PIDs belong to `srun` launchers. Slurm forwards +`TERM` sent to those launchers as `KILL` to tasks, bypassing the proxy drain +handler. Endpoint cleanup now uses the endpoint task's recorded numeric +`job.step` identity to send `TERM` directly with `scancel`. Launchers remain alive +during the drain window and are forcibly stopped only after that window expires. +Cleanup never sends a task-directed signal to a different allocation or to a +whole job. Stale step identity files are removed before launch. + +The process regressions now separate the launcher from its endpoint task and +model Slurm's launcher `TERM` → task `KILL` behavior. They verify successful +draining before backend shutdown on completion, failure, and cancellation, plus +numeric/current-allocation identity checks and forced cleanup after the deadline. + +### Live Slurm drain acceptance — 2026-10-29 + +The [sanitized report](live-drain-acceptance.json) records a separate CPU-only +Slurm allocation exercising the production Bash step-launch, endpoint-start, and +cleanup helpers, real `srun` and `scancel`, and the installed production proxy. +A synthetic HTTP backend streamed an initial chunk, then delayed its final chunk +for two seconds while cleanup began. The response completed on normal completion, +failure cleanup, and controller cancellation (statuses 0, 1, and 130). Each proxy +logged a completed response, released its listener marker, and finished its Slurm +task with exit status 0 before cleanup stopped the backend. The allocation +completed with exit status 0 and released all steps. + +The report records hashes of every shell source used and the installed proxy +source; they match this PR's source. This validates the scheduler signal path +separately from GPU overload acceptance. It does not run the full managed +`slurm execute` state/bootstrap flow, and the synthetic backend is not vLLM. diff --git a/scripts/validate_slurm_proxy_overload.py b/scripts/validate_slurm_proxy_overload.py new file mode 100644 index 000000000..c6e46b84a --- /dev/null +++ b/scripts/validate_slurm_proxy_overload.py @@ -0,0 +1,319 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validate real vLLM queue overload through the production proxy and DD client. + +Run inside an allocation with two idle vLLM replicas, each configured with +max-num-seqs=1, DD_VLLM_MAX_WAITING_REQUESTS=0, and DD_VLLM_RETRY_AFTER_SECONDS=3. +Both replicas must load the package's QueueDepthBackpressureMiddleware. No model, +image, scheduler, or cluster-specific paths are embedded in this harness. +""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import importlib.metadata +import json +import os +import re +import socket +import subprocess +import sys +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +from aiohttp import ClientSession, ClientTimeout + +from data_designer.config import DataDesignerConfigBuilder +from data_designer.config.models import ModelProvider +from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode +from data_designer.engine.models.clients.adapters.openai_compatible import OpenAICompatibleClient +from data_designer.engine.models.clients.errors import ProviderError, ProviderErrorKind +from data_designer.engine.models.clients.model_request_executor import ModelRequestExecutor +from data_designer.engine.models.clients.retry import RetryConfig +from data_designer.engine.models.clients.types import ChatCompletionRequest +from data_designer.engine.models.request_admission.config import RequestAdmissionConfig +from data_designer.engine.models.request_admission.controller import AdaptiveRequestAdmissionController +from data_designer.engine.models.request_admission.resources import RequestDomain, RequestResourceKey +from data_designer.interface import DataDesigner + + +async def queue_depth(client: ClientSession, backend: str) -> float: + """Read actual vLLM waiting-request metrics without synthesizing overload.""" + async with client.get(f"{backend}/metrics") as response: + response.raise_for_status() + text = await response.text() + samples = re.findall(r"^vllm:num_requests_waiting(?:\{[^\n]*\})?\s+([\d.eE+-]+)", text, re.MULTILINE) + if not samples: + raise RuntimeError("required vLLM queue metric is unavailable") + return sum(float(value) for value in samples) + + +@asynccontextmanager +async def saturate(client: ClientSession, backends: list[str], model: str) -> AsyncIterator[list[float]]: + """Queue real long completions and cancel only this harness's requests.""" + + async def generate(backend: str) -> int: + async with client.post( + f"{backend}/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": "Continue counting integers, one per line."}], + "max_tokens": 1024, + "min_tokens": 1024, + "ignore_eos": True, + "temperature": 0, + }, + ) as response: + await response.read() + return response.status + + pending = [asyncio.create_task(generate(backend)) for backend in backends for _ in range(8)] + try: + deadline = time.monotonic() + 20 + while True: + depths = [await queue_depth(client, backend) for backend in backends] + if all(depth > 0 for depth in depths): + await asyncio.sleep(0.2) # Allow the production 100 ms sampler to observe the queue. + break + if time.monotonic() >= deadline: + raise RuntimeError("real backend queues did not saturate within the bounded window") + await asyncio.sleep(0.05) + yield depths + finally: + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + deadline = time.monotonic() + 30 + while any([await queue_depth(client, backend) for backend in backends]): + if time.monotonic() >= deadline: + raise RuntimeError("backend queues did not recover after cancelling validation traffic") + await asyncio.sleep(0.1) + await asyncio.sleep(0.3) + + +def proxy_usage(pid: int) -> dict[str, float]: + """Read Linux process CPU and peak resident memory for the proxy alone.""" + status = Path(f"/proc/{pid}/status").read_text() + peak = re.search(r"^VmHWM:\s+(\d+) kB", status, re.MULTILINE) + fields = Path(f"/proc/{pid}/stat").read_text().split(")", 1)[1].split() + return { + "cpu_seconds": (int(fields[11]) + int(fields[12])) / os.sysconf("SC_CLK_TCK"), + "peak_rss_mib": int(peak.group(1)) / 1024 if peak else 0, + } + + +@asynccontextmanager +async def proxy_process(backends: list[str], directory: Path) -> AsyncIterator[tuple[str, subprocess.Popen[bytes]]]: + """Launch the installed production entrypoint with private raw logs.""" + with socket.socket() as reservation: + reservation.bind(("127.0.0.1", 0)) + port = reservation.getsockname()[1] + command = [sys.executable, "-m", "data_designer.slurm.runtime.proxy", "--listen-port", str(port)] + command.extend(argument for backend in backends for argument in ("--backend", backend)) + directory.mkdir(mode=0o700, parents=True, exist_ok=True) + with (directory / "proxy.log").open("wb") as log: + process = subprocess.Popen(command, stdout=log, stderr=log) + try: + origin = f"http://127.0.0.1:{port}" + async with ClientSession(timeout=ClientTimeout(total=8)) as client: + for _ in range(100): + if process.poll() is not None: + raise RuntimeError("production proxy exited before readiness") + try: + async with client.get(f"{origin}/health") as response: + if response.status == 200: + break + except OSError: + pass + await asyncio.sleep(0.1) + else: + raise RuntimeError("production proxy did not become ready") + yield origin, process + finally: + process.terminate() + try: + await asyncio.to_thread(process.wait, timeout=12) + except subprocess.TimeoutExpired: + process.kill() + await asyncio.to_thread(process.wait) + + +def generate_dataset(origin: str, model: str, directory: Path, records: int) -> dict[str, object]: + """Generate and verify actual persisted Data Designer output.""" + builder = DataDesignerConfigBuilder.from_config( + { + "model_configs": [ + { + "alias": "generator", + "model": model, + "provider": "live", + "inference_parameters": { + "max_tokens": 256, + "temperature": 0, + "max_parallel_requests": 1, + "extra_body": {"reasoning_effort": "low"}, + }, + } + ], + "columns": [ + { + "name": "answer", + "column_type": "llm-text", + "model_alias": "generator", + "prompt": "Reply with only the word READY. Do not explain.", + } + ], + } + ) + designer = DataDesigner( + artifact_path=directory, + model_providers=[ModelProvider(name="live", endpoint=f"{origin}/v1")], + auto_configure_logging=False, + ) + started = time.monotonic() + results = designer.create(builder, num_records=records, dataset_name="live-overload-acceptance") + dataset = results.load_dataset() + assert len(dataset) == records + assert dataset["answer"].map(lambda value: isinstance(value, str) and bool(value.strip())).all() + elapsed = time.monotonic() - started + return { + "expected_records": records, + "actual_records": len(dataset), + "nonempty_answers": True, + "elapsed_seconds": elapsed, + "records_per_minute": records * 60 / elapsed, + } + + +async def validate(arguments: argparse.Namespace) -> dict[str, Any]: + """Execute healthy, partial, and total overload plus cooldown recovery.""" + backends: list[str] = arguments.backend + assert len(backends) == 2, "this bounded acceptance scenario requires exactly two replicas" + assert re.fullmatch(r"[0-9a-f]{40}", arguments.source_commit) + report: dict[str, Any] = {"source_commit": arguments.source_commit, "replicas": 2, "cases": {}} + report["wheel_sha256"] = { + path.name: hashlib.sha256(path.read_bytes()).hexdigest() for path in sorted(arguments.wheels.glob("*.whl")) + } + assert len(report["wheel_sha256"]) == 4 + report["installed_versions"] = { + name: importlib.metadata.version(name) + for name in ("data-designer", "data-designer-config", "data-designer-engine", "data-designer-slurm") + } + async with ClientSession(timeout=ClientTimeout(total=180)) as client: + async with client.get(f"{backends[0]}/version") as response: + report["vllm_version"] = (await response.json())["version"] + async with proxy_process(backends, arguments.output.parent / "healthy") as (origin, process): + report["cases"]["healthy_generation"] = await asyncio.to_thread( + generate_dataset, origin, arguments.model, arguments.output.parent / "dataset", arguments.records + ) + async with client.get(f"{origin}/metrics") as response: + metrics = await response.json() + assert metrics["connections"]["reused"] > 0 + report["cases"]["healthy_generation"].update( + {"proxy_metrics": metrics, "proxy_usage": proxy_usage(process.pid)} + ) + + request = ChatCompletionRequest( + model=arguments.model, + messages=[{"role": "user", "content": "Say READY."}], + max_tokens=256, + extra_body={"reasoning_effort": "low"}, + ) + for total in (False, True): + label = "total_overload" if total else "partial_overload" + async with proxy_process(backends, arguments.output.parent / label) as (origin, process): + admission = AdaptiveRequestAdmissionController( + RequestAdmissionConfig(multiplicative_decrease_factor=0.5, cooldown_seconds=0) + ) + admission.register( + provider_name="live", model_id=arguments.model, alias="generator", max_parallel_requests=8 + ) + adapter = OpenAICompatibleClient( + provider_name="live", + endpoint=f"{origin}/v1", + concurrency_mode=ClientConcurrencyMode.ASYNC, + retry_config=RetryConfig(max_retries=0), + ) + executor = ModelRequestExecutor( + adapter, admission, "live", arguments.model, retry_config=RetryConfig(max_retries=0) + ) + try: + async with saturate(client, backends if total else backends[:1], arguments.model) as depths: + async with client.get(f"{origin}/health") as response: + assert response.status == 200, "overload must not cause health failure" + if total: + try: + await executor.acompletion(request) + except ProviderError as error: + assert error.kind is ProviderErrorKind.RATE_LIMIT + assert error.retry_after == 3 + else: + raise AssertionError("real total overload did not reach client admission as 429") + else: + result = await executor.acompletion(request) + assert result.message.content or result.message.reasoning_content + pressure = admission.pressure.snapshot( + RequestResourceKey("live", arguments.model, RequestDomain.CHAT) + ) + assert pressure is not None and pressure.current_limit == (4 if total else 8) + assert ( + pressure.cooldown_remaining_seconds > 2 + if total + else pressure.cooldown_remaining_seconds == 0 + ) + async with client.get(f"{origin}/metrics") as response: + metrics = await response.json() + assert metrics["attempts"] == 2 + assert metrics["retries"] == {"http_429": 1} + report["cases"][label] = { + "observed_waiting_requests": depths, + "admission_limit": pressure.current_limit, + "cooldown_remaining_seconds": pressure.cooldown_remaining_seconds, + "proxy_metrics": metrics, + } + remaining = admission.pressure.snapshot( + RequestResourceKey("live", arguments.model, RequestDomain.CHAT) + ).cooldown_remaining_seconds + started = time.monotonic() + result = await executor.acompletion(request) + elapsed = time.monotonic() - started + assert result.message.content or result.message.reasoning_content + assert elapsed >= remaining - 0.05 + report["cases"][label].update( + { + "recovery_succeeded": True, + "recovery_elapsed_seconds": elapsed, + "cooldown_remaining_before_recovery_seconds": remaining, + "proxy_usage": proxy_usage(process.pid), + } + ) + finally: + await adapter.aclose() + report["passed"] = True + return report + + +def main() -> None: + """Run explicitly configured live acceptance and write a sanitized summary.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend", action="append", required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--source-commit", required=True) + parser.add_argument("--wheels", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--records", type=int, default=32) + arguments = parser.parse_args() + arguments.output.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + report = asyncio.run(validate(arguments)) + arguments.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps({"passed": report["passed"], "cases": list(report["cases"])})) + + +if __name__ == "__main__": + main()