From fbd559d1e406d1b0d8cee6d91e91c88f3858eaf4 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Tue, 8 Sep 2026 19:54:25 -0300 Subject: [PATCH 01/14] fix(slurm): close one-node M2 integration --- .../data_designer/slurm/client/execution.py | 16 +- .../data_designer/slurm/config/__init__.py | 7 +- .../data_designer/slurm/config/environment.py | 24 +- .../data_designer/slurm/launcher/client.py | 70 ++- .../data_designer/slurm/launcher/renderer.py | 9 +- .../data_designer/slurm/launcher/runner.py | 41 +- .../slurm/runtime/backpressure.py | 237 ++++++++++ .../data_designer/slurm/runtime/bootstrap.py | 366 ++++++++++++++++ .../src/data_designer/slurm/runtime/bundle.py | 23 +- .../data_designer/slurm/runtime/cleanup.sh | 69 +++ .../data_designer/slurm/runtime/context.py | 80 ++++ .../data_designer/slurm/runtime/controller.py | 70 ++- .../data_designer/slurm/runtime/entrypoint.py | 410 +++++++++++++++--- .../data_designer/slurm/runtime/entrypoint.sh | 180 ++++++++ .../slurm/runtime/plan_reader.sh | 197 +++++++++ .../data_designer/slurm/runtime/preflight.py | 10 +- .../slurm/runtime/step_runner.sh | 167 +++++++ .../src/data_designer/slurm/runtime/steps.py | 83 ++-- .../data_designer/slurm/services/artifacts.py | 271 ++++++++++++ .../data_designer/slurm/services/wiring.py | 95 +++- .../data_designer/slurm/state/finalization.py | 24 +- .../src/data_designer/slurm/state/storage.py | 8 + .../src/data_designer/slurm/state/store.py | 10 +- .../tests/launcher/test_client.py | 51 ++- .../tests/launcher/test_renderer.py | 1 + .../tests/launcher/test_runner.py | 10 +- .../tests/runtime/conftest.py | 52 ++- .../tests/runtime/test_backpressure.py | 118 +++++ .../tests/runtime/test_bootstrap.py | 38 ++ .../tests/runtime/test_bundle.py | 34 +- .../tests/runtime/test_controller.py | 25 +- .../tests/runtime/test_entrypoint.py | 115 ++++- .../tests/runtime/test_preflight.py | 4 +- .../tests/runtime/test_shell_runtime.py | 283 ++++++++++++ .../tests/runtime/test_steps.py | 18 + .../tests/services/test_wiring.py | 167 ++++++- .../golden/rendered/multi_node.sbatch | 9 +- .../golden/rendered/single_node.sbatch | 9 +- .../tests/slurm_test_fakes/slurm.py | 2 + .../slurm_test_fakes/test_rendered_scripts.py | 4 +- .../tests/state/test_store.py | 22 + 41 files changed, 3194 insertions(+), 235 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_backpressure.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_bootstrap.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_shell_runtime.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index 8a6d3f05c..f0bdd3e8e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -49,7 +49,7 @@ from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef from data_designer.slurm.config.images import InstalledDistribution from data_designer.slurm.config.run import LocalStdioMCPProviderConfig, RemoteMCPProviderConfig -from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import PlannedShard, ResolvedDependencyLock, ResolvedSlurmRunPlan from data_designer.slurm.state import CandidateOutcome, CandidateOutputFile, CandidateOutputManifest @@ -599,18 +599,6 @@ def _build_candidate_manifest( ), ) created_at = self._clock() - provenance_digest = compute_canonical_json_sha256( - { - "builder_sha256": context.plan.builder.content_sha256, - "client_image_sha256": prepared.client_image_sha256, - "dependency_lock_sha256": prepared.dependency_lock.sha256, - "files": [file.model_dump(mode="json") for file in files], - "attempt_id": prepared.attempt_id, - "resolved_plan_sha256": context.plan.compute_sha256(), - "run_id": context.plan.run_id, - "shard_id": context.shard.shard_id, - } - ) return CandidateOutputManifest( schema_version=1, run_id=context.plan.run_id, @@ -630,7 +618,7 @@ def _build_candidate_manifest( ), files=files, dataset_schema_digest=schema_digest, - provenance_digest=provenance_digest, + provenance_digest=context.plan.compute_sha256(), ) def _publish_success( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index c08456652..00b7f9bf3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -15,7 +15,11 @@ FixedRecordPolicy, ) from data_designer.slurm.config.builder import DataDesignerSlurmConfigBuilder -from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.config.environment import ( + LiteralEnvironmentBinding, + SecretRef, + collect_secret_environment_names, +) from data_designer.slurm.config.errors import SlurmConfigBuilderError, SlurmConfigLoadError from data_designer.slurm.config.images import ( ClientImageInspection, @@ -104,6 +108,7 @@ "RemoteMCPProviderConfig", "SchedulerProfile", "SecretRef", + "collect_secret_environment_names", "SelectedSlurmProfile", "ServerDeploymentConfig", "ServingImageInspection", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py index c835c3983..006effec8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py @@ -9,7 +9,7 @@ from collections.abc import Mapping from typing import Annotated, Literal -from pydantic import Field, JsonValue, StringConstraints, field_validator +from pydantic import BaseModel, Field, JsonValue, StringConstraints, field_validator from data_designer.slurm.contracts import AuthoredConfig, validate_plain_text from data_designer.slurm.types import EnvironmentName @@ -18,6 +18,7 @@ "EnvironmentBinding", "LiteralEnvironmentBinding", "SecretRef", + "collect_secret_environment_names", ] @@ -59,6 +60,13 @@ class SecretRef(AuthoredConfig): ] +def collect_secret_environment_names(value: object) -> tuple[str, ...]: + """Return every external environment variable referenced below a config value.""" + names: set[str] = set() + _collect_secret_environment_names(value, names) + return tuple(sorted(names)) + + def is_secret_bearing_name(value: str) -> bool: """Return whether a Slurm-owned name conventionally carries secret material.""" segments = _secret_name_segments(value) @@ -112,3 +120,17 @@ def _secret_name_segments(value: str) -> list[str]: snake_case = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", value) normalized = re.sub(r"[^a-z0-9]+", "_", snake_case.casefold()).strip("_") return normalized.split("_") + + +def _collect_secret_environment_names(value: object, names: set[str]) -> None: + if isinstance(value, SecretRef): + names.add(value.environment) + elif isinstance(value, BaseModel): + for field_name in type(value).model_fields: + _collect_secret_environment_names(getattr(value, field_name), names) + elif isinstance(value, Mapping): + for child in value.values(): + _collect_secret_environment_names(child, names) + elif isinstance(value, (tuple, list)): + for child in value: + _collect_secret_environment_names(child, names) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py index 47d061ad1..2d822a7f1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -8,7 +8,7 @@ import re import subprocess import unicodedata -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import TypeAlias @@ -32,6 +32,7 @@ _JobSelector: TypeAlias = int | SchedulerIdentity _IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_ENVIRONMENT_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") _MAX_SLURM_INTEGER = (1 << 32) - 1 @@ -43,10 +44,11 @@ class SlurmExecutables: squeue: str = "squeue" sacct: str = "sacct" scancel: str = "scancel" + scontrol: str = "scontrol" sinfo: str = "sinfo" def __post_init__(self) -> None: - for executable in (self.sbatch, self.squeue, self.sacct, self.scancel, self.sinfo): + for executable in (self.sbatch, self.squeue, self.sacct, self.scancel, self.scontrol, self.sinfo): _validate_argument(executable, field_name="Slurm executable") if any(character.isspace() for character in executable): raise ValueError("Slurm executable must be one argument-vector token") @@ -67,22 +69,42 @@ def __init__( self._runner = runner if runner is not None else SubprocessRunner() self._executables = executables if executables is not None else SlurmExecutables() - def submit(self, script_path: str | Path) -> SlurmJobSubmissionReceipt: + def submit( + self, + script_path: str | Path, + *, + hold: bool = False, + export_environment: Mapping[str, str] | None = None, + ) -> SlurmJobSubmissionReceipt: """Submit one rendered batch script and return its assigned job ID.""" path = str(script_path) _validate_argument(path, field_name="batch script path") if path.startswith("-"): raise ValueError("batch script path must not begin with '-'; prefix relative paths with './'") - output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) + hold_arguments = ("--hold",) if hold else () + export_argument, environment = _format_export_environment(export_environment) + output = self._run( + (self._executables.sbatch, "--parsable", *hold_arguments, export_argument, path), + environment=environment, + ) return parse_submission(output) - def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + def submit_script( + self, + script: str, + *, + hold: bool = False, + export_environment: Mapping[str, str] | None = None, + ) -> SlurmJobSubmissionReceipt: """Submit verified batch-script text through standard input.""" if type(script) is not str or not script or "\0" in script: raise ValueError("batch script text must be non-empty UTF-8 text without NUL") + hold_arguments = ("--hold",) if hold else () + export_argument, environment = _format_export_environment(export_environment) output = self._run( - (self._executables.sbatch, "--parsable", "--export=NIL"), + (self._executables.sbatch, "--parsable", *hold_arguments, export_argument), input_text=script, + environment=environment, ) return parse_submission(output) @@ -134,6 +156,10 @@ def cancel(self, selector: _JobSelector) -> None: """Cancel one managed Slurm job, array, or array task.""" self._run((self._executables.scancel, _format_selector(selector))) + def release(self, job_id: int) -> None: + """Release one held managed Slurm job or array.""" + self._run((self._executables.scontrol, "release", _format_job_id(job_id))) + def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, ...]: """Return configured GPU counts reported for eligible node groups.""" command = [self._executables.sinfo, "--noheader", "--format=%G"] @@ -143,12 +169,23 @@ def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, command.append(f"--partition={partition}") return parse_gpu_counts(self._run(command)) - def _run(self, command: Sequence[str], *, input_text: str | None = None) -> str: + def _run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> str: command_name = Path(command[0]).name try: - completed = ( - self._runner.run(command) if input_text is None else self._runner.run(command, input_text=input_text) - ) + if environment is None: + completed = ( + self._runner.run(command) + if input_text is None + else self._runner.run(command, input_text=input_text) + ) + else: + completed = self._runner.run(command, input_text=input_text, environment=environment) except (OSError, subprocess.SubprocessError) as error: raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") from error returncode = getattr(completed, "returncode", None) @@ -162,6 +199,19 @@ def _run(self, command: Sequence[str], *, input_text: str | None = None) -> str: return stdout +def _format_export_environment(environment: Mapping[str, str] | None) -> tuple[str, Mapping[str, str] | None]: + if not environment: + return "--export=NIL", None + names = tuple(sorted(environment)) + for name in names: + if _ENVIRONMENT_NAME_PATTERN.fullmatch(name) is None: + raise ValueError("exported environment names must be valid identifiers") + value = environment[name] + if type(value) is not str or "\0" in value: + raise ValueError("exported environment values must be strings without NUL") + return f"--export={','.join(names)}", environment + + def _format_selectors(selectors: Sequence[_JobSelector]) -> str: if not selectors: raise ValueError("at least one managed Slurm job selector is required") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index e076f5cb1..609e4fbf4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -50,12 +50,11 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi printf -v DD_SHARD_ID 'shard-%05d' "${{DD_ARRAY_TASK_ID}}" readonly DD_SHARD_ID readonly DD_ATTEMPT_DIR="${{DD_RUN_ROOT}}/shards/${{DD_SHARD_ID}}/attempts/attempt-${{DD_ATTEMPT_ORDINAL}}" -install -d -m 0700 "${{DD_ATTEMPT_DIR}}" -DD_RUNTIME_DIR="$(mktemp -d "${{DD_ATTEMPT_DIR}}/runtime.${{DD_RUNTIME_SHA256}}.XXXXXX")" -readonly DD_RUNTIME_DIR -tar -xzf "${{DD_RUNTIME_ARCHIVE}}" -C "${{DD_RUNTIME_DIR}}" +readonly DD_RUNTIME_ROOT="${{DD_ATTEMPT_DIR}}/runtime" +[[ -d ${{DD_RUNTIME_ROOT}} && ! -L ${{DD_RUNTIME_ROOT}} ]] +tar -xzf "${{DD_RUNTIME_ARCHIVE}}" -C "${{DD_RUNTIME_ROOT}}" -source "${{DD_RUNTIME_DIR}}/entrypoint.sh" +source "${{DD_RUNTIME_ROOT}}/entrypoint.sh" dd_slurm_run_allocation "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" """ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py index 9e2e7839f..de0ceb1d5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py @@ -16,7 +16,13 @@ class CommandRunner(Protocol): """Minimal command boundary implemented by production and fake runners.""" - def run(self, command: Sequence[str], *, input_text: str | None = None) -> subprocess.CompletedProcess[str]: + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: """Execute one argument-vector command.""" ... @@ -35,9 +41,12 @@ def __init__( ) -> None: if type(timeout_seconds) not in {int, float} or not math.isfinite(timeout_seconds) or timeout_seconds <= 0: raise ValueError("timeout_seconds must be a finite positive number") - explicit_environment = ( - dict(environment) if environment is not None else {"PATH": os.environ.get("PATH") or os.defpath} - ) + if environment is None: + explicit_environment = {"PATH": os.environ.get("PATH") or os.defpath} + if "SLURM_CONF" in os.environ: + explicit_environment["SLURM_CONF"] = os.environ["SLURM_CONF"] + else: + explicit_environment = dict(environment) for name, value in explicit_environment.items(): if type(name) is not str or not name or "=" in name or "\0" in name: raise ValueError("environment names must be non-empty and must not contain '=' or NUL") @@ -51,8 +60,18 @@ def environment(self) -> Mapping[str, str]: """Return the allowlisted environment forwarded to child processes.""" return self._environment - def run(self, command: Sequence[str], *, input_text: str | None = None) -> subprocess.CompletedProcess[str]: + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: """Execute an argument vector with captured text output.""" + process_environment = dict(self._environment) + if environment is not None: + _validate_environment(environment) + process_environment.update(environment) if input_text is not None: return subprocess.run( tuple(command), @@ -62,7 +81,7 @@ def run(self, command: Sequence[str], *, input_text: str | None = None) -> subpr text=True, encoding="utf-8", errors="replace", - env=dict(self._environment), + env=process_environment, timeout=self._timeout_seconds, ) return subprocess.run( @@ -73,6 +92,14 @@ def run(self, command: Sequence[str], *, input_text: str | None = None) -> subpr text=True, encoding="utf-8", errors="replace", - env=dict(self._environment), + env=process_environment, timeout=self._timeout_seconds, ) + + +def _validate_environment(environment: Mapping[str, str]) -> None: + for name, value in environment.items(): + if type(name) is not str or not name or "=" in name or "\0" in name: + raise ValueError("environment names must be non-empty and must not contain '=' or NUL") + if type(value) is not str or "\0" in value: + raise ValueError("environment values must not contain NUL") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py new file mode 100644 index 000000000..95ff51045 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM middleware that exposes bounded queue admission to client AIMD.""" + +from __future__ import annotations + +import importlib +import json +import math +import os +import threading +import time +from collections.abc import Awaitable, Callable, Iterable, Mapping +from dataclasses import dataclass +from http import HTTPStatus +from typing import Any, Protocol + +MAX_WAITING_REQUESTS_ENVIRONMENT = "DD_VLLM_MAX_WAITING_REQUESTS" +RETRY_AFTER_SECONDS_ENVIRONMENT = "DD_VLLM_RETRY_AFTER_SECONDS" +_METRIC_NAME = "vllm:num_requests_waiting" +_EXEMPT_PATHS = ("/health", "/ready", "/metrics", "/version", "/v1/models", "/ping") + +AsgiMessage = dict[str, Any] +AsgiScope = dict[str, Any] +AsgiReceive = Callable[[], Awaitable[AsgiMessage]] +AsgiSend = Callable[[AsgiMessage], Awaitable[None]] +AsgiApp = Callable[[AsgiScope, AsgiReceive, AsgiSend], Awaitable[None]] + + +class QueueDepthReader(Protocol): + """Read the latest aggregate waiting-request count.""" + + def __call__(self) -> int | None: + """Return queue depth, or ``None`` when metrics are unavailable.""" + ... + + +@dataclass(frozen=True, slots=True) +class QueueBackpressureSettings: + """Validated worker-owned queue admission settings.""" + + max_waiting_requests: int + retry_after_seconds: int | None + poll_interval_seconds: float = 0.1 + stale_after_seconds: float = 1.0 + + def __post_init__(self) -> None: + if type(self.max_waiting_requests) is not int or self.max_waiting_requests < 0: + raise ValueError("maximum waiting requests must be non-negative") + if self.retry_after_seconds is not None and ( + type(self.retry_after_seconds) is not int or self.retry_after_seconds <= 0 + ): + raise ValueError("retry-after seconds must be positive or absent") + if self.poll_interval_seconds <= 0 or self.stale_after_seconds <= 0: + raise ValueError("queue sampler intervals must be positive") + + @classmethod + def from_environment(cls, environment: Mapping[str, str] | None = None) -> QueueBackpressureSettings: + """Load the policy transported by the structured runtime step.""" + source = os.environ if environment is None else environment + maximum = _parse_non_negative_integer(source.get(MAX_WAITING_REQUESTS_ENVIRONMENT), default=128) + retry_value = source.get(RETRY_AFTER_SECONDS_ENVIRONMENT) + retry_after = None if retry_value == "" else _parse_positive_integer(retry_value, default=1) + return cls(maximum, retry_after) + + +@dataclass(frozen=True, slots=True) +class QueueSnapshot: + """One sampled queue depth and its monotonic observation time.""" + + depth: int | None + observed_at: float + + +class QueueBackpressureController: + """Cache queue metrics away from the request path and decide admission.""" + + def __init__( + self, + settings: QueueBackpressureSettings | None = None, + reader: QueueDepthReader | None = None, + *, + start_sampler: bool = True, + ) -> None: + self.settings = settings or QueueBackpressureSettings.from_environment() + self._reader = reader or read_vllm_queue_depth + self._start_sampler = start_sampler + self._snapshot = QueueSnapshot(None, 0.0) + self._lock = threading.Lock() + self._thread: threading.Thread | None = None + + def sample_once(self) -> QueueSnapshot: + """Refresh the cached queue depth, failing open on reader errors.""" + try: + depth = self._reader() + except Exception: + depth = None + if depth is not None: + depth = max(0, depth) + snapshot = QueueSnapshot(depth, time.monotonic()) + with self._lock: + self._snapshot = snapshot + return snapshot + + def should_reject(self) -> tuple[bool, QueueSnapshot]: + """Return the fail-open admission decision and supporting snapshot.""" + self._ensure_sampler() + with self._lock: + snapshot = self._snapshot + stale = time.monotonic() - snapshot.observed_at > self.settings.stale_after_seconds + reject = snapshot.depth is not None and not stale and snapshot.depth > self.settings.max_waiting_requests + return reject, snapshot + + def _ensure_sampler(self) -> None: + if not self._start_sampler or self._thread is not None: + return + with self._lock: + if self._thread is None: + self._thread = threading.Thread(target=self._sample_forever, name="dd-vllm-queue-depth", daemon=True) + self._thread.start() + + def _sample_forever(self) -> None: + while True: + self.sample_once() + time.sleep(self.settings.poll_interval_seconds) + + +class QueueDepthBackpressureMiddleware: + """Reject non-health HTTP requests with 429 above the resolved queue threshold.""" + + def __init__(self, app: AsgiApp, controller: QueueBackpressureController | None = None) -> None: + self.app = app + self.controller = controller or QueueBackpressureController() + + async def __call__(self, scope: AsgiScope, receive: AsgiReceive, send: AsgiSend) -> None: + """Apply queue admission without changing exempt or accepted requests.""" + path = str(scope.get("path", "")) + if scope.get("type") != "http" or path in _EXEMPT_PATHS: + await self.app(scope, receive, send) + return + reject, snapshot = self.controller.should_reject() + if not reject: + await self.app(scope, receive, send) + return + await _send_overload(send, self.controller.settings, snapshot) + + +def read_vllm_queue_depth() -> int | None: + """Return aggregate vLLM queue depth from supported metrics registries.""" + for source in (_read_vllm_metrics, _read_prometheus_metrics): + values = _collect_metric_values(source()) + if values: + return max(0, int(sum(values))) + return None + + +async def _send_overload( + send: AsgiSend, + settings: QueueBackpressureSettings, + snapshot: QueueSnapshot, +) -> None: + body = json.dumps( + { + "error": { + "message": "serving queue admission threshold exceeded", + "type": "rate_limit_exceeded", + "code": HTTPStatus.TOO_MANY_REQUESTS.value, + "queue_depth": snapshot.depth, + "max_waiting_requests": settings.max_waiting_requests, + } + }, + separators=(",", ":"), + ).encode() + headers = [(b"content-type", b"application/json")] + if settings.retry_after_seconds is not None: + headers.append((b"retry-after", str(settings.retry_after_seconds).encode())) + await send({"type": "http.response.start", "status": HTTPStatus.TOO_MANY_REQUESTS.value, "headers": headers}) + await send({"type": "http.response.body", "body": body}) + + +def _read_vllm_metrics() -> Iterable[object]: + try: + module = importlib.import_module("vllm.v1.metrics.reader") + return module.get_metrics_snapshot() + except (ImportError, AttributeError, RuntimeError): + return () + + +def _read_prometheus_metrics() -> Iterable[object]: + try: + module = importlib.import_module("prometheus_client") + return tuple(sample for family in module.REGISTRY.collect() for sample in getattr(family, "samples", ())) + except (ImportError, AttributeError, RuntimeError): + return () + + +def _collect_metric_values(metrics: Iterable[object]) -> tuple[float, ...]: + names = {_METRIC_NAME, _METRIC_NAME.replace(":", "_")} + values: list[float] = [] + for metric in metrics: + if str(getattr(metric, "name", "")) not in names: + continue + try: + value = float(getattr(metric, "value")) + if math.isfinite(value) and value >= 0: + values.append(value) + except (AttributeError, TypeError, ValueError): + continue + return tuple(values) + + +def _parse_non_negative_integer(value: str | None, *, default: int) -> int: + if value is None: + return default + if not value.isascii() or not value.isdigit(): + raise ValueError("queue threshold is invalid") + return int(value) + + +def _parse_positive_integer(value: str | None, *, default: int) -> int: + parsed = _parse_non_negative_integer(value, default=default) + if parsed <= 0: + raise ValueError("retry-after seconds must be positive") + return parsed + + +__all__ = [ + "MAX_WAITING_REQUESTS_ENVIRONMENT", + "QueueBackpressureController", + "QueueBackpressureSettings", + "QueueDepthBackpressureMiddleware", + "QueueDepthReader", + "QueueSnapshot", + "RETRY_AFTER_SECONDS_ENVIRONMENT", + "read_vllm_queue_depth", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py new file mode 100644 index 000000000..8419eb815 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed one-node step manifest produced inside the sealed client image.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Literal + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.config.environment import ( + LiteralEnvironmentBinding, + SecretRef, + collect_secret_environment_names, +) +from data_designer.slurm.contracts import ContractRecord, ContractValue, validate_absolute_path +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStepRole +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.steps import ( + build_client_command, + build_endpoint_command, + build_vllm_command, +) +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.resolver import resolve_vllm_server +from data_designer.slurm.serving.vllm import ResolvedVllmProcess +from data_designer.slurm.types import EnvironmentName, Identifier, NetworkPort, Sha256Digest + + +class RuntimeProbeSpec(ContractValue): + """One loopback readiness target monitored by the Bash controller.""" + + host: Literal["127.0.0.1"] = "127.0.0.1" + port: NetworkPort + path: str + deadline_seconds: PositiveInt + + @field_validator("path") + @classmethod + def validate_path(cls, value: str) -> str: + if not value.startswith("/") or any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("runtime probe path is invalid") + return value + + +class RuntimeStepSpec(ContractValue): + """Container command and environment consumed by the Bash step runner.""" + + step_id: Identifier + role: RuntimeStepRole + image_path: str + command: tuple[str, ...] = Field(min_length=1) + cpus: PositiveInt + gpu_indices: tuple[NonNegativeInt, ...] = () + literal_environment: dict[EnvironmentName, str] = Field(default_factory=dict) + secret_environment: dict[EnvironmentName, EnvironmentName] = Field(default_factory=dict) + environment_prefixes: dict[EnvironmentName, str] = Field(default_factory=dict) + container_environment: tuple[EnvironmentName, ...] = () + stdout_path: str + stderr_path: str + launch_delay_seconds: NonNegativeInt = 0 + readiness: RuntimeProbeSpec | None = None + + _image_path_is_absolute = field_validator("image_path")(validate_absolute_path) + _stdout_path_is_absolute = field_validator("stdout_path")(validate_absolute_path) + _stderr_path_is_absolute = field_validator("stderr_path")(validate_absolute_path) + + @model_validator(mode="after") + def validate_step(self) -> RuntimeStepSpec: + if any(not argument or "\0" in argument for argument in self.command): + raise ValueError("runtime command is invalid") + if self.gpu_indices != tuple(sorted(set(self.gpu_indices))): + raise ValueError("runtime GPU indices must be sorted and unique") + if self.stdout_path == self.stderr_path or Path(self.stdout_path).parent != Path(self.stderr_path).parent: + raise ValueError("runtime log paths must be distinct siblings") + if set(self.environment_prefixes) - (set(self.literal_environment) | set(self.secret_environment)): + raise ValueError("environment prefixes require a materialized variable") + container_names = set(self.container_environment) + if container_names - (set(self.literal_environment) | set(self.secret_environment)): + raise ValueError("container environment contains an unavailable variable") + if self.role is RuntimeStepRole.SERVER and not self.gpu_indices: + raise ValueError("server runtime steps require GPUs") + if self.role is not RuntimeStepRole.SERVER and self.gpu_indices: + raise ValueError("non-server runtime steps cannot request GPUs") + return self + + +class RuntimeBootstrapManifest(ContractRecord): + """Secret-free one-node allocation command manifest.""" + + run_id: Identifier + shard_id: Identifier + attempt_id: Identifier + plan_sha256: Sha256Digest + all_secret_environment_names: tuple[EnvironmentName, ...] + steps: tuple[RuntimeStepSpec, ...] = Field(min_length=4) + + @model_validator(mode="after") + def validate_steps(self) -> RuntimeBootstrapManifest: + step_ids = tuple(step.step_id for step in self.steps) + if len(step_ids) != len(set(step_ids)): + raise ValueError("runtime step identifiers must be unique") + roles = tuple(step.role for step in self.steps) + if roles.count(RuntimeStepRole.CLIENT_PREFLIGHT) != 1 or roles.count(RuntimeStepRole.CLIENT) != 1: + raise ValueError("runtime manifest requires one preflight and generation step") + if RuntimeStepRole.SERVER not in roles or RuntimeStepRole.ENDPOINT not in roles: + raise ValueError("runtime manifest requires server and endpoint steps") + return self + + +def build_runtime_manifest( + context: AllocationContext, + *, + runtime_root: Path, + log_directory: Path, +) -> RuntimeBootstrapManifest: + """Build the secret-free one-node command handoff for the Bash controller.""" + plan = context.plan + runtime_container_root = get_container_path(plan, runtime_root.as_posix(), require_writable=True) + deployments = tuple(resolve_vllm_server(plan, item.deployment_id) for item in plan.deployments) + endpoints = tuple( + RuntimeEndpoint( + model_alias=deployment.model_alias, + served_model_name=deployment.served_model_name, + host="127.0.0.1", + port=deployment.logical_endpoint.port, + ) + for deployment in deployments + ) + steps: list[RuntimeStepSpec] = [ + _build_client_step( + RuntimeStepRole.CLIENT_PREFLIGHT, + "client-preflight", + "preflight", + context, + endpoints, + runtime_container_root, + log_directory, + ) + ] + for deployment in deployments: + steps.extend( + _build_server_step( + deployment, + process, + context, + runtime_root, + runtime_container_root, + log_directory, + ) + for process in deployment.processes + ) + steps.extend(_build_endpoint_step(deployment, context, runtime_root, log_directory) for deployment in deployments) + steps.append( + _build_client_step( + RuntimeStepRole.CLIENT, + "client-generation", + "client", + context, + endpoints, + runtime_container_root, + log_directory, + ) + ) + secret_names = set(collect_secret_environment_names(plan)) + secret_names.update( + name + for deployment in deployments + for name, binding in deployment.launch_policy.environment.items() + if isinstance(binding, SecretRef) + ) + return RuntimeBootstrapManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + plan_sha256=plan.compute_sha256(), + all_secret_environment_names=tuple(sorted(secret_names)), + steps=tuple(steps), + ) + + +def _build_client_step( + role: RuntimeStepRole, + step_id: str, + operation: str, + context: AllocationContext, + endpoints: tuple[RuntimeEndpoint, ...], + runtime_container_root: str, + log_directory: Path, +) -> RuntimeStepSpec: + plan = context.plan + command = build_client_command( + "preflight" if operation == "preflight" else "run", + plan, + context.shard, + context.attempt, + context.attempt_directory, + endpoints, + ) + if operation == "client": + command = ( + "python3", + "-m", + "data_designer.slurm.runtime.entrypoint", + "client", + *command[4:], + ) + secret_names = collect_secret_environment_names( + (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) + ) + return _step( + step_id=step_id, + role=role, + image_path=plan.client.image.path, + command=command, + cpus=plan.client.authored.cpus, + gpu_indices=(), + literal_environment={"LC_ALL": "C", "PYTHONPATH": runtime_container_root}, + secret_environment={name: name for name in secret_names}, + environment_prefixes={}, + container_environment=tuple(sorted((*secret_names, "PYTHONPATH"))), + log_directory=log_directory, + ) + + +def _build_server_step( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + context: AllocationContext, + runtime_root: Path, + runtime_container_root: str, + log_directory: Path, +) -> RuntimeStepSpec: + if process.pipeline_parallel != 1 or process.node_index != 0 or process.http_port is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "one-node runtime received a distributed vLLM process", + ) + literal_environment: dict[str, str] = {"LC_ALL": "C", "PYTHONPATH": runtime_container_root} + secret_environment: dict[str, str] = {} + environment_prefixes: dict[str, str] = {} + for name, binding in deployment.launch_policy.environment.items(): + if isinstance(binding, LiteralEnvironmentBinding): + literal_environment[name] = binding.value + elif isinstance(binding, SecretRef): + secret_environment[name] = binding.environment + else: # pragma: no cover - persisted contracts reject unknown bindings + raise AssertionError(f"unhandled environment binding: {type(binding)!r}") + if "PYTHONPATH" in secret_environment: + literal_environment.pop("PYTHONPATH") + environment_prefixes["PYTHONPATH"] = runtime_container_root + elif "PYTHONPATH" in deployment.launch_policy.environment: + literal_environment["PYTHONPATH"] = os.pathsep.join((runtime_container_root, literal_environment["PYTHONPATH"])) + policy = deployment.launch_policy.queue_backpressure + literal_environment[MAX_WAITING_REQUESTS_ENVIRONMENT] = str(policy.max_waiting_requests) + literal_environment[RETRY_AFTER_SECONDS_ENVIRONMENT] = ( + "" if policy.retry_after_seconds is None else str(policy.retry_after_seconds) + ) + probe = next(item for item in deployment.readiness_probes if item.port == process.http_port) + return _step( + step_id=process.process_id, + role=RuntimeStepRole.SERVER, + image_path=deployment.image.path, + command=build_vllm_command(deployment, process), + cpus=context.plan.client.authored.cpus, + gpu_indices=tuple(process.gpu_indices), + literal_environment=literal_environment, + secret_environment=secret_environment, + environment_prefixes=environment_prefixes, + container_environment=tuple( + sorted( + { + *deployment.launch_policy.environment, + "PYTHONPATH", + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, + } + ) + ), + log_directory=log_directory, + launch_delay_seconds=process.launch_delay_seconds, + readiness=RuntimeProbeSpec( + port=probe.port, + path=probe.path, + deadline_seconds=probe.deadline_seconds, + ), + ) + + +def _build_endpoint_step( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, + runtime_root: Path, + log_directory: Path, +) -> RuntimeStepSpec: + proxy_path = runtime_root / "data_designer/slurm/runtime/proxy.py" + command = build_endpoint_command( + deployment, + context.plan, + proxy_path, + deployment.logical_endpoint.port, + ) + return _step( + step_id=f"{deployment.deployment_id}-endpoint", + role=RuntimeStepRole.ENDPOINT, + image_path=context.plan.client.image.path, + command=command, + cpus=context.plan.client.authored.cpus, + gpu_indices=(), + literal_environment={"LC_ALL": "C"}, + secret_environment={}, + environment_prefixes={}, + container_environment=(), + log_directory=log_directory, + readiness=RuntimeProbeSpec( + port=deployment.logical_endpoint.port, + path="/health", + deadline_seconds=deployment.launch_policy.startup_timeout_seconds, + ), + ) + + +def _step( + *, + step_id: str, + role: RuntimeStepRole, + image_path: str, + command: tuple[str, ...], + cpus: int, + gpu_indices: tuple[int, ...], + literal_environment: dict[str, str], + secret_environment: dict[str, str], + environment_prefixes: dict[str, str], + container_environment: tuple[str, ...], + log_directory: Path, + launch_delay_seconds: int = 0, + readiness: RuntimeProbeSpec | None = None, +) -> RuntimeStepSpec: + return RuntimeStepSpec( + step_id=step_id, + role=role, + image_path=image_path, + command=command, + cpus=cpus, + gpu_indices=gpu_indices, + literal_environment=literal_environment, + secret_environment=secret_environment, + environment_prefixes=environment_prefixes, + container_environment=container_environment, + stdout_path=(log_directory / f"{step_id}.out").as_posix(), + stderr_path=(log_directory / f"{step_id}.err").as_posix(), + launch_delay_seconds=launch_delay_seconds, + readiness=readiness, + ) + + +__all__ = ["RuntimeBootstrapManifest", "RuntimeProbeSpec", "RuntimeStepSpec", "build_runtime_manifest"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py index 7027c496c..7cf08db2f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py @@ -25,23 +25,9 @@ _SOURCE_MODE = 0o400 _MAXIMUM_SOURCE_SIZE = 16 * 1024 * 1024 _TEMPORARY_NAME_PATTERN = re.compile(r"^\.runtime\.[0-9a-f]{16}\.tmp$") -_ENTRYPOINT_NAME = "entrypoint.sh" +_SHELL_SOURCE_NAMES = ("entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh") _SLURM_PACKAGE_ROOT = "data_designer/slurm" _SOURCE_MANIFEST_NAME = f"{_SLURM_PACKAGE_ROOT}/runtime/slurm-sources.txt" -_ENTRYPOINT = b"""#!/usr/bin/env bash -set -Eeuo pipefail - -dd_slurm_run_allocation() { - if [[ $# -ne 2 ]]; then - printf '%s\\n' 'allocation runtime requires plan and attempt directory arguments' >&2 - return 64 - fi - local runtime_root - runtime_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" - PYTHONNOUSERSITE=1 PYTHONSAFEPATH=1 PYTHONPATH="${runtime_root}" \ - python3 -m data_designer.slurm.runtime.entrypoint --plan "$1" --attempt-dir "$2" -} -""" def stage_runtime_bundle(workspace_root: str | Path) -> ArtifactReference: @@ -69,13 +55,16 @@ def stage_runtime_bundle(workspace_root: str | Path) -> ArtifactReference: def _build_runtime_archive() -> bytes: - sources = _collect_slurm_sources(Path(__file__).parents[1]) + source_root = Path(__file__).parents[1] + sources = _collect_slurm_sources(source_root) output = io.BytesIO() with ( gzip.GzipFile(fileobj=output, mode="wb", filename="", mtime=0) as compressed, tarfile.open(fileobj=compressed, mode="w", format=tarfile.PAX_FORMAT) as archive, ): - _add_archive_file(archive, _ENTRYPOINT_NAME, _ENTRYPOINT, mode=_ENTRYPOINT_MODE) + for name in _SHELL_SOURCE_NAMES: + mode = _ENTRYPOINT_MODE if name == "entrypoint.sh" else _SOURCE_MODE + _add_archive_file(archive, name, _read_runtime_source(source_root / "runtime" / name), mode=mode) manifest = "".join(f"{archive_name}\n" for archive_name, _ in sources).encode() _add_archive_file(archive, _SOURCE_MANIFEST_NAME, manifest, mode=_SOURCE_MODE) for archive_name, source_path in sources: 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 new file mode 100644 index 000000000..c545d4e85 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +DD_MANAGED_PIDS=() +DD_REQUIRED_PIDS=() +DD_CLEANUP_COMPLETE=0 + +dd_start_runtime_timer() { + coproc DD_RUNTIME_TIMER { while IFS= read -r _; do :; done; } +} + +dd_sleep() { + read -r -t "$1" -u "${DD_RUNTIME_TIMER[0]}" _ || true +} + +dd_register_required_pid() { + DD_MANAGED_PIDS+=("$1") + DD_REQUIRED_PIDS+=("$1") +} + +dd_require_running() { + local pid + for pid in "${DD_REQUIRED_PIDS[@]+"${DD_REQUIRED_PIDS[@]}"}"; do + kill -0 "${pid}" 2>/dev/null || return 1 + done +} + +dd_wait_for_client() { + local client_pid=$1 + local status=0 + while kill -0 "${client_pid}" 2>/dev/null; do + wait -n || status=$? + if ! kill -0 "${client_pid}" 2>/dev/null; then + return "${status}" + fi + dd_require_running || return 70 + done + wait "${client_pid}" +} + +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]} + 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]} + 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 + done + DD_CLEANUP_COMPLETE=1 +} + +dd_stop_runtime_timer() { + if [[ ${DD_RUNTIME_TIMER_PID:-} =~ ^[0-9]+$ ]]; then + kill -TERM "${DD_RUNTIME_TIMER_PID}" 2>/dev/null || true + wait "${DD_RUNTIME_TIMER_PID}" 2>/dev/null || true + fi +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py new file mode 100644 index 000000000..c0af8c453 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load allocation identity from container-visible persisted state.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path + +from data_designer.slurm.planning import PlannedShard +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.state import SlurmStateWriter + + +def load_allocation_context( + plan_path: Path, + attempt_directory: Path, + environment: Mapping[str, str], +) -> tuple[AllocationContext, SlurmStateWriter]: + """Load one scheduler-selected shard attempt through its container paths.""" + writer = _load_state_writer(plan_path, attempt_directory) + plan = writer.load_resolved_plan() + expected_plan_path = Path(plan.authored_config.path).with_name("resolved-plan.json") + if plan_path.as_posix() != get_container_path(plan, expected_plan_path.as_posix()): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime plan path does not match persisted run intent", + ) + shard = _select_shard(plan.shards, _scheduler_task_id(environment.get("SLURM_ARRAY_TASK_ID"))) + host_attempt_directory = expected_plan_path.parent / "shards" / shard.shard_id / "attempts" / attempt_directory.name + if attempt_directory.as_posix() != get_container_path( + plan, host_attempt_directory.as_posix(), require_writable=True + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "attempt directory does not match the scheduler array task", + ) + attempt = writer.load_attempt(shard.shard_id, attempt_directory.name) + array_job_id = _scheduler_task_id(environment.get("SLURM_ARRAY_JOB_ID")) + if attempt.scheduler is None or attempt.scheduler.array_job_id != array_job_id: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "scheduler array job does not match the persisted attempt", + ) + return AllocationContext(plan, shard, attempt, host_attempt_directory), writer + + +def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWriter: + if not plan_path.is_absolute() or not attempt_directory.is_absolute(): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime paths must be absolute") + if plan_path.name != "resolved-plan.json" or plan_path.parent.parent.name != "runs": + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") + workspace_root = plan_path.parent.parent.parent + run_id = plan_path.parent.name + return SlurmStateWriter(workspace_root, run_id) + + +def _select_shard(shards: tuple[PlannedShard, ...], task_id: int) -> PlannedShard: + selected = tuple(shard for shard in shards if shard.array_task_index == task_id) + if len(selected) != 1: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "scheduler array task does not identify exactly one planned shard", + ) + return selected[0] + + +def _scheduler_task_id(value: str | None) -> int: + if value is None or not value.isascii() or not value.isdigit(): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "SLURM_ARRAY_TASK_ID must be a non-negative integer", + ) + return int(value) + + +__all__ = ["load_allocation_context"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index 2341d617b..7a77db03a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -7,10 +7,11 @@ import subprocess from collections.abc import Mapping +from contextlib import AbstractContextManager from dataclasses import dataclass from datetime import datetime, timedelta from pathlib import Path -from typing import Protocol +from typing import Literal, Protocol from data_designer.slurm.client import ClientResult from data_designer.slurm.contracts import ArtifactReference @@ -39,6 +40,7 @@ ProbeEvidence, ProbeOutcome, ReadinessState, + ShardWinner, StateNotFoundError, ) @@ -65,6 +67,25 @@ def publish_attempt_result( """Publish one complete result pair and bind it to its running attempt.""" ... + def acquire_dataset_workspace( + self, + shard_id: str, + attempt_id: str, + resume_mode: Literal["never", "always", "if_possible"], + ) -> AbstractContextManager[Path]: + """Hold exclusive ownership of the dataset workspace for generation.""" + ... + + def finalize_winner( + self, + shard_id: str, + attempt_id: str, + *, + published_at: datetime, + ) -> ShardWinner: + """Publish the immutable winning candidate for a successful attempt.""" + ... + def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: """Persist the exact next readiness revision.""" ... @@ -106,6 +127,7 @@ def __init__( context: AllocationContext, *, runtime_proxy_path: Path, + runtime_root: Path | None = None, state: RuntimeStateStore, supervisor: StepSupervisor, preflight: AllocationPreflight, @@ -120,6 +142,7 @@ def __init__( raise ValueError("runtime probe intervals must be positive") self._context = context self._runtime_proxy_path = runtime_proxy_path + self._runtime_root = runtime_root if runtime_root is not None else runtime_proxy_path.parent self._state = state self._supervisor = supervisor self._preflight = preflight @@ -145,6 +168,17 @@ def run(self) -> AttemptManifest: if outcome.failure_cause is outcome.failure: raise outcome.failure raise outcome.failure from outcome.failure_cause + try: + self._state.finalize_winner( + terminal.shard_id, + terminal.attempt_id, + published_at=self._now(), + ) + except BaseException as error: + failure = _normalize_failure(error) + if failure is error: + raise + raise failure from error return terminal def _capture_execution(self) -> _RunOutcome: @@ -297,20 +331,25 @@ def _run_client( self._environment, ) ) - generation_started_at = self._now() - self._supervisor.wait(self._supervisor.start(generation), required=required_processes) - self._supervisor.require_running(required_processes) - client_result, candidate = load_complete_client_candidate(self._context, self._attempt) - self._validate_client_timestamps(candidate.created_at, client_result.completed_at, generation_started_at) - candidate_reference = client_result.candidate_output_manifest - if candidate_reference is None: # pragma: no cover - the record contract requires this for complete results - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.FINALIZATION_FAILED, - "complete client result has no candidate reference", - ) - self._state.publish_attempt_result(client_result, candidate) - self._attempt = _copy_attempt(self._attempt, candidate_output=candidate_reference) - return candidate_reference, client_result.completed_at + with self._state.acquire_dataset_workspace( + self._attempt.shard_id, + self._attempt.attempt_id, + self._context.plan.invocation.authored.resume, + ): + generation_started_at = self._now() + self._supervisor.wait(self._supervisor.start(generation), required=required_processes) + self._supervisor.require_running(required_processes) + client_result, candidate = load_complete_client_candidate(self._context, self._attempt) + self._validate_client_timestamps(candidate.created_at, client_result.completed_at, generation_started_at) + candidate_reference = client_result.candidate_output_manifest + if candidate_reference is None: # pragma: no cover - the record contract requires this for complete results + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "complete client result has no candidate reference", + ) + self._state.publish_attempt_result(client_result, candidate) + self._attempt = _copy_attempt(self._attempt, candidate_output=candidate_reference) + return candidate_reference, client_result.completed_at def _validate_client_timestamps( self, @@ -379,6 +418,7 @@ def _start_servers( self._context.plan, self._context.attempt_directory, self._environment, + self._runtime_root, ) ) for process, step in zip(deployment.processes, steps, strict=True): diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index 30e2fd077..45b28e06c 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 @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Allocation-local command entrypoint loaded from the checksummed runtime bundle.""" +"""Typed control phases executed only inside the sealed client image.""" from __future__ import annotations @@ -9,28 +9,50 @@ import os import sys from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta, timezone from pathlib import Path -from data_designer.slurm.planning import PlannedShard -from data_designer.slurm.runtime.controller import OneNodeAllocationController +from data_designer.slurm.client.filesystem import ensure_private_directory, replace_private_text +from data_designer.slurm.client.worker import main as client_worker_main +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.logs import execution_log_directory from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.paths import get_container_path from data_designer.slurm.runtime.preflight import SystemAllocationPreflight -from data_designer.slurm.runtime.probes import HttpReadinessProber -from data_designer.slurm.runtime.signals import TerminationSignalCoordinator -from data_designer.slurm.runtime.steps import DefaultClientStepBuilder -from data_designer.slurm.runtime.supervisor import StepSupervisor, SubprocessStepRunner, SystemRuntimeClock -from data_designer.slurm.state import SlurmStateWriter +from data_designer.slurm.runtime.records import load_complete_client_candidate +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.resolver import resolve_vllm_server +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + DeploymentReadiness, + EndpointPublicationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, + SlurmStateWriter, + StateNotFoundError, +) def main(arguments: Sequence[str] | None = None) -> int: - """Execute one allocation and return a bounded process status.""" - parser = argparse.ArgumentParser(prog="data-designer-slurm-runtime") - parser.add_argument("--plan", required=True) - parser.add_argument("--attempt-dir", required=True) - parsed = parser.parse_args(arguments) + """Execute one container-only allocation control phase.""" + parsed = _parse_arguments(arguments) try: - _run(Path(parsed.plan), Path(parsed.attempt_dir), os.environ) + if parsed.operation == "prepare": + _prepare(parsed, os.environ) + elif parsed.operation == "ready": + _ready(parsed, os.environ) + elif parsed.operation == "client": + _client(parsed, os.environ) + elif parsed.operation == "succeed": + _succeed(parsed, os.environ) + else: + _fail(parsed, os.environ) except SlurmRuntimeError as error: print(f"allocation runtime failed ({error.code.value}): {error}", file=sys.stderr) return ( @@ -45,77 +67,321 @@ def main(arguments: Sequence[str] | None = None) -> int: return 0 -def _run(plan_path: Path, attempt_directory: Path, environment: Mapping[str, str]) -> None: - context, writer = _load_allocation_context(plan_path, attempt_directory, environment) - clock = SystemRuntimeClock() - signals = TerminationSignalCoordinator() - supervisor = StepSupervisor(SubprocessStepRunner(), signals=signals, clock=clock) - controller = OneNodeAllocationController( +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"): + subparser = subparsers.add_parser(operation) + _add_context_arguments(subparser) + prepare = subparsers.choices["prepare"] + prepare.add_argument("--runtime-root", required=True, type=Path) + prepare.add_argument("--manifest", required=True, type=Path) + client = subparsers.add_parser("client") + _add_context_arguments(client) + client.add_argument("--endpoint", action="append", default=[]) + return parser.parse_args(arguments) + + +def _add_context_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--plan", required=True, type=Path) + parser.add_argument("--attempt-dir", required=True, type=Path) + + +def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + _validate_attempt_is_executable(context.attempt) + SystemAllocationPreflight.verify_attempt_directory(context.attempt_directory) + SystemAllocationPreflight.verify_ports(context) + readiness = _begin_attempt(context, writer) + log_directory = execution_log_directory(context.attempt_directory, readiness.revision) + container_log_directory = Path(get_container_path(context.plan, log_directory.as_posix(), require_writable=True)) + ensure_private_directory(container_log_directory) + manifest = build_runtime_manifest( context, - runtime_proxy_path=Path(__file__).with_name("proxy.py"), - state=writer, - supervisor=supervisor, - preflight=SystemAllocationPreflight(), - client_steps=DefaultClientStepBuilder(), - prober=HttpReadinessProber(), - clock=clock, - environment=environment, + runtime_root=arguments.runtime_root, + log_directory=log_directory, ) - with signals.interrupt_on_termination(supervisor.cleanup): - controller.run() - - -def _load_allocation_context( - plan_path: Path, - attempt_directory: Path, - environment: Mapping[str, str], -) -> tuple[AllocationContext, SlurmStateWriter]: - writer = _load_state_writer(plan_path, attempt_directory) - plan = writer.load_resolved_plan() - expected_plan_path = Path(plan.authored_config.path).with_name("resolved-plan.json") - if plan_path != expected_plan_path: - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.INVALID_CONTEXT, - "runtime plan path does not match persisted run intent", + expected_manifest = context.attempt_directory / "runtime-manifest.json" + if arguments.manifest.as_posix() != get_container_path( + context.plan, + expected_manifest.as_posix(), + require_writable=True, + ): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime manifest path is invalid") + replace_private_text(arguments.manifest, manifest.serialize_json()) + + +def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + previous = writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) + timestamp = _now(context.attempt, previous) + deployments = _resolve_deployments(context) + writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + revision=previous.revision + 1, + updated_at=timestamp, + state=ReadinessState.READY, + deployments=tuple( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.model_alias, + state=ReadinessState.READY, + expected_backends=len(deployment.backend_endpoints), + ready_backends=len(deployment.backend_endpoints), + endpoint_publication=EndpointPublicationState.PUBLISHED, + last_probe=_probe(timestamp, ProbeOutcome.SUCCESS, "endpoint_ready", "endpoint ready"), + ) + for deployment in deployments + ), + ) + ) + + +def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + generation_started_at = _now(context.attempt, _load_optional_readiness(context, writer)) + with writer.acquire_dataset_workspace( + context.shard.shard_id, + context.attempt.attempt_id, + context.plan.invocation.authored.resume, + ): + return_code = client_worker_main( + ( + "run", + "--plan", + arguments.plan.as_posix(), + "--shard-id", + context.shard.shard_id, + "--attempt-id", + context.attempt.attempt_id, + "--attempt-dir", + arguments.attempt_dir.as_posix(), + *(argument for endpoint in arguments.endpoint for argument in ("--endpoint", endpoint)), + ) ) - shard = _select_shard(plan.shards, _scheduler_task_id(environment.get("SLURM_ARRAY_TASK_ID"))) - expected_attempt_root = plan_path.parent / "shards" / shard.shard_id / "attempts" - if attempt_directory.parent != expected_attempt_root: + if return_code != 0: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.CLIENT_FAILED, "client generation failed") + client_result, candidate = load_complete_client_candidate(context, context.attempt) + completed_at = client_result.completed_at + if candidate.created_at < generation_started_at or completed_at < generation_started_at: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client result predates the current generation step", + ) + if completed_at > datetime.now(timezone.utc): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client completion timestamp is later than the allocation clock", + ) + writer.publish_attempt_result(client_result, candidate) + + +def _succeed(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + stopped_at = _write_stopped_readiness(context, writer) + attempt = writer.load_attempt(context.shard.shard_id, context.attempt.attempt_id) + if attempt.candidate_output is None: raise SlurmRuntimeError( - SlurmRuntimeErrorCode.INVALID_CONTEXT, - "attempt directory does not match the scheduler array task", + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "successful allocation has no candidate reference", + ) + terminal = writer.update_attempt( + attempt.model_copy( + update={ + "state": AttemptLifecycleState.SUCCEEDED, + "terminal_classification": AttemptTerminalClassification.SUCCEEDED, + "updated_at": max(stopped_at, attempt.updated_at), + } ) - attempt = writer.load_attempt(shard.shard_id, attempt_directory.name) - return AllocationContext(plan, shard, attempt, attempt_directory), writer + ) + writer.finalize_winner( + terminal.shard_id, + terminal.attempt_id, + published_at=max(datetime.now(timezone.utc), terminal.updated_at), + ) -def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWriter: - if not plan_path.is_absolute() or not attempt_directory.is_absolute(): - raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime paths must be absolute") - if plan_path.name != "resolved-plan.json" or plan_path.parent.parent.name != "runs": - raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") - workspace_root = plan_path.parent.parent.parent - run_id = plan_path.parent.name - return SlurmStateWriter(workspace_root, run_id) +def _fail(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + attempt = writer.load_attempt(context.shard.shard_id, context.attempt.attempt_id) + if attempt.state in {AttemptLifecycleState.SUCCEEDED, AttemptLifecycleState.FAILED}: + return + timestamp = _write_failed_and_stopped_readiness(context, writer) + writer.update_attempt( + attempt.model_copy( + update={ + "state": AttemptLifecycleState.FAILED, + "terminal_classification": AttemptTerminalClassification.FAILED, + "updated_at": max(timestamp, attempt.updated_at), + } + ) + ) -def _select_shard(shards: tuple[PlannedShard, ...], task_id: int) -> PlannedShard: - selected = tuple(shard for shard in shards if shard.array_task_index == task_id) - if len(selected) != 1: - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.INVALID_CONTEXT, - "scheduler array task does not identify exactly one planned shard", +def _begin_attempt(context: AllocationContext, writer: SlurmStateWriter) -> AttemptReadiness: + attempt = context.attempt + previous = _load_optional_readiness(context, writer) + timestamp = _now(attempt, previous) + if attempt.state is not AttemptLifecycleState.RUNNING: + attempt = writer.update_attempt( + attempt.model_copy(update={"state": AttemptLifecycleState.RUNNING, "updated_at": timestamp}) ) - return selected[0] + deployments = _resolve_deployments(context) + initial_state = ReadinessState.RESTARTING if previous is not None else ReadinessState.PENDING + initial = writer.write_readiness(_readiness(context, deployments, previous, initial_state, timestamp)) + return writer.write_readiness( + _readiness(context, deployments, initial, ReadinessState.STARTING, _now(attempt, initial)) + ) -def _scheduler_task_id(value: str | None) -> int: - if value is None or not value.isascii() or not value.isdigit(): - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.INVALID_CONTEXT, - "SLURM_ARRAY_TASK_ID must be a non-negative integer", +def _readiness( + context: AllocationContext, + deployments: tuple[ResolvedVllmServerDeployment, ...], + previous: AttemptReadiness | None, + state: ReadinessState, + timestamp: datetime, +) -> AttemptReadiness: + return AttemptReadiness( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + revision=1 if previous is None else previous.revision + 1, + updated_at=timestamp, + state=state, + deployments=tuple( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.model_alias, + state=state, + expected_backends=len(deployment.backend_endpoints), + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + for deployment in deployments + ), + ) + + +def _write_failed_and_stopped_readiness(context: AllocationContext, writer: SlurmStateWriter) -> datetime: + previous = _load_optional_readiness(context, writer) + if previous is None: + return max(datetime.now(timezone.utc), context.attempt.updated_at) + failed_at = _now(context.attempt, previous) + failed = writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=previous.run_id, + shard_id=previous.shard_id, + attempt_id=previous.attempt_id, + revision=previous.revision + 1, + updated_at=failed_at, + state=ReadinessState.FAILED, + deployments=tuple( + deployment.model_copy( + update={ + "state": ReadinessState.FAILED, + "ready_backends": 0, + "endpoint_publication": ( + EndpointPublicationState.FAILED + if deployment.endpoint_publication is EndpointPublicationState.PENDING + else deployment.endpoint_publication + ), + "last_probe": _probe( + failed_at, + ProbeOutcome.FAILURE, + "allocation_failed", + "allocation runtime failed", + ), + } + ) + for deployment in previous.deployments + ), ) - return int(value) + ) + return _write_stopped_readiness(context, writer, previous=failed) + + +def _write_stopped_readiness( + context: AllocationContext, + writer: SlurmStateWriter, + *, + previous: AttemptReadiness | None = None, +) -> datetime: + previous = previous or writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) + timestamp = _now(context.attempt, previous) + writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=previous.run_id, + shard_id=previous.shard_id, + attempt_id=previous.attempt_id, + revision=previous.revision + 1, + updated_at=timestamp, + state=ReadinessState.STOPPED, + deployments=tuple( + deployment.model_copy( + update={ + "state": ReadinessState.STOPPED, + "ready_backends": 0, + "last_probe": _probe( + timestamp, + ProbeOutcome.SUCCESS, + "runtime_stopped", + "allocation processes stopped", + ), + } + ) + for deployment in previous.deployments + ), + ) + ) + return timestamp + + +def _load_optional_readiness(context: AllocationContext, writer: SlurmStateWriter) -> AttemptReadiness | None: + try: + return writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) + except StateNotFoundError: + return None + + +def _resolve_deployments(context: AllocationContext) -> tuple[ResolvedVllmServerDeployment, ...]: + return tuple(resolve_vllm_server(context.plan, item.deployment_id) for item in context.plan.deployments) + + +def _validate_attempt_is_executable(attempt: AttemptManifest) -> None: + if attempt.state not in { + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + AttemptLifecycleState.RUNNING, + }: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation attempt is not executable") + + +def _probe(observed_at: datetime, outcome: ProbeOutcome, reason_code: str, message: str) -> ProbeEvidence: + return ProbeEvidence( + observed_at=observed_at, + outcome=outcome, + reason_code=reason_code, + redacted_message=message, + ) + + +def _now(attempt: AttemptManifest, readiness: AttemptReadiness | None) -> datetime: + value = datetime.now(timezone.utc) + minimum = attempt.updated_at + if readiness is not None and readiness.updated_at > minimum: + minimum = readiness.updated_at + if value < minimum: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime clock moved backward") + if value.tzinfo is None or value.utcoffset() != timedelta(0): # pragma: no cover - system clock is UTC + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime clock is not UTC") + return value if __name__ == "__main__": # pragma: no cover - exercised through the installed module 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 new file mode 100644 index 000000000..ee5c40100 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -Eeuo pipefail + +readonly DD_RUNTIME_DIR=$(cd -- "${BASH_SOURCE[0]%/*}" && pwd -P) +source "${DD_RUNTIME_DIR}/plan_reader.sh" +source "${DD_RUNTIME_DIR}/step_runner.sh" +source "${DD_RUNTIME_DIR}/cleanup.sh" + +DD_RUNTIME_PREPARED=0 +DD_RUNTIME_FINALIZED=0 + +dd_slurm_run_allocation() { + if [[ $# -ne 2 ]]; then + printf '%s\n' 'allocation runtime requires plan and attempt directory arguments' >&2 + return 64 + fi + DD_PLAN_PATH=$1 + DD_ATTEMPT_PATH=$2 + DD_RUNTIME_MANIFEST=${DD_ATTEMPT_PATH}/runtime-manifest.json + dd_read_control_plan "${DD_PLAN_PATH}" + dd_read_plan_secret_names "${DD_PLAN_PATH}" + dd_read_container_path "${DD_PLAN_PATH}" "${DD_PLAN_PATH}" false + DD_PLAN_CONTAINER_PATH=${DD_CONTAINER_PATH} + dd_read_container_path "${DD_PLAN_PATH}" "${DD_ATTEMPT_PATH}" true + DD_ATTEMPT_CONTAINER_DIR=${DD_CONTAINER_PATH} + dd_read_container_path "${DD_PLAN_PATH}" "${DD_RUNTIME_DIR}" true + DD_RUNTIME_CONTAINER_ROOT=${DD_CONTAINER_PATH} + dd_read_container_path "${DD_PLAN_PATH}" "${DD_RUNTIME_MANIFEST}" true + DD_RUNTIME_MANIFEST_CONTAINER_PATH=${DD_CONTAINER_PATH} + export DD_PLAN_PATH DD_ATTEMPT_PATH DD_RUNTIME_MANIFEST + export DD_PLAN_CONTAINER_PATH DD_ATTEMPT_CONTAINER_DIR DD_RUNTIME_CONTAINER_ROOT + + dd_verify_host_context + dd_start_runtime_timer + trap dd_runtime_exit EXIT + trap 'exit 130' INT TERM + + DD_RUNTIME_PREPARED=1 + dd_run_control_phase prepare \ + --runtime-root "${DD_RUNTIME_DIR}" \ + --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" + dd_verify_runtime_manifest \ + "${DD_RUNTIME_MANIFEST}" \ + "${DD_PLAN_SHA256}" \ + "${DD_SHARD_ID}" \ + "attempt-${DD_ATTEMPT_ORDINAL}" + dd_require_plan_secrets + + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" client_preflight + ((${#DD_STEP_IDS[@]} == 1)) + dd_run_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" + + dd_start_servers + dd_wait_for_role_readiness server + dd_start_endpoints + dd_wait_for_role_readiness endpoint + dd_require_running + dd_run_control_phase ready + + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" client + ((${#DD_STEP_IDS[@]} == 1)) + dd_start_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" + local client_pid=${DD_LAST_PID} + DD_MANAGED_PIDS+=("${client_pid}") + dd_wait_for_client "${client_pid}" + dd_require_running + + dd_cleanup_steps + dd_run_control_phase succeed + DD_RUNTIME_FINALIZED=1 +} + +dd_verify_host_context() { + local tool + for tool in bash sha256sum tar jq srun scontrol getent curl; do + command -v "${tool}" >/dev/null || { + printf 'required allocation tool %q is unavailable\n' "${tool}" >&2 + return 69 + } + done + [[ -d ${DD_ATTEMPT_PATH} && ! -L ${DD_ATTEMPT_PATH} ]] + [[ ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]] + [[ ${SLURM_JOB_NUM_NODES:-} == 1 && ${SLURM_NODEID:-} == 0 ]] + dd_verify_gpu_count + dd_read_artifacts "${DD_PLAN_PATH}" "${SLURM_ARRAY_TASK_ID}" + local index path digest actual + for ((index = 0; index < ${#DD_ARTIFACT_FIELDS[@]}; index += 2)); do + path=${DD_ARTIFACT_FIELDS[index]} + digest=${DD_ARTIFACT_FIELDS[index + 1]} + actual=$(sha256sum < "${path}") + [[ ${actual%% *} == "${digest}" ]] || { + printf 'allocation artifact digest mismatch: %s\n' "${path}" >&2 + return 65 + } + done +} + +dd_verify_gpu_count() { + local visible=${CUDA_VISIBLE_DEVICES:-${SLURM_JOB_GPUS:-}} + local count=0 item + local -A devices=() + if [[ -z ${visible//[[:space:]]/} ]]; then + count=0 + elif [[ ${visible} =~ ^gpu(:[^:]+)?:([0-9]+)$ ]]; then + count=${BASH_REMATCH[2]} + else + IFS=, read -r -a device_values <<<"${visible}" + for item in "${device_values[@]+"${device_values[@]}"}"; do + item=${item//[[:space:]]/} + [[ ${item} =~ ^[A-Za-z0-9_.:-]+$ ]] || return 65 + [[ -n ${item} && ! ${devices[${item}]+_} ]] || return 65 + devices[${item}]=1 + ((count += 1)) + done + fi + [[ ${count} == "${DD_EXPECTED_GPUS}" ]] || { + printf '%s\n' 'allocation GPU visibility does not match the resolved plan' >&2 + return 65 + } +} + +dd_start_servers() { + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" server + local step_id deployment delay remaining + local -A deployment_started=() + for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do + dd_read_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + deployment=${step_id%%-replica-*} + deployment_started[${deployment}]=${deployment_started[${deployment}]:-${SECONDS}} + delay=${DD_STEP_DELAY} + remaining=$((delay - (SECONDS - deployment_started[${deployment}]))) + ((remaining <= 0)) || dd_sleep "${remaining}" + dd_start_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + dd_register_required_pid "${DD_LAST_PID}" + dd_require_running + done +} + +dd_start_endpoints() { + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" endpoint + local step_id + for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do + dd_start_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + dd_register_required_pid "${DD_LAST_PID}" + dd_require_running + done +} + +dd_wait_for_role_readiness() { + local role=$1 + local step_id deadline + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" "${role}" + for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do + dd_read_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + deadline=$((SECONDS + DD_STEP_PROBE_DEADLINE)) + until curl --fail --silent --max-time 1 \ + "http://${DD_STEP_PROBE_HOST}:${DD_STEP_PROBE_PORT}${DD_STEP_PROBE_PATH}" >/dev/null 2>&1; do + dd_require_running + ((SECONDS < deadline)) || { + printf 'runtime step %q readiness timed out\n' "${step_id}" >&2 + return 70 + } + dd_sleep 0.5 + done + done +} + +dd_runtime_exit() { + local status=$? + trap - EXIT INT TERM + set +e + dd_cleanup_steps + if ((DD_RUNTIME_PREPARED == 1 && DD_RUNTIME_FINALIZED == 0)); then + dd_run_control_phase fail >/dev/null + fi + dd_stop_runtime_timer + exit "${status}" +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh new file mode 100644 index 000000000..d9424183d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +dd_read_null_values() { + local array_name=$1 + local index=0 value + unset "${array_name}" + declare -g -a "${array_name}" + printf -v "${array_name}[0]" '%s' "" + unset "${array_name}[0]" + while IFS= read -r -d '' value; do + printf -v "${array_name}[${index}]" '%s' "${value}" + ((index += 1)) + done +} + +dd_read_control_plan() { + local plan=$1 + jq -e ' + .schema_version == 1 + and .client.host_node_index == 0 + and ([.deployments[].node_indices] | all(. == [0])) + and ([.container_mounts[] | (.source + .target)] | all(test("[,:]") | not)) + ' "${plan}" >/dev/null + DD_CLIENT_IMAGE=$(jq -er '.client.image.path' "${plan}") + DD_CLIENT_CPUS=$(jq -er '.client.authored.cpus | tostring' "${plan}") + DD_EXPECTED_GPUS=$(jq -er '.resolved_gpus_per_node | tostring' "${plan}") + DD_GPU_REQUEST_MODE=$(jq -er '.selected_profile.profile.gpu_request_mode' "${plan}") + DD_CONTAINER_MOUNTS=$(jq -jr ' + [.container_mounts[] | .source + ":" + .target + (if .read_only then ":ro" else "" end)] + | join(",") + ' "${plan}") +} + +dd_read_container_path() { + local plan=$1 + local host_path=$2 + local require_writable=$3 + DD_CONTAINER_PATH=$(jq -er \ + --arg path "${host_path}" \ + --argjson require_writable "${require_writable}" ' + [ + .container_mounts[] + | . as $mount + | select(($path == $mount.source) or ($path | startswith($mount.source + "/"))) + | select(($require_writable | not) or ($mount.read_only | not)) + ] + | sort_by(.source | length) + | last + | select(. != null) + | . as $mount + | if $path == $mount.source then + $mount.target + else + $mount.target + ($path | ltrimstr($mount.source)) + end + ' "${plan}") +} + +dd_read_artifacts() { + local plan=$1 + local task_id=$2 + dd_read_null_values DD_ARTIFACT_FIELDS < <( + jq -j --argjson task_id "${task_id}" ' + [ + .runtime_bundle, + .client.dependency_lock, + {path: .client.image.path, sha256: .client.image.sha256}, + (.deployments[] | {path: .image.path, sha256: .image.sha256}), + .builder.source, + (.shards[] | select(.array_task_index == $task_id) | .input_partition) + ] + | map(select(. != null)) + | unique_by([.path, .sha256]) + | .[] + | .path, "\u0000", .sha256, "\u0000" + ' "${plan}" + ) +} + +dd_read_plan_secret_names() { + local plan=$1 + dd_read_null_values DD_REQUIRED_SECRET_NAMES < <( + jq -j '[.. | objects | select(.type? == "secret") | .environment] | unique | .[] | ., "\u0000"' "${plan}" + ) + dd_read_null_values DD_ALL_SECRET_NAMES < <( + jq -j ' + [ + (.. | objects | select(.type? == "secret") | .environment), + ( + .deployments[].authored.server.environment + | to_entries[] + | select(.value.type == "secret") + | .key + ) + ] + | unique + | .[] + | ., "\u0000" + ' "${plan}" + ) +} + +dd_verify_runtime_manifest() { + local manifest=$1 + local plan_sha256=$2 + local shard_id=$3 + local attempt_id=$4 + jq -e \ + --arg plan_sha256 "${plan_sha256}" \ + --arg shard_id "${shard_id}" \ + --arg attempt_id "${attempt_id}" ' + .schema_version == 1 + and .plan_sha256 == $plan_sha256 + and .shard_id == $shard_id + and .attempt_id == $attempt_id + and ([.steps[].step_id] | length == (unique | length)) + ' "${manifest}" >/dev/null +} + +dd_read_step_ids() { + local manifest=$1 + local role=$2 + dd_read_null_values DD_STEP_IDS < <( + jq -j --arg role "${role}" '.steps[] | select(.role == $role) | .step_id, "\u0000"' "${manifest}" + ) +} + +dd_read_step() { + local manifest=$1 + local step_id=$2 + jq -e --arg step_id "${step_id}" '[.steps[] | select(.step_id == $step_id)] | length == 1' \ + "${manifest}" >/dev/null + dd_read_null_values DD_STEP_FIELDS < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .image_path, "\u0000", + (.cpus | tostring), "\u0000", + .stdout_path, "\u0000", + .stderr_path, "\u0000", + (.launch_delay_seconds | tostring), "\u0000", + (.readiness.host // ""), "\u0000", + (.readiness.port // "" | tostring), "\u0000", + (.readiness.path // ""), "\u0000", + (.readiness.deadline_seconds // "" | tostring), "\u0000" + ' "${manifest}" + ) + DD_STEP_IMAGE=${DD_STEP_FIELDS[0]} + DD_STEP_CPUS=${DD_STEP_FIELDS[1]} + DD_STEP_STDOUT=${DD_STEP_FIELDS[2]} + DD_STEP_STDERR=${DD_STEP_FIELDS[3]} + DD_STEP_DELAY=${DD_STEP_FIELDS[4]} + DD_STEP_PROBE_HOST=${DD_STEP_FIELDS[5]} + DD_STEP_PROBE_PORT=${DD_STEP_FIELDS[6]} + DD_STEP_PROBE_PATH=${DD_STEP_FIELDS[7]} + DD_STEP_PROBE_DEADLINE=${DD_STEP_FIELDS[8]} + dd_read_null_values DD_STEP_COMMAND < <( + jq -j --arg step_id "${step_id}" '.steps[] | select(.step_id == $step_id) | .command[] | ., "\u0000"' \ + "${manifest}" + ) + dd_read_null_values DD_STEP_GPU_INDICES < <( + jq -j --arg step_id "${step_id}" \ + '.steps[] | select(.step_id == $step_id) | .gpu_indices[] | tostring, "\u0000"' "${manifest}" + ) + dd_read_null_values DD_STEP_CONTAINER_ENV < <( + jq -j --arg step_id "${step_id}" \ + '.steps[] | select(.step_id == $step_id) | .container_environment[] | ., "\u0000"' "${manifest}" + ) + dd_read_null_values DD_STEP_LITERAL_ENV < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .literal_environment + | to_entries[] + | .key, "\u0000", .value, "\u0000" + ' "${manifest}" + ) + dd_read_null_values DD_STEP_SECRET_ENV < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .secret_environment + | to_entries[] + | .key, "\u0000", .value, "\u0000" + ' "${manifest}" + ) + dd_read_null_values DD_STEP_ENV_PREFIXES < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .environment_prefixes + | to_entries[] + | .key, "\u0000", .value, "\u0000" + ' "${manifest}" + ) +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py index 476d62a07..e4c5d8872 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -38,10 +38,10 @@ def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> """Verify every launch-critical fact before model services start.""" try: self._verify_scheduler(context, environment) - self._verify_attempt_directory(context.attempt_directory) + self.verify_attempt_directory(context.attempt_directory) get_container_path(context.plan, context.attempt_directory.as_posix(), require_writable=True) self._verify_artifacts(context) - self._verify_ports(context) + self.verify_ports(context) except SlurmRuntimeError: raise except (OSError, ValueError) as error: @@ -81,7 +81,8 @@ def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str] ) @staticmethod - def _verify_attempt_directory(attempt_directory: Path) -> None: + def verify_attempt_directory(attempt_directory: Path) -> None: + """Require an attempt workspace accessible only to its owner.""" status = attempt_directory.lstat() if not stat.S_ISDIR(status.st_mode) or status.st_mode & 0o077: raise SlurmRuntimeError( @@ -115,7 +116,8 @@ def _verify_artifacts(context: AllocationContext) -> None: _verify_artifact(reference) @staticmethod - def _verify_ports(context: AllocationContext) -> None: + def verify_ports(context: AllocationContext) -> None: + """Verify that every planned one-node port is currently bindable.""" ports = tuple(port.port for port in context.plan.client.ports) + tuple( port.port for deployment in context.plan.deployments for port in deployment.ports ) 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 new file mode 100644 index 000000000..4cd7bf045 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +dd_validate_environment_name() { + [[ $1 =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]] +} + +dd_require_plan_secrets() { + local name + for name in "${DD_REQUIRED_SECRET_NAMES[@]+"${DD_REQUIRED_SECRET_NAMES[@]}"}"; do + dd_validate_environment_name "${name}" || return 64 + [[ ${!name+x} ]] || { + printf 'required secret environment %q is unavailable\n' "${name}" >&2 + return 78 + } + done +} + +dd_scope_control_environment() { + local name + for name in "${DD_ALL_SECRET_NAMES[@]+"${DD_ALL_SECRET_NAMES[@]}"}"; do + unset "${name}" + done + unset CUDA_VISIBLE_DEVICES + export LC_ALL=C + export PYTHONPATH=${DD_RUNTIME_CONTAINER_ROOT} +} + +dd_materialize_step_environment() { + local index name source value prefix current + local -a secret_values=() + for ((index = 0; index < ${#DD_STEP_SECRET_ENV[@]}; index += 2)); do + name=${DD_STEP_SECRET_ENV[index]} + source=${DD_STEP_SECRET_ENV[index + 1]} + dd_validate_environment_name "${name}" && dd_validate_environment_name "${source}" || return 64 + [[ ${!source+x} ]] || return 78 + secret_values+=("${!source}") + done + for name in "${DD_ALL_SECRET_NAMES[@]+"${DD_ALL_SECRET_NAMES[@]}"}"; do + unset "${name}" + done + unset CUDA_VISIBLE_DEVICES + for ((index = 0; index < ${#DD_STEP_LITERAL_ENV[@]}; index += 2)); do + name=${DD_STEP_LITERAL_ENV[index]} + value=${DD_STEP_LITERAL_ENV[index + 1]} + dd_validate_environment_name "${name}" || return 64 + printf -v "${name}" '%s' "${value}" + export "${name}" + done + for ((index = 0; index < ${#DD_STEP_SECRET_ENV[@]}; index += 2)); do + name=${DD_STEP_SECRET_ENV[index]} + printf -v "${name}" '%s' "${secret_values[index / 2]}" + export "${name}" + done + for ((index = 0; index < ${#DD_STEP_ENV_PREFIXES[@]}; index += 2)); do + name=${DD_STEP_ENV_PREFIXES[index]} + prefix=${DD_STEP_ENV_PREFIXES[index + 1]} + dd_validate_environment_name "${name}" || return 64 + current=${!name} + printf -v "${name}" '%s' "${prefix}:${current}" + export "${name}" + done + if [[ -n ${DD_STEP_VISIBLE_GPUS} ]]; then + export CUDA_VISIBLE_DEVICES=${DD_STEP_VISIBLE_GPUS} + fi +} + +dd_build_srun_command() { + local index gpu_mask=0 container_names= + DD_STEP_VISIBLE_GPUS= + DD_SRUN_COMMAND=( + srun + --nodes=1 + --ntasks=1 + --exact + --overlap + --unbuffered + --export=ALL + "--cpus-per-task=${DD_STEP_CPUS}" + "--container-image=${DD_STEP_IMAGE}" + ) + if ((${#DD_STEP_GPU_INDICES[@]})); then + if [[ ${DD_GPU_REQUEST_MODE} == gres ]]; then + for index in "${DD_STEP_GPU_INDICES[@]+"${DD_STEP_GPU_INDICES[@]}"}"; do + ((gpu_mask |= 1 << index)) + done + printf -v gpu_mask '0x%x' "${gpu_mask}" + DD_SRUN_COMMAND+=( + "--gpus-per-task=${#DD_STEP_GPU_INDICES[@]}" + "--gpu-bind=mask_gpu:${gpu_mask}" + ) + else + local visible_gpus + printf -v visible_gpus '%s,' "${DD_STEP_GPU_INDICES[@]}" + DD_STEP_VISIBLE_GPUS=${visible_gpus%,} + DD_STEP_CONTAINER_ENV+=(CUDA_VISIBLE_DEVICES) + fi + else + DD_SRUN_COMMAND+=(--gres=none) + fi + [[ -z ${DD_CONTAINER_MOUNTS} ]] || DD_SRUN_COMMAND+=("--container-mounts=${DD_CONTAINER_MOUNTS}") + if ((${#DD_STEP_CONTAINER_ENV[@]})); then + printf -v container_names '%s,' "${DD_STEP_CONTAINER_ENV[@]}" + DD_SRUN_COMMAND+=("--container-env=${container_names%,}") + fi +} + +dd_start_step() { + local manifest=$1 + local step_id=$2 + dd_read_step "${manifest}" "${step_id}" + dd_build_srun_command + ( + dd_materialize_step_environment + exec "${DD_SRUN_COMMAND[@]}" -- "${DD_STEP_COMMAND[@]}" + ) >"${DD_STEP_STDOUT}" 2>"${DD_STEP_STDERR}" & + DD_LAST_PID=$! +} + +dd_run_step() { + local manifest=$1 + local step_id=$2 + dd_start_step "${manifest}" "${step_id}" + local pid=${DD_LAST_PID} + local index=${#DD_MANAGED_PIDS[@]} + local status=0 + DD_MANAGED_PIDS+=("${pid}") + wait "${pid}" || status=$? + unset "DD_MANAGED_PIDS[${index}]" + DD_MANAGED_PIDS=("${DD_MANAGED_PIDS[@]+"${DD_MANAGED_PIDS[@]}"}") + return "${status}" +} + +dd_run_control_phase() { + local operation=$1 + shift + local -a command=( + srun + --nodes=1 + --ntasks=1 + --exact + --overlap + --unbuffered + --export=ALL + "--cpus-per-task=${DD_CLIENT_CPUS}" + "--container-image=${DD_CLIENT_IMAGE}" + --gres=none + ) + [[ -z ${DD_CONTAINER_MOUNTS} ]] || command+=("--container-mounts=${DD_CONTAINER_MOUNTS}") + command+=( + --container-env=PYTHONPATH + -- + python3 + -m + data_designer.slurm.runtime.entrypoint + "${operation}" + --plan + "${DD_PLAN_CONTAINER_PATH}" + --attempt-dir + "${DD_ATTEMPT_CONTAINER_DIR}" + "$@" + ) + ( + dd_scope_control_environment + exec "${command[@]}" + ) +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py index ed998331a..77ee2346b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -10,10 +10,16 @@ from pathlib import Path from typing import Protocol -from pydantic import BaseModel - -from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.config.environment import ( + LiteralEnvironmentBinding, + SecretRef, + collect_secret_environment_names, +) from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStep, RuntimeStepRole from data_designer.slurm.runtime.paths import get_container_path @@ -125,7 +131,7 @@ def _build_step( endpoints: tuple[RuntimeEndpoint, ...], source_environment: Mapping[str, str], ) -> RuntimeStep: - command = _build_client_command(operation, plan, shard, attempt, attempt_directory, endpoints) + command = build_client_command(operation, plan, shard, attempt, attempt_directory, endpoints) secret_names, environment = _build_client_environment(plan, source_environment) return _build_srun_step( step_id=step_id, @@ -139,7 +145,7 @@ def _build_step( ) -def _build_client_command( +def build_client_command( operation: str, plan: ResolvedSlurmRunPlan, shard: PlannedShard, @@ -191,10 +197,11 @@ def build_vllm_steps( plan: ResolvedSlurmRunPlan, attempt_directory: Path, source_environment: Mapping[str, str], + runtime_root: Path, ) -> tuple[RuntimeStep, ...]: """Build one structured server step per resolved vLLM process.""" return tuple( - _build_vllm_step(deployment, process, plan, attempt_directory, source_environment) + _build_vllm_step(deployment, process, plan, attempt_directory, source_environment, runtime_root) for process in deployment.processes ) @@ -226,7 +233,7 @@ def _build_endpoint_step( host="127.0.0.1", port=deployment.logical_endpoint.port, ) - command = _build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port) + command = build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port) step = _build_srun_step( step_id=f"{deployment.deployment_id}-endpoint", role=RuntimeStepRole.ENDPOINT, @@ -240,7 +247,7 @@ def _build_endpoint_step( return step, endpoint -def _build_endpoint_command( +def build_endpoint_command( deployment: ResolvedVllmServerDeployment, plan: ResolvedSlurmRunPlan, runtime_proxy_path: Path, @@ -270,10 +277,16 @@ def _build_vllm_step( plan: ResolvedSlurmRunPlan, attempt_directory: Path, source_environment: Mapping[str, str], + runtime_root: Path, ) -> RuntimeStep: _validate_local_vllm_process(process) - command = _build_vllm_command(deployment, process) - environment, container_environment = _build_vllm_environment(deployment, source_environment) + command = build_vllm_command(deployment, process) + environment, container_environment = _build_vllm_environment( + deployment, + source_environment, + plan, + runtime_root, + ) return _build_srun_step( step_id=process.process_id, role=RuntimeStepRole.SERVER, @@ -295,7 +308,7 @@ def _validate_local_vllm_process(process: ResolvedVllmProcess) -> None: ) -def _build_vllm_command( +def build_vllm_command( deployment: ResolvedVllmServerDeployment, process: ResolvedVllmProcess, ) -> tuple[str, ...]: @@ -313,6 +326,8 @@ def _build_vllm_command( str(process.http_port), "--tensor-parallel-size", str(process.tensor_parallel), + "--middleware", + "data_designer.slurm.runtime.backpressure.QueueDepthBackpressureMiddleware", ) if deployment.launch_policy.enable_expert_parallel: command += ("--enable-expert-parallel",) @@ -322,6 +337,8 @@ def _build_vllm_command( def _build_vllm_environment( deployment: ResolvedVllmServerDeployment, source_environment: Mapping[str, str], + plan: ResolvedSlurmRunPlan, + runtime_root: Path, ) -> tuple[dict[str, str], tuple[str, ...]]: environment = _base_environment(source_environment) container_environment: list[str] = [] @@ -339,6 +356,17 @@ def _build_vllm_environment( else: # pragma: no cover - persisted contracts reject unknown bindings raise AssertionError(f"unhandled environment binding: {type(binding)!r}") container_environment.append(name) + runtime_pythonpath = get_container_path(plan, runtime_root.as_posix()) + configured_pythonpath = environment.get("PYTHONPATH") + environment["PYTHONPATH"] = ( + os.pathsep.join((runtime_pythonpath, configured_pythonpath)) if configured_pythonpath else runtime_pythonpath + ) + queue_policy = deployment.launch_policy.queue_backpressure + environment[MAX_WAITING_REQUESTS_ENVIRONMENT] = str(queue_policy.max_waiting_requests) + environment[RETRY_AFTER_SECONDS_ENVIRONMENT] = ( + "" if queue_policy.retry_after_seconds is None else str(queue_policy.retry_after_seconds) + ) + container_environment.extend(("PYTHONPATH", MAX_WAITING_REQUESTS_ENVIRONMENT, RETRY_AFTER_SECONDS_ENVIRONMENT)) return environment, tuple(container_environment) @@ -438,21 +466,18 @@ def _base_environment(source_environment: Mapping[str, str]) -> dict[str, str]: def _collect_client_secret_environment_names(plan: ResolvedSlurmRunPlan) -> tuple[str, ...]: - names: set[str] = set() - _collect_secret_environment_names(plan.client.authored.dependencies.index_credentials, names) - _collect_secret_environment_names(plan.invocation.authored.mcp_providers, names) - return tuple(sorted(names)) - - -def _collect_secret_environment_names(value: object, names: set[str]) -> None: - if isinstance(value, SecretRef): - names.add(value.environment) - elif isinstance(value, BaseModel): - for field_name in type(value).model_fields: - _collect_secret_environment_names(getattr(value, field_name), names) - elif isinstance(value, Mapping): - for child in value.values(): - _collect_secret_environment_names(child, names) - elif isinstance(value, (tuple, list)): - for child in value: - _collect_secret_environment_names(child, names) + return collect_secret_environment_names( + (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) + ) + + +__all__ = [ + "ClientStepBuilder", + "DefaultClientStepBuilder", + "build_client_command", + "build_endpoint_command", + "build_endpoint_steps", + "build_vllm_command", + "build_vllm_steps", + "plan_path", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py new file mode 100644 index 000000000..9211898cd --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Publish immutable inputs and initial state for one submitted run.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Callable +from datetime import datetime +from pathlib import Path + +from pydantic import JsonValue + +from data_designer.slurm.client.dependencies import ResolvedClientDependencies +from data_designer.slurm.client.filesystem import ensure_private_directory +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import ArtifactReference, pretty_json +from data_designer.slurm.filesystem import create_restrictive_temporary_file, get_file_facts +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + SchedulerIdentity, + ShardManifest, + SlurmStateError, + SlurmStateWriter, + StateConflictError, +) +from data_designer.slurm.state.filesystem import ( + open_verified_directory, + open_verified_regular_file, + publish_immutable_text, + sync_directory, +) + +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 +_TEMPORARY_PREFIX = ".artifact." +_TEMPORARY_SUFFIX = ".tmp" + + +class StateRunArtifactPublisher: + """Persist submission inputs through the package-owned state workspace.""" + + def __init__(self, workspace_root: str | Path, clock: Callable[[], datetime]) -> None: + self._workspace_root = Path(workspace_root) + self._clock = clock + + def initialize( + self, + authored: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + dependencies: ResolvedClientDependencies, + builder_payload: dict[str, JsonValue] | None, + *, + force: bool, + ) -> None: + if force: + raise StateConflictError("force cannot replace durable run state") + created_at = self._clock() + writer = SlurmStateWriter(self._workspace_root, plan.run_id) + plan_reference = ArtifactReference( + path=(writer.run_root / "resolved-plan.json").as_posix(), + sha256=plan.compute_sha256(), + ) + run = RunManifest( + schema_version=1, + run_id=plan.run_id, + created_at=created_at, + authored_config=plan.authored_config, + resolved_plan=plan_reference, + shard_count=len(plan.shards), + ) + shards = tuple( + ShardManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + shard_index=shard.shard_index, + record_range=shard.record_range, + input_partition=shard.input_partition, + resume_workspace=shard.resume_workspace, + created_at=created_at, + ) + for shard in plan.shards + ) + writer.initialize_run(authored, plan, run, shards) + try: + self._publish_inputs(writer.run_root, plan, dependencies, builder_payload) + except SlurmStateError: + raise + except Exception as error: + raise SlurmStateError(f"cannot publish immutable inputs for run {plan.run_id!r}") from error + + def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitted_at: datetime) -> None: + writer = SlurmStateWriter(self._workspace_root, plan.run_id) + plan_reference = ArtifactReference( + path=(writer.run_root / "resolved-plan.json").as_posix(), + sha256=plan.compute_sha256(), + ) + for shard in plan.shards: + writer.create_attempt( + AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=plan_reference, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=job_id, array_task_id=shard.array_task_index), + created_at=submitted_at, + updated_at=submitted_at, + ) + ) + + def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: + """Mark every initial attempt failed after its held job is cancelled.""" + writer = SlurmStateWriter(self._workspace_root, plan.run_id) + for shard in plan.shards: + attempt = writer.load_attempt(shard.shard_id, "attempt-0001") + writer.update_attempt( + attempt.model_copy( + update={ + "state": AttemptLifecycleState.FAILED, + "terminal_classification": AttemptTerminalClassification.CANCELLED, + "updated_at": failed_at, + } + ) + ) + + @staticmethod + def _publish_inputs( + run_root: Path, + plan: ResolvedSlurmRunPlan, + dependencies: ResolvedClientDependencies, + builder_payload: dict[str, JsonValue] | None, + ) -> None: + _publish_text(run_root, plan.client.dependency_lock, dependencies.lock.serialize_json()) + if (plan.builder.source is None) != (builder_payload is None): + raise SlurmStateError("resolved builder source does not match its staged payload") + if plan.builder.source is not None: + assert builder_payload is not None + _publish_text(run_root, plan.builder.source, pretty_json(builder_payload)) + if (dependencies.lock.source is None) != (dependencies.lock_source is None): + raise SlurmStateError("dependency lock source does not match its staged artifact") + if dependencies.lock.source is not None: + assert dependencies.lock_source is not None + _publish_file(run_root, dependencies.lock.source, dependencies.lock_source) + packages = dependencies.lock.overlay_packages + if len(packages) != len(dependencies.wheel_sources): + raise SlurmStateError("dependency wheel sources do not match the resolved lock") + for package, source in zip(packages, dependencies.wheel_sources, strict=True): + _publish_file(run_root, package.artifact, source) + seed_path = plan.invocation.effective_input_bindings.seed_path + for shard in plan.shards: + if shard.input_partition is None: + continue + _publish_text( + run_root, + shard.input_partition, + pretty_json( + { + "record_range": shard.record_range.model_dump(mode="json"), + "seed_path": seed_path, + } + ), + ) + + +def _publish_text(run_root: Path, reference: ArtifactReference, content: str) -> None: + target = _validate_target(run_root, reference) + if hashlib.sha256(content.encode()).hexdigest() != reference.sha256: + raise SlurmStateError(f"staged artifact {target.name!r} does not match its resolved digest") + ensure_private_directory(target.parent) + with open_verified_directory(target.parent, require_private=True) as descriptor: + publish_immutable_text( + descriptor, + target.name, + content, + target, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + +def _publish_file(run_root: Path, reference: ArtifactReference, source: Path) -> None: + target = _validate_target(run_root, reference) + ensure_private_directory(target.parent) + source_descriptor: int | None = None + try: + source_before = source.lstat() + source_descriptor = os.open( + source, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + ) + source_opened = os.fstat(source_descriptor) + if not stat.S_ISREG(source_opened.st_mode) or get_file_facts(source_before) != get_file_facts(source_opened): + raise OSError(f"source artifact {source} is not a stable regular file") + with open_verified_directory(target.parent, require_private=True) as target_descriptor: + output_descriptor, temporary_name = create_restrictive_temporary_file( + target_descriptor, + prefix=_TEMPORARY_PREFIX, + suffix=_TEMPORARY_SUFFIX, + ) + try: + digest = hashlib.sha256() + try: + while chunk := os.read(source_descriptor, 1024 * 1024): + digest.update(chunk) + remaining = memoryview(chunk) + while remaining: + written = os.write(output_descriptor, remaining) + if written == 0: + raise OSError("artifact copy made no progress") + remaining = remaining[written:] + os.fsync(output_descriptor) + finally: + os.close(output_descriptor) + source_after = os.fstat(source_descriptor) + source_current = source.lstat() + if ( + get_file_facts(source_opened) != get_file_facts(source_after) + or get_file_facts(source_after) != get_file_facts(source_current) + or digest.hexdigest() != reference.sha256 + ): + raise OSError(f"source artifact {source} changed or has an unexpected digest") + try: + os.link( + temporary_name, + target.name, + src_dir_fd=target_descriptor, + dst_dir_fd=target_descriptor, + follow_symlinks=False, + ) + except FileExistsError: + with open_verified_regular_file( + target_descriptor, + target.name, + target, + expected_size=source_after.st_size, + expected_sha256=reference.sha256, + ): + pass + os.unlink(temporary_name, dir_fd=target_descriptor) + temporary_name = None + sync_directory(target_descriptor) + finally: + if temporary_name is not None: + try: + os.unlink(temporary_name, dir_fd=target_descriptor) + except OSError: + pass + finally: + if source_descriptor is not None: + os.close(source_descriptor) + + +def _validate_target(run_root: Path, reference: ArtifactReference) -> Path: + target = Path(reference.path) + try: + target.relative_to(run_root) + except ValueError: + raise SlurmStateError("resolved artifact path is outside the run workspace") from None + return target + + +__all__ = ["StateRunArtifactPublisher"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index ddda16ea7..a2ac391f1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -6,7 +6,8 @@ from __future__ import annotations import importlib.metadata -from collections.abc import Callable, Iterator +import os +from collections.abc import Callable, Iterator, Mapping from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import UTC, datetime @@ -30,6 +31,7 @@ SlurmConfigLoadError, SlurmProfile, SlurmProfileCatalog, + collect_secret_environment_names, load_builder_payload, resolve_profile, ) @@ -46,6 +48,7 @@ from data_designer.slurm.planning.resolution import resolve_slurm_config from data_designer.slurm.runtime.bundle import stage_runtime_bundle from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.services.artifacts import StateRunArtifactPublisher from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation from data_designer.slurm.services.images import SlurmImageService from data_designer.slurm.services.results import ( @@ -88,6 +91,9 @@ def initialize( def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitted_at: datetime) -> None: """Persist the submitted scheduler identity for every initial attempt.""" + def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: + """Mark initial attempts failed after a held submission is cancelled.""" + @dataclass(frozen=True, slots=True) class _PreparedRun: @@ -237,14 +243,16 @@ def __init__( preparer: _RunPreparer, selected_profile: SelectedSlurmProfile, launcher: SlurmCommandClient, - publisher: SlurmRunArtifactPublisher | None, + publisher: SlurmRunArtifactPublisher, clock: Clock, + source_environment: Mapping[str, str], ) -> None: self._preparer = preparer self._profile = selected_profile self._launcher = launcher self._publisher = publisher self._clock = clock + self._source_environment = source_environment def execute( self, @@ -268,12 +276,7 @@ def execute( shard_count=len(plan.shards), batch_script=prepared.batch_script, ) - if self._publisher is None: - raise SlurmServiceError( - SlurmServiceErrorCode.UNAVAILABLE, - SlurmServiceOperation.EXECUTE_RUN, - "run submission is not available; use --dry-run", - ) + export_environment = self._build_export_environment(config) publisher = self._publisher self._initialize_run( publisher, @@ -282,7 +285,11 @@ def execute( force=force, ) try: - receipt = self._launcher.submit_script(prepared.batch_script) + receipt = self._launcher.submit_script( + prepared.batch_script, + hold=True, + export_environment=export_environment, + ) except SlurmLauncherError: raise SlurmServiceError( SlurmServiceErrorCode.UNAVAILABLE, @@ -307,6 +314,30 @@ def execute( except Exception: pass raise + try: + self._launcher.release(receipt.job_id) + except SlurmLauncherError as error: + try: + self._launcher.cancel(receipt.job_id) + except Exception: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + f"held Slurm job {receipt.job_id} could not be released or cancelled", + ) from error + try: + self._record_submission_failure(publisher, plan) + except SlurmServiceError as state_error: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.EXECUTE_RUN, + f"held Slurm job {receipt.job_id} was cancelled but run {plan.run_id!r} could not be updated", + ) from state_error + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + f"held Slurm job {receipt.job_id} could not be released and was cancelled", + ) from None return SlurmRunExecution( run_id=plan.run_id, state="submitted", @@ -315,6 +346,35 @@ def execute( job_id=receipt.job_id, ) + def _build_export_environment(self, config: DataDesignerSlurmConfig) -> dict[str, str]: + environment = {"SLURM_EXPORT_ENV": "ALL"} + if "SLURM_CONF" in self._source_environment: + slurm_conf = self._source_environment["SLURM_CONF"] + if type(slurm_conf) is not str or "\0" in slurm_conf: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.EXECUTE_RUN, + "SLURM_CONF is invalid", + ) + environment["SLURM_CONF"] = slurm_conf + for name in collect_secret_environment_names(config): + try: + value = self._source_environment[name] + except KeyError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.EXECUTE_RUN, + f"required secret environment {name!r} is unavailable", + ) from None + if type(value) is not str or "\0" in value: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.EXECUTE_RUN, + f"required secret environment {name!r} is invalid", + ) + environment[name] = value + return environment + def _initialize_run( self, publisher: SlurmRunArtifactPublisher, @@ -377,6 +437,16 @@ def _record_submission( "submission state cannot be recorded", ) from None + def _record_submission_failure(self, publisher: SlurmRunArtifactPublisher, plan: ResolvedSlurmRunPlan) -> None: + try: + publisher.record_submission_failure(plan, failed_at=self._clock()) + except (StateNotFoundError, StateConflictError, SlurmStateError): + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.EXECUTE_RUN, + "cancelled submission state cannot be recorded", + ) from None + def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: operation = SlurmServiceOperation.STATUS_RUN try: @@ -502,10 +572,12 @@ def create_slurm_run_service( run_id_factory: RunIdFactory | None = None, clock: Clock | None = None, package_version: str | None = None, + source_environment: Mapping[str, str] | None = None, ) -> SlurmRunService: """Create the production run service for one selected cluster profile.""" selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) command_client = launcher or SlurmCommandClient() + selected_clock = clock or _utc_now preparer = _RunPreparer( selected, VerifiedImageRegistry(selected.profile.workspace_root), @@ -518,8 +590,9 @@ def create_slurm_run_service( preparer, selected, command_client, - artifact_publisher, - clock or _utc_now, + artifact_publisher or StateRunArtifactPublisher(selected.profile.workspace_root, selected_clock), + selected_clock, + dict(os.environ if source_environment is None else source_environment), ) return SlurmRunService(_SystemRunPlanner(preparer), render_generation_attempt_script, backend) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py index 474faaeea..83f0576a4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py @@ -53,10 +53,10 @@ def acquire_dataset_workspace( self, shard_id: ShardId, attempt_id: AttemptId, - effective_resume_mode: Literal["never", "always"], + resume_mode: Literal["never", "always", "if_possible"], ) -> Iterator[Path]: with self._storage.acquire_resume_lock(shard_id): - dataset_path = self._prepare_workspace_or_normalize(shard_id, attempt_id, effective_resume_mode) + dataset_path = self._prepare_workspace_or_normalize(shard_id, attempt_id, resume_mode) yield dataset_path def finalize_winner(self, shard_id: ShardId, attempt_id: AttemptId, published_at: datetime) -> ShardWinner: @@ -196,18 +196,20 @@ def _prepare_dataset_workspace( self, shard_id: ShardId, attempt_id: AttemptId, - effective_resume_mode: Literal["never", "always"], + resume_mode: Literal["never", "always", "if_possible"], ) -> Path: with self._storage.acquire_shard_lock(shard_id): run, plan, shard = self._reader.load_shard_context(shard_id) attempts = self._reader.load_validated_shard_attempts(run, plan, shard) attempt = self._reader.get_attempt(attempts, attempt_id) self.require_no_winner(run, plan, shard, attempts) - self._validate_workspace_mode(plan, attempt, effective_resume_mode) - dataset_path = self._storage.ensure_dataset_directory(shard_id, attempt_id, effective_resume_mode) + self._validate_workspace_mode(plan, attempt, resume_mode) + if resume_mode == "if_possible": + return Path(shard.resume_workspace.path) + dataset_path = self._storage.ensure_dataset_directory(shard_id, attempt_id, resume_mode) expected_path = ( Path(shard.resume_workspace.path) - if effective_resume_mode == "always" + if resume_mode == "always" else self._storage.get_attempt_path(shard_id, attempt_id) / "dataset" ) if dataset_path != expected_path: @@ -218,10 +220,10 @@ def _prepare_workspace_or_normalize( self, shard_id: ShardId, attempt_id: AttemptId, - effective_resume_mode: Literal["never", "always"], + resume_mode: Literal["never", "always", "if_possible"], ) -> Path: try: - return self._prepare_dataset_workspace(shard_id, attempt_id, effective_resume_mode) + return self._prepare_dataset_workspace(shard_id, attempt_id, resume_mode) except (StateConflictError, StateCorruptionError, StateNotFoundError, SlurmStateError): raise except (PlanStateContractError, StateContractError) as error: @@ -309,7 +311,7 @@ def _validate_artifact_metadata( def _validate_workspace_mode( plan: ResolvedSlurmRunPlan, attempt: AttemptManifest, - effective_resume_mode: Literal["never", "always"], + resume_mode: Literal["never", "always", "if_possible"], ) -> None: if attempt.state not in { AttemptLifecycleState.SUBMITTED, @@ -318,8 +320,8 @@ def _validate_workspace_mode( }: raise StateContractError("dataset workspace requires an active submitted attempt") requested = plan.invocation.authored.resume - if requested != "if_possible" and effective_resume_mode != requested: - raise StateContractError("effective resume mode does not match the resolved plan") + if requested != "if_possible" and resume_mode != requested: + raise StateContractError("resume mode does not match the resolved plan") def _get_winner_attempt( self, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index 618c156ba..d22ef8be2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -48,6 +48,7 @@ _CANDIDATE_OUTPUT_FILENAME = "output-manifest.json" _WINNER_FILENAME = "winner.json" _DATASET_DIRECTORY_NAME = "dataset" +_RUNTIME_DIRECTORY_NAME = "runtime" _RESUME_LOCK_FILENAME = "resume.lock" _LOCK_DIRECTORY_NAME = ".locks" _MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 @@ -249,6 +250,13 @@ def publish_attempt(self, attempt: AttemptManifest) -> None: ) as attempt_descriptor: self._publish_immutable_record(attempt_descriptor, _ATTEMPT_FILENAME, attempt) + def ensure_runtime_directory(self, shard_id: ShardId, attempt_id: AttemptId) -> None: + attempt_root = self.get_attempt_path(shard_id, attempt_id) + runtime_root = attempt_root / _RUNTIME_DIRECTORY_NAME + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + ensure_private_child_directory(attempt_descriptor, _RUNTIME_DIRECTORY_NAME, runtime_root) + sync_directory(attempt_descriptor) + def replace_attempt(self, attempt: AttemptManifest) -> None: with self.open_attempt_directory(attempt.shard_id, attempt.attempt_id) as attempt_descriptor: self._replace_record(attempt_descriptor, _ATTEMPT_FILENAME, attempt) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py index 95766a8e1..138b0a1e3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -177,17 +177,17 @@ def acquire_dataset_workspace( self, shard_id: ShardId, attempt_id: AttemptId, - effective_resume_mode: Literal["never", "always"], + resume_mode: Literal["never", "always", "if_possible"], ) -> Iterator[Path]: """Yield one validated dataset path while holding its shard lease.""" normalized_shard_id = self._validate_shard_id(shard_id) normalized_attempt_id = self._validate_attempt_id(attempt_id) - if type(effective_resume_mode) is not str or effective_resume_mode not in {"never", "always"}: - raise StateConflictError("effective resume mode must be 'never' or 'always'") + if type(resume_mode) is not str or resume_mode not in {"never", "always", "if_possible"}: + raise StateConflictError("resume mode must be 'never', 'always', or 'if_possible'") with self._finalizer.acquire_dataset_workspace( normalized_shard_id, normalized_attempt_id, - effective_resume_mode, + resume_mode, ) as dataset_path: yield dataset_path @@ -221,6 +221,7 @@ def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifes if existing is not None: if existing != attempt: raise StateConflictError(f"attempt {attempt.attempt_id!r} already contains different state") + self._storage.ensure_runtime_directory(attempt.shard_id, attempt.attempt_id) self._storage.sync_attempt_directory(attempt.shard_id, attempt.attempt_id) return existing self._finalizer.require_no_winner(run, plan, shard, shard_attempts) @@ -230,6 +231,7 @@ def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifes self._reader.validate_attempt_against_plan(run, plan, shard, attempt) validate_shard_attempt_set(run, shard, shard_attempts + (attempt,)) self._storage.publish_attempt(attempt) + self._storage.ensure_runtime_directory(attempt.shard_id, attempt.attempt_id) return attempt def _update_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifest: diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 183680d34..55b8258ff 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -4,7 +4,7 @@ from __future__ import annotations import subprocess -from collections.abc import Sequence +from collections.abc import Mapping, Sequence import pytest from slurm_test_fakes import FakeCommandResponse, FakeSlurmJob, FakeSlurmRunner @@ -40,6 +40,38 @@ def test_client_submits_verified_script_text_through_standard_input() -> None: assert runner.inputs == [script] +def test_client_holds_submission_until_initial_state_is_recorded() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + runner.script_next("scontrol", FakeCommandResponse()) + client = SlurmCommandClient(runner) + + submission = client.submit_script("#!/usr/bin/env bash\n", hold=True) + client.release(submission.job_id) + + assert runner.calls == [ + ("sbatch", "--parsable", "--hold", "--export=NIL"), + ("scontrol", "release", "5101"), + ] + + +def test_client_exports_only_explicit_environment_names_without_values_in_argv() -> None: + runner = _EnvironmentRunner() + client = SlurmCommandClient(runner) + + client.submit_script( + "#!/usr/bin/env bash\n", + export_environment={"HF_TOKEN": "secret-value", "SLURM_EXPORT_ENV": "ALL"}, + ) + + assert runner.command == ( + "sbatch", + "--parsable", + "--export=HF_TOKEN,SLURM_EXPORT_ENV", + ) + assert runner.environment == {"HF_TOKEN": "secret-value", "SLURM_EXPORT_ENV": "ALL"} + assert "secret-value" not in " ".join(runner.command) + + def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) client.submit("run.sbatch") @@ -281,6 +313,23 @@ def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: raise FileNotFoundError("missing executable") +class _EnvironmentRunner: + def __init__(self) -> None: + self.command: tuple[str, ...] = () + self.environment: dict[str, str] = {} + + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + self.command = tuple(command) + self.environment = dict(environment or {}) + return subprocess.CompletedProcess(command, 0, stdout="42\n", stderr="") + + class _TimeoutRunner: def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: raise subprocess.TimeoutExpired(command, 30.0) diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 08d93d747..90cff5b07 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -192,6 +192,7 @@ def test_rendered_script_verifies_exact_persisted_plan_bytes_before_sourcing_run ) -> None: run_root = tmp_path / "run" run_root.mkdir() + (run_root / "shards/shard-00000/attempts/attempt-0001/runtime").mkdir(parents=True) captured_plan_path = tmp_path / "captured-plan.json" entrypoint_path = tmp_path / "entrypoint.sh" entrypoint_path.write_text( diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index b41811178..98c74b636 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -61,13 +61,18 @@ def fake_run( } -def test_subprocess_runner_default_environment_forwards_only_search_path(monkeypatch: pytest.MonkeyPatch) -> None: +def test_subprocess_runner_default_environment_forwards_only_slurm_bootstrap(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("PATH", "/workspace/slurm/bin:/usr/bin") + monkeypatch.setenv("SLURM_CONF", "/etc/slurm/slurm.conf") monkeypatch.setenv("SECRET", "must-not-leak") runner = SubprocessRunner() - assert runner.environment == {"LC_ALL": "C", "PATH": "/workspace/slurm/bin:/usr/bin"} + assert runner.environment == { + "LC_ALL": "C", + "PATH": "/workspace/slurm/bin:/usr/bin", + "SLURM_CONF": "/etc/slurm/slurm.conf", + } def test_subprocess_runner_forwards_explicit_standard_input(monkeypatch: pytest.MonkeyPatch) -> None: @@ -110,6 +115,7 @@ def fake_run( def test_subprocess_runner_default_environment_replaces_empty_search_path(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("PATH", "") + monkeypatch.delenv("SLURM_CONF", raising=False) runner = SubprocessRunner() diff --git a/packages/data-designer-slurm/tests/runtime/conftest.py b/packages/data-designer-slurm/tests/runtime/conftest.py index 9dc83da6d..c2be41a0d 100644 --- a/packages/data-designer-slurm/tests/runtime/conftest.py +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -4,10 +4,12 @@ from __future__ import annotations import json +from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import cast +from typing import Literal, cast import pytest @@ -22,6 +24,7 @@ AttemptReadiness, CandidateOutputManifest, SchedulerIdentity, + ShardWinner, StateNotFoundError, validate_attempt_transition, validate_readiness_transition, @@ -39,6 +42,9 @@ class RuntimeCase: class FakeStateStore: attempt: AttemptManifest readiness: list[AttemptReadiness] = field(default_factory=list) + dataset_workspace_lease_active: bool = False + dataset_workspace_modes: list[str] = field(default_factory=list) + winners: list[ShardWinner] = field(default_factory=list) def update_attempt(self, attempt: AttemptManifest) -> AttemptManifest: validate_attempt_transition(self.attempt, attempt) @@ -50,6 +56,7 @@ def publish_attempt_result( client_result: ClientResult, candidate: CandidateOutputManifest, ) -> tuple[ClientResult, CandidateOutputManifest]: + assert self.dataset_workspace_lease_active reference = client_result.candidate_output_manifest assert reference is not None assert candidate.compute_sha256() == reference.sha256 @@ -59,6 +66,44 @@ def publish_attempt_result( self.attempt = bound_attempt return client_result, candidate + @contextmanager + def acquire_dataset_workspace( + self, + shard_id: str, + attempt_id: str, + resume_mode: Literal["never", "always", "if_possible"], + ) -> Iterator[Path]: + assert shard_id == self.attempt.shard_id + assert attempt_id == self.attempt.attempt_id + assert not self.dataset_workspace_lease_active + self.dataset_workspace_modes.append(resume_mode) + self.dataset_workspace_lease_active = True + try: + yield Path(self.attempt.resolved_plan.path).parent / "dataset" + finally: + self.dataset_workspace_lease_active = False + + def finalize_winner( + self, + shard_id: str, + attempt_id: str, + *, + published_at: datetime, + ) -> ShardWinner: + assert self.attempt.state is AttemptLifecycleState.SUCCEEDED + assert self.attempt.candidate_output is not None + winner = ShardWinner( + schema_version=1, + run_id=self.attempt.run_id, + shard_id=shard_id, + attempt_id=attempt_id, + attempt_ordinal=self.attempt.attempt_ordinal, + candidate_manifest=self.attempt.candidate_output, + published_at=published_at, + ) + self.winners.append(winner) + return winner + def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: if self.readiness: validate_readiness_transition(self.readiness[-1], readiness) @@ -75,6 +120,11 @@ def load_readiness(self, shard_id: str, attempt_id: str) -> AttemptReadiness: assert readiness.attempt_id == attempt_id return readiness + def load_attempt(self, shard_id: str, attempt_id: str) -> AttemptManifest: + assert self.attempt.shard_id == shard_id + assert self.attempt.attempt_id == attempt_id + return self.attempt + class FakePreflight: def __init__(self, failure: BaseException | None = None) -> None: diff --git a/packages/data-designer-slurm/tests/runtime/test_backpressure.py b/packages/data-designer-slurm/tests/runtime/test_backpressure.py new file mode 100644 index 000000000..4478a3690 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_backpressure.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, + AsgiMessage, + QueueBackpressureController, + QueueBackpressureSettings, + QueueDepthBackpressureMiddleware, +) + + +def test_backpressure_settings_are_strict_and_support_absent_retry_header() -> None: + settings = QueueBackpressureSettings.from_environment( + { + MAX_WAITING_REQUESTS_ENVIRONMENT: "7", + RETRY_AFTER_SECONDS_ENVIRONMENT: "", + } + ) + assert settings.max_waiting_requests == 7 + assert settings.retry_after_seconds is None + with pytest.raises(ValueError, match="threshold"): + QueueBackpressureSettings.from_environment({MAX_WAITING_REQUESTS_ENVIRONMENT: "-1"}) + + +@pytest.mark.asyncio +async def test_middleware_rejects_fresh_overload_and_preserves_health() -> None: + downstream_calls: list[str] = [] + messages: list[AsgiMessage] = [] + + async def app(scope: dict[str, object], receive: object, send: object) -> None: + del receive, send + downstream_calls.append(str(scope["path"])) + + async def receive() -> AsgiMessage: + return {"type": "http.request"} + + async def send(message: AsgiMessage) -> None: + messages.append(message) + + controller = QueueBackpressureController( + QueueBackpressureSettings(2, 3), + reader=lambda: 4, + start_sampler=False, + ) + controller.sample_once() + middleware = QueueDepthBackpressureMiddleware(app, controller) + + await middleware({"type": "http", "path": "/v1/chat/completions"}, receive, send) + await middleware({"type": "http", "path": "/health"}, receive, send) + + assert messages[0]["status"] == 429 + assert (b"retry-after", b"3") in messages[0]["headers"] + assert downstream_calls == ["/health"] + + +@pytest.mark.asyncio +async def test_middleware_fails_open_when_metrics_are_unavailable() -> None: + called = False + + async def app(scope: dict[str, object], receive: object, send: object) -> None: + nonlocal called + del scope, receive, send + called = True + + async def receive() -> AsgiMessage: + return {"type": "http.request"} + + async def send(message: AsgiMessage) -> None: + raise AssertionError(f"unexpected middleware response: {message}") + + controller = QueueBackpressureController(reader=lambda: None, start_sampler=False) + controller.sample_once() + middleware = QueueDepthBackpressureMiddleware(app, controller) + await middleware({"type": "http", "path": "/v1/completions"}, receive, send) + + assert called + + +def test_sampler_fails_open_after_reader_error_and_recovers() -> None: + readings: list[Exception | int] = [RuntimeError("metrics unavailable"), 4] + + def read_queue_depth() -> int: + reading = readings.pop(0) + if isinstance(reading, Exception): + raise reading + return reading + + controller = QueueBackpressureController( + QueueBackpressureSettings(2, 3), + reader=read_queue_depth, + start_sampler=False, + ) + + failed_snapshot = controller.sample_once() + failed_reject, _ = controller.should_reject() + recovered_snapshot = controller.sample_once() + recovered_reject, _ = controller.should_reject() + + assert failed_snapshot.depth is None + assert not failed_reject + assert recovered_snapshot.depth == 4 + assert recovered_reject + + +def test_sampler_preserves_process_control_exceptions() -> None: + def interrupt() -> int: + raise KeyboardInterrupt + + controller = QueueBackpressureController(reader=interrupt, start_sampler=False) + + with pytest.raises(KeyboardInterrupt): + controller.sample_once() diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py new file mode 100644 index 000000000..eb9a767ff --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from conftest import RuntimeCase + +from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, build_runtime_manifest +from data_designer.slurm.runtime.models import RuntimeStepRole + + +def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + runtime_root = context.attempt_directory / "runtime" + log_directory = context.attempt_directory / "logs/execution-00000002" + + manifest = build_runtime_manifest( + context, + runtime_root=runtime_root, + log_directory=log_directory, + ) + reloaded = RuntimeBootstrapManifest.model_validate_json(manifest.serialize_json()) + + assert reloaded == manifest + assert [step.role for step in manifest.steps] == [ + RuntimeStepRole.CLIENT_PREFLIGHT, + RuntimeStepRole.SERVER, + RuntimeStepRole.ENDPOINT, + RuntimeStepRole.CLIENT, + ] + assert all(step.command[0] != "srun" for step in manifest.steps) + assert all(step.stdout_path.startswith(context.attempt_directory.as_posix()) for step in manifest.steps) + assert manifest.steps[-1].command[:4] == ( + "python3", + "-m", + "data_designer.slurm.runtime.entrypoint", + "client", + ) diff --git a/packages/data-designer-slurm/tests/runtime/test_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py index a1d67bde8..40174aba1 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bundle.py +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -36,8 +36,9 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p with tarfile.open(fileobj=io.BytesIO(content), mode="r:gz") as archive: names = archive.getnames() assert names[0] == "entrypoint.sh" - assert names[1] == "data_designer/slurm/runtime/slurm-sources.txt" - assert names[2] == "data_designer/slurm/__init__.py" + assert names[1:4] == ["plan_reader.sh", "step_runner.sh", "cleanup.sh"] + assert names[4] == "data_designer/slurm/runtime/slurm-sources.txt" + assert names[5] == "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/state/store.py" in names @@ -46,7 +47,7 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p assert all(archive.getmember(name).uid == 0 for name in names) entrypoint = archive.extractfile("entrypoint.sh") assert entrypoint is not None - assert b'PYTHONPATH="${runtime_root}"' in entrypoint.read() + assert b"python3 -m data_designer.slurm.runtime.entrypoint" not in entrypoint.read() def test_runtime_bundle_recursively_collects_and_imports_nested_packages( @@ -62,6 +63,8 @@ def test_runtime_bundle_recursively_collects_and_imports_nested_packages( (source_root / "__init__.py").write_text("") (runtime_root / "__init__.py").write_text("") (runtime_root / "bundle.py").write_text("") + for name in ("entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh"): + (runtime_root / name).write_text("") (nested_root / "__init__.py").write_text("") (nested_root / "worker.py").write_text("VALUE = 42\n") monkeypatch.setattr(runtime_bundle, "__file__", (runtime_root / "bundle.py").as_posix()) @@ -233,6 +236,31 @@ def test_extracted_bundle_runtime_takes_precedence_over_installed_sources(tmp_pa assert completed.stdout == (extracted / "data_designer/slurm/runtime/__init__.py").as_posix() +def test_backpressure_module_imports_from_the_bundle_without_site_packages(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + extracted = tmp_path / "extracted" + workspace.mkdir(mode=0o700) + reference = stage_runtime_bundle(workspace) + with tarfile.open(reference.path, mode="r:gz") as archive: + archive.extractall(extracted, filter="data") + + completed = subprocess.run( + ( + sys.executable, + "-S", + "-c", + "import sys; " + f"sys.path.insert(0, {extracted.as_posix()!r}); " + "from data_designer.slurm.runtime.backpressure import QueueDepthBackpressureMiddleware", + ), + capture_output=True, + text=True, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + + def _get_module_name(source_name: str) -> str: module_name = source_name.removesuffix(".py").replace("/", ".") return module_name.removesuffix(".__init__") diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index de4b363c5..524878bf7 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib from dataclasses import dataclass, replace from datetime import timedelta from pathlib import Path @@ -12,6 +13,7 @@ from conftest import FakeClientStepBuilder, FakePreflight, FakeStateStore, RuntimeCase from slurm_test_fakes import FakeClock +import data_designer.lazy_heavy_imports as lazy from data_designer.slurm.client import ClientOutcome, ClientResult from data_designer.slurm.config import DataDesignerSlurmConfig from data_designer.slurm.contracts import ArtifactReference @@ -34,6 +36,7 @@ ShardManifest, SlurmStateWriter, ) +from data_designer.slurm.state.artifacts import compute_candidate_schema_digest @dataclass(slots=True) @@ -149,7 +152,12 @@ def _supervisor( def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_case: RuntimeCase) -> None: clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) state = FakeStateStore(runtime_case.context.attempt) - runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + + def write_result_under_dataset_lease() -> None: + assert state.dataset_workspace_lease_active + _write_complete_result(runtime_case, clock) + + runner = _FakeRunner(generation_hook=write_result_under_dataset_lease) supervisor = _supervisor(runner, clock, poll_interval_seconds=0.1) controller = OneNodeAllocationController( runtime_case.context, @@ -168,6 +176,9 @@ def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_c assert result.state is AttemptLifecycleState.SUCCEEDED assert result.terminal_classification is AttemptTerminalClassification.SUCCEEDED assert result.candidate_output is not None + assert state.dataset_workspace_modes == ["never"] + assert not state.dataset_workspace_lease_active + assert state.winners[0].attempt_id == result.attempt_id assert [step.role for step in runner.steps] == [ RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.SERVER, @@ -207,6 +218,7 @@ def test_controller_publishes_result_before_success_with_real_state_writer( assert persisted == result assert persisted.state is AttemptLifecycleState.SUCCEEDED assert persisted.candidate_output is not None + assert state.load_winner(result.shard_id).attempt_id == result.attempt_id def test_preflight_failure_starts_no_process_and_fails_attempt(runtime_case: RuntimeCase) -> None: @@ -598,6 +610,11 @@ def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: context = runtime_case.context requested = context.shard.requested_records dataset_path = (context.attempt_directory / "dataset").as_posix() + output_path = Path(dataset_path) / "part-00000.parquet" + output_path.parent.mkdir(exist_ok=True) + table = lazy.pa.table({"record_id": range(requested)}) + lazy.pq.write_table(table, output_path) + output_bytes = output_path.read_bytes() candidate = CandidateOutputManifest( schema_version=1, run_id=context.plan.run_id, @@ -612,12 +629,12 @@ def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: files=( CandidateOutputFile( relative_path="part-00000.parquet", - sha256="a" * 64, - byte_size=128, + sha256=hashlib.sha256(output_bytes).hexdigest(), + byte_size=len(output_bytes), record_count=requested, ), ), - dataset_schema_digest="b" * 64, + dataset_schema_digest=compute_candidate_schema_digest(table.schema), provenance_digest=context.plan.compute_sha256(), ) candidate_path = context.attempt_directory / "output-manifest.json" diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index b83f83c34..8927732cd 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -3,13 +3,124 @@ from __future__ import annotations +from pathlib import Path +from types import SimpleNamespace + import pytest +from conftest import FakeStateStore, RuntimeCase -from data_designer.slurm.runtime.entrypoint import main +import data_designer.slurm.runtime.entrypoint as entrypoint +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.state import AttemptLifecycleState, ReadinessState def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.CaptureFixture[str]) -> None: - assert main(("--plan", "resolved-plan.json", "--attempt-dir", "attempt-0001")) == 64 + assert ( + entrypoint.main( + ( + "prepare", + "--plan", + "resolved-plan.json", + "--attempt-dir", + "attempt-0001", + "--runtime-root", + "runtime", + "--manifest", + "runtime-manifest.json", + ) + ) + == 64 + ) captured = capsys.readouterr() assert "runtime paths must be absolute" in captured.err assert "Traceback" not in captured.err + + +def test_control_phases_record_running_ready_and_failed( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, +) -> None: + state = FakeStateStore(runtime_case.context.attempt) + runtime_root = runtime_case.context.attempt_directory / "runtime" + manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" + _patch_runtime_context(monkeypatch, runtime_case, state) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda context: None) + monkeypatch.setattr( + entrypoint, + "build_runtime_manifest", + lambda *args, **kwargs: SimpleNamespace(serialize_json=lambda: "{}"), + ) + + assert entrypoint.main(_phase_arguments("prepare", runtime_case, runtime_root, manifest_path)) == 0 + assert state.attempt.state is AttemptLifecycleState.RUNNING + assert [item.state for item in state.readiness] == [ReadinessState.PENDING, ReadinessState.STARTING] + assert manifest_path.read_text() == "{}" + + assert entrypoint.main(_phase_arguments("ready", runtime_case)) == 0 + assert state.readiness[-1].state is ReadinessState.READY + + assert entrypoint.main(_phase_arguments("fail", runtime_case)) == 0 + assert state.attempt.state is AttemptLifecycleState.FAILED + assert [item.state for item in state.readiness[-2:]] == [ReadinessState.FAILED, ReadinessState.STOPPED] + + +def test_succeed_phase_stops_runtime_and_finalizes_winner( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, +) -> None: + state = FakeStateStore(runtime_case.context.attempt) + runtime_root = runtime_case.context.attempt_directory / "runtime" + manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" + _patch_runtime_context(monkeypatch, runtime_case, state) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda context: None) + monkeypatch.setattr( + entrypoint, + "build_runtime_manifest", + lambda *args, **kwargs: SimpleNamespace(serialize_json=lambda: "{}"), + ) + assert entrypoint.main(_phase_arguments("prepare", runtime_case, runtime_root, manifest_path)) == 0 + assert entrypoint.main(_phase_arguments("ready", runtime_case)) == 0 + state.attempt = state.attempt.model_copy( + update={ + "candidate_output": ArtifactReference( + path=(runtime_case.context.attempt_directory / "output-manifest.json").as_posix(), + sha256="a" * 64, + ) + } + ) + + assert entrypoint.main(_phase_arguments("succeed", runtime_case)) == 0 + + assert state.attempt.state is AttemptLifecycleState.SUCCEEDED + assert state.readiness[-1].state is ReadinessState.STOPPED + assert state.winners[0].attempt_id == state.attempt.attempt_id + + +def _patch_runtime_context( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, + state: FakeStateStore, +) -> None: + monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args: (runtime_case.context, state)) + monkeypatch.setattr(entrypoint, "get_container_path", lambda plan, path, **kwargs: path) + + +def _phase_arguments( + operation: str, + runtime_case: RuntimeCase, + runtime_root: Path | None = None, + manifest_path: Path | None = None, +) -> tuple[str, ...]: + arguments = ( + operation, + "--plan", + (runtime_case.workspace / "runs/run-single/resolved-plan.json").as_posix(), + "--attempt-dir", + runtime_case.context.attempt_directory.as_posix(), + ) + if operation == "prepare": + assert runtime_root is not None and manifest_path is not None + return (*arguments, "--runtime-root", runtime_root.as_posix(), "--manifest", manifest_path.as_posix()) + return arguments diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index e4be62e92..1766840b0 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -127,7 +127,7 @@ def close(self) -> None: monkeypatch.setattr("data_designer.slurm.runtime.preflight.socket.socket", _UnavailableSocket) with pytest.raises(SlurmRuntimeError, match="ports are unavailable"): - SystemAllocationPreflight._verify_ports(runtime_case.context) + SystemAllocationPreflight.verify_ports(runtime_case.context) assert _UnavailableSocket.closed @@ -135,4 +135,4 @@ def test_attempt_directory_must_be_restrictive(runtime_case: RuntimeCase) -> Non runtime_case.context.attempt_directory.chmod(0o755) with pytest.raises(SlurmRuntimeError, match="restrictive directory"): - SystemAllocationPreflight._verify_attempt_directory(runtime_case.context.attempt_directory) + SystemAllocationPreflight.verify_attempt_directory(runtime_case.context.attempt_directory) diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py new file mode 100644 index 000000000..04ed4f208 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import os +import shlex +import subprocess +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) +def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( + tmp_path: Path, + gpu_request_mode: str, +) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + attempt_directory = tmp_path / "runs/run-shell/shards/shard-00000/attempts/attempt-0001" + log_directory = attempt_directory / "logs/execution-00000002" + fake_bin = tmp_path / "bin" + log_directory.mkdir(parents=True) + fake_bin.mkdir() + artifacts = tuple(_artifact(tmp_path, name) for name in ("runtime", "lock", "client", "server")) + plan_path = tmp_path / "runs/run-shell/resolved-plan.json" + plan_path.parent.mkdir(parents=True, exist_ok=True) + plan = _plan(tmp_path, runtime_root, artifacts, gpu_request_mode) + plan_path.write_text(json.dumps(plan)) + manifest_path = tmp_path / "manifest-source.json" + manifest_path.write_text(json.dumps(_manifest(attempt_directory, artifacts, "a" * 64))) + marker_path = tmp_path / "host-python-ran" + _write_executable(fake_bin / "python3", f"#!/usr/bin/env bash\nprintf ran > {marker_path}\nexit 99\n") + _write_executable(fake_bin / "curl", "#!/usr/bin/env bash\nexit 0\n") + _write_executable(fake_bin / "getent", "#!/usr/bin/env bash\nexit 0\n") + _write_executable(fake_bin / "scontrol", "#!/usr/bin/env bash\nexit 0\n") + _write_executable(fake_bin / "srun", _fake_srun()) + command = f""" + set -Eeuo pipefail +DD_PLAN_SHA256={"a" * 64} +DD_SHARD_ID=shard-00000 +DD_ATTEMPT_ORDINAL=0001 +readonly DD_PLAN={shlex.quote(plan_path.as_posix())} +readonly DD_ATTEMPT_DIR={shlex.quote(attempt_directory.as_posix())} +source {shlex.quote((runtime_root / "entrypoint.sh").as_posix())} +dd_slurm_run_allocation "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" +""" + environment = { + **os.environ, + "PATH": f"{fake_bin}:{os.environ['PATH']}", + "FAKE_MANIFEST_SOURCE": manifest_path.as_posix(), + "FAKE_GPU_MODE": gpu_request_mode, + "SOURCE_TOKEN": "supersecret", + "CUDA_VISIBLE_DEVICES": "0", + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_NUM_NODES": "1", + "SLURM_NODEID": "0", + } + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + check=False, + env=environment, + timeout=10, + ) + + assert completed.returncode == 0, completed.stderr + assert not marker_path.exists() + assert "supersecret" not in completed.stdout + completed.stderr + manifest_path.read_text() + + +def test_staged_shell_modules_parse_as_bash(tmp_path: Path) -> None: + del tmp_path + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + scripts = tuple(runtime_root / name for name in ("entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh")) + + completed = subprocess.run(("bash", "-n", *(path.as_posix() for path in scripts)), capture_output=True, text=True) + + assert completed.returncode == 0, completed.stderr + + +def test_shell_helpers_handle_empty_and_sparse_arrays() -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "plan_reader.sh").as_posix())} +source {shlex.quote((runtime_root / "cleanup.sh").as_posix())} +dd_read_null_values DD_VALUES < <(printf '') +(( ${{#DD_VALUES[@]}} == 0 )) +dd_start_runtime_timer +sleep 30 & +pid=$! +DD_MANAGED_PIDS[2]=${{pid}} +dd_cleanup_steps +! kill -0 "${{pid}}" 2>/dev/null +dd_stop_runtime_timer +""" + + completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=5) + + assert completed.returncode == 0, completed.stderr + + +def _artifact(root: Path, name: str) -> tuple[str, str]: + path = root / f"{name}.artifact" + content = name.encode() + path.write_bytes(content) + return path.as_posix(), hashlib.sha256(content).hexdigest() + + +def _plan( + root: Path, + runtime_root: Path, + artifacts: tuple[tuple[str, str], ...], + gpu_request_mode: str, +) -> dict[str, object]: + runtime, lock, client, server = artifacts + mounts = [ + {"source": root.as_posix(), "target": root.as_posix(), "read_only": False}, + {"source": runtime_root.as_posix(), "target": runtime_root.as_posix(), "read_only": False}, + ] + return { + "schema_version": 1, + "runtime_bundle": {"path": runtime[0], "sha256": runtime[1]}, + "client": { + "host_node_index": 0, + "authored": {"cpus": 1}, + "dependency_lock": {"path": lock[0], "sha256": lock[1]}, + "image": {"path": client[0], "sha256": client[1]}, + }, + "resolved_gpus_per_node": 1, + "selected_profile": {"profile": {"gpu_request_mode": gpu_request_mode}}, + "container_mounts": mounts, + "deployments": [ + { + "node_indices": [0], + "image": {"path": server[0], "sha256": server[1]}, + "authored": { + "server": {"environment": {"SERVER_TOKEN": {"type": "secret", "environment": "SOURCE_TOKEN"}}} + }, + } + ], + "builder": {"source": None}, + "shards": [{"array_task_index": 0, "input_partition": None}], + } + + +def _manifest( + attempt_directory: Path, + artifacts: tuple[tuple[str, str], ...], + plan_sha256: str, +) -> dict[str, object]: + _, _, client, server = artifacts + return { + "schema_version": 1, + "run_id": "run-shell", + "shard_id": "shard-00000", + "attempt_id": "attempt-0001", + "plan_sha256": plan_sha256, + "all_secret_environment_names": ["SERVER_TOKEN", "SOURCE_TOKEN"], + "steps": [ + _step(attempt_directory, "client-preflight", "client_preflight", client[0], "fake-preflight"), + _step( + attempt_directory, + "deployment-00000-replica-00000-rank-00000", + "server", + server[0], + "fake-server", + gpu_indices=[0], + secret_environment={"SERVER_TOKEN": "SOURCE_TOKEN"}, + container_environment=["SERVER_TOKEN"], + readiness={"host": "127.0.0.1", "port": 18000, "path": "/health", "deadline_seconds": 2}, + ), + _step( + attempt_directory, + "deployment-00000-endpoint", + "endpoint", + client[0], + "fake-endpoint", + readiness={"host": "127.0.0.1", "port": 17000, "path": "/health", "deadline_seconds": 2}, + ), + _step( + attempt_directory, + "client-generation", + "client", + client[0], + "fake-client", + secret_environment={"SOURCE_TOKEN": "SOURCE_TOKEN"}, + container_environment=["SOURCE_TOKEN"], + ), + ], + } + + +def _step( + attempt_directory: Path, + step_id: str, + role: str, + image_path: str, + command: str, + *, + gpu_indices: list[int] | None = None, + secret_environment: dict[str, str] | None = None, + container_environment: list[str] | None = None, + readiness: dict[str, object] | None = None, +) -> dict[str, object]: + log_root = attempt_directory / "logs/execution-00000002" + return { + "step_id": step_id, + "role": role, + "image_path": image_path, + "command": [command], + "cpus": 1, + "gpu_indices": gpu_indices or [], + "literal_environment": {"LC_ALL": "C"}, + "secret_environment": secret_environment or {}, + "environment_prefixes": {}, + "container_environment": container_environment or [], + "stdout_path": (log_root / f"{step_id}.out").as_posix(), + "stderr_path": (log_root / f"{step_id}.err").as_posix(), + "launch_delay_seconds": 0, + "readiness": readiness, + } + + +def _write_executable(path: Path, content: str) -> None: + path.write_text(content) + path.chmod(0o700) + + +def _fake_srun() -> str: + return """#!/usr/bin/env bash +set -Eeuo pipefail +arguments=("$@") +for ((index = 0; index < ${#arguments[@]}; index++)); do + [[ ${arguments[index]} == -- ]] && break +done +command=${arguments[index + 1]} +if [[ ${command} == python3 ]]; then + operation=${arguments[index + 4]} + [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + if [[ ${operation} == prepare ]]; then + for ((position = index + 5; position < ${#arguments[@]}; position++)); do + if [[ ${arguments[position]} == --manifest ]]; then + cp "${FAKE_MANIFEST_SOURCE}" "${arguments[position + 1]}" + break + fi + done + fi + exit 0 +fi +case ${command} in + fake-preflight) + [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + ;; + fake-client) + [[ ${SOURCE_TOKEN} == supersecret && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + ;; + fake-server) + [[ ${SERVER_TOKEN} == supersecret && ! ${SOURCE_TOKEN+x} ]] + if [[ ${FAKE_GPU_MODE} == visible ]]; then + [[ ${CUDA_VISIBLE_DEVICES} == 0 ]] + else + [[ ! ${CUDA_VISIBLE_DEVICES+x} ]] + fi + trap 'exit 0' TERM INT + while :; do :; done + ;; + fake-endpoint) + [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + trap 'exit 0' TERM INT + while :; do :; done + ;; + *) + exit 90 + ;; +esac +""" diff --git a/packages/data-designer-slurm/tests/runtime/test_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py index b01ed41aa..aa400785e 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -12,6 +12,10 @@ from data_designer.slurm.config import QueueBackpressureConfig from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStepRole from data_designer.slurm.runtime.steps import ( @@ -37,6 +41,7 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti context.plan, context.attempt_directory, source_environment, + context.attempt_directory / "runtime", ) endpoint_steps = build_endpoint_steps( (deployment,), @@ -82,6 +87,19 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti assert f"--gpu-bind=mask_gpu:{expected_mask:#x}" in server.command assert "CUDA_VISIBLE_DEVICES" not in server.environment assert all("CUDA_VISIBLE_DEVICES" not in argument for argument in server.command) + assert "--middleware" in server.command + assert "data_designer.slurm.runtime.backpressure.QueueDepthBackpressureMiddleware" in server.command + assert server.environment[MAX_WAITING_REQUESTS_ENVIRONMENT] == "128" + assert server.environment[RETRY_AFTER_SECONDS_ENVIRONMENT] == "1" + assert server.environment["PYTHONPATH"].endswith("/runtime") + assert any( + argument.startswith("--container-env=") + and all( + name in argument + for name in ("PYTHONPATH", MAX_WAITING_REQUESTS_ENVIRONMENT, RETRY_AFTER_SECONDS_ENVIRONMENT) + ) + for argument in server.command + ) assert all("--gpus-per-task=" not in argument for step in client_steps for argument in step.command) assert all("--gres=none" in step.command for step in client_steps) assert all("CUDA_VISIBLE_DEVICES" not in step.environment for step in client_steps) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index cd8a09d2a..20db42325 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -4,6 +4,7 @@ from __future__ import annotations import hashlib +from collections.abc import Mapping from datetime import UTC, datetime from pathlib import Path @@ -14,12 +15,14 @@ BuilderInput, DataDesignerSlurmConfig, ImageBuildRequest, + SecretRef, SlurmProfile, SlurmProfileCatalog, ) from data_designer.slurm.contracts import ArtifactReference, canonical_json from data_designer.slurm.images.records import RegisteredImage from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.launcher.errors import SlurmLauncherError from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.services import ( @@ -31,6 +34,7 @@ from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, + AttemptTerminalClassification, RunManifest, SchedulerIdentity, ShardManifest, @@ -40,14 +44,32 @@ class _Launcher: - def __init__(self, gpu_counts: tuple[int, ...] = (), *, cancel_error: Exception | None = None) -> None: + def __init__( + self, + gpu_counts: tuple[int, ...] = (), + *, + cancel_error: Exception | None = None, + release_error: Exception | None = None, + ) -> None: self.submissions: list[str] = [] self.cancellations: list[int] = [] + self.releases: list[int] = [] + self.held_submissions: list[bool] = [] + self.exported_environments: list[dict[str, str]] = [] self.gpu_counts = gpu_counts self.cancel_error = cancel_error + self.release_error = release_error - def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + def submit_script( + self, + script: str, + *, + hold: bool = False, + export_environment: Mapping[str, str] | None = None, + ) -> SlurmJobSubmissionReceipt: self.submissions.append(script) + self.held_submissions.append(hold) + self.exported_environments.append(dict(export_environment or {})) return SlurmJobSubmissionReceipt(job_id=42) def cancel(self, job_id: int) -> None: @@ -55,6 +77,11 @@ def cancel(self, job_id: int) -> None: if self.cancel_error is not None: raise self.cancel_error + def release(self, job_id: int) -> None: + self.releases.append(job_id) + if self.release_error is not None: + raise self.release_error + def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: assert partition is not None return self.gpu_counts @@ -69,6 +96,7 @@ def __init__( ) -> None: self.initializations: list[tuple[str, bool]] = [] self.submissions: list[tuple[str, int, datetime]] = [] + self.submission_failures: list[tuple[str, datetime]] = [] self.initialization_error = initialization_error self.submission_error = submission_error @@ -93,6 +121,9 @@ def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitte raise self.submission_error self.submissions.append((plan.run_id, job_id, submitted_at)) + def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: + self.submission_failures.append((plan.run_id, failed_at)) + def _profile(tmp_path: Path, profile_catalog: SlurmProfileCatalog) -> SlurmProfile: return profile_catalog.clusters["primary"].model_copy(update={"workspace_root": tmp_path.as_posix()}) @@ -199,9 +230,12 @@ def test_production_wiring_submits_after_publisher_initialization( assert publisher.initializations == [("run-wired", True)] assert publisher.submissions == [("run-wired", 42, submitted_at)] assert len(launcher.submissions) == 1 + assert launcher.held_submissions == [True] + assert launcher.releases == [42] + assert launcher.exported_environments == [{"SLURM_EXPORT_ENV": "ALL"}] -def test_production_wiring_stops_before_submission_without_state_publisher( +def test_production_publisher_rejects_force_before_submission( tmp_path: Path, profile_catalog: SlurmProfileCatalog, authored_run_single: DataDesignerSlurmConfig, @@ -216,14 +250,104 @@ def test_production_wiring_stops_before_submission_without_state_publisher( package_version="0.9.2", ) - with pytest.raises(SlurmServiceError) as caught: - service.execute(authored_run_single, source_root=tmp_path) + with pytest.raises(SlurmServiceError, match="different inputs") as caught: + service.execute(authored_run_single, source_root=tmp_path, force=True) - assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE - assert str(caught.value) == "run submission is not available; use --dry-run" + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert launcher.submissions == [] + + +def test_production_wiring_exports_referenced_secrets_to_the_allocation( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + deployment = authored_run_single.deployments[0] + server = deployment.server.model_copy( + update={"environment": {"HF_TOKEN": SecretRef(type="secret", environment="SOURCE_TOKEN")}} + ) + authored = authored_run_single.model_copy( + update={"deployments": [deployment.model_copy(update={"server": server})]} + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=_Publisher(), # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + source_environment={"SOURCE_TOKEN": "secret-value"}, + ) + + service.execute(authored, source_root=tmp_path) + + assert launcher.exported_environments == [{"SLURM_EXPORT_ENV": "ALL", "SOURCE_TOKEN": "secret-value"}] + + +def test_production_wiring_rejects_missing_secret_before_publishing( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + deployment = authored_run_single.deployments[0] + server = deployment.server.model_copy( + update={"environment": {"HF_TOKEN": SecretRef(type="secret", environment="MISSING_TOKEN")}} + ) + authored = authored_run_single.model_copy( + update={"deployments": [deployment.model_copy(update={"server": server})]} + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + publisher = _Publisher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=publisher, # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + source_environment={}, + ) + + with pytest.raises(SlurmServiceError, match="MISSING_TOKEN") as caught: + service.execute(authored, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert publisher.initializations == [] assert launcher.submissions == [] +def test_production_wiring_publishes_initial_state_before_releasing_submission( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + result = service.execute(authored_run_single, source_root=tmp_path) + + writer = SlurmStateWriter(tmp_path, result.run_id) + attempts = writer.load_attempts("shard-00000") + assert result.state == "submitted" + assert writer.load_run().run_id == result.run_id + dependency_bytes = (tmp_path / "runs/run-wired/dependency-lock.json").read_bytes() + assert hashlib.sha256(dependency_bytes).hexdigest() == writer.load_resolved_plan().client.dependency_lock.sha256 + assert attempts[0].state is AttemptLifecycleState.SUBMITTED + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=42, array_task_id=0) + assert launcher.held_submissions == [True] + assert launcher.releases == [42] + + def test_auto_gpu_resolution_rejects_mixed_node_shapes( tmp_path: Path, profile_catalog: SlurmProfileCatalog, @@ -293,6 +417,35 @@ def test_recording_conflict_cancels_the_accepted_job( assert launcher.cancellations == [42] +def test_release_failure_cancels_the_held_job( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher(release_error=SlurmLauncherError("release failed")) + failed_at = datetime(2026, 9, 8, tzinfo=UTC) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: failed_at, + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError, match="could not be released and was cancelled") as caught: + service.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert launcher.releases == [42] + assert launcher.cancellations == [42] + attempt = SlurmStateWriter(tmp_path, "run-wired").load_attempt("shard-00000", "attempt-0001") + assert attempt.state is AttemptLifecycleState.FAILED + assert attempt.terminal_classification is AttemptTerminalClassification.CANCELLED + assert attempt.updated_at == failed_at + + def test_failed_compensating_cancel_reports_accepted_job_id( tmp_path: Path, profile_catalog: SlurmProfileCatalog, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 3f75a41b6..f9bf180ab 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -33,10 +33,9 @@ readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID}" printf -v DD_SHARD_ID 'shard-%05d' "${DD_ARRAY_TASK_ID}" readonly DD_SHARD_ID readonly DD_ATTEMPT_DIR="${DD_RUN_ROOT}/shards/${DD_SHARD_ID}/attempts/attempt-${DD_ATTEMPT_ORDINAL}" -install -d -m 0700 "${DD_ATTEMPT_DIR}" -DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" -readonly DD_RUNTIME_DIR -tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" +readonly DD_RUNTIME_ROOT="${DD_ATTEMPT_DIR}/runtime" +[[ -d ${DD_RUNTIME_ROOT} && ! -L ${DD_RUNTIME_ROOT} ]] +tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_ROOT}" -source "${DD_RUNTIME_DIR}/entrypoint.sh" +source "${DD_RUNTIME_ROOT}/entrypoint.sh" dd_slurm_run_allocation "${DD_PLAN}" "${DD_ATTEMPT_DIR}" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 553fadd86..60d0ce84e 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -33,10 +33,9 @@ readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID}" printf -v DD_SHARD_ID 'shard-%05d' "${DD_ARRAY_TASK_ID}" readonly DD_SHARD_ID readonly DD_ATTEMPT_DIR="${DD_RUN_ROOT}/shards/${DD_SHARD_ID}/attempts/attempt-${DD_ATTEMPT_ORDINAL}" -install -d -m 0700 "${DD_ATTEMPT_DIR}" -DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" -readonly DD_RUNTIME_DIR -tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" +readonly DD_RUNTIME_ROOT="${DD_ATTEMPT_DIR}/runtime" +[[ -d ${DD_RUNTIME_ROOT} && ! -L ${DD_RUNTIME_ROOT} ]] +tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_ROOT}" -source "${DD_RUNTIME_DIR}/entrypoint.sh" +source "${DD_RUNTIME_ROOT}/entrypoint.sh" dd_slurm_run_allocation "${DD_PLAN}" "${DD_ATTEMPT_DIR}" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py index 5a8494cfc..2350b5ab2 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -110,6 +110,7 @@ def run( *, check: bool = False, input_text: str | None = None, + environment: Mapping[str, str] | None = None, ) -> subprocess.CompletedProcess[str]: """Run one fake Slurm command and optionally raise on failure.""" if not command: @@ -117,6 +118,7 @@ def run( argv = tuple(command) self.calls.append(argv) self.inputs.append(input_text) + del environment command_name = Path(argv[0]).name scripted = self._scripted_responses.get(command_name) if scripted: diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 621d4a33e..357405e01 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="8238dd22393a46538169a83046e91230166aaa0f96e0a91255adf6a6074bfa20", + expected_fixture_sha256="8cbbe5dd355d64f1affbf1aa2875ad3400d504ac47e7c0818a8e2f1736d88130", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="6153c1769d9b90315c4e7acc767bd5c469d0a5e1f5fc7843baddc1a733cb8e9b", + expected_fixture_sha256="bf0de906e6d3f7f2dae54220f13a920488db7eb9e720d350cfe527d1a9320097", ) diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 8557b113e..7006f5eb9 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -187,6 +187,9 @@ def track_read( attempt = _submitted_attempt(case) case.writer.create_attempt(attempt) + runtime_directory = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/runtime" + assert runtime_directory.is_dir() + assert runtime_directory.stat().st_mode & 0o777 == 0o700 readiness = _readiness(case, attempt) case.writer.write_readiness(readiness) record_names.clear() @@ -1963,6 +1966,25 @@ def second_writer() -> Path: pass +def test_if_possible_lease_does_not_create_resume_workspace( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan_payload = single_node_plan.model_dump(mode="python") + plan_payload["invocation"]["authored"]["resume"] = "if_possible" + case = _initialized_case(tmp_path, authored_run_single, ResolvedSlurmRunPlan.model_validate(plan_payload)) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + resume_path = Path(case.shards[0].resume_workspace.path) + + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "if_possible") as path: + assert path == resume_path + assert not path.exists() + + assert not resume_path.exists() + + def test_dataset_lock_rejects_unsafe_files_without_reclassifying_body_errors( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, From 3419ee51ee269ecdd27b3cdf6196b36a52d4e28c Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 02:38:07 -0300 Subject: [PATCH 02/14] fix(slurm): address integration review findings --- .../data_designer/slurm/runtime/controller.py | 46 +++++-- .../data_designer/slurm/runtime/entrypoint.py | 25 ++-- .../data_designer/slurm/runtime/records.py | 7 +- .../data_designer/slurm/services/artifacts.py | 6 +- .../data_designer/slurm/services/wiring.py | 13 ++ .../data_designer/slurm/state/finalization.py | 121 +++++++++++++++--- .../src/data_designer/slurm/state/store.py | 16 ++- .../tests/runtime/conftest.py | 12 ++ .../tests/runtime/test_controller.py | 29 +++++ .../tests/runtime/test_entrypoint.py | 80 +++++++++++- .../tests/services/test_wiring.py | 47 +++++++ .../tests/state/test_store.py | 75 +++++++++++ 12 files changed, 428 insertions(+), 49 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index 7a77db03a..112e85c99 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -59,6 +59,10 @@ def update_attempt(self, attempt: AttemptManifest) -> AttemptManifest: """Persist a monotonic attempt update.""" ... + def load_attempt(self, shard_id: str, attempt_id: str) -> AttemptManifest: + """Load one persisted attempt.""" + ... + def publish_attempt_result( self, client_result: ClientResult, @@ -82,8 +86,9 @@ def finalize_winner( attempt_id: str, *, published_at: datetime, + completed_at: datetime | None = None, ) -> ShardWinner: - """Publish the immutable winning candidate for a successful attempt.""" + """Commit attempt success and publish its immutable winning candidate.""" ... def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: @@ -163,24 +168,39 @@ def run(self) -> AttemptManifest: self._record_outcome_failure(outcome) self._cleanup_runtime(outcome) self._record_stopped_readiness(outcome) - terminal = self._persist_terminal_outcome(outcome) + if outcome.failure is None: + try: + terminal = self._finalize_success(outcome) + except BaseException as error: + outcome.failure = _normalize_failure(error) + outcome.failure_cause = error + terminal = self._persist_terminal_outcome(outcome) + else: + terminal = self._persist_terminal_outcome(outcome) if outcome.failure is not None: if outcome.failure_cause is outcome.failure: raise outcome.failure raise outcome.failure from outcome.failure_cause - try: - self._state.finalize_winner( - terminal.shard_id, - terminal.attempt_id, - published_at=self._now(), - ) - except BaseException as error: - failure = _normalize_failure(error) - if failure is error: - raise - raise failure from error return terminal + def _finalize_success(self, outcome: _RunOutcome) -> AttemptManifest: + if outcome.candidate_reference is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "successful allocation has no candidate reference", + ) + completed_at = self._now() + if outcome.client_completed_at is not None and outcome.client_completed_at > completed_at: + completed_at = outcome.client_completed_at + self._state.finalize_winner( + self._attempt.shard_id, + self._attempt.attempt_id, + completed_at=completed_at, + published_at=max(self._now(), completed_at), + ) + self._attempt = self._state.load_attempt(self._attempt.shard_id, self._attempt.attempt_id) + return self._attempt + def _capture_execution(self) -> _RunOutcome: try: candidate_reference, completed_at = self._execute() 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 45b28e06c..0a00849b1 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 @@ -90,7 +90,7 @@ def _add_context_arguments(parser: argparse.ArgumentParser) -> None: def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) _validate_attempt_is_executable(context.attempt) - SystemAllocationPreflight.verify_attempt_directory(context.attempt_directory) + SystemAllocationPreflight.verify_attempt_directory(arguments.attempt_dir) SystemAllocationPreflight.verify_ports(context) readiness = _begin_attempt(context, writer) log_directory = execution_log_directory(context.attempt_directory, readiness.revision) @@ -165,7 +165,11 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No ) if return_code != 0: raise SlurmRuntimeError(SlurmRuntimeErrorCode.CLIENT_FAILED, "client generation failed") - client_result, candidate = load_complete_client_candidate(context, context.attempt) + client_result, candidate = load_complete_client_candidate( + context, + context.attempt, + attempt_directory=arguments.attempt_dir, + ) completed_at = client_result.completed_at if candidate.created_at < generation_started_at or completed_at < generation_started_at: raise SlurmRuntimeError( @@ -189,19 +193,12 @@ def _succeed(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N SlurmRuntimeErrorCode.FINALIZATION_FAILED, "successful allocation has no candidate reference", ) - terminal = writer.update_attempt( - attempt.model_copy( - update={ - "state": AttemptLifecycleState.SUCCEEDED, - "terminal_classification": AttemptTerminalClassification.SUCCEEDED, - "updated_at": max(stopped_at, attempt.updated_at), - } - ) - ) + completed_at = max(stopped_at, attempt.updated_at) writer.finalize_winner( - terminal.shard_id, - terminal.attempt_id, - published_at=max(datetime.now(timezone.utc), terminal.updated_at), + attempt.shard_id, + attempt.attempt_id, + completed_at=completed_at, + published_at=max(datetime.now(timezone.utc), completed_at), ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py index c41454609..e61f48017 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py @@ -27,15 +27,18 @@ def load_complete_client_candidate( context: AllocationContext, attempt: AttemptManifest, + *, + attempt_directory: Path | None = None, ) -> tuple[ClientResult, CandidateOutputManifest]: """Load a complete semantic client result and its digest-bound candidate.""" - client_result = _read_record(context.attempt_directory, _CLIENT_RESULT_NAME, ClientResult) + record_directory = context.attempt_directory if attempt_directory is None else attempt_directory + client_result = _read_record(record_directory, _CLIENT_RESULT_NAME, ClientResult) if client_result.outcome is not ClientOutcome.COMPLETE: raise SlurmRuntimeError( SlurmRuntimeErrorCode.CLIENT_FAILED, f"client generation finished with outcome {client_result.outcome.value!r}", ) - candidate = _read_record(context.attempt_directory, _CANDIDATE_NAME, CandidateOutputManifest) + candidate = _read_record(record_directory, _CANDIDATE_NAME, CandidateOutputManifest) try: PlanStateValidator(context.plan).validate_client_candidate( context.shard, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py index 9211898cd..5ae28389c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py @@ -30,6 +30,7 @@ SlurmStateError, SlurmStateWriter, StateConflictError, + StateNotFoundError, ) from data_designer.slurm.state.filesystem import ( open_verified_directory, @@ -122,7 +123,10 @@ def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: da """Mark every initial attempt failed after its held job is cancelled.""" writer = SlurmStateWriter(self._workspace_root, plan.run_id) for shard in plan.shards: - attempt = writer.load_attempt(shard.shard_id, "attempt-0001") + try: + attempt = writer.load_attempt(shard.shard_id, "attempt-0001") + except StateNotFoundError: + continue writer.update_attempt( attempt.model_copy( update={ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index a2ac391f1..581c5f4b7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -307,12 +307,25 @@ def execute( SlurmServiceOperation.EXECUTE_RUN, f"Slurm job {receipt.job_id} was submitted but could not be recorded or cancelled", ) from error + try: + self._record_submission_failure(publisher, plan) + except SlurmServiceError as state_error: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.EXECUTE_RUN, + f"Slurm job {receipt.job_id} was cancelled but its partial submission state could not be updated", + ) from state_error raise except BaseException: try: self._launcher.cancel(receipt.job_id) except Exception: pass + else: + try: + self._record_submission_failure(publisher, plan) + except BaseException: + pass raise try: self._launcher.release(receipt.job_id) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py index 83f0576a4..26f8d0895 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py @@ -24,12 +24,22 @@ StateCorruptionError, StateNotFoundError, ) -from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + ShardManifest, +) from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError from data_designer.slurm.state.reader import StateReader from data_designer.slurm.state.storage import StateStorage -from data_designer.slurm.state.validation import StateContractError, validate_shard_winner +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_transition, + validate_shard_winner, +) @dataclass(frozen=True, slots=True) @@ -37,6 +47,8 @@ class _WinnerResolution: winner: ShardWinner candidate: CandidateOutputManifest plan: ResolvedSlurmRunPlan + persisted_attempt: AttemptManifest + successful_attempt: AttemptManifest already_published: bool @@ -59,10 +71,16 @@ def acquire_dataset_workspace( dataset_path = self._prepare_workspace_or_normalize(shard_id, attempt_id, resume_mode) yield dataset_path - def finalize_winner(self, shard_id: ShardId, attempt_id: AttemptId, published_at: datetime) -> ShardWinner: + def finalize_winner( + self, + shard_id: ShardId, + attempt_id: AttemptId, + published_at: datetime, + completed_at: datetime | None = None, + ) -> ShardWinner: try: with self._storage.acquire_resume_lock(shard_id): - return self._finalize_with_dataset_lease(shard_id, attempt_id, published_at) + return self._finalize_with_dataset_lease(shard_id, attempt_id, published_at, completed_at) except (StateConflictError, StateCorruptionError, StateNotFoundError): raise except (PlanStateContractError, StateContractError, ValidationError) as error: @@ -124,14 +142,22 @@ def _finalize_with_dataset_lease( shard_id: ShardId, attempt_id: AttemptId, published_at: datetime, + completed_at: datetime | None, ) -> ShardWinner: - resolution = self._resolve_under_state_locks(shard_id, attempt_id, published_at) + resolution = self._resolve_under_state_locks(shard_id, attempt_id, published_at, completed_at) if resolution.already_published: return resolution.winner with ExitStack() as resources: artifacts = self._open_candidate_artifacts(resources, resolution) self._validate_artifact_metadata(resolution.candidate, artifacts) - return self._publish_verified_resolution(shard_id, attempt_id, published_at, resolution, artifacts) + return self._publish_verified_resolution( + shard_id, + attempt_id, + published_at, + completed_at, + resolution, + artifacts, + ) def _open_candidate_artifacts( self, @@ -148,20 +174,35 @@ def _publish_verified_resolution( shard_id: ShardId, attempt_id: AttemptId, published_at: datetime, + completed_at: datetime | None, expected: _WinnerResolution, artifacts: VerifiedCandidateArtifacts, ) -> ShardWinner: with self._storage.acquire_shard_lock(shard_id): - current = self._resolve_winner(shard_id, attempt_id, published_at) + current = self._resolve_winner(shard_id, attempt_id, published_at, completed_at) if current.already_published: return current.winner - if current.winner != expected.winner or current.candidate != expected.candidate: + if ( + current.winner != expected.winner + or current.candidate != expected.candidate + or current.successful_attempt != expected.successful_attempt + ): raise StateContractError("attempt finalization records changed during verification") try: artifacts.rebind() except OSError as error: raise StateContractError("candidate output paths changed during finalization") from error - self._storage.publish_winner(expected.winner) + if current.persisted_attempt != current.successful_attempt: + self._storage.replace_attempt(current.successful_attempt) + try: + self._storage.publish_winner(expected.winner) + except BaseException: + winner_was_published = self._winner_was_published(expected.winner) + if not winner_was_published: + self._storage.replace_attempt(current.persisted_attempt) + elif completed_at is not None: + return expected.winner + raise return expected.winner def _resolve_under_state_locks( @@ -169,28 +210,76 @@ def _resolve_under_state_locks( shard_id: ShardId, attempt_id: AttemptId, published_at: datetime, + completed_at: datetime | None, ) -> _WinnerResolution: with self._storage.acquire_shard_lock(shard_id): - return self._resolve_winner(shard_id, attempt_id, published_at) + return self._resolve_winner(shard_id, attempt_id, published_at, completed_at) def _resolve_winner( self, shard_id: ShardId, attempt_id: AttemptId, published_at: datetime, + completed_at: datetime | None, ) -> _WinnerResolution: run, plan, shard = self._reader.load_shard_context(shard_id) attempts = self._reader.load_validated_shard_attempts(run, plan, shard) - attempt = self._reader.get_attempt(attempts, attempt_id) + persisted_attempt = self._reader.get_attempt(attempts, attempt_id) existing = self.load_optional_winner(run, plan, shard, attempts) - client_result, candidate = self._load_finalization_records(attempt) + client_result, candidate = self._load_finalization_records(persisted_attempt) if existing is not None: if existing.attempt_id != attempt_id: raise StateConflictError(f"shard {shard_id!r} already has an immutable winner") - return _WinnerResolution(existing, candidate, plan, True) - winner = self._build_winner(run, shard, attempt, client_result, published_at) - self._validate_finalization_chain(run, plan, shard, attempt, client_result, candidate, winner) - return _WinnerResolution(winner, candidate, plan, False) + return _WinnerResolution( + existing, + candidate, + plan, + persisted_attempt, + persisted_attempt, + True, + ) + successful_attempt = self._successful_attempt(persisted_attempt, completed_at) + winner = self._build_winner(run, shard, successful_attempt, client_result, published_at) + self._validate_finalization_chain( + run, + plan, + shard, + successful_attempt, + client_result, + candidate, + winner, + ) + return _WinnerResolution( + winner, + candidate, + plan, + persisted_attempt, + successful_attempt, + False, + ) + + def _winner_was_published(self, expected: ShardWinner) -> bool: + try: + persisted = self._storage.read_winner(expected.shard_id) + except FileNotFoundError: + return False + if persisted != expected: + raise StateCorruptionError("persisted winner changed during finalization") + return True + + @staticmethod + def _successful_attempt(attempt: AttemptManifest, completed_at: datetime | None) -> AttemptManifest: + if completed_at is None or attempt.state is AttemptLifecycleState.SUCCEEDED: + return attempt + successful = attempt.model_copy( + update={ + "state": AttemptLifecycleState.SUCCEEDED, + "terminal_classification": AttemptTerminalClassification.SUCCEEDED, + "updated_at": max(completed_at, attempt.updated_at), + } + ) + validate_attempt_transition(attempt, successful) + return successful def _prepare_dataset_workspace( self, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py index 138b0a1e3..aa1f86d3c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -197,8 +197,9 @@ def finalize_winner( attempt_id: AttemptId, *, published_at: datetime, + completed_at: datetime | None = None, ) -> ShardWinner: - """Validate attempt-local artifacts and publish one immutable winner.""" + """Validate artifacts and optionally commit attempt success with its winner.""" normalized_shard_id = self._validate_shard_id(shard_id) normalized_attempt_id = self._validate_attempt_id(attempt_id) if ( @@ -207,7 +208,18 @@ def finalize_winner( or published_at.utcoffset() != timedelta(0) ): raise StateConflictError("winner publication timestamp must be timezone-aware UTC") - return self._finalizer.finalize_winner(normalized_shard_id, normalized_attempt_id, published_at) + if completed_at is not None and ( + not isinstance(completed_at, datetime) + or completed_at.tzinfo is None + or completed_at.utcoffset() != timedelta(0) + ): + raise StateConflictError("attempt completion timestamp must be timezone-aware UTC") + return self._finalizer.finalize_winner( + normalized_shard_id, + normalized_attempt_id, + published_at, + completed_at, + ) def load_winner(self, shard_id: ShardId) -> ShardWinner: """Load and validate one shard's immutable winner chain.""" diff --git a/packages/data-designer-slurm/tests/runtime/conftest.py b/packages/data-designer-slurm/tests/runtime/conftest.py index c2be41a0d..475a56ff5 100644 --- a/packages/data-designer-slurm/tests/runtime/conftest.py +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -22,6 +22,7 @@ AttemptLifecycleState, AttemptManifest, AttemptReadiness, + AttemptTerminalClassification, CandidateOutputManifest, SchedulerIdentity, ShardWinner, @@ -89,7 +90,18 @@ def finalize_winner( attempt_id: str, *, published_at: datetime, + completed_at: datetime | None = None, ) -> ShardWinner: + if completed_at is not None: + self.update_attempt( + self.attempt.model_copy( + update={ + "state": AttemptLifecycleState.SUCCEEDED, + "terminal_classification": AttemptTerminalClassification.SUCCEEDED, + "updated_at": completed_at, + } + ) + ) assert self.attempt.state is AttemptLifecycleState.SUCCEEDED assert self.attempt.candidate_output is not None winner = ShardWinner( diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index 524878bf7..51638b2c7 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -221,6 +221,35 @@ def test_controller_publishes_result_before_success_with_real_state_writer( assert state.load_winner(result.shard_id).attempt_id == result.attempt_id +def test_winner_finalization_failure_leaves_controller_attempt_retryable(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + + class FailingState(FakeStateStore): + def finalize_winner(self, *args: object, **kwargs: object) -> None: + raise RuntimeError("injected finalization failure") + + state = FailingState(runtime_case.context.attempt) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment={}, + ) + + with pytest.raises(SlurmRuntimeError, match="allocation runtime failed"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + assert state.attempt.candidate_output is not None + assert state.winners == [] + + def test_preflight_failure_starts_no_process_and_fails_attempt(runtime_case: RuntimeCase) -> None: clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) state = FakeStateStore(runtime_case.context.attempt) diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 8927732cd..18ef8020d 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -14,6 +14,10 @@ from data_designer.slurm.state import AttemptLifecycleState, ReadinessState +class _InjectedFailure(Exception): + pass + + def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.CaptureFixture[str]) -> None: assert ( entrypoint.main( @@ -36,6 +40,45 @@ def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.Capt assert "Traceback" not in captured.err +def test_container_phases_use_the_container_attempt_directory( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, +) -> None: + state = FakeStateStore(runtime_case.context.attempt) + container_attempt_directory = Path("/container/workspace/runs/run-single/shards/shard-00000/attempts/attempt-0001") + _patch_runtime_context(monkeypatch, runtime_case, state) + + def verify_attempt_directory(path: Path) -> None: + assert path == container_attempt_directory + raise _InjectedFailure + + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", verify_attempt_directory) + prepare = entrypoint._parse_arguments( + _phase_arguments( + "prepare", + runtime_case, + runtime_case.context.attempt_directory / "runtime", + container_attempt_directory / "runtime-manifest.json", + attempt_directory=container_attempt_directory, + ) + ) + with pytest.raises(_InjectedFailure): + entrypoint._prepare(prepare, {}) + + monkeypatch.setattr(entrypoint, "client_worker_main", lambda arguments: 0) + + def load_candidate(*args: object, attempt_directory: Path | None = None) -> None: + assert attempt_directory == container_attempt_directory + raise _InjectedFailure + + monkeypatch.setattr(entrypoint, "load_complete_client_candidate", load_candidate) + client = entrypoint._parse_arguments( + _phase_arguments("client", runtime_case, attempt_directory=container_attempt_directory) + ) + with pytest.raises(_InjectedFailure): + entrypoint._client(client, {}) + + def test_control_phases_record_running_ready_and_failed( monkeypatch: pytest.MonkeyPatch, runtime_case: RuntimeCase, @@ -98,6 +141,39 @@ def test_succeed_phase_stops_runtime_and_finalizes_winner( assert state.winners[0].attempt_id == state.attempt.attempt_id +def test_succeed_phase_does_not_strand_success_when_winner_finalization_fails( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, +) -> None: + class FailingState(FakeStateStore): + def finalize_winner(self, *args: object, **kwargs: object) -> None: + raise RuntimeError("injected finalization failure") + + state = FailingState(runtime_case.context.attempt) + runtime_root = runtime_case.context.attempt_directory / "runtime" + manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" + _patch_runtime_context(monkeypatch, runtime_case, state) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda context: None) + monkeypatch.setattr( + entrypoint, + "build_runtime_manifest", + lambda *args, **kwargs: SimpleNamespace(serialize_json=lambda: "{}"), + ) + assert entrypoint.main(_phase_arguments("prepare", runtime_case, runtime_root, manifest_path)) == 0 + assert entrypoint.main(_phase_arguments("ready", runtime_case)) == 0 + state.attempt = state.attempt.model_copy( + update={ + "candidate_output": ArtifactReference( + path=(runtime_case.context.attempt_directory / "output-manifest.json").as_posix(), + sha256="a" * 64, + ), + } + ) + assert entrypoint.main(_phase_arguments("succeed", runtime_case)) == 70 + assert state.attempt.state is AttemptLifecycleState.RUNNING + + def _patch_runtime_context( monkeypatch: pytest.MonkeyPatch, runtime_case: RuntimeCase, @@ -112,13 +188,15 @@ def _phase_arguments( runtime_case: RuntimeCase, runtime_root: Path | None = None, manifest_path: Path | None = None, + *, + attempt_directory: Path | None = None, ) -> tuple[str, ...]: arguments = ( operation, "--plan", (runtime_case.workspace / "runs/run-single/resolved-plan.json").as_posix(), "--attempt-dir", - runtime_case.context.attempt_directory.as_posix(), + (attempt_directory or runtime_case.context.attempt_directory).as_posix(), ) if operation == "prepare": assert runtime_root is not None and manifest_path is not None diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 20db42325..7b4d59270 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -40,6 +40,7 @@ ShardManifest, SlurmStateWriter, StateConflictError, + StateNotFoundError, ) @@ -415,6 +416,52 @@ def test_recording_conflict_cancels_the_accepted_job( assert caught.value.code is SlurmServiceErrorCode.CONFLICT assert launcher.cancellations == [42] + assert len(publisher.submission_failures) == 1 + + +def test_partial_submission_recording_failure_cancels_job_and_fails_created_attempts( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={"array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2})} + ) + _register_images(tmp_path, authored, single_node_plan) + original_create = SlurmStateWriter.create_attempt + created = 0 + + def fail_second_attempt(writer: SlurmStateWriter, attempt: AttemptManifest) -> AttemptManifest: + nonlocal created + created += 1 + if created == 2: + raise StateConflictError("injected partial submission failure") + return original_create(writer, attempt) + + monkeypatch.setattr(SlurmStateWriter, "create_attempt", fail_second_attempt) + launcher = _Launcher() + failed_at = datetime(2026, 9, 8, tzinfo=UTC) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: failed_at, + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert launcher.cancellations == [42] + writer = SlurmStateWriter(tmp_path, "run-wired") + attempt = writer.load_attempt("shard-00000", "attempt-0001") + assert attempt.state is AttemptLifecycleState.FAILED + assert attempt.terminal_classification is AttemptTerminalClassification.CANCELLED + with pytest.raises(StateNotFoundError): + writer.load_attempt("shard-00001", "attempt-0001") def test_release_failure_cancels_the_held_job( diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 7006f5eb9..e46661621 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -1239,6 +1239,81 @@ def test_result_publication_binds_candidate_before_success_transition( case.writer.update_attempt(conflicting_success) +def test_winner_publication_failure_restores_running_attempt_for_retry( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result(case, attempt, dataset_path, complete_attempt=False) + original_publish = case.writer._storage.publish_winner + + def fail_winner_publication(winner: ShardWinner) -> None: + raise OSError("injected winner publication failure") + + monkeypatch.setattr(case.writer._storage, "publish_winner", fail_winner_publication) + completed_at = case.created_at + timedelta(minutes=5) + with pytest.raises(SlurmStateError, match="cannot finalize"): + case.writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=completed_at, + published_at=finalization.published_at, + ) + + persisted = case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) + assert persisted.state is AttemptLifecycleState.RUNNING + assert persisted.candidate_output == finalization.client_result.candidate_output_manifest + + monkeypatch.setattr(case.writer._storage, "publish_winner", original_publish) + winner = case.writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=completed_at, + published_at=finalization.published_at, + ) + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id).state is AttemptLifecycleState.SUCCEEDED + assert case.writer.load_winner(attempt.shard_id) == winner + + +def test_runtime_finalization_converges_after_committed_winner_sync_failure( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result(case, attempt, dataset_path, complete_attempt=False) + winner_path = case.writer.run_root / "shards/shard-00000/winner.json" + original_fsync = state_filesystem.os.fsync + failed = False + + def fail_after_winner_link(descriptor: int) -> None: + nonlocal failed + if winner_path.exists() and not failed: + failed = True + raise OSError("injected winner directory fsync failure") + original_fsync(descriptor) + + monkeypatch.setattr(state_filesystem.os, "fsync", fail_after_winner_link) + winner = case.writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=case.created_at + timedelta(minutes=5), + published_at=finalization.published_at, + ) + + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id).state is AttemptLifecycleState.SUCCEEDED + assert case.writer.load_winner(attempt.shard_id) == winner + + def test_attempt_update_cannot_bind_candidate_before_result_publication( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, From 458dc28bef58a8a5091ea28931cb7e48e241aaf5 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 11:39:10 -0300 Subject: [PATCH 03/14] fix(slurm): address smoke and review findings --- .../data_designer/slurm/client/environment.py | 81 +++++++++--- .../data_designer/slurm/client/execution.py | 46 ++++--- .../data_designer/slurm/config/profiles.py | 11 ++ .../data_designer/slurm/launcher/renderer.py | 8 +- .../data_designer/slurm/runtime/bootstrap.py | 3 +- .../data_designer/slurm/runtime/context.py | 30 ++++- .../src/data_designer/slurm/runtime/paths.py | 40 +++++- .../data_designer/slurm/runtime/preflight.py | 11 +- .../src/data_designer/slurm/services/run.py | 4 +- .../data_designer/slurm/services/wiring.py | 122 +++++++++++++++++- .../data_designer/slurm/state/artifacts.py | 7 +- .../data_designer/slurm/state/finalization.py | 6 +- .../src/data_designer/slurm/state/reader.py | 8 +- .../slurm/state/reconciliation.py | 16 ++- .../src/data_designer/slurm/state/storage.py | 25 +++- .../src/data_designer/slurm/state/store.py | 56 +++++++- .../tests/client/conftest.py | 3 + .../tests/client/test_environment.py | 53 +++++++- .../tests/client/test_worker.py | 46 ++++++- .../contracts/golden/multi_node_plan.json | 3 +- .../contracts/golden/profile_catalog.json | 2 + .../contracts/golden/single_node_plan.json | 3 +- .../tests/contracts/test_profiles.py | 1 + .../golden/finalization_chain.json | 12 +- .../tests/launcher/test_renderer.py | 11 ++ .../tests/runtime/test_bootstrap.py | 4 + .../tests/runtime/test_context.py | 96 ++++++++++++++ .../tests/runtime/test_paths.py | 3 +- .../tests/runtime/test_preflight.py | 40 ++++++ .../tests/services/test_wiring.py | 62 ++++++++- .../golden/rendered/multi_node.sbatch | 3 +- .../golden/rendered/single_node.sbatch | 3 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- .../tests/state/test_state_golden_records.py | 2 +- .../tests/state/test_store.py | 42 ++++++ 35 files changed, 773 insertions(+), 94 deletions(-) create mode 100644 packages/data-designer-slurm/tests/runtime/test_context.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py index aa48076e8..cd339f9e3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -7,6 +7,7 @@ import importlib.metadata import json import os +import posixpath import re import subprocess import sys @@ -45,6 +46,7 @@ class PreparedClientEnvironment: class _BootstrapInputs: run_id: str run_root: Path + logical_run_root: Path shard_id: str attempt_id: str attempt_dir: Path @@ -117,17 +119,25 @@ def _load_bootstrap_inputs( raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "shard identifier is invalid") if not re.fullmatch(r"attempt-[0-9]{4,}", attempt_id) or int(attempt_id.removeprefix("attempt-")) < 1: raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "attempt identifier is invalid") - plan = _load_json_object(plan_path, ClientErrorCode.INVALID_INPUT) - run_id = _require_string(plan, "run_id") + plan_payload = _load_json_object(plan_path, ClientErrorCode.INVALID_INPUT) + run_id = _require_string(plan_payload, "run_id") run_root = plan_path.parent - if plan_path.name != "resolved-plan.json" or run_root.name != run_id or run_root.parent.name != "runs": + authored_reference = _artifact_reference(_require_object(plan_payload, "authored_config")) + logical_run_root = Path(authored_reference.path).parent + if ( + plan_path.name != "resolved-plan.json" + or run_root.name != run_id + or run_root.parent.name != "runs" + or _get_container_path(plan_payload, (logical_run_root / plan_path.name).as_posix()) != plan_path.as_posix() + ): raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "resolved plan path is not canonical") - expected_attempt = run_root / "shards" / shard_id / "attempts" / attempt_id + logical_attempt = logical_run_root / "shards" / shard_id / "attempts" / attempt_id + expected_attempt = Path(_get_container_path(plan_payload, logical_attempt.as_posix(), require_writable=True)) if attempt_dir.as_posix() != expected_attempt.as_posix(): raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "attempt directory does not match the plan") ensure_private_directory(attempt_dir) - client = _require_object(plan, "client") + client = _require_object(plan_payload, "client") image = _require_object(client, "image") image_sha256 = _require_digest(image, "sha256") inspection_record = _require_object(image, "inspection") @@ -142,11 +152,12 @@ def _load_bootstrap_inputs( raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client installer path is invalid") lock_reference = _artifact_reference(_require_object(client, "dependency_lock")) - if Path(lock_reference.path).as_posix() != (run_root / "dependency-lock.json").as_posix(): + if Path(lock_reference.path).as_posix() != (logical_run_root / "dependency-lock.json").as_posix(): raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency lock path is not canonical") return _BootstrapInputs( run_id=run_id, run_root=run_root, + logical_run_root=logical_run_root, shard_id=shard_id, attempt_id=attempt_id, attempt_dir=attempt_dir, @@ -159,7 +170,8 @@ def _load_bootstrap_inputs( def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependencies: lock_bytes = read_regular_bytes( - Path(inputs.dependency_lock.path), missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING + inputs.run_root / "dependency-lock.json", + missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING, ) if _sha256_bytes(lock_bytes) != inputs.dependency_lock.sha256: raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency lock digest differs") @@ -185,7 +197,11 @@ def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependen source = lock.get("source") if source is not None: source_reference = _artifact_reference(_as_object(source)) - _verify_input_artifact(source_reference, inputs.run_root / "inputs") + _verify_input_artifact( + source_reference, + inputs.logical_run_root / "inputs", + inputs.run_root / "inputs", + ) return _VerifiedDependencies( image_distributions=expected_image, overlay_packages=tuple(_as_object_list(lock.get("overlay_packages"))), @@ -199,6 +215,7 @@ def _prepare_overlay( expected_overlay, wheels = _verify_wheels( dependencies.overlay_packages, inputs.run_root / "dependencies", + inputs.logical_run_root / "dependencies", dependencies.image_distributions, ) overlay_path = inputs.attempt_dir / "client-env" / "site-packages" @@ -294,6 +311,7 @@ def _run_installer(command: tuple[str, ...]) -> None: def _verify_wheels( packages: tuple[dict[str, object], ...], dependencies_root: Path, + logical_dependencies_root: Path, image_distributions: tuple[InstalledDistribution, ...], ) -> tuple[tuple[InstalledDistribution, ...], tuple[Path, ...]]: expected: list[InstalledDistribution] = [] @@ -305,12 +323,13 @@ def _verify_wheels( if name in image_names or name in {item.name for item in expected}: raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency distributions overlap") artifact = _artifact_reference(_require_object(package, "artifact")) - wheel = Path(artifact.path) + logical_wheel = Path(artifact.path) try: - wheel.relative_to(dependencies_root) + relative_path = logical_wheel.relative_to(logical_dependencies_root) except ValueError as error: raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency wheel path is not canonical") from error - if wheel.parent != dependencies_root or wheel.suffix != ".whl": + wheel = dependencies_root / relative_path + if logical_wheel.parent != logical_dependencies_root or logical_wheel.suffix != ".whl": raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency wheel path is not canonical") if compute_file_sha256(wheel, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) != artifact.sha256: raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency wheel digest differs") @@ -332,13 +351,14 @@ def _verify_wheels( return tuple(pair[0] for pair in sorted_pairs), tuple(pair[1] for pair in sorted_pairs) -def _verify_input_artifact(reference: ArtifactReference, root: Path) -> None: - path = Path(reference.path) +def _verify_input_artifact(reference: ArtifactReference, logical_root: Path, root: Path) -> None: + logical_path = Path(reference.path) try: - path.relative_to(root) + relative_path = logical_path.relative_to(logical_root) except ValueError as error: raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency source path is not canonical") from error - if path.parent != root or path.suffix != ".json": + path = root / relative_path + if logical_path.parent != logical_root or logical_path.suffix != ".json": raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency source path is not canonical") if compute_file_sha256(path, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) != reference.sha256: raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency source digest differs") @@ -408,6 +428,37 @@ def _artifact_reference(value: dict[str, object]) -> ArtifactReference: raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client artifact reference is invalid") from error +def _get_container_path(plan: dict[str, object], host_path: str, *, require_writable: bool = False) -> str: + if not host_path.startswith("/") or posixpath.normpath(host_path) != host_path: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is invalid") + values = plan.get("container_mounts") + if not isinstance(values, list): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan mounts are invalid") + candidates: list[tuple[str, str, bool]] = [] + for value in values: + mount = _as_object(value) + source = _require_string(mount, "source") + target = _require_string(mount, "target") + read_only = mount.get("read_only") + if ( + not source.startswith("/") + or posixpath.normpath(source) != source + or not target.startswith("/") + or posixpath.normpath(target) != target + or type(read_only) is not bool + ): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan mounts are invalid") + if host_path == source or host_path.startswith(f"{source}/"): + candidates.append((source, target, read_only)) + if not candidates: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is not mounted") + source, target, read_only = max(candidates, key=lambda item: len(item[0])) + if require_writable and read_only: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is not writable") + relative_path = posixpath.relpath(host_path, source) + return target if relative_path == "." else posixpath.join(target, relative_path) + + def _sha256_bytes(value: bytes) -> str: return hashlib.sha256(value).hexdigest() diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index f0bdd3e8e..9b85625d1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -51,6 +51,7 @@ from data_designer.slurm.config.run import LocalStdioMCPProviderConfig, RemoteMCPProviderConfig from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import PlannedShard, ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.paths import get_container_path, get_host_path from data_designer.slurm.state import CandidateOutcome, CandidateOutputFile, CandidateOutputManifest Clock = Callable[[], datetime] @@ -250,7 +251,8 @@ def _build_context( self._validate_prepared(plan, shard, prepared) lock = ResolvedDependencyLock.model_validate_json( read_regular_bytes( - Path(plan.client.dependency_lock.path), missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING + Path(get_container_path(plan, plan.client.dependency_lock.path)), + missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING, ) ) if lock.compute_sha256() != plan.client.dependency_lock.sha256: @@ -275,13 +277,13 @@ def _build_context( self._validate_model_references(builder) self._materialize_seed(plan, shard, builder) mcp_providers = self._materialize_mcp_providers(plan) - self._validate_assets(plan) + managed_assets_path = self._validate_assets(plan) requested_resume = ResumeMode(plan.invocation.authored.resume) - dataset_path = self._dataset_path(shard, prepared, requested_resume) + dataset_path = self._dataset_path(plan, shard, prepared, requested_resume) designer = self._data_designer_factory( artifact_path=dataset_path.parent, model_providers=providers, - managed_assets_path=plan.invocation.effective_input_bindings.managed_assets_path, + managed_assets_path=managed_assets_path.as_posix(), mcp_providers=mcp_providers, auto_configure_logging=False, ) @@ -298,7 +300,8 @@ def _validate_prepared( shard: PlannedShard, prepared: PreparedClientEnvironment, ) -> None: - expected_attempt = Path(shard.resume_workspace.path).parent / "attempts" / prepared.attempt_id + logical_attempt = Path(shard.resume_workspace.path).parent / "attempts" / prepared.attempt_id + expected_attempt = Path(get_container_path(plan, logical_attempt.as_posix(), require_writable=True)) inspection = plan.client.image.inspection_facts if ( plan.run_id != prepared.run_id @@ -314,7 +317,10 @@ def _load_builder(plan: ResolvedSlurmRunPlan) -> dict[str, object]: if plan.builder.inline is not None: return cast(dict[str, object], plan.builder.inline) assert plan.builder.source is not None - payload = read_regular_bytes(Path(plan.builder.source.path), missing_code=ClientErrorCode.INVALID_INPUT) + payload = read_regular_bytes( + Path(get_container_path(plan, plan.builder.source.path)), + missing_code=ClientErrorCode.INVALID_INPUT, + ) if hashlib.sha256(payload).hexdigest() != plan.builder.source.sha256: raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "builder artifact digest differs") try: @@ -390,15 +396,18 @@ def _materialize_seed( seed = builder.get_seed_config() if seed is None or "path" not in type(seed.source).model_fields: raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "seed binding does not match the builder") - path = Path(seed_path) + path = Path(get_container_path(plan, seed_path)) if not path.exists() or not os.access(path, os.R_OK): raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "seed input is unavailable") if shard.input_partition is None: raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "seed partition artifact is missing") - payload = read_regular_bytes(Path(shard.input_partition.path), missing_code=ClientErrorCode.INVALID_INPUT) + payload = read_regular_bytes( + Path(get_container_path(plan, shard.input_partition.path)), + missing_code=ClientErrorCode.INVALID_INPUT, + ) if hashlib.sha256(payload).hexdigest() != shard.input_partition.sha256: raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "seed partition artifact digest differs") - source = seed.source.model_copy(update={"path": seed_path}) + source = seed.source.model_copy(update={"path": path.as_posix()}) builder.with_seed_dataset( source, sampling_strategy=seed.sampling_strategy, @@ -437,20 +446,22 @@ def _materialize_mcp_providers( return providers @staticmethod - def _validate_assets(plan: ResolvedSlurmRunPlan) -> None: + def _validate_assets(plan: ResolvedSlurmRunPlan) -> Path: value = plan.invocation.effective_input_bindings.managed_assets_path assert value is not None - path = Path(value) + path = Path(get_container_path(plan, value)) if not path.is_dir() or not os.access(path, os.R_OK | os.X_OK): raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "managed assets are unavailable") + return path @staticmethod def _dataset_path( + plan: ResolvedSlurmRunPlan, shard: PlannedShard, prepared: PreparedClientEnvironment, resume: ResumeMode, ) -> Path: - resume_path = Path(shard.resume_workspace.path) + resume_path = Path(get_container_path(plan, shard.resume_workspace.path, require_writable=True)) if resume_path.is_symlink(): raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "resume workspace is invalid") if resume is ResumeMode.ALWAYS and (not resume_path.is_dir() or not any(resume_path.iterdir())): @@ -528,7 +539,7 @@ def _validate_creation_result( dataset_path = Path(results.dataset_path) effective_resume = results.effective_resume_mode - shared_path = Path(context.shard.resume_workspace.path) + shared_path = Path(get_container_path(context.plan, context.shard.resume_workspace.path, require_writable=True)) expected_path = shared_path if effective_resume is ResumeMode.ALWAYS else prepared.attempt_dir / "dataset" if ( not dataset_path.is_absolute() @@ -589,7 +600,9 @@ def _build_candidate_manifest( metadata = lazy.pq.read_metadata(exported_path) if metadata.num_rows != creation.actual_records: raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "exported record count differs") - schema_digest = hashlib.sha256(lazy.pq.read_schema(exported_path).serialize().to_pybytes()).hexdigest() + schema_digest = hashlib.sha256( + lazy.pq.read_schema(exported_path).remove_metadata().serialize().to_pybytes() + ).hexdigest() files = ( CandidateOutputFile( relative_path=exported_path.relative_to(dataset_path).as_posix(), @@ -606,7 +619,7 @@ def _build_candidate_manifest( attempt_id=prepared.attempt_id, attempt_ordinal=int(prepared.attempt_id.removeprefix("attempt-")), created_at=created_at, - dataset_path=dataset_path.as_posix(), + dataset_path=get_host_path(context.plan, dataset_path.as_posix(), require_writable=True), requested_records=creation.requested_records, actual_records=creation.actual_records, outcome=( @@ -648,7 +661,8 @@ def _publish_success( requested_resume_mode=context.requested_resume.value, effective_resume_mode=creation.effective_resume.value, candidate_output_manifest=ArtifactReference( - path=candidate_path.as_posix(), sha256=candidate.compute_sha256() + path=get_host_path(context.plan, candidate_path.as_posix(), require_writable=True), + sha256=candidate.compute_sha256(), ), ) publish_private_text(prepared.attempt_dir / "client-result.json", result.serialize_json()) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py index f3fc57987..5567850f0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -30,6 +30,17 @@ class SchedulerProfile(AuthoredConfig): account: Identifier | None = None partition: Identifier | None = None mem_per_gpu: Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:K|M|G|T)$")] | None = None + bin_path: str | None = None + + @field_validator("bin_path") + @classmethod + def validate_bin_path(cls, value: str | None) -> str | None: + if value is None: + return None + validated = validate_absolute_path(value) + if ":" in validated: + raise ValueError("scheduler bin path must name one directory") + return validated class ImageBuildProfile(AuthoredConfig): diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index 609e4fbf4..b145688e7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -11,6 +11,8 @@ from data_designer.slurm.launcher.errors import SlurmBatchRenderError from data_designer.slurm.planning import ResolvedSlurmRunPlan +_SYSTEM_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" + def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: """Render a resolved generation plan as one thin deterministic entrypoint.""" @@ -19,13 +21,15 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi run_root = posixpath.dirname(plan.authored_config.path) plan_path = posixpath.join(run_root, "resolved-plan.json") - directive_text = render_batch_directives(_build_generation_directives(plan)) + directive_text = f"{render_batch_directives(_build_generation_directives(plan))}\n#SBATCH --exclusive" attempt = f"{attempt_ordinal:04d}" + scheduler_bin_path = plan.selected_profile.profile.scheduler.bin_path + command_path = _SYSTEM_PATH if scheduler_bin_path is None else f"{scheduler_bin_path}:{_SYSTEM_PATH}" return f"""#!/usr/bin/env bash {directive_text} set -Eeuo pipefail -export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" +export PATH={quote_shell_value(command_path)} readonly DD_RUNTIME_ARCHIVE={quote_shell_value(plan.runtime_bundle.path)} readonly DD_RUNTIME_SHA256={quote_shell_value(plan.runtime_bundle.sha256)} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index 8419eb815..467312672 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -212,7 +212,8 @@ def _build_client_step( "-m", "data_designer.slurm.runtime.entrypoint", "client", - *command[4:], + *command[4:6], + *command[10:], ) secret_names = collect_secret_environment_names( (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py index c0af8c453..ded559a7f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -8,11 +8,16 @@ from collections.abc import Mapping from pathlib import Path -from data_designer.slurm.planning import PlannedShard +from pydantic import ValidationError + +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.paths import get_container_path from data_designer.slurm.state import SlurmStateWriter +from data_designer.slurm.state.filesystem import open_verified_directory, read_regular_text + +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 def load_allocation_context( @@ -55,7 +60,28 @@ def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWr raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") workspace_root = plan_path.parent.parent.parent run_id = plan_path.parent.name - return SlurmStateWriter(workspace_root, run_id) + try: + with open_verified_directory(plan_path.parent, require_private=True) as descriptor: + content = read_regular_text( + descriptor, + plan_path.name, + plan_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + plan = ResolvedSlurmRunPlan.model_validate_json(content) + except (OSError, UnicodeError, ValueError, ValidationError) as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan is unavailable or invalid" + ) from error + logical_workspace_root = plan.selected_profile.profile.workspace_root + logical_plan_path = Path(logical_workspace_root) / "runs" / run_id / plan_path.name + if plan.run_id != run_id or get_container_path(plan, logical_plan_path.as_posix()) != plan_path.as_posix(): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") + return SlurmStateWriter( + workspace_root, + run_id, + logical_workspace_root=logical_workspace_root, + ) def _select_shard(shards: tuple[PlannedShard, ...], task_id: int) -> PlannedShard: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py index 23315a2c1..1274fe054 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py @@ -44,4 +44,42 @@ def get_container_path( return validate_absolute_path(mapped) -__all__ = ["get_container_path"] +def get_host_path( + plan: ResolvedSlurmRunPlan, + container_path: str, + *, + require_writable: bool = False, +) -> str: + """Map one absolute container path back through the most specific resolved mount.""" + try: + normalized = validate_absolute_path(container_path) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime container path is invalid") from error + candidates = tuple( + mount + for mount in plan.container_mounts + if normalized == mount.target or normalized.startswith(f"{mount.target}/") + ) + if not candidates: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "runtime path is not available through a resolved container mount", + ) + mount = max(candidates, key=lambda value: len(value.target)) + if require_writable and mount.read_only: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "runtime path requires a writable resolved container mount", + ) + relative_path = posixpath.relpath(normalized, mount.target) + mapped = mount.source if relative_path == "." else posixpath.join(mount.source, relative_path) + host_path = validate_absolute_path(mapped) + if get_container_path(plan, host_path, require_writable=require_writable) != normalized: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime path does not map unambiguously through resolved container mounts", + ) + return host_path + + +__all__ = ["get_container_path", "get_host_path"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py index e4c5d8872..5e90ef801 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -38,8 +38,10 @@ def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> """Verify every launch-critical fact before model services start.""" try: self._verify_scheduler(context, environment) - self.verify_attempt_directory(context.attempt_directory) - get_container_path(context.plan, context.attempt_directory.as_posix(), require_writable=True) + attempt_directory = Path( + get_container_path(context.plan, context.attempt_directory.as_posix(), require_writable=True) + ) + self.verify_attempt_directory(attempt_directory) self._verify_artifacts(context) self.verify_ports(context) except SlurmRuntimeError: @@ -113,6 +115,11 @@ def _verify_artifacts(context: AllocationContext) -> None: references.extend(reference for reference in optional_references if reference is not None) unique_references = {(reference.path, reference.sha256): reference for reference in references} for reference in unique_references.values(): + if any( + reference.path == mount.source or reference.path.startswith(f"{mount.source}/") + for mount in plan.container_mounts + ): + reference = reference.model_copy(update={"path": get_container_path(plan, reference.path)}) _verify_artifact(reference) @staticmethod diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py index 9790b877b..68eeaa90a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -66,7 +66,7 @@ def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: """Return the durable M2 status for one run.""" def cancel(self, run_id: Identifier) -> SlurmRunCancellation: - """Request cancellation of active jobs without changing persisted state.""" + """Request cancellation of active jobs.""" class SlurmRunService: @@ -179,7 +179,7 @@ def execute_run() -> SlurmRunExecution: return _invoke_service_backend(operation, execute_run) def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: - """Return persisted M2 records without scheduler reconciliation.""" + """Reconcile scheduler observations and return persisted M2 records.""" operation = SlurmServiceOperation.STATUS_RUN normalized_run_id = _validate_run_id(run_id, operation) backend = self._require_backend(operation) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 581c5f4b7..b4d7dbd69 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -10,7 +10,7 @@ from collections.abc import Callable, Iterator, Mapping from contextlib import ExitStack, contextmanager from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import datetime, timezone from pathlib import Path from typing import Protocol, TypeVar from uuid import uuid4 @@ -22,6 +22,8 @@ ClientDependencyResolver, ResolvedClientDependencies, ) +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.filesystem import ensure_private_directory from data_designer.slurm.config import ( DataDesignerSlurmConfig, ImageBuildRequest, @@ -62,16 +64,34 @@ from data_designer.slurm.serving.resolver import resolve_vllm_server from data_designer.slurm.state import ( AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + EffectiveAttemptState, + SchedulerObservation, + SchedulerState, SlurmStateError, SlurmStateWriter, StateConflictError, StateNotFoundError, + reconcile_attempt_observation, ) RunIdFactory = Callable[[], str] Clock = Callable[[], datetime] _ResultT = TypeVar("_ResultT") _MAX_VISIBLE_JOB_IDS = 16 +_ACTIVE_ATTEMPT_STATES = frozenset( + {AttemptLifecycleState.SUBMITTED, AttemptLifecycleState.PENDING, AttemptLifecycleState.RUNNING} +) +_FAILURE_CLASSIFICATIONS = { + SchedulerState.FAILED: AttemptTerminalClassification.FAILED, + SchedulerState.CANCELLED: AttemptTerminalClassification.CANCELLED, + SchedulerState.TIMED_OUT: AttemptTerminalClassification.TIMED_OUT, + SchedulerState.NODE_FAILED: AttemptTerminalClassification.NODE_FAILED, + SchedulerState.PREEMPTED: AttemptTerminalClassification.PREEMPTED, + SchedulerState.REQUEUED: AttemptTerminalClassification.REQUEUED, + SchedulerState.OUT_OF_MEMORY: AttemptTerminalClassification.OUT_OF_MEMORY, +} class SlurmRunArtifactPublisher(Protocol): @@ -276,6 +296,7 @@ def execute( shard_count=len(plan.shards), batch_script=prepared.batch_script, ) + self._materialize_default_managed_assets(config, plan) export_environment = self._build_export_environment(config) publisher = self._publisher self._initialize_run( @@ -359,6 +380,24 @@ def execute( job_id=receipt.job_id, ) + @staticmethod + def _materialize_default_managed_assets( + config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + ) -> None: + if config.invocation.input_bindings.managed_assets_path is not None: + return + path = plan.invocation.effective_input_bindings.managed_assets_path + assert path is not None + try: + ensure_private_directory(Path(path)) + except (ClientWorkerError, OSError): + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + "managed assets workspace cannot be prepared", + ) from None + def _build_export_environment(self, config: DataDesignerSlurmConfig) -> dict[str, str]: environment = {"SLURM_EXPORT_ENV": "ALL"} if "SLURM_CONF" in self._source_environment: @@ -465,8 +504,14 @@ def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: try: writer = SlurmStateWriter(self._profile.profile.workspace_root, run_id) run = writer.load_run() + persisted_shards = writer.load_shards() + self._reconcile_attempts( + writer, + tuple(attempt for shard in persisted_shards for attempt in writer.load_attempts(shard.shard_id)), + ) shards = [] - for shard in writer.load_shards(): + for shard in persisted_shards: + writer.resume_incomplete_finalization(shard.shard_id, published_at=self._clock()) attempts = tuple( SlurmPersistedAttemptStatus( attempt=attempt, @@ -493,16 +538,83 @@ def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: except SlurmStateError: raise SlurmServiceError(SlurmServiceErrorCode.INTERNAL, operation, "run state cannot be read") from None + def _reconcile_attempts( + self, + writer: SlurmStateWriter, + attempts: tuple[AttemptManifest, ...], + ) -> None: + active = tuple( + attempt for attempt in attempts if attempt.state in _ACTIVE_ATTEMPT_STATES and attempt.scheduler is not None + ) + if not active: + return + identities = tuple(attempt.scheduler for attempt in active if attempt.scheduler is not None) + try: + queue = {entry.job_identity: entry.state for entry in self._launcher.query_queue(identities)} + missing = tuple(identity for identity in identities if identity not in queue) + accounting = ( + {entry.job_identity: entry.state for entry in self._launcher.query_accounting(missing)} + if missing + else {} + ) + except SlurmLauncherError: + return + for attempt in active: + scheduler = attempt.scheduler + assert scheduler is not None + scheduler_state = queue.get(scheduler, accounting.get(scheduler)) + if scheduler_state is None: + continue + readiness = _load_optional(lambda: writer.load_readiness(attempt.shard_id, attempt.attempt_id)) + observed_at = max( + self._clock(), + attempt.updated_at, + readiness.updated_at if readiness is not None else attempt.updated_at, + ) + observation = SchedulerObservation( + schema_version=1, + scheduler=scheduler, + observed_at=observed_at, + state=scheduler_state, + ) + effective = reconcile_attempt_observation( + attempt, + readiness, + observation, + current_time=observed_at, + ) + update: dict[str, object] = {"updated_at": observed_at} + if effective is EffectiveAttemptState.PENDING and attempt.state is AttemptLifecycleState.SUBMITTED: + update["state"] = AttemptLifecycleState.PENDING + elif effective is EffectiveAttemptState.RUNNING and attempt.state in { + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + }: + update["state"] = AttemptLifecycleState.RUNNING + elif effective is EffectiveAttemptState.FAILED: + update.update( + state=AttemptLifecycleState.FAILED, + terminal_classification=_FAILURE_CLASSIFICATIONS.get( + scheduler_state, + AttemptTerminalClassification.UNKNOWN, + ), + ) + else: + continue + try: + writer.update_attempt(attempt.model_copy(update=update)) + except StateConflictError: + continue + def cancel(self, run_id: Identifier) -> SlurmRunCancellation: status = self.status(run_id) - active = {AttemptLifecycleState.SUBMITTED, AttemptLifecycleState.PENDING, AttemptLifecycleState.RUNNING} job_ids = tuple( sorted( { attempt.attempt.scheduler.array_job_id for shard in status.shards for attempt in shard.attempts - if attempt.attempt.state in active and attempt.attempt.scheduler is not None + if attempt.attempt.state in _ACTIVE_ATTEMPT_STATES and attempt.attempt.scheduler is not None } ) ) @@ -644,7 +756,7 @@ def _new_run_id() -> str: def _utc_now() -> datetime: - return datetime.now(UTC) + return datetime.now(timezone.utc) def _format_job_ids(job_ids: list[int]) -> str: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py index 3a3b6b466..c397e50f6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py @@ -12,7 +12,7 @@ from contextlib import ExitStack, contextmanager from dataclasses import dataclass from pathlib import Path, PurePosixPath -from typing import Protocol +from typing import Callable, Protocol import data_designer.lazy_heavy_imports as lazy from data_designer.slurm.state.filesystem import ( @@ -104,9 +104,12 @@ def rebind(self) -> None: class CandidateArtifactVerifier: """Verify one manifest-bounded candidate and lease its files through publication.""" + def __init__(self, path_resolver: Callable[[str], Path] = Path) -> None: + self._path_resolver = path_resolver + @contextmanager def verify(self, candidate: CandidateOutputManifest) -> Iterator[VerifiedCandidateArtifacts]: - dataset_path = Path(candidate.dataset_path) + dataset_path = self._path_resolver(candidate.dataset_path) with ExitStack() as resources: dataset_descriptor = resources.enter_context(open_verified_directory(dataset_path, require_private=True)) dataset = _DirectoryBinding(None, None, dataset_descriptor, dataset_path) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py index 26f8d0895..441ffe13a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py @@ -58,7 +58,7 @@ class WinnerFinalizer: def __init__(self, storage: StateStorage, reader: StateReader) -> None: self._storage = storage self._reader = reader - self._artifacts = CandidateArtifactVerifier() + self._artifacts = CandidateArtifactVerifier(storage.get_local_path) @contextmanager def acquire_dataset_workspace( @@ -294,10 +294,10 @@ def _prepare_dataset_workspace( self.require_no_winner(run, plan, shard, attempts) self._validate_workspace_mode(plan, attempt, resume_mode) if resume_mode == "if_possible": - return Path(shard.resume_workspace.path) + return self._storage.get_local_path(shard.resume_workspace.path) dataset_path = self._storage.ensure_dataset_directory(shard_id, attempt_id, resume_mode) expected_path = ( - Path(shard.resume_workspace.path) + self._storage.get_local_path(shard.resume_workspace.path) if resume_mode == "always" else self._storage.get_attempt_path(shard_id, attempt_id) / "dataset" ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py index 57f0fe4a5..67e23c2a9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py @@ -33,8 +33,8 @@ def load_run(self) -> RunManifest: run = self._storage.read_run() if ( run.run_id != self._run_id - or run.authored_config.path != self._storage.authored_config_path.as_posix() - or run.resolved_plan.path != self._storage.resolved_plan_path.as_posix() + or run.authored_config.path != self._storage.logical_authored_config_path.as_posix() + or run.resolved_plan.path != self._storage.logical_resolved_plan_path.as_posix() ): raise StateCorruptionError(f"run {self._run_id!r} manifest does not match its persisted location") return run @@ -52,7 +52,7 @@ def load_authored_config(self, run: RunManifest | None = None) -> DataDesignerSl except (FileNotFoundError, OSError) as error: raise StateCorruptionError(f"run {self._run_id!r} has no valid authored config") from error if ( - bound_run.authored_config.path != self._storage.authored_config_path.as_posix() + bound_run.authored_config.path != self._storage.logical_authored_config_path.as_posix() or bound_run.authored_config.sha256 != authored_config.compute_sha256() ): raise StateCorruptionError(f"run {self._run_id!r} authored config does not match its manifest") @@ -65,7 +65,7 @@ def load_resolved_plan(self, run: RunManifest | None = None) -> ResolvedSlurmRun except (FileNotFoundError, OSError) as error: raise StateCorruptionError(f"run {self._run_id!r} has no valid resolved plan") from error if ( - bound_run.resolved_plan.path != self._storage.resolved_plan_path.as_posix() + bound_run.resolved_plan.path != self._storage.logical_resolved_plan_path.as_posix() or bound_run.resolved_plan.sha256 != plan.compute_sha256() ): raise StateCorruptionError(f"run {self._run_id!r} resolved plan does not match its manifest") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index 0d00e7b3b..63c0d2a1a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -145,7 +145,7 @@ def validate_readiness_transition( def reconcile_attempt_observation( attempt: AttemptManifest, - readiness: AttemptReadiness, + readiness: AttemptReadiness | None, scheduler: SchedulerObservation, *, current_time: datetime, @@ -154,14 +154,16 @@ def reconcile_attempt_observation( _require_utc(current_time, "current_time") _require(current_time >= scheduler.observed_at, "current_time cannot precede scheduler observation") _require(current_time >= attempt.updated_at, "current_time cannot precede attempt update") - _require(current_time >= readiness.updated_at, "current_time cannot precede readiness update") - _require(readiness.run_id == attempt.run_id, "readiness run_id does not match attempt") - _require(readiness.shard_id == attempt.shard_id, "readiness shard_id does not match attempt") - _require(readiness.attempt_id == attempt.attempt_id, "readiness attempt_id does not match attempt") + if readiness is not None: + _require(current_time >= readiness.updated_at, "current_time cannot precede readiness update") + _require(readiness.run_id == attempt.run_id, "readiness run_id does not match attempt") + _require(readiness.shard_id == attempt.shard_id, "readiness shard_id does not match attempt") + _require(readiness.attempt_id == attempt.attempt_id, "readiness attempt_id does not match attempt") _require(attempt.scheduler is not None, "attempt has no scheduler identity") _require(scheduler.scheduler == attempt.scheduler, "scheduler identity does not match attempt") _require(scheduler.observed_at >= attempt.created_at, "scheduler observation cannot precede attempt creation") - _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") + if readiness is not None: + _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") if scheduler.state in _SCHEDULER_FAILURE_STATES: return EffectiveAttemptState.FAILED @@ -180,7 +182,7 @@ def reconcile_attempt_observation( return EffectiveAttemptState.UNKNOWN if scheduler.state is SchedulerState.UNKNOWN: return EffectiveAttemptState.UNKNOWN - if readiness.state is ReadinessState.FAILED: + if readiness is not None and readiness.state is ReadinessState.FAILED: return EffectiveAttemptState.FAILED if scheduler.state is SchedulerState.PENDING: return EffectiveAttemptState.PENDING diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index d22ef8be2..28c04438e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -59,12 +59,20 @@ class StateStorage: """Own descriptor-bound paths, locking, and record serialization.""" - def __init__(self, workspace_root: Path, run_id: Identifier) -> None: + def __init__( + self, + workspace_root: Path, + run_id: Identifier, + *, + logical_workspace_root: Path | None = None, + ) -> None: self.workspace_root = workspace_root + self.logical_workspace_root = logical_workspace_root or workspace_root self.run_id = run_id self.runs_root = workspace_root / "runs" self.locks_root = self.runs_root / _LOCK_DIRECTORY_NAME self.run_root = self.runs_root / run_id + self.logical_run_root = self.logical_workspace_root / "runs" / run_id @property def authored_config_path(self) -> Path: @@ -74,6 +82,21 @@ def authored_config_path(self) -> Path: def resolved_plan_path(self) -> Path: return self.run_root / _RESOLVED_PLAN_FILENAME + @property + def logical_authored_config_path(self) -> Path: + return self.logical_run_root / _AUTHORED_CONFIG_FILENAME + + @property + def logical_resolved_plan_path(self) -> Path: + return self.logical_run_root / _RESOLVED_PLAN_FILENAME + + def get_local_path(self, logical_path: str | Path) -> Path: + try: + relative_path = Path(logical_path).relative_to(self.logical_workspace_root) + except ValueError as error: + raise StateCorruptionError("persisted path is outside the selected workspace") from error + return self.workspace_root / relative_path + def get_shard_path(self, shard_id: str) -> Path: return self.run_root / _SHARDS_DIRECTORY_NAME / shard_id diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py index aa1f86d3c..52d6efb1f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -23,7 +23,7 @@ StateCorruptionError, StateNotFoundError, ) -from data_designer.slurm.state.execution import AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest from data_designer.slurm.state.finalization import WinnerFinalizer from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError @@ -53,17 +53,31 @@ class SlurmStateWriter: audit snapshots. Args: - workspace_root: Selected compute-visible workspace root. + workspace_root: Workspace root visible to this process. run_id: Stable application-owned run identity. + logical_workspace_root: Optional host-side root persisted in plan records. """ - def __init__(self, workspace_root: str | Path, run_id: Identifier) -> None: + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + *, + logical_workspace_root: str | Path | None = None, + ) -> None: try: normalized_root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_logical_root = validate_absolute_path( + Path(logical_workspace_root if logical_workspace_root is not None else workspace_root).as_posix() + ) normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) except (ValidationError, ValueError) as error: raise SlurmStateError("invalid persisted run location") from error - self._storage = StateStorage(Path(normalized_root), normalized_run_id) + self._storage = StateStorage( + Path(normalized_root), + normalized_run_id, + logical_workspace_root=Path(normalized_logical_root), + ) self._reader = StateReader(self._storage, normalized_run_id) self._results = AttemptResultPublisher(self._storage, self._reader) self._finalizer = WinnerFinalizer(self._storage, self._reader) @@ -225,6 +239,34 @@ def load_winner(self, shard_id: ShardId) -> ShardWinner: """Load and validate one shard's immutable winner chain.""" return self._finalizer.load_winner(self._validate_shard_id(shard_id)) + def resume_incomplete_finalization( + self, + shard_id: ShardId, + *, + published_at: datetime, + ) -> ShardWinner | None: + """Finish winner publication for a successful attempt after process interruption.""" + normalized_shard_id = self._validate_shard_id(shard_id) + try: + return self._finalizer.load_winner(normalized_shard_id) + except StateNotFoundError: + pass + successful = tuple( + attempt + for attempt in self._reader.load_attempts(normalized_shard_id) + if attempt.state is AttemptLifecycleState.SUCCEEDED + ) + if not successful: + return None + if len(successful) != 1: + raise StateCorruptionError(f"shard {normalized_shard_id!r} has multiple successful attempts") + attempt = successful[0] + return self.finalize_winner( + normalized_shard_id, + attempt.attempt_id, + published_at=max(published_at, attempt.updated_at), + ) + def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifest: with self._storage.acquire_resume_and_shard_locks(attempt.shard_id): run, plan, shard = self._reader.load_shard_context(attempt.shard_id) @@ -339,16 +381,16 @@ def _validate_initial_bindings( ) -> None: if run.run_id != self._run_id or resolved_plan.run_id != self._run_id: raise StateContractError("run identity does not match the state writer") - if resolved_plan.selected_profile.profile.workspace_root != self._storage.workspace_root.as_posix(): + if resolved_plan.selected_profile.profile.workspace_root != self._storage.logical_workspace_root.as_posix(): raise StateContractError("resolved plan workspace does not match the state writer") if resolved_plan.authored_config.sha256 != authored_config.compute_sha256(): raise StateContractError("authored config digest does not match the resolved plan") if run.authored_config != resolved_plan.authored_config: raise StateContractError("run authored config does not match the resolved plan") - if run.authored_config.path != self._storage.authored_config_path.as_posix(): + if run.authored_config.path != self._storage.logical_authored_config_path.as_posix(): raise StateContractError("run authored config reference does not match its persisted location") if ( - run.resolved_plan.path != self._storage.resolved_plan_path.as_posix() + run.resolved_plan.path != self._storage.logical_resolved_plan_path.as_posix() or run.resolved_plan.sha256 != resolved_plan.compute_sha256() ): raise StateContractError("run resolved plan reference does not match persisted plan bytes") diff --git a/packages/data-designer-slurm/tests/client/conftest.py b/packages/data-designer-slurm/tests/client/conftest.py index 343225a5d..a47eeed2f 100644 --- a/packages/data-designer-slurm/tests/client/conftest.py +++ b/packages/data-designer-slurm/tests/client/conftest.py @@ -118,6 +118,9 @@ def client_worker_case(tmp_path: Path) -> ClientWorkerCase: lock_payload["python_abi"] = python_abi lock = ResolvedDependencyLock.model_validate_json(json.dumps(lock_payload)) payload["client"]["dependency_lock"]["sha256"] = lock.compute_sha256() + mount = {"source": workspace.as_posix(), "target": workspace.as_posix(), "read_only": False} + payload["selected_profile"]["profile"]["container_mounts"] = [mount] + payload["container_mounts"] = [mount] payload["selected_profile"]["profile_sha256"] = compute_canonical_json_sha256( payload["selected_profile"]["profile"] ) diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 9578adc7d..3eae6f502 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -9,16 +9,19 @@ import subprocess import sys from pathlib import Path +from typing import cast from unittest.mock import Mock import pytest -from conftest import ClientWorkerCase +from conftest import ClientWorkerCase, FakeDataDesigner from data_designer.slurm.client.environment import ClientEnvironmentBuilder, inspect_distributions from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.execution import ClientWorker from data_designer.slurm.client.plugins import discover_plugins from data_designer.slurm.client.records import ClientErrorCode, ClientInstallerOutcome -from data_designer.slurm.contracts import InstalledDistribution +from data_designer.slurm.config import SlurmProfile +from data_designer.slurm.contracts import InstalledDistribution, compute_canonical_json_sha256 from data_designer.slurm.planning import ResolvedSlurmRunPlan @@ -37,6 +40,52 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: assert prepared.installed_distributions == client_worker_case.lock.image_distributions +def test_client_runs_through_non_identity_workspace_mount(client_worker_case: ClientWorkerCase) -> None: + physical_workspace = client_worker_case.plan_path.parents[2] + logical_workspace = "/host/workspace" + payload = cast( + dict[str, object], + json.loads(client_worker_case.plan.serialize_json().replace(physical_workspace.as_posix(), logical_workspace)), + ) + selected = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected["profile"]) + mount = {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False} + profile_payload["container_mounts"] = [mount] + payload["container_mounts"] = [mount] + profile = SlurmProfile.model_validate(profile_payload) + selected["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + return client_worker_case.lock.image_distributions if path is None else () + + prepared = ClientEnvironmentBuilder(inventory=inventory).prepare( + client_worker_case.plan_path, + shard_id=plan.shards[0].shard_id, + attempt_id="attempt-0001", + attempt_dir=client_worker_case.attempt_dir, + ) + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + worker.preflight( + client_worker_case.plan_path, + prepared=prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + result = worker.run( + client_worker_case.plan_path, + prepared=prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert result.dataset_path.startswith(logical_workspace) + assert result.candidate_output_manifest is not None + assert result.candidate_output_manifest.path.startswith(logical_workspace) + + def test_inspect_distributions_omits_path_for_active_environment(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[dict[str, object]] = [] diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index 7d4089df6..0ddaf4916 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -12,8 +12,10 @@ from types import SimpleNamespace from unittest.mock import Mock +import pyarrow as pa +import pyarrow.parquet as pq import pytest -from conftest import ClientWorkerCase, FakeDataDesigner +from conftest import ClientWorkerCase, FakeCreationResults, FakeDataDesigner import data_designer.slurm.client.worker as worker_module from data_designer.config import ResumeMode @@ -30,7 +32,7 @@ ) from data_designer.slurm.contracts import compute_serialized_json_sha256 from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan -from data_designer.slurm.state import CandidateOutputManifest +from data_designer.slurm.state import CandidateOutputManifest, compute_candidate_schema_digest def test_preflight_materializes_endpoint_and_ready_environment(client_worker_case: ClientWorkerCase) -> None: @@ -167,6 +169,46 @@ def test_run_persists_semantic_result_and_candidate( assert progress.phase is ClientProgressPhase.COMPLETE +def test_run_normalizes_parquet_schema_metadata( + client_worker_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def export_with_metadata(results: FakeCreationResults, path: Path, *, format: str) -> Path: + assert format == "parquet" + schema = pa.schema( + [("value", pa.int64())], + metadata={b"producer": b"data-designer"}, + ) + table = pa.Table.from_arrays( + [pa.array(range(results.actual_num_records))], + schema=schema, + ) + pq.write_table(table, path) + return path + + monkeypatch.setattr(FakeCreationResults, "export", export_with_metadata) + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + candidate = CandidateOutputManifest.model_validate_json( + (client_worker_case.attempt_dir / "output-manifest.json").read_text() + ) + schema = pq.read_schema(client_worker_case.attempt_dir / "dataset/part-00000.parquet") + assert candidate.dataset_schema_digest == compute_candidate_schema_digest(schema) + + @pytest.mark.parametrize( ("failure", "expected_code"), ( diff --git a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json index 4d257a8fb..a16835e3f 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -419,13 +419,14 @@ }, "scheduler": { "account": "research", + "bin_path": null, "mem_per_gpu": null, "partition": "batch" }, "schema_version": 1, "workspace_root": "/workspace/primary" }, - "profile_sha256": "80e50266bc5748469c0429533731e1973856f6f58259d89ab6cde6d4dc0e5bc0", + "profile_sha256": "657ae8b87f9f6d271b059e0987b571a83f06622fab1633447dae1326f28f232d", "schema_version": 1, "selection_source": "explicit" }, diff --git a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json index 7fc47e93d..baf71e3b7 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json +++ b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json @@ -15,6 +15,7 @@ }, "scheduler": { "account": "lab", + "bin_path": null, "mem_per_gpu": null, "partition": "gpu" }, @@ -43,6 +44,7 @@ }, "scheduler": { "account": "research", + "bin_path": null, "mem_per_gpu": null, "partition": "batch" }, diff --git a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json index 19b781d06..2ac3d4def 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -265,13 +265,14 @@ }, "scheduler": { "account": "research", + "bin_path": null, "mem_per_gpu": null, "partition": "batch" }, "schema_version": 1, "workspace_root": "/workspace/primary" }, - "profile_sha256": "80e50266bc5748469c0429533731e1973856f6f58259d89ab6cde6d4dc0e5bc0", + "profile_sha256": "657ae8b87f9f6d271b059e0987b571a83f06622fab1633447dae1326f28f232d", "schema_version": 1, "selection_source": "explicit" }, diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py index c0c67c00a..29a7f2dda 100644 --- a/packages/data-designer-slurm/tests/contracts/test_profiles.py +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -98,6 +98,7 @@ def test_explicit_selection_rejects_unknown_cluster(profile_catalog: SlurmProfil lambda payload: payload["clusters"]["lab"].update(host_patterns=["primary-login-*"]), lambda payload: payload["clusters"]["primary"].update(extra="unknown"), lambda payload: payload["clusters"]["primary"].update(workspace_root="relative"), + lambda payload: payload["clusters"]["primary"]["scheduler"].update(bin_path="/opt/slurm:/usr/bin"), lambda payload: payload["clusters"]["primary"].update(host_patterns=["login[broken"]), lambda payload: payload["clusters"]["primary"]["image_build"].update(cpus_per_task=0), lambda payload: payload["clusters"]["primary"]["image_build"].update(memory="0G"), diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index e9933b74c..6798f7507 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -4,12 +4,12 @@ "attempt_ordinal": 1, "candidate_output": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" + "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" }, "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" + "sha256": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" }, "run_id": "run-single", "scheduler": { @@ -38,7 +38,7 @@ } ], "outcome": "complete", - "provenance_digest": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604", + "provenance_digest": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3", "requested_records": 8, "run_id": "run-single", "schema_version": 1, @@ -49,7 +49,7 @@ "attempt_id": "attempt-0001", "candidate_output_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" + "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" }, "completed_at": "2026-08-19T12:05:01Z", "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" + "sha256": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" }, "run_id": "run-single", "schema_version": 1, @@ -118,7 +118,7 @@ "attempt_ordinal": 1, "candidate_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" + "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" }, "published_at": "2026-08-19T12:05:03Z", "run_id": "run-single", diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 90cff5b07..c62a3c391 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -85,6 +85,17 @@ def test_renderer_emits_mem_per_gpu_for_gres_mode(single_node_plan: ResolvedSlur assert "#SBATCH --mem-per-gpu=80G\n" in render_generation_attempt_script(plan, attempt_ordinal=1) +def test_renderer_uses_profile_slurm_bin_path(single_node_plan: ResolvedSlurmRunPlan) -> None: + scheduler = single_node_plan.selected_profile.profile.scheduler.model_copy(update={"bin_path": "/opt/slurm/bin"}) + profile = single_node_plan.selected_profile.profile.model_copy(update={"scheduler": scheduler}) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert 'export PATH="/opt/slurm/bin:/usr/local/sbin:' in script + assert "#SBATCH --exclusive\n" in script + + @pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) def test_renderer_reserves_client_cpus_for_each_gpu_request_mode( single_node_plan: ResolvedSlurmRunPlan, diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index eb9a767ff..8d7de9aac 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -36,3 +36,7 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru "data_designer.slurm.runtime.entrypoint", "client", ) + assert "--shard-id" not in manifest.steps[-1].command + assert "--attempt-id" not in manifest.steps[-1].command + assert "--plan" in manifest.steps[-1].command + assert "--attempt-dir" in manifest.steps[-1].command diff --git a/packages/data-designer-slurm/tests/runtime/test_context.py b/packages/data-designer-slurm/tests/runtime/test_context.py new file mode 100644 index 000000000..731aec670 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_context.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import cast + +from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.context import load_allocation_context +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + RunManifest, + SchedulerIdentity, + ShardManifest, + SlurmStateWriter, +) + + +def test_allocation_context_reads_and_updates_state_through_remapped_workspace( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + physical_workspace = tmp_path / "workspace" + physical_workspace.mkdir() + logical_workspace = single_node_plan.selected_profile.profile.workspace_root + payload = cast(dict[str, object], json.loads(single_node_plan.serialize_json())) + selected = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected["profile"]) + mount = {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False} + profile_payload["container_mounts"] = [mount] + payload["container_mounts"] = [mount] + profile = SlurmProfile.model_validate(profile_payload) + selected["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + created_at = datetime(2026, 9, 9, tzinfo=timezone.utc) + plan_reference = ArtifactReference( + path=f"{logical_workspace}/runs/{plan.run_id}/resolved-plan.json", + sha256=plan.compute_sha256(), + ) + run = RunManifest( + schema_version=1, + run_id=plan.run_id, + created_at=created_at, + authored_config=plan.authored_config, + resolved_plan=plan_reference, + shard_count=1, + ) + shard = ShardManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=plan.shards[0].shard_id, + shard_index=0, + record_range=plan.shards[0].record_range, + input_partition=plan.shards[0].input_partition, + resume_workspace=plan.shards[0].resume_workspace, + created_at=created_at, + ) + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=0) + attempt = AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=plan_reference, + state=AttemptLifecycleState.SUBMITTED, + scheduler=scheduler, + created_at=created_at, + updated_at=created_at, + ) + host_writer = SlurmStateWriter( + physical_workspace, + plan.run_id, + logical_workspace_root=logical_workspace, + ) + host_writer.initialize_run(authored_run_single, plan, run, (shard,)) + host_writer.create_attempt(attempt) + plan_path = physical_workspace / "runs" / plan.run_id / "resolved-plan.json" + attempt_directory = plan_path.parent / "shards" / shard.shard_id / "attempts" / attempt.attempt_id + + context, runtime_writer = load_allocation_context( + plan_path, + attempt_directory, + {"SLURM_ARRAY_TASK_ID": "0", "SLURM_ARRAY_JOB_ID": "4101"}, + ) + runtime_writer.update_attempt(context.attempt.model_copy(update={"state": AttemptLifecycleState.RUNNING})) + + assert context.attempt_directory.as_posix().startswith(logical_workspace) + assert host_writer.load_attempt(shard.shard_id, attempt.attempt_id).state is AttemptLifecycleState.RUNNING diff --git a/packages/data-designer-slurm/tests/runtime/test_paths.py b/packages/data-designer-slurm/tests/runtime/test_paths.py index ebdd820db..7e55a3806 100644 --- a/packages/data-designer-slurm/tests/runtime/test_paths.py +++ b/packages/data-designer-slurm/tests/runtime/test_paths.py @@ -8,7 +8,7 @@ from data_designer.slurm.config import ContainerMount from data_designer.slurm.runtime.errors import SlurmRuntimeError -from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.paths import get_container_path, get_host_path def test_container_path_uses_most_specific_mount_and_preserves_relative_path(runtime_case: RuntimeCase) -> None: @@ -26,6 +26,7 @@ def test_container_path_uses_most_specific_mount_and_preserves_relative_path(run mapped = get_container_path(plan, f"{nested}/run-single/resolved-plan.json") assert mapped == "/container/runs/run-single/resolved-plan.json" + assert get_host_path(plan, mapped) == f"{nested}/run-single/resolved-plan.json" def test_container_path_rejects_unmounted_and_read_only_writes(runtime_case: RuntimeCase) -> None: diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index 1766840b0..570309622 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -10,9 +10,11 @@ import pytest from conftest import RuntimeCase +from data_designer.slurm.config import ContainerMount from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.runtime import preflight as runtime_preflight from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.preflight import SystemAllocationPreflight, _verify_artifact @@ -136,3 +138,41 @@ def test_attempt_directory_must_be_restrictive(runtime_case: RuntimeCase) -> Non with pytest.raises(SlurmRuntimeError, match="restrictive directory"): SystemAllocationPreflight.verify_attempt_directory(runtime_case.context.attempt_directory) + + +def test_preflight_translates_host_paths_into_container( + runtime_case: RuntimeCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + logical_attempt = runtime_case.context.attempt_directory + container_workspace = tmp_path / "container-workspace" + container_attempt = container_workspace / logical_attempt.relative_to(runtime_case.workspace) + logical_attempt.rmdir() + container_attempt.mkdir(parents=True, mode=0o700) + plan = runtime_case.context.plan.model_copy( + update={ + "container_mounts": ( + ContainerMount( + source=runtime_case.workspace.as_posix(), + target=container_workspace.as_posix(), + ), + ) + } + ) + context = AllocationContext( + plan=plan, + shard=runtime_case.context.shard, + attempt=runtime_case.context.attempt, + attempt_directory=logical_attempt, + ) + verified_paths: list[str] = [] + monkeypatch.setattr(SystemAllocationPreflight, "_verify_scheduler", lambda *args: None) + monkeypatch.setattr(SystemAllocationPreflight, "verify_ports", lambda *args: None) + monkeypatch.setattr(runtime_preflight, "_verify_artifact", lambda reference: verified_paths.append(reference.path)) + + SystemAllocationPreflight().verify(context, {}) + + assert verified_paths + assert any(path.startswith(container_workspace.as_posix()) for path in verified_paths) + assert not any(path.startswith(runtime_case.workspace.as_posix()) for path in verified_paths) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 7b4d59270..8e5c1976a 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -5,7 +5,7 @@ import hashlib from collections.abc import Mapping -from datetime import UTC, datetime +from datetime import datetime, timezone from pathlib import Path import pytest @@ -23,7 +23,12 @@ from data_designer.slurm.images.records import RegisteredImage from data_designer.slurm.images.registry import ImageRegistryStore from data_designer.slurm.launcher.errors import SlurmLauncherError -from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmProcessExitCode, + SlurmQueueEntry, +) from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.services import ( SlurmServiceError, @@ -37,6 +42,7 @@ AttemptTerminalClassification, RunManifest, SchedulerIdentity, + SchedulerState, ShardManifest, SlurmStateWriter, StateConflictError, @@ -57,6 +63,8 @@ def __init__( self.releases: list[int] = [] self.held_submissions: list[bool] = [] self.exported_environments: list[dict[str, str]] = [] + self.queue_entries: tuple[SlurmQueueEntry, ...] = () + self.accounting_entries: tuple[SlurmAccountingEntry, ...] = () self.gpu_counts = gpu_counts self.cancel_error = cancel_error self.release_error = release_error @@ -87,6 +95,14 @@ def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: assert partition is not None return self.gpu_counts + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + del selectors + return self.queue_entries + + def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ...]: + del selectors + return self.accounting_entries + class _Publisher: def __init__( @@ -214,7 +230,7 @@ def test_production_wiring_submits_after_publisher_initialization( _register_images(tmp_path, authored_run_single, single_node_plan) launcher = _Launcher() publisher = _Publisher() - submitted_at = datetime(2026, 9, 8, tzinfo=UTC) + submitted_at = datetime(2026, 9, 8, tzinfo=timezone.utc) service = create_slurm_run_service( profile=_profile(tmp_path, profile_catalog), artifact_publisher=publisher, # type: ignore[arg-type] @@ -234,6 +250,7 @@ def test_production_wiring_submits_after_publisher_initialization( assert launcher.held_submissions == [True] assert launcher.releases == [42] assert launcher.exported_environments == [{"SLURM_EXPORT_ENV": "ALL"}] + assert (tmp_path / "managed-assets").is_dir() def test_production_publisher_rejects_force_before_submission( @@ -349,6 +366,39 @@ def test_production_wiring_publishes_initial_state_before_releasing_submission( assert launcher.releases == [42] +def test_status_reconciles_cancelled_scheduler_attempt( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: datetime(2026, 9, 8, tzinfo=timezone.utc), + package_version="0.9.2", + ) + result = service.execute(authored_run_single, source_root=tmp_path) + service.cancel(result.run_id) + scheduler = SchedulerIdentity(array_job_id=42, array_task_id=0) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=scheduler, + state=SchedulerState.CANCELLED, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=15), + ), + ) + + status = service.status(result.run_id) + + attempt = status.shards[0].attempts[0].attempt + assert attempt.state is AttemptLifecycleState.FAILED + assert attempt.terminal_classification is AttemptTerminalClassification.CANCELLED + + def test_auto_gpu_resolution_rejects_mixed_node_shapes( tmp_path: Path, profile_catalog: SlurmProfileCatalog, @@ -442,7 +492,7 @@ def fail_second_attempt(writer: SlurmStateWriter, attempt: AttemptManifest) -> A monkeypatch.setattr(SlurmStateWriter, "create_attempt", fail_second_attempt) launcher = _Launcher() - failed_at = datetime(2026, 9, 8, tzinfo=UTC) + failed_at = datetime(2026, 9, 8, tzinfo=timezone.utc) service = create_slurm_run_service( profile=_profile(tmp_path, profile_catalog), launcher=launcher, # type: ignore[arg-type] @@ -472,7 +522,7 @@ def test_release_failure_cancels_the_held_job( ) -> None: _register_images(tmp_path, authored_run_single, single_node_plan) launcher = _Launcher(release_error=SlurmLauncherError("release failed")) - failed_at = datetime(2026, 9, 8, tzinfo=UTC) + failed_at = datetime(2026, 9, 8, tzinfo=timezone.utc) service = create_slurm_run_service( profile=_profile(tmp_path, profile_catalog), launcher=launcher, # type: ignore[arg-type] @@ -552,7 +602,7 @@ def test_production_status_and_cancel_use_only_persisted_m2_records( package_version="0.9.2", ) plan = service.plan(authored_run_single) - now = datetime(2026, 9, 8, tzinfo=UTC) + now = datetime(2026, 9, 8, tzinfo=timezone.utc) plan_reference = ArtifactReference( path=(tmp_path / "runs/run-wired/resolved-plan.json").as_posix(), sha256=plan.compute_sha256(), diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index f9bf180ab..1a6ac1c3c 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -7,13 +7,14 @@ #SBATCH --time=03:55:00 #SBATCH --array=0-1%2 #SBATCH --gres=gpu:8 +#SBATCH --exclusive set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="902919292da35aca426a191bee775acf140b0f8b7440b55887fb2ba60ed98a52" +readonly DD_PLAN_SHA256="69ff0f0e29331a513d732f9ad9f5fc979f5ed833b75061d8fdf769feff96ab11" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 60d0ce84e..f6ba56a7e 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -7,13 +7,14 @@ #SBATCH --time=03:55:00 #SBATCH --array=0 #SBATCH --gres=gpu:8 +#SBATCH --exclusive set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" +readonly DD_PLAN_SHA256="b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 357405e01..b38ab401f 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="8cbbe5dd355d64f1affbf1aa2875ad3400d504ac47e7c0818a8e2f1736d88130", + expected_fixture_sha256="391ca71f5d1f1a15808d66f2ebeb2242b1aeb248c72ee0b7dea26a561669defb", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="bf0de906e6d3f7f2dae54220f13a920488db7eb9e720d350cfe527d1a9320097", + expected_fixture_sha256="5e2dff54c1d3f534f4ae43e93063067681f8762b8995a996b9dd74bce2ecb5ef", ) diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index 6599cf992..057ba5bf9 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -44,7 +44,7 @@ CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json", CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json", ) -_COMPATIBILITY_FIXTURE_DIGEST = "1a51941bbb2d8aac7114c197f51105104297de1c6a367cfab0f4111b5d7eebaa" +_COMPATIBILITY_FIXTURE_DIGEST = "16460710f383bb9f4650f17e3bac95e4e4a495ca3d80873d7282aa3953533f97" @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index e46661621..026a94cf4 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -1280,6 +1280,48 @@ def fail_winner_publication(winner: ShardWinner) -> None: assert case.writer.load_winner(attempt.shard_id) == winner +def test_fresh_writer_resumes_finalization_interrupted_after_success_commit( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result(case, attempt, dataset_path, complete_attempt=False) + original_replace = case.writer._storage.replace_attempt + + def interrupt_after_success_commit(updated: AttemptManifest) -> None: + original_replace(updated) + if updated.state is AttemptLifecycleState.SUCCEEDED: + raise KeyboardInterrupt("injected process interruption") + + monkeypatch.setattr(case.writer._storage, "replace_attempt", interrupt_after_success_commit) + with pytest.raises(KeyboardInterrupt, match="process interruption"): + case.writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=case.created_at + timedelta(minutes=5), + published_at=finalization.published_at, + ) + + resumed = SlurmStateWriter(case.workspace, case.plan.run_id) + persisted = resumed.load_attempt(attempt.shard_id, attempt.attempt_id) + assert persisted.state is AttemptLifecycleState.SUCCEEDED + with pytest.raises(StateNotFoundError): + resumed.load_winner(attempt.shard_id) + + winner = resumed.resume_incomplete_finalization( + attempt.shard_id, + published_at=finalization.published_at, + ) + + assert winner is not None + assert resumed.load_winner(attempt.shard_id) == winner + + def test_runtime_finalization_converges_after_committed_winner_sync_failure( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, From 26a8f7eebf969b95e771201ec37b16b829cd52ae Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 11:56:27 -0300 Subject: [PATCH 04/14] docs(slurm): align status reconciliation help --- packages/data-designer-slurm/src/data_designer/slurm/cli.py | 4 ++-- .../src/data_designer/slurm/services/results.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 7315f4d3b..f2c9e12fd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -70,7 +70,7 @@ def status_command( profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), cluster: str | None = typer.Option(None, "--cluster"), ) -> None: - """Show persisted M2 run status without scheduler reconciliation.""" + """Reconcile scheduler observations and show persisted M2 run status.""" operation = SlurmServiceOperation.STATUS_RUN result = _invoke( operation, @@ -85,7 +85,7 @@ def cancel_command( profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), cluster: str | None = typer.Option(None, "--cluster"), ) -> None: - """Request job cancellation; status changes after reconciliation.""" + """Request cancellation of active jobs.""" operation = SlurmServiceOperation.CANCEL_RUN result = _invoke( operation, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/results.py b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py index 51950bdcf..2417b000d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/results.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py @@ -40,7 +40,7 @@ def validate_state(self) -> SlurmRunExecution: class SlurmPersistedAttemptStatus(ContractValue): - """Persisted attempt state available without scheduler reconciliation.""" + """Persisted attempt state returned after scheduler reconciliation.""" attempt: AttemptManifest readiness: AttemptReadiness | None = None From 89e38d751449c875cb801963c8a71b1b3a8d0310 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 13:09:26 -0300 Subject: [PATCH 05/14] fix(slurm): honor nested runtime mounts --- .../data_designer/slurm/runtime/context.py | 1 + .../data_designer/slurm/state/finalization.py | 4 +- .../src/data_designer/slurm/state/storage.py | 28 +++++- .../src/data_designer/slurm/state/store.py | 5 +- .../tests/runtime/test_context.py | 23 ++++- .../tests/state/test_store.py | 94 ++++++++++++++++++- 6 files changed, 141 insertions(+), 14 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py index ded559a7f..0230a1e52 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -81,6 +81,7 @@ def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWr workspace_root, run_id, logical_workspace_root=logical_workspace_root, + local_path_resolver=lambda path: get_container_path(plan, path, require_writable=True), ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py index 441ffe13a..38fb75138 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py @@ -299,7 +299,9 @@ def _prepare_dataset_workspace( expected_path = ( self._storage.get_local_path(shard.resume_workspace.path) if resume_mode == "always" - else self._storage.get_attempt_path(shard_id, attempt_id) / "dataset" + else self._storage.get_local_path( + self._storage.logical_run_root / "shards" / shard_id / "attempts" / attempt_id / "dataset" + ) ) if dataset_path != expected_path: raise StateContractError("dataset workspace path does not match the resolved plan") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index 28c04438e..f940c3238 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -11,13 +11,13 @@ from collections.abc import Iterator from contextlib import ExitStack, contextmanager from pathlib import Path -from typing import Literal, TypeVar +from typing import Callable, Literal, TypeVar from pydantic import ValidationError from data_designer.slurm.client import ClientResult from data_designer.slurm.config import DataDesignerSlurmConfig -from data_designer.slurm.contracts import AttemptId, ContractRecord, Identifier, ShardId +from data_designer.slurm.contracts import AttemptId, ContractRecord, Identifier, ShardId, validate_absolute_path from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.state.errors import SlurmStateError, StateCorruptionError, StateNotFoundError from data_designer.slurm.state.execution import AttemptManifest, RunManifest, ShardManifest @@ -54,6 +54,7 @@ _MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 _ATTEMPT_NAME_PATTERN = re.compile(r"^attempt-[0-9]{4,}$") _RecordT = TypeVar("_RecordT", bound=ContractRecord) +_LocalPathResolver = Callable[[str], str | Path] class StateStorage: @@ -65,6 +66,7 @@ def __init__( run_id: Identifier, *, logical_workspace_root: Path | None = None, + local_path_resolver: _LocalPathResolver | None = None, ) -> None: self.workspace_root = workspace_root self.logical_workspace_root = logical_workspace_root or workspace_root @@ -73,6 +75,7 @@ def __init__( self.locks_root = self.runs_root / _LOCK_DIRECTORY_NAME self.run_root = self.runs_root / run_id self.logical_run_root = self.logical_workspace_root / "runs" / run_id + self._local_path_resolver = local_path_resolver @property def authored_config_path(self) -> Path: @@ -91,10 +94,16 @@ def logical_resolved_plan_path(self) -> Path: return self.logical_run_root / _RESOLVED_PLAN_FILENAME def get_local_path(self, logical_path: str | Path) -> Path: + logical_path = Path(logical_path) try: - relative_path = Path(logical_path).relative_to(self.logical_workspace_root) + relative_path = logical_path.relative_to(self.logical_workspace_root) except ValueError as error: raise StateCorruptionError("persisted path is outside the selected workspace") from error + if self._local_path_resolver is not None: + try: + return Path(validate_absolute_path(Path(self._local_path_resolver(logical_path.as_posix())).as_posix())) + except Exception as error: + raise StateCorruptionError("persisted path has no valid local mapping") from error return self.workspace_root / relative_path def get_shard_path(self, shard_id: str) -> Path: @@ -315,11 +324,20 @@ def ensure_dataset_directory( dataset_path = self.get_shard_path(shard_id) / _DATASET_DIRECTORY_NAME with self.open_shard_directory(shard_id) as descriptor: ensure_private_child_directory(descriptor, _DATASET_DIRECTORY_NAME, dataset_path) - return dataset_path + logical_path = self.logical_run_root / _SHARDS_DIRECTORY_NAME / shard_id / _DATASET_DIRECTORY_NAME + return self.get_local_path(logical_path) dataset_path = self.get_attempt_path(shard_id, attempt_id) / _DATASET_DIRECTORY_NAME with self.open_attempt_directory(shard_id, attempt_id) as descriptor: ensure_private_child_directory(descriptor, _DATASET_DIRECTORY_NAME, dataset_path) - return dataset_path + logical_path = ( + self.logical_run_root + / _SHARDS_DIRECTORY_NAME + / shard_id + / _ATTEMPTS_DIRECTORY_NAME + / attempt_id + / _DATASET_DIRECTORY_NAME + ) + return self.get_local_path(logical_path) def read_finalization_records( self, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py index 52d6efb1f..9076e1955 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -5,7 +5,7 @@ from __future__ import annotations -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager from datetime import datetime, timedelta from pathlib import Path @@ -56,6 +56,7 @@ class SlurmStateWriter: workspace_root: Workspace root visible to this process. run_id: Stable application-owned run identity. logical_workspace_root: Optional host-side root persisted in plan records. + local_path_resolver: Optional mapping for persisted paths with nested mounts. """ def __init__( @@ -64,6 +65,7 @@ def __init__( run_id: Identifier, *, logical_workspace_root: str | Path | None = None, + local_path_resolver: Callable[[str], str | Path] | None = None, ) -> None: try: normalized_root = validate_absolute_path(Path(workspace_root).as_posix()) @@ -77,6 +79,7 @@ def __init__( Path(normalized_root), normalized_run_id, logical_workspace_root=Path(normalized_logical_root), + local_path_resolver=local_path_resolver, ) self._reader = StateReader(self._storage, normalized_run_id) self._results = AttemptResultPublisher(self._storage, self._reader) diff --git a/packages/data-designer-slurm/tests/runtime/test_context.py b/packages/data-designer-slurm/tests/runtime/test_context.py index 731aec670..52d7fc3fa 100644 --- a/packages/data-designer-slurm/tests/runtime/test_context.py +++ b/packages/data-designer-slurm/tests/runtime/test_context.py @@ -30,12 +30,24 @@ def test_allocation_context_reads_and_updates_state_through_remapped_workspace( physical_workspace = tmp_path / "workspace" physical_workspace.mkdir() logical_workspace = single_node_plan.selected_profile.profile.workspace_root + logical_attempts = ( + Path(logical_workspace) + / "runs" + / single_node_plan.run_id + / "shards" + / single_node_plan.shards[0].shard_id + / "attempts" + ) + fast_attempts = tmp_path / "fast-attempts" payload = cast(dict[str, object], json.loads(single_node_plan.serialize_json())) selected = cast(dict[str, object], payload["selected_profile"]) profile_payload = cast(dict[str, object], selected["profile"]) - mount = {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False} - profile_payload["container_mounts"] = [mount] - payload["container_mounts"] = [mount] + mounts = [ + {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False}, + {"source": logical_attempts.as_posix(), "target": fast_attempts.as_posix(), "read_only": False}, + ] + profile_payload["container_mounts"] = mounts + payload["container_mounts"] = mounts profile = SlurmProfile.model_validate(profile_payload) selected["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) @@ -83,7 +95,8 @@ def test_allocation_context_reads_and_updates_state_through_remapped_workspace( host_writer.initialize_run(authored_run_single, plan, run, (shard,)) host_writer.create_attempt(attempt) plan_path = physical_workspace / "runs" / plan.run_id / "resolved-plan.json" - attempt_directory = plan_path.parent / "shards" / shard.shard_id / "attempts" / attempt.attempt_id + attempt_directory = fast_attempts / attempt.attempt_id + attempt_directory.mkdir(parents=True, mode=0o700) context, runtime_writer = load_allocation_context( plan_path, @@ -91,6 +104,8 @@ def test_allocation_context_reads_and_updates_state_through_remapped_workspace( {"SLURM_ARRAY_TASK_ID": "0", "SLURM_ARRAY_JOB_ID": "4101"}, ) runtime_writer.update_attempt(context.attempt.model_copy(update={"state": AttemptLifecycleState.RUNNING})) + with runtime_writer.acquire_dataset_workspace(shard.shard_id, attempt.attempt_id, "never") as dataset_path: + assert dataset_path == attempt_directory / "dataset" assert context.attempt_directory.as_posix().startswith(logical_workspace) assert host_writer.load_attempt(shard.shard_id, attempt.attempt_id).state is AttemptLifecycleState.RUNNING diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 026a94cf4..9a59c86cd 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -23,9 +23,10 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.slurm import filesystem as slurm_filesystem from data_designer.slurm.client import ClientOutcome, ClientResult -from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile +from data_designer.slurm.config import ContainerMount, DataDesignerSlurmConfig, SlurmProfile from data_designer.slurm.contracts import ArtifactReference, ContractValue, compute_canonical_json_sha256, pretty_json from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.paths import get_container_path from data_designer.slurm.state import ( AttemptId, AttemptLifecycleState, @@ -1322,6 +1323,91 @@ def interrupt_after_success_commit(updated: AttemptManifest) -> None: assert resumed.load_winner(attempt.shard_id) == winner +def test_interrupted_finalization_resumes_through_nested_attempt_mount( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + base_case = _build_case(tmp_path, authored_run_single, single_node_plan) + logical_attempts = base_case.workspace / "runs" / base_case.plan.run_id / "shards" / "shard-00000" / "attempts" + fast_attempts = tmp_path / "fast-attempts" + mounts = ( + *base_case.plan.container_mounts, + ContainerMount(source=logical_attempts.as_posix(), target=fast_attempts.as_posix()), + ) + profile = SlurmProfile.model_validate( + base_case.plan.selected_profile.profile.model_dump(mode="python") | {"container_mounts": list(mounts)} + ) + selected_profile = base_case.plan.selected_profile.model_copy( + update={ + "profile": profile, + "profile_sha256": compute_canonical_json_sha256(profile.model_dump(mode="json")), + } + ) + plan = ResolvedSlurmRunPlan.model_validate( + base_case.plan.model_dump(mode="python") | {"selected_profile": selected_profile, "container_mounts": mounts} + ) + plan_reference = base_case.run.resolved_plan.model_copy(update={"sha256": plan.compute_sha256()}) + run = base_case.run.model_copy(update={"resolved_plan": plan_reference}) + writer = SlurmStateWriter( + base_case.workspace, + plan.run_id, + local_path_resolver=lambda path: get_container_path(plan, path, require_writable=True), + ) + case = _StateCase( + workspace=base_case.workspace, + authored_config=base_case.authored_config, + plan=plan, + run=run, + shards=base_case.shards, + writer=writer, + created_at=base_case.created_at, + ) + writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + attempt = _submitted_attempt(case).model_copy(update={"resolved_plan": plan_reference}) + writer.create_attempt(attempt) + logical_dataset = logical_attempts / attempt.attempt_id / "dataset" + with writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + assert dataset_path == fast_attempts / attempt.attempt_id / "dataset" + dataset_path.mkdir(parents=True, mode=0o700) + finalization = _persist_complete_result( + case, + attempt, + dataset_path, + manifest_dataset_path=logical_dataset, + complete_attempt=False, + ) + original_replace = writer._storage.replace_attempt + + def interrupt_after_success_commit(updated: AttemptManifest) -> None: + original_replace(updated) + if updated.state is AttemptLifecycleState.SUCCEEDED: + raise KeyboardInterrupt("injected process interruption") + + monkeypatch.setattr(writer._storage, "replace_attempt", interrupt_after_success_commit) + with pytest.raises(KeyboardInterrupt, match="process interruption"): + writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=case.created_at + timedelta(minutes=5), + published_at=finalization.published_at, + ) + + resumed = SlurmStateWriter( + case.workspace, + plan.run_id, + local_path_resolver=lambda path: get_container_path(plan, path, require_writable=True), + ) + winner = resumed.resume_incomplete_finalization( + attempt.shard_id, + published_at=finalization.published_at, + ) + + assert winner is not None + assert resumed.load_winner(attempt.shard_id) == winner + + def test_runtime_finalization_converges_after_committed_winner_sync_failure( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -2356,6 +2442,7 @@ def _persist_complete_result( relative_path: str = "part-00000.parquet", physical_records: int | None = None, reported_schema_digest: str | None = None, + manifest_dataset_path: Path | None = None, complete_attempt: bool = True, ) -> _FinalizationCase: running_attempt = _validated_copy( @@ -2375,6 +2462,7 @@ def _persist_complete_result( lazy.pq.write_table(table, output_path) output_path.chmod(0o644) content = output_path.read_bytes() + persisted_dataset_path = manifest_dataset_path or dataset_path candidate = CandidateOutputManifest( schema_version=1, run_id=case.plan.run_id, @@ -2382,7 +2470,7 @@ def _persist_complete_result( attempt_id=attempt.attempt_id, attempt_ordinal=attempt.attempt_ordinal, created_at=case.created_at + timedelta(minutes=3), - dataset_path=dataset_path.as_posix(), + dataset_path=persisted_dataset_path.as_posix(), requested_records=requested_records, actual_records=requested_records, outcome=CandidateOutcome.COMPLETE, @@ -2411,7 +2499,7 @@ def _persist_complete_result( requested_records=requested_records, actual_records=requested_records, outcome=ClientOutcome.COMPLETE, - dataset_path=dataset_path.as_posix(), + dataset_path=persisted_dataset_path.as_posix(), early_shutdown=False, requested_resume_mode=case.plan.invocation.authored.resume, effective_resume_mode="never", From 0eafebecf053105a0bb0523b9cd51ae0f0f311a5 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 13:17:26 -0300 Subject: [PATCH 06/14] fix(slurm): allow exclusive GPU allocations --- .../src/data_designer/slurm/runtime/preflight.py | 2 +- .../tests/runtime/test_preflight.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py index 5e90ef801..6d1e4f1e1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -76,7 +76,7 @@ def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str] f"scheduler environment {name!r} does not match the resolved plan", ) visible_gpus = environment.get("CUDA_VISIBLE_DEVICES") or environment.get("SLURM_JOB_GPUS") - if _parse_gpu_count(visible_gpus) != context.plan.resolved_gpus_per_node: + if _parse_gpu_count(visible_gpus) < context.plan.resolved_gpus_per_node: raise SlurmRuntimeError( SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "allocation GPU visibility does not match the resolved plan", diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index 570309622..2f510a1f6 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -30,6 +30,18 @@ def test_scheduler_preflight_accepts_exact_one_node_gpu_shape(runtime_case: Runt SystemAllocationPreflight._verify_scheduler(runtime_case.context, environment) +def test_scheduler_preflight_accepts_exclusive_allocation_gpu_superset(runtime_case: RuntimeCase) -> None: + environment = { + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_NUM_NODES": "1", + "SLURM_NODEID": "0", + "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7,8", + } + + SystemAllocationPreflight._verify_scheduler(runtime_case.context, environment) + + @pytest.mark.parametrize( ("name", "value"), ( From c70dfd9fa743aad835a2aa6435d7c0212a345a03 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 13:43:10 -0300 Subject: [PATCH 07/14] fix(slurm): accept exclusive job GPU visibility --- .../data_designer/slurm/runtime/entrypoint.sh | 2 +- .../tests/runtime/test_shell_runtime.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) 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 ee5c40100..dc15efd34 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 @@ -115,7 +115,7 @@ dd_verify_gpu_count() { ((count += 1)) done fi - [[ ${count} == "${DD_EXPECTED_GPUS}" ]] || { + [[ ${DD_EXPECTED_GPUS} =~ ^[1-9][0-9]*$ && ${count} -ge ${DD_EXPECTED_GPUS} ]] || { printf '%s\n' 'allocation GPU visibility does not match the resolved plan' >&2 return 65 } 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 04ed4f208..fb29b92aa 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -84,6 +84,33 @@ def test_staged_shell_modules_parse_as_bash(tmp_path: Path) -> None: assert completed.returncode == 0, completed.stderr +@pytest.mark.parametrize( + ("visible_gpus", "expected_gpus", "expected_status"), + (("0,1,2,3,4,5,6,7", 1, 0), ("0", 2, 65)), +) +def test_bash_gpu_count_requires_planned_minimum( + visible_gpus: str, + expected_gpus: int, + expected_status: int, +) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "entrypoint.sh").as_posix())} +DD_EXPECTED_GPUS={expected_gpus} +dd_verify_gpu_count +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={**os.environ, "CUDA_VISIBLE_DEVICES": visible_gpus}, + ) + + assert completed.returncode == expected_status + + def test_shell_helpers_handle_empty_and_sparse_arrays() -> None: runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" command = f""" From 49ce1b72b966dd774d6d3ab83596bd47bc906379 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 14:07:25 -0300 Subject: [PATCH 08/14] fix(slurm): load plans through nested mounts --- .../src/data_designer/slurm/runtime/context.py | 8 ++++---- .../data-designer-slurm/tests/runtime/test_context.py | 8 +++++++- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py index 0230a1e52..6817300b4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -56,10 +56,8 @@ def load_allocation_context( def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWriter: if not plan_path.is_absolute() or not attempt_directory.is_absolute(): raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime paths must be absolute") - if plan_path.name != "resolved-plan.json" or plan_path.parent.parent.name != "runs": + if plan_path.name != "resolved-plan.json": raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") - workspace_root = plan_path.parent.parent.parent - run_id = plan_path.parent.name try: with open_verified_directory(plan_path.parent, require_private=True) as descriptor: content = read_regular_text( @@ -73,10 +71,12 @@ def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWr raise SlurmRuntimeError( SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan is unavailable or invalid" ) from error + run_id = plan.run_id logical_workspace_root = plan.selected_profile.profile.workspace_root logical_plan_path = Path(logical_workspace_root) / "runs" / run_id / plan_path.name - if plan.run_id != run_id or get_container_path(plan, logical_plan_path.as_posix()) != plan_path.as_posix(): + if get_container_path(plan, logical_plan_path.as_posix()) != plan_path.as_posix(): raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") + workspace_root = Path(get_container_path(plan, logical_workspace_root, require_writable=True)) return SlurmStateWriter( workspace_root, run_id, diff --git a/packages/data-designer-slurm/tests/runtime/test_context.py b/packages/data-designer-slurm/tests/runtime/test_context.py index 52d7fc3fa..a0d0e7d19 100644 --- a/packages/data-designer-slurm/tests/runtime/test_context.py +++ b/packages/data-designer-slurm/tests/runtime/test_context.py @@ -30,6 +30,8 @@ def test_allocation_context_reads_and_updates_state_through_remapped_workspace( physical_workspace = tmp_path / "workspace" physical_workspace.mkdir() logical_workspace = single_node_plan.selected_profile.profile.workspace_root + logical_runs = Path(logical_workspace) / "runs" + fast_runs = tmp_path / "fast-runs" logical_attempts = ( Path(logical_workspace) / "runs" @@ -44,6 +46,7 @@ def test_allocation_context_reads_and_updates_state_through_remapped_workspace( profile_payload = cast(dict[str, object], selected["profile"]) mounts = [ {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False}, + {"source": logical_runs.as_posix(), "target": fast_runs.as_posix(), "read_only": False}, {"source": logical_attempts.as_posix(), "target": fast_attempts.as_posix(), "read_only": False}, ] profile_payload["container_mounts"] = mounts @@ -94,7 +97,10 @@ def test_allocation_context_reads_and_updates_state_through_remapped_workspace( ) host_writer.initialize_run(authored_run_single, plan, run, (shard,)) host_writer.create_attempt(attempt) - plan_path = physical_workspace / "runs" / plan.run_id / "resolved-plan.json" + plan_path = fast_runs / plan.run_id / "resolved-plan.json" + plan_path.parent.mkdir(parents=True, mode=0o700) + plan_path.write_text(plan.serialize_json()) + plan_path.chmod(0o600) attempt_directory = fast_attempts / attempt.attempt_id attempt_directory.mkdir(parents=True, mode=0o700) From da96e4e57599dd33840c6efa64b4680a3d959dd1 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 14:24:24 -0300 Subject: [PATCH 09/14] fix(slurm): validate logical client paths --- .../data_designer/slurm/client/environment.py | 4 +- .../tests/client/test_environment.py | 46 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py index cd339f9e3..c3ee4880c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -126,8 +126,8 @@ def _load_bootstrap_inputs( logical_run_root = Path(authored_reference.path).parent if ( plan_path.name != "resolved-plan.json" - or run_root.name != run_id - or run_root.parent.name != "runs" + or logical_run_root.name != run_id + or logical_run_root.parent.name != "runs" or _get_container_path(plan_payload, (logical_run_root / plan_path.name).as_posix()) != plan_path.as_posix() ): raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "resolved plan path is not canonical") diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 3eae6f502..ec7f9f643 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -86,6 +86,52 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: assert result.candidate_output_manifest.path.startswith(logical_workspace) +def test_environment_prepares_through_nested_runs_mount( + client_worker_case: ClientWorkerCase, + tmp_path: Path, +) -> None: + physical_workspace = client_worker_case.plan_path.parents[2] + physical_run = tmp_path / "fast-run" + logical_workspace = "/host/workspace" + payload = cast( + dict[str, object], + json.loads(client_worker_case.plan.serialize_json().replace(physical_workspace.as_posix(), logical_workspace)), + ) + selected = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected["profile"]) + mounts = [ + {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False}, + { + "source": f"{logical_workspace}/runs/{payload['run_id']}", + "target": physical_run.as_posix(), + "read_only": False, + }, + ] + profile_payload["container_mounts"] = mounts + payload["container_mounts"] = mounts + profile = SlurmProfile.model_validate(profile_payload) + selected["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + plan_path = physical_run / "resolved-plan.json" + plan_path.parent.mkdir(parents=True) + plan_path.write_text(plan.serialize_json()) + (plan_path.parent / "dependency-lock.json").write_text(client_worker_case.lock.serialize_json()) + shard_id = plan.shards[0].shard_id + attempt_dir = plan_path.parent / "shards" / shard_id / "attempts" / "attempt-0001" + + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + return client_worker_case.lock.image_distributions if path is None else () + + prepared = ClientEnvironmentBuilder(inventory=inventory).prepare( + plan_path, + shard_id=shard_id, + attempt_id="attempt-0001", + attempt_dir=attempt_dir, + ) + + assert prepared.attempt_dir == attempt_dir + + def test_inspect_distributions_omits_path_for_active_environment(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[dict[str, object]] = [] From 5017c7e57a2aa7cf3882c73ae3c098f0381aa2a5 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 14:32:45 -0300 Subject: [PATCH 10/14] fix(slurm): resolve client artifacts by mount --- .../data_designer/slurm/client/environment.py | 38 +++++++++---------- .../tests/client/test_environment.py | 19 +++++----- 2 files changed, 28 insertions(+), 29 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py index c3ee4880c..6e96d15ea 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -45,7 +45,6 @@ class PreparedClientEnvironment: @dataclass(frozen=True) class _BootstrapInputs: run_id: str - run_root: Path logical_run_root: Path shard_id: str attempt_id: str @@ -55,6 +54,7 @@ class _BootstrapInputs: installer_path: Path inspection: dict[str, object] dependency_lock: ArtifactReference + plan: dict[str, object] @dataclass(frozen=True) @@ -121,7 +121,6 @@ def _load_bootstrap_inputs( raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "attempt identifier is invalid") plan_payload = _load_json_object(plan_path, ClientErrorCode.INVALID_INPUT) run_id = _require_string(plan_payload, "run_id") - run_root = plan_path.parent authored_reference = _artifact_reference(_require_object(plan_payload, "authored_config")) logical_run_root = Path(authored_reference.path).parent if ( @@ -156,7 +155,6 @@ def _load_bootstrap_inputs( raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency lock path is not canonical") return _BootstrapInputs( run_id=run_id, - run_root=run_root, logical_run_root=logical_run_root, shard_id=shard_id, attempt_id=attempt_id, @@ -166,11 +164,12 @@ def _load_bootstrap_inputs( installer_path=installer_path, inspection=inspection, dependency_lock=lock_reference, + plan=plan_payload, ) def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependencies: lock_bytes = read_regular_bytes( - inputs.run_root / "dependency-lock.json", + Path(_get_container_path(inputs.plan, inputs.dependency_lock.path)), missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING, ) if _sha256_bytes(lock_bytes) != inputs.dependency_lock.sha256: @@ -200,7 +199,7 @@ def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependen _verify_input_artifact( source_reference, inputs.logical_run_root / "inputs", - inputs.run_root / "inputs", + inputs.plan, ) return _VerifiedDependencies( image_distributions=expected_image, @@ -214,11 +213,20 @@ def _prepare_overlay( ) -> tuple[Path, ClientInstallerOutcome, tuple[InstalledDistribution, ...]]: expected_overlay, wheels = _verify_wheels( dependencies.overlay_packages, - inputs.run_root / "dependencies", inputs.logical_run_root / "dependencies", dependencies.image_distributions, + inputs.plan, ) - overlay_path = inputs.attempt_dir / "client-env" / "site-packages" + logical_overlay = ( + inputs.logical_run_root + / "shards" + / inputs.shard_id + / "attempts" + / inputs.attempt_id + / "client-env" + / "site-packages" + ) + overlay_path = Path(_get_container_path(inputs.plan, logical_overlay.as_posix(), require_writable=True)) outcome = self._install_overlay(inputs.installer_path, wheels, expected_overlay, overlay_path) installed = tuple(sorted((*dependencies.image_distributions, *expected_overlay), key=lambda item: item.name)) return overlay_path, outcome, installed @@ -310,9 +318,9 @@ def _run_installer(command: tuple[str, ...]) -> None: def _verify_wheels( packages: tuple[dict[str, object], ...], - dependencies_root: Path, logical_dependencies_root: Path, image_distributions: tuple[InstalledDistribution, ...], + plan: dict[str, object], ) -> tuple[tuple[InstalledDistribution, ...], tuple[Path, ...]]: expected: list[InstalledDistribution] = [] wheels: list[Path] = [] @@ -324,13 +332,9 @@ def _verify_wheels( raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency distributions overlap") artifact = _artifact_reference(_require_object(package, "artifact")) logical_wheel = Path(artifact.path) - try: - relative_path = logical_wheel.relative_to(logical_dependencies_root) - except ValueError as error: - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency wheel path is not canonical") from error - wheel = dependencies_root / relative_path if logical_wheel.parent != logical_dependencies_root or logical_wheel.suffix != ".whl": raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency wheel path is not canonical") + wheel = Path(_get_container_path(plan, logical_wheel.as_posix())) if compute_file_sha256(wheel, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) != artifact.sha256: raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency wheel digest differs") try: @@ -351,15 +355,11 @@ def _verify_wheels( return tuple(pair[0] for pair in sorted_pairs), tuple(pair[1] for pair in sorted_pairs) -def _verify_input_artifact(reference: ArtifactReference, logical_root: Path, root: Path) -> None: +def _verify_input_artifact(reference: ArtifactReference, logical_root: Path, plan: dict[str, object]) -> None: logical_path = Path(reference.path) - try: - relative_path = logical_path.relative_to(logical_root) - except ValueError as error: - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency source path is not canonical") from error - path = root / relative_path if logical_path.parent != logical_root or logical_path.suffix != ".json": raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency source path is not canonical") + path = Path(_get_container_path(plan, logical_path.as_posix())) if compute_file_sha256(path, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) != reference.sha256: raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency source digest differs") diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index ec7f9f643..c42e284fd 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -86,24 +86,25 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: assert result.candidate_output_manifest.path.startswith(logical_workspace) -def test_environment_prepares_through_nested_runs_mount( +def test_environment_maps_each_artifact_through_nested_mounts( client_worker_case: ClientWorkerCase, tmp_path: Path, ) -> None: physical_workspace = client_worker_case.plan_path.parents[2] - physical_run = tmp_path / "fast-run" + physical_plan = tmp_path / "fast-plan" / "resolved-plan.json" logical_workspace = "/host/workspace" payload = cast( dict[str, object], json.loads(client_worker_case.plan.serialize_json().replace(physical_workspace.as_posix(), logical_workspace)), ) + logical_plan = f"{logical_workspace}/runs/{payload['run_id']}/resolved-plan.json" selected = cast(dict[str, object], payload["selected_profile"]) profile_payload = cast(dict[str, object], selected["profile"]) mounts = [ {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False}, { - "source": f"{logical_workspace}/runs/{payload['run_id']}", - "target": physical_run.as_posix(), + "source": logical_plan, + "target": physical_plan.as_posix(), "read_only": False, }, ] @@ -112,18 +113,16 @@ def test_environment_prepares_through_nested_runs_mount( profile = SlurmProfile.model_validate(profile_payload) selected["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) - plan_path = physical_run / "resolved-plan.json" - plan_path.parent.mkdir(parents=True) - plan_path.write_text(plan.serialize_json()) - (plan_path.parent / "dependency-lock.json").write_text(client_worker_case.lock.serialize_json()) + physical_plan.parent.mkdir(parents=True) + physical_plan.write_text(plan.serialize_json()) shard_id = plan.shards[0].shard_id - attempt_dir = plan_path.parent / "shards" / shard_id / "attempts" / "attempt-0001" + attempt_dir = physical_workspace / "runs" / plan.run_id / "shards" / shard_id / "attempts" / "attempt-0001" def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: return client_worker_case.lock.image_distributions if path is None else () prepared = ClientEnvironmentBuilder(inventory=inventory).prepare( - plan_path, + physical_plan, shard_id=shard_id, attempt_id="attempt-0001", attempt_dir=attempt_dir, From d772ec8b05086dcc74053e49ef21b96ad9eab5c5 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 14:48:41 -0300 Subject: [PATCH 11/14] fix(slurm): map serving model paths --- .../data_designer/slurm/runtime/bootstrap.py | 2 +- .../src/data_designer/slurm/runtime/steps.py | 6 +++-- .../tests/runtime/test_steps.py | 27 +++++++++++++++++-- 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index 467312672..d53d5c67a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -271,7 +271,7 @@ def _build_server_step( step_id=process.process_id, role=RuntimeStepRole.SERVER, image_path=deployment.image.path, - command=build_vllm_command(deployment, process), + command=build_vllm_command(deployment, process, context.plan), cpus=context.plan.client.authored.cpus, gpu_indices=tuple(process.gpu_indices), literal_environment=literal_environment, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py index 77ee2346b..387520a0b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -280,7 +280,7 @@ def _build_vllm_step( runtime_root: Path, ) -> RuntimeStep: _validate_local_vllm_process(process) - command = build_vllm_command(deployment, process) + command = build_vllm_command(deployment, process, plan) environment, container_environment = _build_vllm_environment( deployment, source_environment, @@ -311,13 +311,15 @@ def _validate_local_vllm_process(process: ResolvedVllmProcess) -> None: def build_vllm_command( deployment: ResolvedVllmServerDeployment, process: ResolvedVllmProcess, + plan: ResolvedSlurmRunPlan, ) -> tuple[str, ...]: if process.http_port is None: # pragma: no cover - validated before command construction raise AssertionError("vLLM HTTP port is unavailable") + model = get_container_path(plan, deployment.model) if deployment.model.startswith("/") else deployment.model command: tuple[str, ...] = ( deployment.executable_path, "serve", - deployment.model, + model, "--served-model-name", deployment.served_model_name, "--host", diff --git a/packages/data-designer-slurm/tests/runtime/test_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py index aa400785e..a9aeefcbe 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -3,14 +3,16 @@ from __future__ import annotations +import json from datetime import datetime, timezone from pathlib import Path +from typing import cast import pytest from conftest import RuntimeCase, relocate_plan -from data_designer.slurm.config import QueueBackpressureConfig -from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.config import QueueBackpressureConfig, SlurmProfile +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.runtime.backpressure import ( MAX_WAITING_REQUESTS_ENVIRONMENT, @@ -21,6 +23,7 @@ from data_designer.slurm.runtime.steps import ( DefaultClientStepBuilder, build_endpoint_steps, + build_vllm_command, build_vllm_steps, ) from data_designer.slurm.serving.resolver import resolve_vllm_server @@ -105,6 +108,26 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti assert all("CUDA_VISIBLE_DEVICES" not in step.environment for step in client_steps) +def test_vllm_command_maps_absolute_model_path(runtime_case: RuntimeCase) -> None: + plan = runtime_case.context.plan + host_model = runtime_case.workspace / "models" / "generator" + payload = cast(dict[str, object], json.loads(plan.serialize_json())) + deployment_payload = cast(dict[str, object], cast(list[object], payload["deployments"])[0]) + authored = cast(dict[str, object], deployment_payload["authored"]) + authored["model"] = host_model.as_posix() + authored["served_model_name"] = deployment_payload["served_model_name"] + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected_profile["profile"]) + profile = SlurmProfile.model_validate(profile_payload) + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + deployment = resolve_vllm_server(plan, plan.deployments[0].deployment_id) + + command = build_vllm_command(deployment, deployment.processes[0], plan) + + assert command[2] == "/workspace/primary/models/generator" + + def test_client_worker_receives_only_persisted_identity_and_logical_endpoint(runtime_case: RuntimeCase) -> None: context = runtime_case.context deployment = resolve_vllm_server(context.plan, context.plan.deployments[0].deployment_id) From 7cbd678dfae06a42b1d62d4681e5bc1ec2899b9b Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 15:30:29 -0300 Subject: [PATCH 12/14] fix(slurm): avoid idle GPU reaper --- .../data_designer/slurm/launcher/renderer.py | 2 +- .../data_designer/slurm/runtime/bootstrap.py | 4 +- .../data_designer/slurm/runtime/controller.py | 7 +- .../data_designer/slurm/runtime/entrypoint.py | 4 +- .../src/data_designer/slurm/runtime/ports.py | 75 +++++++++++++++++++ .../data_designer/slurm/runtime/preflight.py | 5 +- .../tests/launcher/test_renderer.py | 2 +- .../tests/runtime/test_ports.py | 54 +++++++++++++ .../golden/rendered/multi_node.sbatch | 1 - .../golden/rendered/single_node.sbatch | 1 - .../slurm_test_fakes/test_rendered_scripts.py | 4 +- 11 files changed, 141 insertions(+), 18 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_ports.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index b145688e7..8fdae3522 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -21,7 +21,7 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi run_root = posixpath.dirname(plan.authored_config.path) plan_path = posixpath.join(run_root, "resolved-plan.json") - directive_text = f"{render_batch_directives(_build_generation_directives(plan))}\n#SBATCH --exclusive" + directive_text = render_batch_directives(_build_generation_directives(plan)) attempt = f"{attempt_ordinal:04d}" scheduler_bin_path = plan.selected_profile.profile.scheduler.bin_path command_path = _SYSTEM_PATH if scheduler_bin_path is None else f"{scheduler_bin_path}:{_SYSTEM_PATH}" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index d53d5c67a..8a7e35745 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -24,13 +24,13 @@ from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStepRole from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.ports import resolve_allocation_deployments from data_designer.slurm.runtime.steps import ( build_client_command, build_endpoint_command, build_vllm_command, ) from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment -from data_designer.slurm.serving.resolver import resolve_vllm_server from data_designer.slurm.serving.vllm import ResolvedVllmProcess from data_designer.slurm.types import EnvironmentName, Identifier, NetworkPort, Sha256Digest @@ -125,7 +125,7 @@ def build_runtime_manifest( """Build the secret-free one-node command handoff for the Bash controller.""" plan = context.plan runtime_container_root = get_container_path(plan, runtime_root.as_posix(), require_writable=True) - deployments = tuple(resolve_vllm_server(plan, item.deployment_id) for item in plan.deployments) + deployments = resolve_allocation_deployments(context) endpoints = tuple( RuntimeEndpoint( model_alias=deployment.model_alias, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index 112e85c99..c78d93173 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -18,6 +18,7 @@ from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.logs import bind_execution_logs, execution_log_directory from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStep +from data_designer.slurm.runtime.ports import resolve_allocation_deployments from data_designer.slurm.runtime.preflight import AllocationPreflight from data_designer.slurm.runtime.probes import ReadinessProber from data_designer.slurm.runtime.records import load_complete_client_candidate @@ -28,7 +29,6 @@ ) from data_designer.slurm.runtime.supervisor import ManagedStep, RuntimeClock, StepSupervisor from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment -from data_designer.slurm.serving.resolver import resolve_vllm_server from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, @@ -268,10 +268,7 @@ def _execute(self) -> tuple[ArtifactReference, datetime]: def _prepare_runtime(self) -> _RuntimeTopology: self._validate_attempt_state() - deployments = tuple( - resolve_vllm_server(self._context.plan, deployment.deployment_id) - for deployment in self._context.plan.deployments - ) + deployments = resolve_allocation_deployments(self._context) restarting = self._readiness is not None if restarting: self._begin_execution(deployments, ReadinessState.RESTARTING) 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 0a00849b1..7a68f8198 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 @@ -20,10 +20,10 @@ from data_designer.slurm.runtime.logs import execution_log_directory from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.ports import resolve_allocation_deployments from data_designer.slurm.runtime.preflight import SystemAllocationPreflight from data_designer.slurm.runtime.records import load_complete_client_candidate from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment -from data_designer.slurm.serving.resolver import resolve_vllm_server from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, @@ -348,7 +348,7 @@ def _load_optional_readiness(context: AllocationContext, writer: SlurmStateWrite def _resolve_deployments(context: AllocationContext) -> tuple[ResolvedVllmServerDeployment, ...]: - return tuple(resolve_vllm_server(context.plan, item.deployment_id) for item in context.plan.deployments) + return resolve_allocation_deployments(context) def _validate_attempt_is_executable(attempt: AttemptManifest) -> None: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py new file mode 100644 index 000000000..d0dc3693c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-local port isolation for shared Slurm nodes.""" + +from __future__ import annotations + +import hashlib + +from data_designer.slurm.planning import PortClaim, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.resolver import resolve_vllm_server + +_PORT_RANGE_START = 10000 +_PORT_RANGE_SIZE = 20000 + + +def resolve_allocation_deployments(context: AllocationContext) -> tuple[ResolvedVllmServerDeployment, ...]: + """Resolve deployments with ports isolated to one scheduler array element.""" + plan = _remap_plan_ports(context) + return tuple(resolve_vllm_server(plan, item.deployment_id) for item in plan.deployments) + + +def allocation_ports(context: AllocationContext) -> tuple[int, ...]: + """Return every allocation-local port in deterministic claim order.""" + plan = _remap_plan_ports(context) + return tuple(port.port for port in plan.client.ports) + tuple( + port.port for deployment in plan.deployments for port in deployment.ports + ) + + +def _remap_plan_ports(context: AllocationContext) -> ResolvedSlurmRunPlan: + scheduler = context.attempt.scheduler + if scheduler is None: # pragma: no cover - AllocationContext rejects this state + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation has no scheduler identity") + claims = context.plan.client.ports + tuple( + port for deployment in context.plan.deployments for port in deployment.ports + ) + claims_by_node: dict[int, set[int]] = {} + for claim in claims: + claims_by_node.setdefault(claim.node_index, set()).add(claim.port) + reserved_by_node: dict[int, set[int]] = {} + otel_port = context.plan.invocation.effective_run_config.get("otel_metrics_port") + if type(otel_port) is int and _PORT_RANGE_START <= otel_port < _PORT_RANGE_START + _PORT_RANGE_SIZE: + reserved_by_node[context.plan.client.host_node_index] = {otel_port} + mapping: dict[tuple[int, int], int] = {} + for node_index, planned_ports in claims_by_node.items(): + seed = f"{scheduler.array_job_id}:{scheduler.array_task_id}:{node_index}".encode() + offset = int.from_bytes(hashlib.sha256(seed).digest()[:8], "big") % _PORT_RANGE_SIZE + reserved = reserved_by_node.get(node_index, set()) + for planned_port in sorted(planned_ports): + for _ in range(_PORT_RANGE_SIZE): + port = _PORT_RANGE_START + offset + offset = (offset + 1) % _PORT_RANGE_SIZE + if port not in reserved: + mapping[(node_index, planned_port)] = port + reserved.add(port) + break + else: # pragma: no cover - plan contracts bound claims below the allocation range + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation port range is exhausted") + + def remap(port: PortClaim) -> PortClaim: + return port.model_copy(update={"port": mapping[(port.node_index, port.port)]}) + + client = context.plan.client.model_copy(update={"ports": tuple(remap(port) for port in context.plan.client.ports)}) + deployments = tuple( + deployment.model_copy(update={"ports": tuple(remap(port) for port in deployment.ports)}) + for deployment in context.plan.deployments + ) + return context.plan.model_copy(update={"client": client, "deployments": deployments}) + + +__all__ = ["allocation_ports", "resolve_allocation_deployments"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py index 6d1e4f1e1..e9c89f335 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -18,6 +18,7 @@ from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.ports import allocation_ports _DIGEST_CHUNK_SIZE = 1024 * 1024 _GPU_COUNT_PATTERN = re.compile(r"^(?:gpu(?::[^:]+)?):([0-9]+)$") @@ -125,9 +126,7 @@ def _verify_artifacts(context: AllocationContext) -> None: @staticmethod def verify_ports(context: AllocationContext) -> None: """Verify that every planned one-node port is currently bindable.""" - ports = tuple(port.port for port in context.plan.client.ports) + tuple( - port.port for deployment in context.plan.deployments for port in deployment.ports - ) + ports = allocation_ports(context) reservations: list[socket.socket] = [] try: for port in ports: diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index c62a3c391..5d8f41ef0 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -93,7 +93,7 @@ def test_renderer_uses_profile_slurm_bin_path(single_node_plan: ResolvedSlurmRun script = render_generation_attempt_script(plan, attempt_ordinal=1) assert 'export PATH="/opt/slurm/bin:/usr/local/sbin:' in script - assert "#SBATCH --exclusive\n" in script + assert "#SBATCH --exclusive" not in script @pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) diff --git a/packages/data-designer-slurm/tests/runtime/test_ports.py b/packages/data-designer-slurm/tests/runtime/test_ports.py new file mode 100644 index 000000000..1f6fed915 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_ports.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from conftest import RuntimeCase + +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.ports import allocation_ports, resolve_allocation_deployments + + +def test_allocation_ports_are_deterministic_and_isolated_by_scheduler_identity(runtime_case: RuntimeCase) -> None: + first = allocation_ports(runtime_case.context) + scheduler = runtime_case.context.attempt.scheduler + assert scheduler is not None + alternate_attempt = runtime_case.context.attempt.model_copy( + update={"scheduler": scheduler.model_copy(update={"array_job_id": 4102})} + ) + alternate = AllocationContext( + plan=runtime_case.context.plan, + shard=runtime_case.context.shard, + attempt=alternate_attempt, + attempt_directory=runtime_case.context.attempt_directory, + ) + + assert first == allocation_ports(runtime_case.context) + assert set(first).isdisjoint(allocation_ports(alternate)) + assert all(10000 <= port < 30000 for port in first) + + +def test_allocation_deployment_uses_remapped_ports(runtime_case: RuntimeCase) -> None: + deployment = resolve_allocation_deployments(runtime_case.context)[0] + ports = set(allocation_ports(runtime_case.context)) + + assert deployment.logical_endpoint.port in ports + assert {backend.port for backend in deployment.backend_endpoints} <= ports + assert {probe.port for probe in deployment.readiness_probes} <= ports + assert {process.http_port for process in deployment.processes if process.http_port is not None} <= ports + + +def test_allocation_ports_skip_client_otel_port(runtime_case: RuntimeCase) -> None: + otel_port = allocation_ports(runtime_case.context)[0] + invocation = runtime_case.context.plan.invocation.model_copy( + update={"effective_run_config": {"otel_metrics_port": otel_port}} + ) + plan = runtime_case.context.plan.model_copy(update={"invocation": invocation}) + context = AllocationContext( + plan=plan, + shard=runtime_case.context.shard, + attempt=runtime_case.context.attempt, + attempt_directory=runtime_case.context.attempt_directory, + ) + + assert otel_port not in allocation_ports(context) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 1a6ac1c3c..c38572c6f 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -7,7 +7,6 @@ #SBATCH --time=03:55:00 #SBATCH --array=0-1%2 #SBATCH --gres=gpu:8 -#SBATCH --exclusive set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index f6ba56a7e..72a565bd4 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -7,7 +7,6 @@ #SBATCH --time=03:55:00 #SBATCH --array=0 #SBATCH --gres=gpu:8 -#SBATCH --exclusive set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index b38ab401f..6c5578a70 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="391ca71f5d1f1a15808d66f2ebeb2242b1aeb248c72ee0b7dea26a561669defb", + expected_fixture_sha256="3523fe78e80edd4b4acd1e090d9888b3a9879430f1460fecd3576755a4ca923e", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="5e2dff54c1d3f534f4ae43e93063067681f8762b8995a996b9dd74bce2ecb5ef", + expected_fixture_sha256="e69a4303113764670742d8d0ce12b1ebe4cc80751734bf68341b535b578108bf", ) From bf3202c7b1742733c358cb08de03c1bfc1223a97 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 15:55:30 -0300 Subject: [PATCH 13/14] fix(slurm): isolate shared-node service ports --- .../data_designer/slurm/launcher/renderer.py | 2 + .../data_designer/slurm/runtime/bootstrap.py | 4 +- .../data_designer/slurm/runtime/controller.py | 2 +- .../data_designer/slurm/runtime/entrypoint.py | 22 +++++-- .../src/data_designer/slurm/runtime/ports.py | 66 +++++++++++++------ .../data_designer/slurm/runtime/preflight.py | 6 +- .../slurm/runtime/step_runner.sh | 2 +- .../tests/launcher/test_renderer.py | 1 + .../tests/runtime/test_bootstrap.py | 1 + .../tests/runtime/test_controller.py | 32 ++++----- .../tests/runtime/test_entrypoint.py | 7 +- .../tests/runtime/test_ports.py | 61 +++++++++++------ .../tests/runtime/test_preflight.py | 2 +- .../tests/runtime/test_shell_runtime.py | 2 + 14 files changed, 137 insertions(+), 73 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index 8fdae3522..ad71b0af5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -22,6 +22,8 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi run_root = posixpath.dirname(plan.authored_config.path) plan_path = posixpath.join(run_root, "resolved-plan.json") directive_text = render_batch_directives(_build_generation_directives(plan)) + if plan.selected_profile.profile.gpu_request_mode == "visible": + directive_text = f"{directive_text}\n#SBATCH --exclusive" attempt = f"{attempt_ordinal:04d}" scheduler_bin_path = plan.selected_profile.profile.scheduler.bin_path command_path = _SYSTEM_PATH if scheduler_bin_path is None else f"{scheduler_bin_path}:{_SYSTEM_PATH}" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index 8a7e35745..9b0c1d658 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -6,6 +6,7 @@ from __future__ import annotations import os +from collections.abc import Mapping from pathlib import Path from typing import Literal @@ -118,6 +119,7 @@ def validate_steps(self) -> RuntimeBootstrapManifest: def build_runtime_manifest( context: AllocationContext, + environment: Mapping[str, str], *, runtime_root: Path, log_directory: Path, @@ -125,7 +127,7 @@ def build_runtime_manifest( """Build the secret-free one-node command handoff for the Bash controller.""" plan = context.plan runtime_container_root = get_container_path(plan, runtime_root.as_posix(), require_writable=True) - deployments = resolve_allocation_deployments(context) + deployments = resolve_allocation_deployments(context, environment) endpoints = tuple( RuntimeEndpoint( model_alias=deployment.model_alias, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index c78d93173..e90cbc555 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -268,7 +268,7 @@ def _execute(self) -> tuple[ArtifactReference, datetime]: def _prepare_runtime(self) -> _RuntimeTopology: self._validate_attempt_state() - deployments = resolve_allocation_deployments(self._context) + deployments = resolve_allocation_deployments(self._context, self._environment) restarting = self._readiness is not None if restarting: self._begin_execution(deployments, ReadinessState.RESTARTING) 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 7a68f8198..66ee59f21 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 @@ -91,13 +91,14 @@ def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) _validate_attempt_is_executable(context.attempt) SystemAllocationPreflight.verify_attempt_directory(arguments.attempt_dir) - SystemAllocationPreflight.verify_ports(context) - readiness = _begin_attempt(context, writer) + SystemAllocationPreflight.verify_ports(context, environment) + readiness = _begin_attempt(context, writer, environment) log_directory = execution_log_directory(context.attempt_directory, readiness.revision) container_log_directory = Path(get_container_path(context.plan, log_directory.as_posix(), require_writable=True)) ensure_private_directory(container_log_directory) manifest = build_runtime_manifest( context, + environment, runtime_root=arguments.runtime_root, log_directory=log_directory, ) @@ -115,7 +116,7 @@ def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> Non context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) previous = writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) timestamp = _now(context.attempt, previous) - deployments = _resolve_deployments(context) + deployments = _resolve_deployments(context, environment) writer.write_readiness( AttemptReadiness( schema_version=1, @@ -219,7 +220,11 @@ def _fail(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None ) -def _begin_attempt(context: AllocationContext, writer: SlurmStateWriter) -> AttemptReadiness: +def _begin_attempt( + context: AllocationContext, + writer: SlurmStateWriter, + environment: Mapping[str, str], +) -> AttemptReadiness: attempt = context.attempt previous = _load_optional_readiness(context, writer) timestamp = _now(attempt, previous) @@ -227,7 +232,7 @@ def _begin_attempt(context: AllocationContext, writer: SlurmStateWriter) -> Atte attempt = writer.update_attempt( attempt.model_copy(update={"state": AttemptLifecycleState.RUNNING, "updated_at": timestamp}) ) - deployments = _resolve_deployments(context) + deployments = _resolve_deployments(context, environment) initial_state = ReadinessState.RESTARTING if previous is not None else ReadinessState.PENDING initial = writer.write_readiness(_readiness(context, deployments, previous, initial_state, timestamp)) return writer.write_readiness( @@ -347,8 +352,11 @@ def _load_optional_readiness(context: AllocationContext, writer: SlurmStateWrite return None -def _resolve_deployments(context: AllocationContext) -> tuple[ResolvedVllmServerDeployment, ...]: - return resolve_allocation_deployments(context) +def _resolve_deployments( + context: AllocationContext, + environment: Mapping[str, str], +) -> tuple[ResolvedVllmServerDeployment, ...]: + return resolve_allocation_deployments(context, environment) def _validate_attempt_is_executable(attempt: AttemptManifest) -> None: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py index d0dc3693c..b5a03591d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py @@ -5,7 +5,7 @@ from __future__ import annotations -import hashlib +from collections.abc import Mapping from data_designer.slurm.planning import PortClaim, ResolvedSlurmRunPlan from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode @@ -14,27 +14,38 @@ from data_designer.slurm.serving.resolver import resolve_vllm_server _PORT_RANGE_START = 10000 -_PORT_RANGE_SIZE = 20000 +_PORT_RANGE_END = 65536 +_PORTS_PER_GPU = 256 -def resolve_allocation_deployments(context: AllocationContext) -> tuple[ResolvedVllmServerDeployment, ...]: +def resolve_allocation_deployments( + context: AllocationContext, + environment: Mapping[str, str], +) -> tuple[ResolvedVllmServerDeployment, ...]: """Resolve deployments with ports isolated to one scheduler array element.""" - plan = _remap_plan_ports(context) + plan = _remap_plan_ports(context, environment) return tuple(resolve_vllm_server(plan, item.deployment_id) for item in plan.deployments) -def allocation_ports(context: AllocationContext) -> tuple[int, ...]: +def allocation_ports(context: AllocationContext, environment: Mapping[str, str]) -> tuple[int, ...]: """Return every allocation-local port in deterministic claim order.""" - plan = _remap_plan_ports(context) + plan = _remap_plan_ports(context, environment) return tuple(port.port for port in plan.client.ports) + tuple( port.port for deployment in plan.deployments for port in deployment.ports ) -def _remap_plan_ports(context: AllocationContext) -> ResolvedSlurmRunPlan: - scheduler = context.attempt.scheduler - if scheduler is None: # pragma: no cover - AllocationContext rejects this state - raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation has no scheduler identity") +def _remap_plan_ports(context: AllocationContext, environment: Mapping[str, str]) -> ResolvedSlurmRunPlan: + if context.plan.selected_profile.profile.gpu_request_mode == "visible": + return context.plan + gpu_ids = _parse_gpu_ids(environment.get("SLURM_JOB_GPUS")) + block_start = _PORT_RANGE_START + min(gpu_ids) * _PORTS_PER_GPU + block_end = block_start + _PORTS_PER_GPU + if block_end > _PORT_RANGE_END: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocated GPU index exceeds the supported allocation port range", + ) claims = context.plan.client.ports + tuple( port for deployment in context.plan.deployments for port in deployment.ports ) @@ -43,23 +54,20 @@ def _remap_plan_ports(context: AllocationContext) -> ResolvedSlurmRunPlan: claims_by_node.setdefault(claim.node_index, set()).add(claim.port) reserved_by_node: dict[int, set[int]] = {} otel_port = context.plan.invocation.effective_run_config.get("otel_metrics_port") - if type(otel_port) is int and _PORT_RANGE_START <= otel_port < _PORT_RANGE_START + _PORT_RANGE_SIZE: + if type(otel_port) is int and block_start <= otel_port < block_end: reserved_by_node[context.plan.client.host_node_index] = {otel_port} mapping: dict[tuple[int, int], int] = {} for node_index, planned_ports in claims_by_node.items(): - seed = f"{scheduler.array_job_id}:{scheduler.array_task_id}:{node_index}".encode() - offset = int.from_bytes(hashlib.sha256(seed).digest()[:8], "big") % _PORT_RANGE_SIZE reserved = reserved_by_node.get(node_index, set()) + next_port = block_start for planned_port in sorted(planned_ports): - for _ in range(_PORT_RANGE_SIZE): - port = _PORT_RANGE_START + offset - offset = (offset + 1) % _PORT_RANGE_SIZE - if port not in reserved: - mapping[(node_index, planned_port)] = port - reserved.add(port) - break - else: # pragma: no cover - plan contracts bound claims below the allocation range + while next_port in reserved: + next_port += 1 + if next_port >= block_end: raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation port range is exhausted") + mapping[(node_index, planned_port)] = next_port + reserved.add(next_port) + next_port += 1 def remap(port: PortClaim) -> PortClaim: return port.model_copy(update={"port": mapping[(port.node_index, port.port)]}) @@ -72,4 +80,20 @@ def remap(port: PortClaim) -> PortClaim: return context.plan.model_copy(update={"client": client, "deployments": deployments}) +def _parse_gpu_ids(value: str | None) -> tuple[int, ...]: + if value is None: + gpu_ids = () + else: + fields = value.split(",") + gpu_ids = tuple(int(field) for field in fields if field.isascii() and field.isdigit()) + if len(gpu_ids) != len(fields): + gpu_ids = () + if not gpu_ids or len(gpu_ids) != len(set(gpu_ids)): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "scheduler environment 'SLURM_JOB_GPUS' is unavailable or invalid", + ) + return gpu_ids + + __all__ = ["allocation_ports", "resolve_allocation_deployments"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py index e9c89f335..78366c390 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -44,7 +44,7 @@ def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> ) self.verify_attempt_directory(attempt_directory) self._verify_artifacts(context) - self.verify_ports(context) + self.verify_ports(context, environment) except SlurmRuntimeError: raise except (OSError, ValueError) as error: @@ -124,9 +124,9 @@ def _verify_artifacts(context: AllocationContext) -> None: _verify_artifact(reference) @staticmethod - def verify_ports(context: AllocationContext) -> None: + def verify_ports(context: AllocationContext, environment: Mapping[str, str]) -> None: """Verify that every planned one-node port is currently bindable.""" - ports = allocation_ports(context) + ports = allocation_ports(context, environment) reservations: list[socket.socket] = [] try: for port in ports: 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 4cd7bf045..7dc5aeae3 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 @@ -148,7 +148,7 @@ dd_run_control_phase() { ) [[ -z ${DD_CONTAINER_MOUNTS} ]] || command+=("--container-mounts=${DD_CONTAINER_MOUNTS}") command+=( - --container-env=PYTHONPATH + --container-env=PYTHONPATH,SLURM_JOB_GPUS -- python3 -m diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 5d8f41ef0..f4ed66d29 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -73,6 +73,7 @@ def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fiel assert "#SBATCH --gres=" not in script assert "#SBATCH --account=" not in script assert "#SBATCH --partition=" not in script + assert "#SBATCH --exclusive\n" in script assert '#SBATCH --comment="safe test run"\n' in script diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 8d7de9aac..51a56aa2c 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -16,6 +16,7 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru manifest = build_runtime_manifest( context, + {"SLURM_JOB_GPUS": "0"}, runtime_root=runtime_root, log_directory=log_directory, ) diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index 51638b2c7..45ff94b19 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -38,6 +38,8 @@ ) from data_designer.slurm.state.artifacts import compute_candidate_schema_digest +_ALLOCATION_ENVIRONMENT = {"SLURM_JOB_GPUS": "0"} + @dataclass(slots=True) class _FakeProcess: @@ -168,7 +170,7 @@ def write_result_under_dataset_lease() -> None: client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) result = controller.run() @@ -209,7 +211,7 @@ def test_controller_publishes_result_before_success_with_real_state_writer( client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) result = controller.run() @@ -239,7 +241,7 @@ def finalize_winner(self, *args: object, **kwargs: object) -> None: client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="allocation runtime failed"): @@ -264,7 +266,7 @@ def test_preflight_failure_starts_no_process_and_fails_attempt(runtime_case: Run client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="injected preflight failure") as raised: @@ -329,7 +331,7 @@ def test_requeued_running_attempt_publishes_restart_epoch_and_uses_fresh_logs( client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) result = controller.run() @@ -363,7 +365,7 @@ def test_required_server_exit_fails_and_cleans_partial_start(runtime_case: Runti client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="required runtime step"): @@ -394,7 +396,7 @@ def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="required runtime step") as raised: @@ -418,7 +420,7 @@ def test_readiness_timeout_fails_and_terminates_server(runtime_case: RuntimeCase client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=False, clock=clock, advance_on_failure=10_000), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="readiness timed out"): @@ -454,7 +456,7 @@ def test_managed_step_failure_fails_attempt_and_cleans_started_services( client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="status 23"): @@ -491,7 +493,7 @@ def test_cleanup_failure_prevents_false_success(runtime_case: RuntimeCase) -> No client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="cleanup failed"): @@ -520,7 +522,7 @@ def test_cleanup_failure_is_retained_as_a_note_on_the_primary_failure(runtime_ca client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="status 23") as raised: @@ -551,7 +553,7 @@ def test_incomplete_cleanup_does_not_publish_stopped_readiness(runtime_case: Run client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="cleanup failed"): @@ -580,7 +582,7 @@ def write_future_result() -> None: client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="later than the allocation clock"): @@ -602,7 +604,7 @@ def test_partial_client_result_is_classified_before_candidate_loading(runtime_ca client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError) as raised: @@ -626,7 +628,7 @@ def test_stale_candidate_from_an_earlier_generation_cannot_succeed(runtime_case: client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="predates the current generation"): diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 18ef8020d..86439bcc8 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -88,7 +88,7 @@ def test_control_phases_record_running_ready_and_failed( manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" _patch_runtime_context(monkeypatch, runtime_case, state) monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) - monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda context: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda *args: None) monkeypatch.setattr( entrypoint, "build_runtime_manifest", @@ -117,7 +117,7 @@ def test_succeed_phase_stops_runtime_and_finalizes_winner( manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" _patch_runtime_context(monkeypatch, runtime_case, state) monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) - monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda context: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda *args: None) monkeypatch.setattr( entrypoint, "build_runtime_manifest", @@ -154,7 +154,7 @@ def finalize_winner(self, *args: object, **kwargs: object) -> None: manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" _patch_runtime_context(monkeypatch, runtime_case, state) monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) - monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda context: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda *args: None) monkeypatch.setattr( entrypoint, "build_runtime_manifest", @@ -181,6 +181,7 @@ def _patch_runtime_context( ) -> None: monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args: (runtime_case.context, state)) monkeypatch.setattr(entrypoint, "get_container_path", lambda plan, path, **kwargs: path) + monkeypatch.setenv("SLURM_JOB_GPUS", "0") def _phase_arguments( diff --git a/packages/data-designer-slurm/tests/runtime/test_ports.py b/packages/data-designer-slurm/tests/runtime/test_ports.py index 1f6fed915..337d7cee5 100644 --- a/packages/data-designer-slurm/tests/runtime/test_ports.py +++ b/packages/data-designer-slurm/tests/runtime/test_ports.py @@ -3,34 +3,28 @@ from __future__ import annotations +import pytest from conftest import RuntimeCase +from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.ports import allocation_ports, resolve_allocation_deployments -def test_allocation_ports_are_deterministic_and_isolated_by_scheduler_identity(runtime_case: RuntimeCase) -> None: - first = allocation_ports(runtime_case.context) - scheduler = runtime_case.context.attempt.scheduler - assert scheduler is not None - alternate_attempt = runtime_case.context.attempt.model_copy( - update={"scheduler": scheduler.model_copy(update={"array_job_id": 4102})} - ) - alternate = AllocationContext( - plan=runtime_case.context.plan, - shard=runtime_case.context.shard, - attempt=alternate_attempt, - attempt_directory=runtime_case.context.attempt_directory, - ) +def test_allocation_ports_are_deterministic_and_isolated_by_gpu(runtime_case: RuntimeCase) -> None: + first = allocation_ports(runtime_case.context, {"SLURM_JOB_GPUS": "0"}) + second = allocation_ports(runtime_case.context, {"SLURM_JOB_GPUS": "1"}) - assert first == allocation_ports(runtime_case.context) - assert set(first).isdisjoint(allocation_ports(alternate)) - assert all(10000 <= port < 30000 for port in first) + assert first == allocation_ports(runtime_case.context, {"SLURM_JOB_GPUS": "0"}) + assert set(first).isdisjoint(second) + assert all(10000 <= port < 10256 for port in first) + assert all(10256 <= port < 10512 for port in second) def test_allocation_deployment_uses_remapped_ports(runtime_case: RuntimeCase) -> None: - deployment = resolve_allocation_deployments(runtime_case.context)[0] - ports = set(allocation_ports(runtime_case.context)) + environment = {"SLURM_JOB_GPUS": "0"} + deployment = resolve_allocation_deployments(runtime_case.context, environment)[0] + ports = set(allocation_ports(runtime_case.context, environment)) assert deployment.logical_endpoint.port in ports assert {backend.port for backend in deployment.backend_endpoints} <= ports @@ -39,7 +33,8 @@ def test_allocation_deployment_uses_remapped_ports(runtime_case: RuntimeCase) -> def test_allocation_ports_skip_client_otel_port(runtime_case: RuntimeCase) -> None: - otel_port = allocation_ports(runtime_case.context)[0] + environment = {"SLURM_JOB_GPUS": "0"} + otel_port = allocation_ports(runtime_case.context, environment)[0] invocation = runtime_case.context.plan.invocation.model_copy( update={"effective_run_config": {"otel_metrics_port": otel_port}} ) @@ -51,4 +46,30 @@ def test_allocation_ports_skip_client_otel_port(runtime_case: RuntimeCase) -> No attempt_directory=runtime_case.context.attempt_directory, ) - assert otel_port not in allocation_ports(context) + assert otel_port not in allocation_ports(context, environment) + + +@pytest.mark.parametrize("environment", ({}, {"SLURM_JOB_GPUS": ""}, {"SLURM_JOB_GPUS": "0,0"})) +def test_allocation_ports_reject_invalid_gpu_ids( + runtime_case: RuntimeCase, + environment: dict[str, str], +) -> None: + with pytest.raises(SlurmRuntimeError, match="SLURM_JOB_GPUS"): + allocation_ports(runtime_case.context, environment) + + +def test_visible_gpu_mode_keeps_planned_ports(runtime_case: RuntimeCase) -> None: + profile = runtime_case.context.plan.selected_profile.profile.model_copy(update={"gpu_request_mode": "visible"}) + selected_profile = runtime_case.context.plan.selected_profile.model_copy(update={"profile": profile}) + plan = runtime_case.context.plan.model_copy(update={"selected_profile": selected_profile}) + context = AllocationContext( + plan=plan, + shard=runtime_case.context.shard, + attempt=runtime_case.context.attempt, + attempt_directory=runtime_case.context.attempt_directory, + ) + planned = tuple(port.port for port in plan.client.ports) + tuple( + port.port for deployment in plan.deployments for port in deployment.ports + ) + + assert allocation_ports(context, {}) == planned diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index 2f510a1f6..4cc0890e6 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -141,7 +141,7 @@ def close(self) -> None: monkeypatch.setattr("data_designer.slurm.runtime.preflight.socket.socket", _UnavailableSocket) with pytest.raises(SlurmRuntimeError, match="ports are unavailable"): - SystemAllocationPreflight.verify_ports(runtime_case.context) + SystemAllocationPreflight.verify_ports(runtime_case.context, {"SLURM_JOB_GPUS": "0"}) assert _UnavailableSocket.closed 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 fb29b92aa..ef877ff6a 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -57,6 +57,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( "SLURM_ARRAY_JOB_ID": "4101", "SLURM_ARRAY_TASK_ID": "0", "SLURM_JOB_NUM_NODES": "1", + "SLURM_JOB_GPUS": "0", "SLURM_NODEID": "0", } @@ -271,6 +272,7 @@ def _fake_srun() -> str: if [[ ${command} == python3 ]]; then operation=${arguments[index + 4]} [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + [[ ${arguments[*]} == *--container-env=PYTHONPATH,SLURM_JOB_GPUS* ]] if [[ ${operation} == prepare ]]; then for ((position = index + 5; position < ${#arguments[@]}; position++)); do if [[ ${arguments[position]} == --manifest ]]; then From ae26a95f6179bf3c9a0861062e826be878906ef9 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 9 Sep 2026 16:12:04 -0300 Subject: [PATCH 14/14] fix(slurm): validate allocation-local endpoints --- .../data_designer/slurm/client/execution.py | 6 +++- .../data_designer/slurm/runtime/bootstrap.py | 12 ++++++-- .../src/data_designer/slurm/runtime/ports.py | 30 ++++++++++--------- .../src/data_designer/slurm/runtime/steps.py | 3 +- .../tests/client/conftest.py | 13 +++++++- .../tests/client/test_worker.py | 14 ++++++++- .../tests/runtime/test_steps.py | 3 +- 7 files changed, 60 insertions(+), 21 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index 9b85625d1..5f3fa3184 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -52,6 +52,7 @@ from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import PlannedShard, ResolvedDependencyLock, ResolvedSlurmRunPlan from data_designer.slurm.runtime.paths import get_container_path, get_host_path +from data_designer.slurm.runtime.ports import resolve_allocation_plan from data_designer.slurm.state import CandidateOutcome, CandidateOutputFile, CandidateOutputManifest Clock = Callable[[], datetime] @@ -155,9 +156,11 @@ def __init__( *, data_designer_factory: DataDesignerFactory = DataDesigner, clock: Clock | None = None, + environment: Mapping[str, str] | None = None, ) -> None: self._data_designer_factory = data_designer_factory self._clock = clock or (lambda: datetime.now(timezone.utc)) + self._environment = os.environ if environment is None else environment def preflight( self, @@ -273,7 +276,8 @@ def _build_context( raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "client environment differs from the lock") builder_payload = self._load_builder(plan) builder = DataDesignerConfigBuilder.from_config(builder_payload) - providers = self._materialize_model_endpoints(plan, builder, endpoints) + allocation_plan = resolve_allocation_plan(plan, self._environment) + providers = self._materialize_model_endpoints(allocation_plan, builder, endpoints) self._validate_model_references(builder) self._materialize_seed(plan, shard, builder) mcp_providers = self._materialize_mcp_providers(plan) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index 9b0c1d658..54a053231 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -143,6 +143,7 @@ def build_runtime_manifest( "client-preflight", "preflight", context, + environment, endpoints, runtime_container_root, log_directory, @@ -167,6 +168,7 @@ def build_runtime_manifest( "client-generation", "client", context, + environment, endpoints, runtime_container_root, log_directory, @@ -195,6 +197,7 @@ def _build_client_step( step_id: str, operation: str, context: AllocationContext, + environment: Mapping[str, str], endpoints: tuple[RuntimeEndpoint, ...], runtime_container_root: str, log_directory: Path, @@ -220,6 +223,11 @@ def _build_client_step( secret_names = collect_secret_environment_names( (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) ) + allocation_environment = ( + {"SLURM_JOB_GPUS": environment["SLURM_JOB_GPUS"]} + if plan.selected_profile.profile.gpu_request_mode == "gres" + else {} + ) return _step( step_id=step_id, role=role, @@ -227,10 +235,10 @@ def _build_client_step( command=command, cpus=plan.client.authored.cpus, gpu_indices=(), - literal_environment={"LC_ALL": "C", "PYTHONPATH": runtime_container_root}, + literal_environment={"LC_ALL": "C", "PYTHONPATH": runtime_container_root, **allocation_environment}, secret_environment={name: name for name in secret_names}, environment_prefixes={}, - container_environment=tuple(sorted((*secret_names, "PYTHONPATH"))), + container_environment=tuple(sorted((*secret_names, *allocation_environment, "PYTHONPATH"))), log_directory=log_directory, ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py index b5a03591d..53d6f049a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py @@ -23,21 +23,25 @@ def resolve_allocation_deployments( environment: Mapping[str, str], ) -> tuple[ResolvedVllmServerDeployment, ...]: """Resolve deployments with ports isolated to one scheduler array element.""" - plan = _remap_plan_ports(context, environment) + plan = resolve_allocation_plan(context.plan, environment) return tuple(resolve_vllm_server(plan, item.deployment_id) for item in plan.deployments) def allocation_ports(context: AllocationContext, environment: Mapping[str, str]) -> tuple[int, ...]: """Return every allocation-local port in deterministic claim order.""" - plan = _remap_plan_ports(context, environment) + plan = resolve_allocation_plan(context.plan, environment) return tuple(port.port for port in plan.client.ports) + tuple( port.port for deployment in plan.deployments for port in deployment.ports ) -def _remap_plan_ports(context: AllocationContext, environment: Mapping[str, str]) -> ResolvedSlurmRunPlan: - if context.plan.selected_profile.profile.gpu_request_mode == "visible": - return context.plan +def resolve_allocation_plan( + plan: ResolvedSlurmRunPlan, + environment: Mapping[str, str], +) -> ResolvedSlurmRunPlan: + """Return a plan with ports bound to the allocation's assigned GPUs.""" + if plan.selected_profile.profile.gpu_request_mode == "visible": + return plan gpu_ids = _parse_gpu_ids(environment.get("SLURM_JOB_GPUS")) block_start = _PORT_RANGE_START + min(gpu_ids) * _PORTS_PER_GPU block_end = block_start + _PORTS_PER_GPU @@ -46,16 +50,14 @@ def _remap_plan_ports(context: AllocationContext, environment: Mapping[str, str] SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "allocated GPU index exceeds the supported allocation port range", ) - claims = context.plan.client.ports + tuple( - port for deployment in context.plan.deployments for port in deployment.ports - ) + claims = plan.client.ports + tuple(port for deployment in plan.deployments for port in deployment.ports) claims_by_node: dict[int, set[int]] = {} for claim in claims: claims_by_node.setdefault(claim.node_index, set()).add(claim.port) reserved_by_node: dict[int, set[int]] = {} - otel_port = context.plan.invocation.effective_run_config.get("otel_metrics_port") + otel_port = plan.invocation.effective_run_config.get("otel_metrics_port") if type(otel_port) is int and block_start <= otel_port < block_end: - reserved_by_node[context.plan.client.host_node_index] = {otel_port} + reserved_by_node[plan.client.host_node_index] = {otel_port} mapping: dict[tuple[int, int], int] = {} for node_index, planned_ports in claims_by_node.items(): reserved = reserved_by_node.get(node_index, set()) @@ -72,12 +74,12 @@ def _remap_plan_ports(context: AllocationContext, environment: Mapping[str, str] def remap(port: PortClaim) -> PortClaim: return port.model_copy(update={"port": mapping[(port.node_index, port.port)]}) - client = context.plan.client.model_copy(update={"ports": tuple(remap(port) for port in context.plan.client.ports)}) + client = plan.client.model_copy(update={"ports": tuple(remap(port) for port in plan.client.ports)}) deployments = tuple( deployment.model_copy(update={"ports": tuple(remap(port) for port in deployment.ports)}) - for deployment in context.plan.deployments + for deployment in plan.deployments ) - return context.plan.model_copy(update={"client": client, "deployments": deployments}) + return plan.model_copy(update={"client": client, "deployments": deployments}) def _parse_gpu_ids(value: str | None) -> tuple[int, ...]: @@ -96,4 +98,4 @@ def _parse_gpu_ids(value: str | None) -> tuple[int, ...]: return gpu_ids -__all__ = ["allocation_ports", "resolve_allocation_deployments"] +__all__ = ["allocation_ports", "resolve_allocation_deployments", "resolve_allocation_plan"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py index 387520a0b..8aeb96eb2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -133,13 +133,14 @@ def _build_step( ) -> RuntimeStep: command = build_client_command(operation, plan, shard, attempt, attempt_directory, endpoints) secret_names, environment = _build_client_environment(plan, source_environment) + allocation_environment = ("SLURM_JOB_GPUS",) if plan.selected_profile.profile.gpu_request_mode == "gres" else () return _build_srun_step( step_id=step_id, role=role, image_path=plan.client.image.path, command=command, environment=environment, - container_environment=secret_names, + container_environment=(*secret_names, *allocation_environment), plan=plan, attempt_directory=attempt_directory, ) diff --git a/packages/data-designer-slurm/tests/client/conftest.py b/packages/data-designer-slurm/tests/client/conftest.py index a47eeed2f..95220553e 100644 --- a/packages/data-designer-slurm/tests/client/conftest.py +++ b/packages/data-designer-slurm/tests/client/conftest.py @@ -18,6 +18,7 @@ from data_designer.slurm.client.records import ClientInstallerOutcome from data_designer.slurm.contracts import InstalledDistribution, compute_canonical_json_sha256 from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.ports import resolve_allocation_plan GOLDEN_DIRECTORY = Path(__file__).parents[1] / "contracts" / "golden" @@ -104,6 +105,11 @@ def create( return results +@pytest.fixture(autouse=True) +def allocation_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SLURM_JOB_GPUS", "0") + + @pytest.fixture def client_worker_case(tmp_path: Path) -> ClientWorkerCase: workspace = tmp_path / "workspace" @@ -150,5 +156,10 @@ def client_worker_case(tmp_path: Path) -> ClientWorkerCase: for distribution in lock.image_distributions ), ) - endpoints = {plan.deployments[0].authored.model_alias: f"http://127.0.0.1:{plan.client.ports[0].port}/v1"} + allocation_plan = resolve_allocation_plan(plan, {"SLURM_JOB_GPUS": "0"}) + endpoints = { + allocation_plan.deployments[ + 0 + ].authored.model_alias: f"http://127.0.0.1:{allocation_plan.client.ports[0].port}/v1" + } return ClientWorkerCase(plan, plan_path, lock, attempt_dir, prepared, endpoints) diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index 0ddaf4916..c2175577b 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -60,6 +60,18 @@ def factory(**kwargs: object) -> FakeDataDesigner: ) +def test_preflight_rejects_endpoint_from_another_gpu_allocation(client_worker_case: ClientWorkerCase) -> None: + alias = client_worker_case.plan.deployments[0].authored.model_alias + + with pytest.raises(ClientWorkerError, match="runtime model endpoint is invalid"): + ClientWorker(data_designer_factory=FakeDataDesigner, environment={"SLURM_JOB_GPUS": "0"}).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints={alias: "http://127.0.0.1:10256/v1"}, + plugins=(), + ) + + def test_preflight_rejects_missing_managed_assets(client_worker_case: ClientWorkerCase) -> None: Path(client_worker_case.plan.invocation.effective_input_bindings.managed_assets_path).rmdir() @@ -480,7 +492,7 @@ def test_preflight_rejects_plugin_secondary_model_alias( ClientWorker().preflight( plan_path, prepared=prepared, - endpoints={"generator": "http://127.0.0.1:17000/v1"}, + endpoints={"generator": "http://127.0.0.1:10000/v1"}, plugins=plugins, ) except ClientWorkerError as error: diff --git a/packages/data-designer-slurm/tests/runtime/test_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py index a9aeefcbe..bda517d92 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -244,12 +244,13 @@ def test_client_receives_client_secrets_without_server_only_secrets( { "PACKAGE_INDEX_TOKEN": "client-secret", "HF_TOKEN": "server-secret", + "SLURM_JOB_GPUS": "0", }, ) assert step.environment["PACKAGE_INDEX_TOKEN"] == "client-secret" assert "HF_TOKEN" not in step.environment - assert "--container-env=PACKAGE_INDEX_TOKEN" in step.command + assert "--container-env=PACKAGE_INDEX_TOKEN,SLURM_JOB_GPUS" in step.command with pytest.raises(SlurmRuntimeError, match="PACKAGE_INDEX_TOKEN"): DefaultClientStepBuilder().build_preflight_step(