From d7a05421f01e2c42ca2fd5c3a3ffd6bb420f1a9a Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:26:00 -0600 Subject: [PATCH 01/23] feat: add Slurm command client and renderer Part of #868 Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/__init__.py | 39 ++++ .../data_designer/slurm/launcher/client.py | 153 ++++++++++++++++ .../data_designer/slurm/launcher/errors.py | 22 +++ .../data_designer/slurm/launcher/models.py | 42 +++++ .../data_designer/slurm/launcher/parsing.py | 168 ++++++++++++++++++ .../data_designer/slurm/launcher/renderer.py | 119 +++++++++++++ .../data_designer/slurm/launcher/runner.py | 60 +++++++ .../tests/launcher/test_client.py | 120 +++++++++++++ .../tests/launcher/test_parsing.py | 133 ++++++++++++++ .../tests/launcher/test_renderer.py | 127 +++++++++++++ .../tests/launcher/test_runner.py | 72 ++++++++ .../golden/rendered/multi_node.sbatch | 9 +- .../golden/rendered/single_node.sbatch | 9 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- 14 files changed, 1071 insertions(+), 6 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_client.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_parsing.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_renderer.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_runner.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py new file mode 100644 index 000000000..c3e29dce3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Structured Slurm submission, observation, and batch rendering.""" + +from __future__ import annotations + +from data_designer.slurm.launcher.client import SlurmCommandClient, SlurmExecutables +from data_designer.slurm.launcher.errors import ( + BatchRenderError, + SlurmCommandError, + SlurmLauncherError, + SlurmParseError, +) +from data_designer.slurm.launcher.models import ( + AccountingRecord, + QueueRecord, + SlurmExitCode, + SlurmSubmission, +) +from data_designer.slurm.launcher.renderer import BatchDirective, render_batch_script +from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner + +__all__ = [ + "AccountingRecord", + "BatchDirective", + "BatchRenderError", + "CommandRunner", + "QueueRecord", + "SlurmCommandClient", + "SlurmCommandError", + "SlurmExecutables", + "SlurmExitCode", + "SlurmLauncherError", + "SlurmParseError", + "SlurmSubmission", + "SubprocessRunner", + "render_batch_script", +] 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 new file mode 100644 index 000000000..29b9717e0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed argument-vector client for Slurm command-line tools.""" + +from __future__ import annotations + +import re +import subprocess +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path + +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.launcher.errors import SlurmCommandError +from data_designer.slurm.launcher.models import AccountingRecord, QueueRecord, SlurmSubmission +from data_designer.slurm.launcher.parsing import ( + parse_accounting, + parse_gpu_counts, + parse_queue, + parse_submission, +) +from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner +from data_designer.slurm.state import SchedulerIdentity + +JobSelector = int | SchedulerIdentity +_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") + + +@dataclass(frozen=True, slots=True) +class SlurmExecutables: + """Executable paths used for bounded Slurm operations.""" + + sbatch: str = "sbatch" + squeue: str = "squeue" + sacct: str = "sacct" + scancel: str = "scancel" + sinfo: str = "sinfo" + + def __post_init__(self) -> None: + for executable in (self.sbatch, self.squeue, self.sacct, self.scancel, 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") + + +class SlurmCommandClient: + """Submit, observe, and cancel Slurm jobs through structured commands.""" + + _executables: SlurmExecutables + _runner: CommandRunner + + def __init__( + self, + runner: CommandRunner | None = None, + *, + executables: SlurmExecutables | None = None, + ) -> None: + 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) -> SlurmSubmission: + """Submit one rendered batch script and return its assigned job ID.""" + path = str(script_path) + _validate_argument(path, field_name="batch script path") + output = self._run((self._executables.sbatch, "--parsable", path)) + return parse_submission(output) + + def query_queue(self, selectors: Sequence[JobSelector]) -> tuple[QueueRecord, ...]: + """Return normalized active-queue rows for explicit managed jobs.""" + jobs = _format_selectors(selectors) + output = self._run( + ( + self._executables.squeue, + "--noheader", + "--array", + "--format=%i|%T", + f"--jobs={jobs}", + ) + ) + return parse_queue(output) + + def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[AccountingRecord, ...]: + """Return normalized accounting rows for explicit managed jobs.""" + jobs = _format_selectors(selectors) + output = self._run( + ( + self._executables.sacct, + "--noheader", + "--parsable2", + "--format=%i|%State|%ExitCode", + f"--jobs={jobs}", + ) + ) + return parse_accounting(output) + + def cancel(self, selector: JobSelector) -> None: + """Cancel one managed Slurm array or array task.""" + self._run((self._executables.scancel, _format_selector(selector))) + + 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"] + if partition is not None: + if type(partition) is not str or _IDENTIFIER_PATTERN.fullmatch(partition) is None: + raise ValueError("Slurm partition must be a valid identifier") + command.append(f"--partition={partition}") + return parse_gpu_counts(self._run(command)) + + def _run(self, command: Sequence[str]) -> str: + command_name = Path(command[0]).name + try: + completed = self._runner.run(command) + except (OSError, subprocess.SubprocessError) as error: + raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") from error + if completed.returncode: + detail = _normalize_bounded_text(completed.stderr) or "no diagnostic output" + raise SlurmCommandError(f"{command_name} failed with exit code {completed.returncode}: {detail}") + if not isinstance(completed.stdout, str): + raise SlurmCommandError(f"{command_name} did not return text output") + return completed.stdout + + +def _format_selectors(selectors: Sequence[JobSelector]) -> str: + if not selectors: + raise ValueError("at least one managed Slurm job selector is required") + return ",".join(dict.fromkeys(_format_selector(selector) for selector in selectors)) + + +def _format_selector(selector: JobSelector) -> str: + if isinstance(selector, SchedulerIdentity): + return f"{selector.array_job_id}_{selector.array_task_id}" + if type(selector) is not int or selector <= 0: + raise ValueError("Slurm job IDs must be positive integers") + return str(selector) + + +def _validate_argument(value: str, *, field_name: str) -> None: + if type(value) is not str or not value: + raise ValueError(f"{field_name} must not be empty") + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError(f"{field_name} must not contain control characters") + + +def _normalize_bounded_text(value: str, *, limit: int = 512) -> str: + normalized = " ".join(value.split()) + return normalized if len(normalized) <= limit else f"{normalized[:limit]}..." + + +def _format_error_detail(error: BaseException) -> str: + if isinstance(error, subprocess.TimeoutExpired): + return "command timed out" + return _normalize_bounded_text(str(error)) or error.__class__.__name__ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py new file mode 100644 index 000000000..791ae69b7 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Canonical errors for the Slurm launcher boundary.""" + +from __future__ import annotations + + +class SlurmLauncherError(RuntimeError): + """Base error for structured Slurm launcher operations.""" + + +class SlurmCommandError(SlurmLauncherError): + """A Slurm command could not be executed successfully.""" + + +class SlurmParseError(SlurmLauncherError, ValueError): + """Slurm returned output that violates the requested format.""" + + +class BatchRenderError(SlurmLauncherError, ValueError): + """A resolved plan cannot be rendered as a safe batch script.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py new file mode 100644 index 000000000..2574ef204 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transient typed values returned by Slurm commands.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + + +@dataclass(frozen=True, slots=True) +class SlurmSubmission: + """Identity assigned by Slurm to one accepted array submission.""" + + array_job_id: int + + +@dataclass(frozen=True, slots=True) +class SlurmExitCode: + """Slurm's process status and terminating signal pair.""" + + status: int + signal: int + + +@dataclass(frozen=True, slots=True) +class QueueRecord: + """One normalized active-queue row.""" + + scheduler: SchedulerIdentity + state: SchedulerState + + +@dataclass(frozen=True, slots=True) +class AccountingRecord: + """One normalized accounting row.""" + + scheduler: SchedulerIdentity + state: SchedulerState + exit_code: SlurmExitCode diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py new file mode 100644 index 000000000..67937dd75 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict parsers for bounded, machine-readable Slurm output.""" + +from __future__ import annotations + +import re + +from data_designer.slurm.launcher.errors import SlurmParseError +from data_designer.slurm.launcher.models import ( + AccountingRecord, + QueueRecord, + SlurmExitCode, + SlurmSubmission, +) +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + +_ARRAY_ID_PATTERN = re.compile(r"^(?P[1-9][0-9]*)_(?P[0-9]+)$") +_ARRAY_STEP_ID_PATTERN = re.compile(r"^[1-9][0-9]*_[0-9]+\.[^\s|]+$") +_JOB_ID_PATTERN = re.compile(r"^[1-9][0-9]*$") +_CLUSTER_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_EXIT_CODE_PATTERN = re.compile(r"^(?P[0-9]+):(?P[0-9]+)$") +_GRES_GPU_PATTERN = re.compile(r"^gpu:(?:(?:[^:,()]+):)*(?P[1-9][0-9]*)(?:\([^\r\n]*\))?$") + +_STATE_MAP = { + "BOOT_FAIL": SchedulerState.FAILED, + "CANCELLED": SchedulerState.CANCELLED, + "COMPLETED": SchedulerState.COMPLETED, + "COMPLETING": SchedulerState.RUNNING, + "CONFIGURING": SchedulerState.PENDING, + "DEADLINE": SchedulerState.FAILED, + "FAILED": SchedulerState.FAILED, + "NODE_FAIL": SchedulerState.NODE_FAILED, + "OUT_OF_MEMORY": SchedulerState.OUT_OF_MEMORY, + "PENDING": SchedulerState.PENDING, + "PREEMPTED": SchedulerState.PREEMPTED, + "REQUEUED": SchedulerState.REQUEUED, + "REQUEUE_FED": SchedulerState.PENDING, + "REQUEUE_HOLD": SchedulerState.PENDING, + "RESIZING": SchedulerState.RUNNING, + "REVOKED": SchedulerState.FAILED, + "RUNNING": SchedulerState.RUNNING, + "SIGNALING": SchedulerState.RUNNING, + "SPECIAL_EXIT": SchedulerState.FAILED, + "STAGE_OUT": SchedulerState.RUNNING, + "STOPPED": SchedulerState.RUNNING, + "SUSPENDED": SchedulerState.RUNNING, + "TIMEOUT": SchedulerState.TIMED_OUT, +} + + +def parse_submission(output: str) -> SlurmSubmission: + """Parse ``sbatch --parsable`` output.""" + value = output.strip() + job_id, separator, cluster_name = value.partition(";") + if not job_id.isascii() or not job_id.isdecimal() or int(job_id) <= 0: + raise SlurmParseError("sbatch returned an invalid job ID") + if separator and _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: + raise SlurmParseError("sbatch returned an invalid cluster name") + return SlurmSubmission(array_job_id=int(job_id)) + + +def parse_queue(output: str) -> tuple[QueueRecord, ...]: + """Parse ``squeue --format=%i|%T`` rows.""" + records: list[QueueRecord] = [] + identities: set[SchedulerIdentity] = set() + for line_number, line in _collect_nonempty_lines(output): + fields = line.split("|") + if len(fields) != 2: + raise SlurmParseError(f"squeue line {line_number} must contain two fields") + scheduler = _parse_array_identity(fields[0], command="squeue", line_number=line_number) + _reject_duplicate(scheduler, identities, command="squeue", line_number=line_number) + records.append(QueueRecord(scheduler=scheduler, state=parse_state(fields[1]))) + return tuple(records) + + +def parse_accounting(output: str) -> tuple[AccountingRecord, ...]: + """Parse array-task rows from ``sacct --format=%i|%State|%ExitCode``.""" + records: list[AccountingRecord] = [] + identities: set[SchedulerIdentity] = set() + for line_number, line in _collect_nonempty_lines(output): + fields = line.split("|") + if len(fields) != 3: + raise SlurmParseError(f"sacct line {line_number} must contain three fields") + if _JOB_ID_PATTERN.fullmatch(fields[0]) is not None or _ARRAY_STEP_ID_PATTERN.fullmatch(fields[0]) is not None: + continue + scheduler = _parse_array_identity(fields[0], command="sacct", line_number=line_number) + _reject_duplicate(scheduler, identities, command="sacct", line_number=line_number) + records.append( + AccountingRecord( + scheduler=scheduler, + state=parse_state(fields[1]), + exit_code=_parse_exit_code(fields[2], line_number=line_number), + ) + ) + return tuple(records) + + +def parse_gpu_counts(output: str) -> tuple[int, ...]: + """Parse configured per-node GPU counts from ``sinfo --format=%G`` rows.""" + counts: list[int] = [] + for line_number, line in _collect_nonempty_lines(output): + if line in {"(null)", "N/A"}: + continue + line_counts: list[int] = [] + for gres in line.split(","): + if not gres.startswith("gpu:"): + continue + match = _GRES_GPU_PATTERN.fullmatch(gres) + if match is None: + raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + line_counts.append(int(match.group("count"))) + if line_counts: + counts.append(sum(line_counts)) + return tuple(counts) + + +def parse_state(value: str) -> SchedulerState: + """Normalize one Slurm long state spelling without guessing unknown states.""" + normalized = value.strip().upper().removesuffix("+") + if not normalized: + raise SlurmParseError("scheduler state must not be empty") + if normalized.startswith("CANCELLED BY "): + canceller = normalized.removeprefix("CANCELLED BY ") + if not canceller.isascii() or not canceller.isdecimal(): + raise SlurmParseError("cancelled scheduler state has an invalid owner") + normalized = "CANCELLED" + elif any(character.isspace() for character in normalized): + raise SlurmParseError("scheduler state contains unexpected whitespace") + return _STATE_MAP.get(normalized, SchedulerState.UNKNOWN) + + +def _collect_nonempty_lines(output: str) -> tuple[tuple[int, str], ...]: + return tuple( + (line_number, line) + for line_number, raw_line in enumerate(output.splitlines(), start=1) + if (line := raw_line.strip()) + ) + + +def _parse_array_identity(value: str, *, command: str, line_number: int) -> SchedulerIdentity: + match = _ARRAY_ID_PATTERN.fullmatch(value) + if match is None: + raise SlurmParseError(f"{command} line {line_number} contains an invalid array-task ID") + return SchedulerIdentity( + array_job_id=int(match.group("job")), + array_task_id=int(match.group("task")), + ) + + +def _parse_exit_code(value: str, *, line_number: int) -> SlurmExitCode: + match = _EXIT_CODE_PATTERN.fullmatch(value) + if match is None: + raise SlurmParseError(f"sacct line {line_number} contains an invalid exit code") + return SlurmExitCode(status=int(match.group("status")), signal=int(match.group("signal"))) + + +def _reject_duplicate( + scheduler: SchedulerIdentity, + identities: set[SchedulerIdentity], + *, + command: str, + line_number: int, +) -> None: + if scheduler in identities: + raise SlurmParseError(f"{command} line {line_number} duplicates an array-task ID") + identities.add(scheduler) 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 new file mode 100644 index 000000000..d63532dd8 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Safe deterministic rendering for thin Slurm batch entrypoints.""" + +from __future__ import annotations + +import posixpath +import re +from dataclasses import dataclass + +from data_designer.slurm.launcher.errors import BatchRenderError +from data_designer.slurm.planning import ResolvedSlurmRunPlan + +_DIRECTIVE_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9-]*$") +_DIRECTIVE_TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/,%+-]*$") + + +@dataclass(frozen=True, slots=True) +class BatchDirective: + """One validated ``#SBATCH`` option.""" + + name: str + value: str + + def render(self) -> str: + """Render the directive as one non-executable scheduler line.""" + if type(self.name) is not str or _DIRECTIVE_NAME_PATTERN.fullmatch(self.name) is None: + raise BatchRenderError("batch directive name is invalid") + if type(self.value) is not str: + raise BatchRenderError("batch directive value must be text") + _reject_control_characters(self.value, field_name=f"--{self.name} value") + value = self.value if _DIRECTIVE_TOKEN_PATTERN.fullmatch(self.value) else _quote_double_value(self.value) + return f"#SBATCH --{self.name}={value}" + + +def render_batch_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int = 1) -> str: + """Render a resolved generation plan as one thin deterministic entrypoint.""" + if type(attempt_ordinal) is not int or attempt_ordinal <= 0: + raise BatchRenderError("attempt_ordinal must be a positive integer") + + run_root = posixpath.dirname(plan.authored_config.path) + plan_path = posixpath.join(run_root, "resolved-plan.json") + directives = _build_generation_directives(plan) + directive_text = "\n".join(directive.render() for directive in directives) + attempt = f"{attempt_ordinal:04d}" + + return f"""#!/usr/bin/env bash +{directive_text} +set -Eeuo pipefail + +readonly DD_RUNTIME_ARCHIVE={_quote_double_value(plan.runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={_quote_double_value(plan.runtime_bundle.sha256)} +readonly DD_PLAN={_quote_double_value(plan_path)} +readonly DD_PLAN_SHA256={_quote_double_value(plan.compute_sha256())} +readonly DD_RUN_ROOT={_quote_double_value(run_root)} +readonly DD_ATTEMPT_ORDINAL={_quote_double_value(attempt)} + +verify_sha256() {{ + printf '%s %s\\n' "$1" "$2" | sha256sum --check --status - +}} + +verify_sha256 "${{DD_RUNTIME_SHA256}}" "${{DD_RUNTIME_ARCHIVE}}" +verify_sha256 "${{DD_PLAN_SHA256}}" "${{DD_PLAN}}" +if [[ ! ${{SLURM_ARRAY_TASK_ID:-}} =~ ^[0-9]+$ ]]; then + printf '%s\\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 + exit 64 +fi +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}}" + +source "${{DD_RUNTIME_DIR}}/entrypoint.sh" +dd_slurm_run_allocation "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" +""" + + +def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[BatchDirective, ...]: + node_indices = ( + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + ) + node_count = max(node_indices) + 1 + array = "0" + if plan.array_tasks.count > 1: + array = f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + + values: list[tuple[str, str | None]] = [ + ("job-name", plan.submission.job_name), + ("account", plan.submission.account), + ("partition", plan.submission.partition), + ("nodes", str(node_count)), + ("time", plan.submission.time_limit), + ("array", array), + ] + profile = plan.selected_profile.profile + if profile.gpu_request_mode == "gres": + values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) + if profile.scheduler.mem_per_gpu is not None: + values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) + if plan.submission.comment is not None: + values.append(("comment", plan.submission.comment)) + return tuple(BatchDirective(name=name, value=value) for name, value in values if value is not None) + + +def _quote_double_value(value: str) -> str: + _reject_control_characters(value, field_name="shell value") + escaped = value.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`") + return f'"{escaped}"' + + +def _reject_control_characters(value: str, *, field_name: str) -> None: + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise BatchRenderError(f"{field_name} must not contain control characters") 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 new file mode 100644 index 000000000..44f6b2f84 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Injectable process execution for Slurm command-line tools.""" + +from __future__ import annotations + +import subprocess +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Protocol + + +class CommandRunner(Protocol): + """Minimal command boundary implemented by production and fake runners.""" + + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + """Execute one argument-vector command.""" + ... + + +class SubprocessRunner: + """Run commands without a shell or unrestricted ambient environment.""" + + _environment: Mapping[str, str] + _timeout_seconds: float + + def __init__( + self, + *, + environment: Mapping[str, str] | None = None, + timeout_seconds: float = 30.0, + ) -> None: + if timeout_seconds <= 0: + raise ValueError("timeout_seconds must be positive") + explicit_environment = dict(environment or {}) + 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") + if type(value) is not str or "\0" in value: + raise ValueError("environment values must not contain NUL") + self._environment = MappingProxyType({**explicit_environment, "LC_ALL": "C"}) + self._timeout_seconds = timeout_seconds + + @property + def environment(self) -> Mapping[str, str]: + """Return the explicit environment forwarded to child processes.""" + return self._environment + + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + """Execute an argument vector with captured text output.""" + return subprocess.run( + tuple(command), + check=False, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + env=dict(self._environment), + timeout=self._timeout_seconds, + ) diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py new file mode 100644 index 000000000..49cbe220c --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +from collections.abc import Sequence + +import pytest +from slurm_test_fakes import FakeCommandResponse, FakeSlurmRunner + +from data_designer.slurm.launcher import SlurmCommandClient, SlurmCommandError +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + + +def test_client_submits_and_observes_one_managed_array(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + submission = client.submit("/workspace/run.sbatch") + queue = client.query_queue((submission.array_job_id,)) + + assert submission.array_job_id == 4101 + assert tuple(record.state for record in queue) == (SchedulerState.PENDING, SchedulerState.RUNNING) + assert fake_slurm_runner.calls == [ + ("sbatch", "--parsable", "/workspace/run.sbatch"), + ("squeue", "--noheader", "--array", "--format=%i|%T", "--jobs=4101"), + ] + + +def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + client.submit("run.sbatch") + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=1) + + client.cancel(scheduler) + accounting = client.query_accounting((scheduler,)) + + assert len(accounting) == 1 + assert accounting[0].scheduler == scheduler + assert accounting[0].state is SchedulerState.CANCELLED + assert fake_slurm_runner.calls[-2:] == [ + ("scancel", "4101_1"), + ("sacct", "--noheader", "--parsable2", "--format=%i|%State|%ExitCode", "--jobs=4101_1"), + ] + + +def test_client_deduplicates_explicit_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + client.submit("run.sbatch") + + client.query_queue((4101, 4101, SchedulerIdentity(array_job_id=4101, array_task_id=0))) + + assert fake_slurm_runner.calls[-1][-1] == "--jobs=4101,4101_0" + + +def test_client_rejects_unbounded_or_invalid_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="at least one"): + client.query_queue(()) + with pytest.raises(ValueError, match="positive integers"): + client.query_accounting((0,)) + with pytest.raises(ValueError, match="positive integers"): + client.cancel(True) + + assert fake_slurm_runner.calls == [] + + +def test_client_queries_bounded_gpu_inventory(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + assert client.query_gpu_counts() == (2,) + assert fake_slurm_runner.calls == [("sinfo", "--noheader", "--format=%G")] + + +def test_client_rejects_invalid_gpu_partition_without_running_command(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="valid identifier"): + client.query_gpu_counts(partition="batch,other") + + assert fake_slurm_runner.calls == [] + + +def test_client_normalizes_command_failures(fake_slurm_runner: FakeSlurmRunner) -> None: + fake_slurm_runner.script_next( + "sacct", + FakeCommandResponse(stderr="accounting\nservice unavailable\n", returncode=2), + ) + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(SlurmCommandError, match="sacct failed with exit code 2: accounting service unavailable"): + client.query_accounting((4101,)) + + +def test_client_normalizes_execution_errors() -> None: + client = SlurmCommandClient(_FailingRunner()) + + with pytest.raises(SlurmCommandError, match="squeue could not be executed") as error: + client.query_queue((4101,)) + + assert isinstance(error.value.__cause__, FileNotFoundError) + + +def test_script_path_is_one_argument_vector_token(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + client.submit("/workspace/run; touch injected.sbatch") + + assert fake_slurm_runner.calls[0] == ( + "sbatch", + "--parsable", + "/workspace/run; touch injected.sbatch", + ) + + +class _FailingRunner: + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + del command + raise FileNotFoundError("missing executable") diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py new file mode 100644 index 000000000..f0f4f0c0f --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path + +import pytest + +from data_designer.slurm.launcher import QueueRecord, SlurmParseError +from data_designer.slurm.launcher.parsing import ( + parse_accounting, + parse_gpu_counts, + parse_queue, + parse_state, + parse_submission, +) +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "slurm" + + +@pytest.mark.parametrize( + ("output", "expected_job_id"), + (("4101\n", 4101), ("4101;primary\n", 4101)), +) +def test_parse_submission_accepts_parsable_sbatch_output(output: str, expected_job_id: int) -> None: + assert parse_submission(output).array_job_id == expected_job_id + + +@pytest.mark.parametrize("output", ("", "0", "Submitted batch job 4101", "٤١٠١", "4101;", "4101;bad name")) +def test_parse_submission_rejects_malformed_output(output: str) -> None: + with pytest.raises(SlurmParseError, match="invalid"): + parse_submission(output) + + +def test_parse_queue_normalizes_active_array_tasks() -> None: + records = parse_queue((GOLDEN_DIRECTORY / "squeue_active.txt").read_text()) + + assert records == ( + _make_queue_record(0, SchedulerState.PENDING), + _make_queue_record(1, SchedulerState.RUNNING), + ) + + +@pytest.mark.parametrize( + ("raw_state", "expected"), + ( + ("CONFIGURING", SchedulerState.PENDING), + ("COMPLETING", SchedulerState.RUNNING), + ("COMPLETED+", SchedulerState.COMPLETED), + ("CANCELLED by 1234", SchedulerState.CANCELLED), + ("TIMEOUT", SchedulerState.TIMED_OUT), + ("NODE_FAIL", SchedulerState.NODE_FAILED), + ("PREEMPTED", SchedulerState.PREEMPTED), + ("REQUEUED", SchedulerState.REQUEUED), + ("OUT_OF_MEMORY", SchedulerState.OUT_OF_MEMORY), + ("A_NEW_STATE", SchedulerState.UNKNOWN), + ), +) +def test_parse_state_normalizes_long_slurm_spellings(raw_state: str, expected: SchedulerState) -> None: + assert parse_state(raw_state) is expected + + +def test_parse_accounting_normalizes_terminal_rows_and_ignores_step_rows() -> None: + output = ( + "4101|RUNNING|0:0\n" + (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() + "4101_0.batch|FAILED|1:0\n" + ) + + records = parse_accounting(output) + + assert tuple(record.state for record in records) == ( + SchedulerState.TIMED_OUT, + SchedulerState.NODE_FAILED, + SchedulerState.PREEMPTED, + SchedulerState.REQUEUED, + SchedulerState.OUT_OF_MEMORY, + SchedulerState.CANCELLED, + ) + assert records[0].exit_code.status == 0 + assert records[0].exit_code.signal == 125 + + +def test_empty_scheduler_output_preserves_absent_evidence_for_reconciliation() -> None: + assert parse_queue("") == () + assert parse_accounting("\n") == () + + +@pytest.mark.parametrize( + ("parser", "output", "message"), + ( + (parse_queue, "malformed scheduler output\n", "two fields"), + (parse_queue, "4101|RUNNING\n", "array-task ID"), + (parse_queue, "4101_0|RUNNING\n4101_0|PENDING\n", "duplicates"), + (parse_accounting, "4101_0|FAILED\n", "three fields"), + (parse_accounting, "4101_0|FAILED|not-an-exit-code\n", "exit code"), + (parse_accounting, "garbage.step|FAILED|1:0\n", "array-task ID"), + (parse_queue, "4101_0|COMPLETED unexpectedly\n", "unexpected whitespace"), + ), +) +def test_scheduler_parsers_reject_malformed_or_ambiguous_rows( + parser: Callable[[str], object], + output: str, + message: str, +) -> None: + with pytest.raises(SlurmParseError, match=message): + parser(output) + + +@pytest.mark.parametrize( + ("output", "expected"), + ( + ("gpu:2\n", (2,)), + ("gpu:a100:8(S:0-7)\n", (8,)), + ("gpu:a100:4,gpu:h100:4\n(null)\n", (8,)), + ("(null)\nN/A\n", ()), + ), +) +def test_parse_gpu_counts_normalizes_configured_gres(output: str, expected: tuple[int, ...]) -> None: + assert parse_gpu_counts(output) == expected + + +def test_parse_gpu_counts_rejects_malformed_gpu_resources() -> None: + with pytest.raises(SlurmParseError, match="invalid GPU resource"): + parse_gpu_counts("gpu:a100:not-a-count\n") + + +def _make_queue_record(array_task_id: int, state: SchedulerState) -> QueueRecord: + return QueueRecord( + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=array_task_id), + state=state, + ) diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py new file mode 100644 index 000000000..738e70228 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from data_designer.slurm.config import SchedulerProfile, injected_profile +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.launcher import BatchDirective, BatchRenderError, render_batch_script +from data_designer.slurm.planning import ResolvedSlurmRunPlan, ResolvedSubmission + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "rendered" + + +@pytest.mark.parametrize( + ("fixture_name", "plan_fixture"), + (("single_node.sbatch", "single_node_plan"), ("multi_node.sbatch", "multi_node_plan")), +) +def test_renderer_matches_contract_bound_goldens( + fixture_name: str, + plan_fixture: str, + request: pytest.FixtureRequest, +) -> None: + plan = request.getfixturevalue(plan_fixture) + + assert render_batch_script(plan) == (GOLDEN_DIRECTORY / fixture_name).read_text() + + +def test_renderer_omits_gres_for_visible_mode_and_emits_optional_fields( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + profile = single_node_plan.selected_profile.profile.model_copy( + update={ + "gpu_request_mode": "visible", + "scheduler": SchedulerProfile(account="research", partition="batch", mem_per_gpu="80G"), + } + ) + plan = single_node_plan.model_copy( + update={ + "selected_profile": injected_profile(profile), + "submission": ResolvedSubmission( + job_name="data-designer", + account=None, + partition=None, + time_limit="01:00:00", + comment="safe test run", + ), + } + ) + + script = render_batch_script(plan) + + assert "#SBATCH --gres=" not in script + assert "#SBATCH --account=" not in script + assert "#SBATCH --partition=" not in script + assert "#SBATCH --mem-per-gpu=80G\n" in script + assert '#SBATCH --comment="safe test run"\n' in script + + +def test_renderer_escapes_shell_expansion_in_structured_paths( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = single_node_plan.model_copy( + update={ + "runtime_bundle": ArtifactReference( + path='/workspace/runtime/$(touch owned)-`whoami`-"bundle".tar.gz', + sha256="e" * 64, + ) + } + ) + + script = render_batch_script(plan) + + assert ( + 'readonly DD_RUNTIME_ARCHIVE="/workspace/runtime/\\$(touch owned)-\\`whoami\\`-\\"bundle\\".tar.gz"' in script + ) + completed = subprocess.run(("bash", "-n"), input=script, capture_output=True, text=True, check=False) + assert completed.returncode == 0, completed.stderr + + +def test_renderer_keeps_user_text_on_one_non_executable_directive( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + comment = '$(touch owned)" --output=/tmp/owned; `whoami`' + plan = single_node_plan.model_copy( + update={"submission": single_node_plan.submission.model_copy(update={"comment": comment})} + ) + + script = render_batch_script(plan) + + comment_lines = [line for line in script.splitlines() if line.startswith("#SBATCH --comment=")] + assert len(comment_lines) == 1 + assert "\\$(touch owned)" in comment_lines[0] + assert "\\`whoami\\`" in comment_lines[0] + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +@pytest.mark.parametrize( + "attempt_ordinal", + (0, -1, True), +) +def test_renderer_rejects_invalid_attempt_ordinals( + single_node_plan: ResolvedSlurmRunPlan, + attempt_ordinal: object, +) -> None: + with pytest.raises(BatchRenderError, match="positive integer"): + render_batch_script(single_node_plan, attempt_ordinal=attempt_ordinal) # type: ignore[arg-type] + + +def test_batch_directive_rejects_invalid_names_and_control_characters() -> None: + with pytest.raises(BatchRenderError, match="name is invalid"): + BatchDirective(name="output\n", value="safe").render() + with pytest.raises(BatchRenderError, match="control characters"): + BatchDirective(name="comment", value="first\n#SBATCH --output=owned").render() + + +def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) -> None: + script = render_batch_script(single_node_plan, attempt_ordinal=12) + + assert script.count("dd_slurm_run_allocation") == 1 + assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script + assert len(script.splitlines()) < 40 + assert script.endswith("\n") diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py new file mode 100644 index 000000000..8aa47288e --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +from collections.abc import Mapping, Sequence + +import pytest + +from data_designer.slurm.launcher import SubprocessRunner + + +def test_subprocess_runner_uses_argv_and_only_explicit_environment(monkeypatch: pytest.MonkeyPatch) -> None: + observed: dict[str, object] = {} + + def fake_run( + command: Sequence[str], + *, + check: bool, + stdin: int, + capture_output: bool, + text: bool, + env: Mapping[str, str], + timeout: float, + ) -> subprocess.CompletedProcess[str]: + observed.update( + command=command, + check=check, + stdin=stdin, + capture_output=capture_output, + text=text, + env=env, + timeout=timeout, + ) + return subprocess.CompletedProcess(command, 0, stdout="ok\n", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + supplied_environment = {"PATH": "/usr/bin", "LC_ALL": "fr_FR.UTF-8"} + runner = SubprocessRunner(environment=supplied_environment, timeout_seconds=4.0) + supplied_environment["SECRET"] = "must-not-leak" + + completed = runner.run(("squeue", "--noheader")) + + assert completed.stdout == "ok\n" + assert observed == { + "command": ("squeue", "--noheader"), + "check": False, + "stdin": subprocess.DEVNULL, + "capture_output": True, + "text": True, + "env": {"LC_ALL": "C", "PATH": "/usr/bin"}, + "timeout": 4.0, + } + + +def test_subprocess_runner_environment_is_immutable() -> None: + runner = SubprocessRunner() + + with pytest.raises(TypeError): + runner.environment["SECRET"] = "value" # type: ignore[index] + + +def test_subprocess_runner_rejects_nonpositive_timeout() -> None: + with pytest.raises(ValueError, match="positive"): + SubprocessRunner(timeout_seconds=0) + + +@pytest.mark.parametrize("environment", ({"BAD=NAME": "value"}, {"NAME": "bad\0value"})) +def test_subprocess_runner_rejects_invalid_environment(environment: dict[str, str]) -> None: + with pytest.raises(ValueError, match="environment"): + SubprocessRunner(environment=environment) 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 0cf7a8e85..f2a0b94c0 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 @@ -21,12 +21,17 @@ verify_sha256() { verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" verify_sha256 "${DD_PLAN_SHA256}" "${DD_PLAN}" -readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID:?SLURM_ARRAY_TASK_ID is required}" +if [[ ! ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]]; then + printf '%s\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 + exit 64 +fi +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}" -readonly DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +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}" source "${DD_RUNTIME_DIR}/entrypoint.sh" 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 60751e07c..e0849f783 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 @@ -21,12 +21,17 @@ verify_sha256() { verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" verify_sha256 "${DD_PLAN_SHA256}" "${DD_PLAN}" -readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID:?SLURM_ARRAY_TASK_ID is required}" +if [[ ! ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]]; then + printf '%s\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 + exit 64 +fi +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}" -readonly DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +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}" source "${DD_RUNTIME_DIR}/entrypoint.sh" 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 1d793f500..687b1e32c 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="cdb5204e8ce7a7affda99ea9c121c6b60a16bf42c510c8eb6bc57661969620b9", + expected_fixture_sha256="5e80f485138e4eac7a3280eb3dc7cb19d7ded0be73988a5f2d01ac0d083ad9a2", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="c6708cdbc0a03e095062c0642cfd141f066153b958b7ad3dec779afd6414fa34", + expected_fixture_sha256="965edcb71c34ff55367d9296963c404a29dbedbe944c92d97b4ee0eca90dec97", ) From 3a75c8cb6ff58c25bd57b05e519f31f877ed2795 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:35:22 -0600 Subject: [PATCH 02/23] fix: tighten Slurm launcher boundaries Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/__init__.py | 3 +-- .../src/data_designer/slurm/launcher/client.py | 2 ++ .../src/data_designer/slurm/launcher/models.py | 2 ++ .../src/data_designer/slurm/launcher/parsing.py | 2 +- .../src/data_designer/slurm/launcher/renderer.py | 6 +++--- .../tests/launcher/test_client.py | 10 +++++++++- .../tests/launcher/test_parsing.py | 15 +++++++++++---- .../tests/launcher/test_renderer.py | 14 +++++++++----- 8 files changed, 38 insertions(+), 16 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py index c3e29dce3..cc69332b5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py @@ -18,12 +18,11 @@ SlurmExitCode, SlurmSubmission, ) -from data_designer.slurm.launcher.renderer import BatchDirective, render_batch_script +from data_designer.slurm.launcher.renderer import render_batch_script from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner __all__ = [ "AccountingRecord", - "BatchDirective", "BatchRenderError", "CommandRunner", "QueueRecord", 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 29b9717e0..0bdbcba1e 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 @@ -87,6 +87,8 @@ def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[Accounting ( self._executables.sacct, "--noheader", + "--array", + "--allocations", "--parsable2", "--format=%i|%State|%ExitCode", f"--jobs={jobs}", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py index 2574ef204..7a45ac4e1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -7,6 +7,7 @@ from dataclasses import dataclass +from data_designer.slurm.contracts import Identifier from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -15,6 +16,7 @@ class SlurmSubmission: """Identity assigned by Slurm to one accepted array submission.""" array_job_id: int + cluster_name: Identifier | None = None @dataclass(frozen=True, slots=True) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index 67937dd75..c5b488124 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -58,7 +58,7 @@ def parse_submission(output: str) -> SlurmSubmission: raise SlurmParseError("sbatch returned an invalid job ID") if separator and _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: raise SlurmParseError("sbatch returned an invalid cluster name") - return SlurmSubmission(array_job_id=int(job_id)) + return SlurmSubmission(array_job_id=int(job_id), cluster_name=cluster_name or None) def parse_queue(output: str) -> tuple[QueueRecord, ...]: 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 d63532dd8..87c2984ae 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 @@ -17,7 +17,7 @@ @dataclass(frozen=True, slots=True) -class BatchDirective: +class _BatchDirective: """One validated ``#SBATCH`` option.""" name: str @@ -80,7 +80,7 @@ def render_batch_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int = 1) """ -def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[BatchDirective, ...]: +def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDirective, ...]: node_indices = ( plan.client.host_node_index, *(index for deployment in plan.deployments for index in deployment.node_indices), @@ -105,7 +105,7 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[BatchDirec values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) if plan.submission.comment is not None: values.append(("comment", plan.submission.comment)) - return tuple(BatchDirective(name=name, value=value) for name, value in values if value is not None) + return tuple(_BatchDirective(name=name, value=value) for name, value in values if value is not None) def _quote_double_value(value: str) -> str: diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 49cbe220c..a7898ed65 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -40,7 +40,15 @@ def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: assert accounting[0].state is SchedulerState.CANCELLED assert fake_slurm_runner.calls[-2:] == [ ("scancel", "4101_1"), - ("sacct", "--noheader", "--parsable2", "--format=%i|%State|%ExitCode", "--jobs=4101_1"), + ( + "sacct", + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=%i|%State|%ExitCode", + "--jobs=4101_1", + ), ] diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index f0f4f0c0f..609cbcc2c 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -22,11 +22,18 @@ @pytest.mark.parametrize( - ("output", "expected_job_id"), - (("4101\n", 4101), ("4101;primary\n", 4101)), + ("output", "expected_job_id", "expected_cluster"), + (("4101\n", 4101, None), ("4101;primary\n", 4101, "primary")), ) -def test_parse_submission_accepts_parsable_sbatch_output(output: str, expected_job_id: int) -> None: - assert parse_submission(output).array_job_id == expected_job_id +def test_parse_submission_accepts_parsable_sbatch_output( + output: str, + expected_job_id: int, + expected_cluster: str | None, +) -> None: + submission = parse_submission(output) + + assert submission.array_job_id == expected_job_id + assert submission.cluster_name == expected_cluster @pytest.mark.parametrize("output", ("", "0", "Submitted batch job 4101", "٤١٠١", "4101;", "4101;bad name")) diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 738e70228..12e98f0bf 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -10,7 +10,7 @@ from data_designer.slurm.config import SchedulerProfile, injected_profile from data_designer.slurm.contracts import ArtifactReference -from data_designer.slurm.launcher import BatchDirective, BatchRenderError, render_batch_script +from data_designer.slurm.launcher import BatchRenderError, render_batch_script from data_designer.slurm.planning import ResolvedSlurmRunPlan, ResolvedSubmission GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "rendered" @@ -111,11 +111,15 @@ def test_renderer_rejects_invalid_attempt_ordinals( render_batch_script(single_node_plan, attempt_ordinal=attempt_ordinal) # type: ignore[arg-type] -def test_batch_directive_rejects_invalid_names_and_control_characters() -> None: - with pytest.raises(BatchRenderError, match="name is invalid"): - BatchDirective(name="output\n", value="safe").render() +def test_renderer_rejects_control_characters_from_unvalidated_plan_copies( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = single_node_plan.model_copy( + update={"submission": single_node_plan.submission.model_copy(update={"comment": "unsafe\ntext"})} + ) + with pytest.raises(BatchRenderError, match="control characters"): - BatchDirective(name="comment", value="first\n#SBATCH --output=owned").render() + render_batch_script(plan) def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) -> None: From 7c9ab7c55cc0ff9210015d058ec887e3127caa75 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:37:05 -0600 Subject: [PATCH 03/23] fix: use native sacct output fields Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/client.py | 2 +- .../src/data_designer/slurm/launcher/parsing.py | 2 +- .../data-designer-slurm/tests/launcher/test_client.py | 2 +- .../tests/slurm_test_fakes/slurm.py | 8 +++++++- .../tests/slurm_test_fakes/test_slurm.py | 10 ++++++++-- 5 files changed, 18 insertions(+), 6 deletions(-) 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 0bdbcba1e..6cc4e8694 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 @@ -90,7 +90,7 @@ def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[Accounting "--array", "--allocations", "--parsable2", - "--format=%i|%State|%ExitCode", + "--format=JobIDRaw,State,ExitCode", f"--jobs={jobs}", ) ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index c5b488124..919c94e7d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -76,7 +76,7 @@ def parse_queue(output: str) -> tuple[QueueRecord, ...]: def parse_accounting(output: str) -> tuple[AccountingRecord, ...]: - """Parse array-task rows from ``sacct --format=%i|%State|%ExitCode``.""" + """Parse array-task rows from ``sacct --format=JobIDRaw,State,ExitCode``.""" records: list[AccountingRecord] = [] identities: set[SchedulerIdentity] = set() for line_number, line in _collect_nonempty_lines(output): diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index a7898ed65..f3807ca98 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -46,7 +46,7 @@ def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: "--array", "--allocations", "--parsable2", - "--format=%i|%State|%ExitCode", + "--format=JobIDRaw,State,ExitCode", "--jobs=4101_1", ), ] 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 785af0630..15b588e9c 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -15,7 +15,13 @@ _JOB_SELECTOR_PATTERN = re.compile(r"^[0-9]+(?:_[0-9]+)?$") _SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--format=%i|%T") -_SACCT_REQUIRED_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") +_SACCT_REQUIRED_ARGUMENTS = ( + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobIDRaw,State,ExitCode", +) @dataclass(frozen=True) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py index 8bfb3de0e..02a8a807d 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -13,7 +13,13 @@ GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" SQUEUE_ARGUMENTS = ("--noheader", "--format=%i|%T") -SACCT_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") +SACCT_ARGUMENTS = ( + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobIDRaw,State,ExitCode", +) def _submit(runner: FakeSlurmRunner) -> None: @@ -174,7 +180,7 @@ def test_fake_slurm_runner_matches_sbatch_parsable_mode( "command", ( ("squeue", "--noheader"), - ("sacct", "--noheader", "--format=%i|%State|%ExitCode"), + ("sacct", "--noheader", "--format=JobIDRaw,State,ExitCode"), ), ) def test_fake_slurm_runner_rejects_underspecified_state_queries( From 996ac7ff6739155b4aeeabdb89762e81a4ea47f9 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:38:07 -0600 Subject: [PATCH 04/23] test: require expanded Slurm array queries Part of #868 Signed-off-by: Nabin Mulepati --- packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py | 2 +- .../data-designer-slurm/tests/slurm_test_fakes/test_slurm.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 15b588e9c..5f29ca68b 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -14,7 +14,7 @@ from data_designer.slurm.state import SchedulerIdentity _JOB_SELECTOR_PATTERN = re.compile(r"^[0-9]+(?:_[0-9]+)?$") -_SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--format=%i|%T") +_SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") _SACCT_REQUIRED_ARGUMENTS = ( "--noheader", "--array", diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py index 02a8a807d..e13b9f586 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -12,7 +12,7 @@ from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" -SQUEUE_ARGUMENTS = ("--noheader", "--format=%i|%T") +SQUEUE_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") SACCT_ARGUMENTS = ( "--noheader", "--array", From 0bf21421483bbb9996c1ff0c3c2ebcfa6674cc05 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:39:38 -0600 Subject: [PATCH 05/23] fix: sanitize Slurm command diagnostics Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/client.py | 4 +++- .../tests/launcher/test_client.py | 13 +++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) 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 6cc4e8694..dcfb32f49 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 @@ -7,6 +7,7 @@ import re import subprocess +import unicodedata from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path @@ -145,7 +146,8 @@ def _validate_argument(value: str, *, field_name: str) -> None: def _normalize_bounded_text(value: str, *, limit: int = 512) -> str: - normalized = " ".join(value.split()) + sanitized = "".join(" " if unicodedata.category(character).startswith("C") else character for character in value) + normalized = " ".join(sanitized.split()) return normalized if len(normalized) <= limit else f"{normalized[:limit]}..." diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index f3807ca98..63ebe939b 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -101,6 +101,19 @@ def test_client_normalizes_command_failures(fake_slurm_runner: FakeSlurmRunner) client.query_accounting((4101,)) +def test_client_removes_terminal_controls_from_command_failures(fake_slurm_runner: FakeSlurmRunner) -> None: + fake_slurm_runner.script_next( + "squeue", + FakeCommandResponse(stderr="queue unavailable\x1b[31m\n", returncode=2), + ) + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(SlurmCommandError) as error: + client.query_queue((4101,)) + + assert "\x1b" not in str(error.value) + + def test_client_normalizes_execution_errors() -> None: client = SlurmCommandClient(_FailingRunner()) From 65804af315687f5793d4cd76643bc92309f7fe46 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:41:35 -0600 Subject: [PATCH 06/23] fix: reject invalid visible GPU memory Part of #868 Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/renderer.py | 2 ++ .../tests/launcher/test_renderer.py | 29 +++++++++++++++++-- 2 files changed, 28 insertions(+), 3 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 87c2984ae..cec3954c3 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 @@ -101,6 +101,8 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire profile = plan.selected_profile.profile if profile.gpu_request_mode == "gres": values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) + elif profile.scheduler.mem_per_gpu is not None: + raise BatchRenderError("mem_per_gpu requires GRES GPU request mode") if profile.scheduler.mem_per_gpu is not None: values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) if plan.submission.comment is not None: diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 12e98f0bf..455baf688 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -30,13 +30,13 @@ def test_renderer_matches_contract_bound_goldens( assert render_batch_script(plan) == (GOLDEN_DIRECTORY / fixture_name).read_text() -def test_renderer_omits_gres_for_visible_mode_and_emits_optional_fields( +def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fields( single_node_plan: ResolvedSlurmRunPlan, ) -> None: profile = single_node_plan.selected_profile.profile.model_copy( update={ "gpu_request_mode": "visible", - "scheduler": SchedulerProfile(account="research", partition="batch", mem_per_gpu="80G"), + "scheduler": SchedulerProfile(account="research", partition="batch"), } ) plan = single_node_plan.model_copy( @@ -57,10 +57,33 @@ def test_renderer_omits_gres_for_visible_mode_and_emits_optional_fields( assert "#SBATCH --gres=" not in script assert "#SBATCH --account=" not in script assert "#SBATCH --partition=" not in script - assert "#SBATCH --mem-per-gpu=80G\n" in script assert '#SBATCH --comment="safe test run"\n' in script +def test_renderer_emits_mem_per_gpu_for_gres_mode(single_node_plan: ResolvedSlurmRunPlan) -> None: + profile = single_node_plan.selected_profile.profile.model_copy( + update={"scheduler": SchedulerProfile(account="research", partition="batch", mem_per_gpu="80G")} + ) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + + assert "#SBATCH --mem-per-gpu=80G\n" in render_batch_script(plan) + + +def test_renderer_rejects_mem_per_gpu_without_a_slurm_gpu_request( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + profile = single_node_plan.selected_profile.profile.model_copy( + update={ + "gpu_request_mode": "visible", + "scheduler": SchedulerProfile(account="research", partition="batch", mem_per_gpu="80G"), + } + ) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + + with pytest.raises(BatchRenderError, match="requires GRES"): + render_batch_script(plan) + + def test_renderer_escapes_shell_expansion_in_structured_paths( single_node_plan: ResolvedSlurmRunPlan, ) -> None: From 926a2212225eac656f705d52f442ac46394d0b96 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:43:53 -0600 Subject: [PATCH 07/23] fix: hash rendered inputs without manifests Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/renderer.py | 4 +++- packages/data-designer-slurm/tests/launcher/test_renderer.py | 2 +- .../tests/slurm_test_fakes/golden/rendered/multi_node.sbatch | 4 +++- .../tests/slurm_test_fakes/golden/rendered/single_node.sbatch | 4 +++- .../tests/slurm_test_fakes/test_rendered_scripts.py | 4 ++-- 5 files changed, 12 insertions(+), 6 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 cec3954c3..f1e764327 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 @@ -57,7 +57,9 @@ def render_batch_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int = 1) readonly DD_ATTEMPT_ORDINAL={_quote_double_value(attempt)} verify_sha256() {{ - printf '%s %s\\n' "$1" "$2" | sha256sum --check --status - + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${{actual_sha256%% *}}" == "$1" ]] }} verify_sha256 "${{DD_RUNTIME_SHA256}}" "${{DD_RUNTIME_ARCHIVE}}" diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 455baf688..dd9c11e3e 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -150,5 +150,5 @@ def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) - assert script.count("dd_slurm_run_allocation") == 1 assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script - assert len(script.splitlines()) < 40 + assert len(script.splitlines()) <= 40 assert script.endswith("\n") 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 f2a0b94c0..eed9df12c 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 @@ -16,7 +16,9 @@ readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" verify_sha256() { - printf '%s %s\n' "$1" "$2" | sha256sum --check --status - + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${actual_sha256%% *}" == "$1" ]] } verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" 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 e0849f783..833a363f8 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 @@ -16,7 +16,9 @@ readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" verify_sha256() { - printf '%s %s\n' "$1" "$2" | sha256sum --check --status - + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${actual_sha256%% *}" == "$1" ]] } verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" 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 687b1e32c..c3f80c236 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="5e80f485138e4eac7a3280eb3dc7cb19d7ded0be73988a5f2d01ac0d083ad9a2", + expected_fixture_sha256="8bb51021f8b8b1e4144335829c92c70d78948e3010ef6b93bda1c498a0b76ed7", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="965edcb71c34ff55367d9296963c404a29dbedbe944c92d97b4ee0eca90dec97", + expected_fixture_sha256="4d491afe35e815f28367ea694b648fc34779c4584b43cfd89e31de6df01f9f34", ) From c5df7d6cb313ec06cd7ff089b8ebd593c11f9287 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:45:11 -0600 Subject: [PATCH 08/23] fix: make Slurm output decoding stable Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/runner.py | 2 ++ packages/data-designer-slurm/tests/launcher/test_runner.py | 6 ++++++ 2 files changed, 8 insertions(+) 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 44f6b2f84..532b4d7c4 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 @@ -55,6 +55,8 @@ def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: stdin=subprocess.DEVNULL, capture_output=True, text=True, + encoding="utf-8", + errors="replace", env=dict(self._environment), timeout=self._timeout_seconds, ) diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index 8aa47288e..f50ddab40 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -21,6 +21,8 @@ def fake_run( stdin: int, capture_output: bool, text: bool, + encoding: str, + errors: str, env: Mapping[str, str], timeout: float, ) -> subprocess.CompletedProcess[str]: @@ -30,6 +32,8 @@ def fake_run( stdin=stdin, capture_output=capture_output, text=text, + encoding=encoding, + errors=errors, env=env, timeout=timeout, ) @@ -49,6 +53,8 @@ def fake_run( "stdin": subprocess.DEVNULL, "capture_output": True, "text": True, + "encoding": "utf-8", + "errors": "replace", "env": {"LC_ALL": "C", "PATH": "/usr/bin"}, "timeout": 4.0, } From 6d20a8c4db974b741619bd1a7c7250e68767ae90 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:47:55 -0600 Subject: [PATCH 09/23] fix: validate finite command timeouts Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/runner.py | 7 ++++--- packages/data-designer-slurm/tests/launcher/test_runner.py | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) 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 532b4d7c4..07bf1b312 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 @@ -5,6 +5,7 @@ from __future__ import annotations +import math import subprocess from collections.abc import Mapping, Sequence from types import MappingProxyType @@ -31,8 +32,8 @@ def __init__( environment: Mapping[str, str] | None = None, timeout_seconds: float = 30.0, ) -> None: - if timeout_seconds <= 0: - raise ValueError("timeout_seconds must be positive") + if isinstance(timeout_seconds, bool) or not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise ValueError("timeout_seconds must be a finite positive number") explicit_environment = dict(environment or {}) for name, value in explicit_environment.items(): if type(name) is not str or not name or "=" in name or "\0" in name: @@ -40,7 +41,7 @@ def __init__( if type(value) is not str or "\0" in value: raise ValueError("environment values must not contain NUL") self._environment = MappingProxyType({**explicit_environment, "LC_ALL": "C"}) - self._timeout_seconds = timeout_seconds + self._timeout_seconds = float(timeout_seconds) @property def environment(self) -> Mapping[str, str]: diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index f50ddab40..47a9057c9 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -67,9 +67,10 @@ def test_subprocess_runner_environment_is_immutable() -> None: runner.environment["SECRET"] = "value" # type: ignore[index] -def test_subprocess_runner_rejects_nonpositive_timeout() -> None: - with pytest.raises(ValueError, match="positive"): - SubprocessRunner(timeout_seconds=0) +@pytest.mark.parametrize("timeout_seconds", [0, -1, True, float("nan"), float("inf")]) +def test_subprocess_runner_rejects_invalid_timeout(timeout_seconds: float) -> None: + with pytest.raises(ValueError, match="finite positive"): + SubprocessRunner(timeout_seconds=timeout_seconds) @pytest.mark.parametrize("environment", ({"BAD=NAME": "value"}, {"NAME": "bad\0value"})) From e2297b184896e353affdd0efbb5c3f2bfdf4bbdc Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:53:09 -0600 Subject: [PATCH 10/23] fix: harden Slurm boundary inputs Part of #868 Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/client.py | 2 ++ .../data_designer/slurm/launcher/parsing.py | 22 ++++++++++++++++++- .../data_designer/slurm/launcher/runner.py | 2 +- .../tests/launcher/test_client.py | 9 ++++++++ .../tests/launcher/test_parsing.py | 6 +++-- .../tests/launcher/test_runner.py | 6 ++--- 6 files changed, 40 insertions(+), 7 deletions(-) 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 dcfb32f49..69124b27e 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 @@ -64,6 +64,8 @@ def submit(self, script_path: str | Path) -> SlurmSubmission: """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", path)) return parse_submission(output) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index 919c94e7d..dfb9fe9c7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -104,7 +104,7 @@ def parse_gpu_counts(output: str) -> tuple[int, ...]: if line in {"(null)", "N/A"}: continue line_counts: list[int] = [] - for gres in line.split(","): + for gres in _split_gres_fields(line, line_number=line_number): if not gres.startswith("gpu:"): continue match = _GRES_GPU_PATTERN.fullmatch(gres) @@ -116,6 +116,26 @@ def parse_gpu_counts(output: str) -> tuple[int, ...]: return tuple(counts) +def _split_gres_fields(value: str, *, line_number: int) -> tuple[str, ...]: + fields: list[str] = [] + start = 0 + annotation_depth = 0 + for index, character in enumerate(value): + if character == "(": + annotation_depth += 1 + elif character == ")": + annotation_depth -= 1 + if annotation_depth < 0: + raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + elif character == "," and annotation_depth == 0: + fields.append(value[start:index]) + start = index + 1 + if annotation_depth: + raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + fields.append(value[start:]) + return tuple(fields) + + def parse_state(value: str) -> SchedulerState: """Normalize one Slurm long state spelling without guessing unknown states.""" normalized = value.strip().upper().removesuffix("+") 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 07bf1b312..d87781155 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 @@ -32,7 +32,7 @@ def __init__( environment: Mapping[str, str] | None = None, timeout_seconds: float = 30.0, ) -> None: - if isinstance(timeout_seconds, bool) or not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + 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 or {}) for name, value in explicit_environment.items(): diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 63ebe939b..276f07b83 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -135,6 +135,15 @@ def test_script_path_is_one_argument_vector_token(fake_slurm_runner: FakeSlurmRu ) +def test_client_rejects_option_like_script_path(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="must not begin"): + client.submit("--wrap=unexpected") + + assert fake_slurm_runner.calls == [] + + class _FailingRunner: def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: del command diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index 609cbcc2c..96f3b26a9 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -120,6 +120,7 @@ def test_scheduler_parsers_reject_malformed_or_ambiguous_rows( ( ("gpu:2\n", (2,)), ("gpu:a100:8(S:0-7)\n", (8,)), + ("gpu:a100:4(S:0-1,4-5)\n", (4,)), ("gpu:a100:4,gpu:h100:4\n(null)\n", (8,)), ("(null)\nN/A\n", ()), ), @@ -128,9 +129,10 @@ def test_parse_gpu_counts_normalizes_configured_gres(output: str, expected: tupl assert parse_gpu_counts(output) == expected -def test_parse_gpu_counts_rejects_malformed_gpu_resources() -> None: +@pytest.mark.parametrize("output", ("gpu:a100:not-a-count\n", "gpu:a100:4(S:0-1,4-5\n")) +def test_parse_gpu_counts_rejects_malformed_gpu_resources(output: str) -> None: with pytest.raises(SlurmParseError, match="invalid GPU resource"): - parse_gpu_counts("gpu:a100:not-a-count\n") + parse_gpu_counts(output) def _make_queue_record(array_task_id: int, state: SchedulerState) -> QueueRecord: diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index 47a9057c9..55adc2940 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -67,10 +67,10 @@ def test_subprocess_runner_environment_is_immutable() -> None: runner.environment["SECRET"] = "value" # type: ignore[index] -@pytest.mark.parametrize("timeout_seconds", [0, -1, True, float("nan"), float("inf")]) -def test_subprocess_runner_rejects_invalid_timeout(timeout_seconds: float) -> None: +@pytest.mark.parametrize("timeout_seconds", [0, -1, True, "30", float("nan"), float("inf")]) +def test_subprocess_runner_rejects_invalid_timeout(timeout_seconds: object) -> None: with pytest.raises(ValueError, match="finite positive"): - SubprocessRunner(timeout_seconds=timeout_seconds) + SubprocessRunner(timeout_seconds=timeout_seconds) # type: ignore[arg-type] @pytest.mark.parametrize("environment", ({"BAD=NAME": "value"}, {"NAME": "bad\0value"})) From 1ea40f0dd37d2c591a29cb454b4247a563bc92db Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:55:44 -0600 Subject: [PATCH 11/23] fix: reject malformed Slurm GRES lists Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/parsing.py | 4 ++++ .../tests/launcher/test_parsing.py | 13 ++++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index dfb9fe9c7..c6a0a8a9e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -123,6 +123,8 @@ def _split_gres_fields(value: str, *, line_number: int) -> tuple[str, ...]: for index, character in enumerate(value): if character == "(": annotation_depth += 1 + if annotation_depth > 1: + raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") elif character == ")": annotation_depth -= 1 if annotation_depth < 0: @@ -133,6 +135,8 @@ def _split_gres_fields(value: str, *, line_number: int) -> tuple[str, ...]: if annotation_depth: raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") fields.append(value[start:]) + if any(not field for field in fields): + raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") return tuple(fields) diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index 96f3b26a9..c7b5da3a4 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -129,7 +129,18 @@ def test_parse_gpu_counts_normalizes_configured_gres(output: str, expected: tupl assert parse_gpu_counts(output) == expected -@pytest.mark.parametrize("output", ("gpu:a100:not-a-count\n", "gpu:a100:4(S:0-1,4-5\n")) +@pytest.mark.parametrize( + "output", + ( + "gpu:a100:not-a-count\n", + "gpu:a100:4(S:0-1,4-5\n", + "gpu:a100:4)\n", + "gpu:a100:4((S:0-1))\n", + "gpu:a100:4,\n", + ",gpu:a100:4\n", + "gpu:a100:4,,mps:100\n", + ), +) def test_parse_gpu_counts_rejects_malformed_gpu_resources(output: str) -> None: with pytest.raises(SlurmParseError, match="invalid GPU resource"): parse_gpu_counts(output) From 1f0c70be5be6221941f92494a8fa1cb2647445bd Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:57:54 -0600 Subject: [PATCH 12/23] fix: isolate submitted Slurm environments Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/client.py | 2 +- packages/data-designer-slurm/tests/launcher/test_client.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) 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 69124b27e..1eda7ee6f 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 @@ -66,7 +66,7 @@ def submit(self, script_path: str | Path) -> SlurmSubmission: _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", path)) + output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) return parse_submission(output) def query_queue(self, selectors: Sequence[JobSelector]) -> tuple[QueueRecord, ...]: diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 276f07b83..0a69af287 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -22,7 +22,7 @@ def test_client_submits_and_observes_one_managed_array(fake_slurm_runner: FakeSl assert submission.array_job_id == 4101 assert tuple(record.state for record in queue) == (SchedulerState.PENDING, SchedulerState.RUNNING) assert fake_slurm_runner.calls == [ - ("sbatch", "--parsable", "/workspace/run.sbatch"), + ("sbatch", "--parsable", "--export=NIL", "/workspace/run.sbatch"), ("squeue", "--noheader", "--array", "--format=%i|%T", "--jobs=4101"), ] @@ -131,6 +131,7 @@ def test_script_path_is_one_argument_vector_token(fake_slurm_runner: FakeSlurmRu assert fake_slurm_runner.calls[0] == ( "sbatch", "--parsable", + "--export=NIL", "/workspace/run; touch injected.sbatch", ) From edb5140056b67ff0453cef9418c764315b076660 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:02:38 -0600 Subject: [PATCH 13/23] test: verify launcher wheel import Part of #868 Signed-off-by: Nabin Mulepati --- scripts/test_slurm_package_install.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index f80247355..7de100df2 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -130,6 +130,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.contracts import ResumeWorkspace as ContractResumeWorkspace from data_designer.slurm.integration import PlanStateValidator from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.launcher import SlurmCommandClient from data_designer.slurm.planning import ArtifactReference as PlanningArtifactReference from data_designer.slurm.planning import RecordRange as PlanningRecordRange from data_designer.slurm.planning import ResumeWorkspace as PlanningResumeWorkspace @@ -139,6 +140,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.state import RunManifest assert RunManifest.__name__ == "RunManifest" assert ImageRegistryStore.__name__ == "ImageRegistryStore" +assert "SlurmCommandClient" in str(SlurmCommandClient) assert PlanningArtifactReference is ContractArtifactReference assert PlanningRecordRange is ContractRecordRange assert PlanningResumeWorkspace is ContractResumeWorkspace From 15b9dc3b87e7c3e56cae04f1dab39607a4d1702e Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:04:55 -0600 Subject: [PATCH 14/23] fix: pin batch host tool path Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/renderer.py | 1 + packages/data-designer-slurm/tests/launcher/test_renderer.py | 2 +- .../tests/slurm_test_fakes/golden/rendered/multi_node.sbatch | 1 + .../tests/slurm_test_fakes/golden/rendered/single_node.sbatch | 1 + .../tests/slurm_test_fakes/test_rendered_scripts.py | 4 ++-- 5 files changed, 6 insertions(+), 3 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 f1e764327..c489ca382 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 @@ -48,6 +48,7 @@ def render_batch_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int = 1) 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" readonly DD_RUNTIME_ARCHIVE={_quote_double_value(plan.runtime_bundle.path)} readonly DD_RUNTIME_SHA256={_quote_double_value(plan.runtime_bundle.sha256)} diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index dd9c11e3e..5012d4635 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -150,5 +150,5 @@ def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) - assert script.count("dd_slurm_run_allocation") == 1 assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script - assert len(script.splitlines()) <= 40 + assert len(script.splitlines()) <= 41 assert script.endswith("\n") 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 eed9df12c..405a5ec82 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,6 +7,7 @@ #SBATCH --array=0-1%2 #SBATCH --gres=gpu:8 set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" 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 833a363f8..b8daa536b 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,6 +7,7 @@ #SBATCH --array=0 #SBATCH --gres=gpu:8 set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" 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 c3f80c236..d6b4ae84a 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="8bb51021f8b8b1e4144335829c92c70d78948e3010ef6b93bda1c498a0b76ed7", + expected_fixture_sha256="cc6bb2a035a541b01f422cafc661012a9ecfcb751547b921285e8e4eac19d94a", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="4d491afe35e815f28367ea694b648fc34779c4584b43cfd89e31de6df01f9f34", + expected_fixture_sha256="2c29159b25f56250cc5e0b502115940f9354c681b1404d10ed599a82fed244c4", ) From 347484370f5aab25747cb66616d76371efd84632 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:08:55 -0600 Subject: [PATCH 15/23] test: cover Slurm launcher boundaries Part of #868 Signed-off-by: Nabin Mulepati --- .../tests/launcher/test_client.py | 54 ++++++++++++++++++- .../tests/launcher/test_parsing.py | 7 +++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 0a69af287..dc704ebf3 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -9,7 +9,7 @@ import pytest from slurm_test_fakes import FakeCommandResponse, FakeSlurmRunner -from data_designer.slurm.launcher import SlurmCommandClient, SlurmCommandError +from data_designer.slurm.launcher import SlurmCommandClient, SlurmCommandError, SlurmExecutables from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -81,6 +81,14 @@ def test_client_queries_bounded_gpu_inventory(fake_slurm_runner: FakeSlurmRunner assert fake_slurm_runner.calls == [("sinfo", "--noheader", "--format=%G")] +def test_client_queries_partition_scoped_gpu_inventory() -> None: + command = ("sinfo", "--noheader", "--format=%G", "--partition=batch") + runner = FakeSlurmRunner(sinfo_responses={command: FakeCommandResponse(stdout="gpu:a100:8\n")}) + + assert SlurmCommandClient(runner).query_gpu_counts(partition="batch") == (8,) + assert runner.calls == [command] + + def test_client_rejects_invalid_gpu_partition_without_running_command(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) @@ -123,6 +131,22 @@ def test_client_normalizes_execution_errors() -> None: assert isinstance(error.value.__cause__, FileNotFoundError) +def test_client_normalizes_command_timeouts() -> None: + client = SlurmCommandClient(_TimeoutRunner()) + + with pytest.raises(SlurmCommandError, match="command timed out") as error: + client.query_queue((4101,)) + + assert isinstance(error.value.__cause__, subprocess.TimeoutExpired) + + +def test_client_rejects_non_text_runner_output() -> None: + client = SlurmCommandClient(_NonTextRunner()) + + with pytest.raises(SlurmCommandError, match="did not return text output"): + client.query_queue((4101,)) + + def test_script_path_is_one_argument_vector_token(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) @@ -145,7 +169,35 @@ def test_client_rejects_option_like_script_path(fake_slurm_runner: FakeSlurmRunn assert fake_slurm_runner.calls == [] +@pytest.mark.parametrize("script_path", ("", "bad\npath")) +def test_client_rejects_invalid_script_path(fake_slurm_runner: FakeSlurmRunner, script_path: str) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="batch script path"): + client.submit(script_path) + + assert fake_slurm_runner.calls == [] + + +@pytest.mark.parametrize("executable", ("", "sbatch --wait", "sbatch\n")) +def test_executables_reject_invalid_tokens(executable: str) -> None: + with pytest.raises(ValueError, match="Slurm executable"): + SlurmExecutables(sbatch=executable) + + class _FailingRunner: def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: del command raise FileNotFoundError("missing executable") + + +class _TimeoutRunner: + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + raise subprocess.TimeoutExpired(command, 30.0) + + +class _NonTextRunner: + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + completed = subprocess.CompletedProcess(command, 0, stdout="ok", stderr="") + completed.stdout = b"not text" # type: ignore[assignment] + return completed diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index c7b5da3a4..16c9dc383 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -121,6 +121,7 @@ def test_scheduler_parsers_reject_malformed_or_ambiguous_rows( ("gpu:2\n", (2,)), ("gpu:a100:8(S:0-7)\n", (8,)), ("gpu:a100:4(S:0-1,4-5)\n", (4,)), + ("mps:100,gpu:a100:4\n", (4,)), ("gpu:a100:4,gpu:h100:4\n(null)\n", (8,)), ("(null)\nN/A\n", ()), ), @@ -146,6 +147,12 @@ def test_parse_gpu_counts_rejects_malformed_gpu_resources(output: str) -> None: parse_gpu_counts(output) +@pytest.mark.parametrize("state", ("", "CANCELLED by root")) +def test_parse_state_rejects_invalid_spellings(state: str) -> None: + with pytest.raises(SlurmParseError): + parse_state(state) + + def _make_queue_record(array_task_id: int, state: SchedulerState) -> QueueRecord: return QueueRecord( scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=array_task_id), From 3f576413ce6345341554544f53a89fa17ed1e974 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:11:45 -0600 Subject: [PATCH 16/23] fix: correlate Slurm query responses Part of #868 Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/client.py | 39 ++++++++++++++++--- .../tests/launcher/test_client.py | 14 ++++++- 2 files changed, 47 insertions(+), 6 deletions(-) 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 1eda7ee6f..fbb477e06 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 @@ -13,7 +13,7 @@ from pathlib import Path from data_designer.slurm.contracts import Identifier -from data_designer.slurm.launcher.errors import SlurmCommandError +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmParseError from data_designer.slurm.launcher.models import AccountingRecord, QueueRecord, SlurmSubmission from data_designer.slurm.launcher.parsing import ( parse_accounting, @@ -71,7 +71,8 @@ def submit(self, script_path: str | Path) -> SlurmSubmission: def query_queue(self, selectors: Sequence[JobSelector]) -> tuple[QueueRecord, ...]: """Return normalized active-queue rows for explicit managed jobs.""" - jobs = _format_selectors(selectors) + requested = tuple(selectors) + jobs = _format_selectors(requested) output = self._run( ( self._executables.squeue, @@ -81,11 +82,18 @@ def query_queue(self, selectors: Sequence[JobSelector]) -> tuple[QueueRecord, .. f"--jobs={jobs}", ) ) - return parse_queue(output) + records = parse_queue(output) + _validate_selected_schedulers( + tuple(record.scheduler for record in records), + requested, + command="squeue", + ) + return records def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[AccountingRecord, ...]: """Return normalized accounting rows for explicit managed jobs.""" - jobs = _format_selectors(selectors) + requested = tuple(selectors) + jobs = _format_selectors(requested) output = self._run( ( self._executables.sacct, @@ -97,7 +105,13 @@ def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[Accounting f"--jobs={jobs}", ) ) - return parse_accounting(output) + records = parse_accounting(output) + _validate_selected_schedulers( + tuple(record.scheduler for record in records), + requested, + command="sacct", + ) + return records def cancel(self, selector: JobSelector) -> None: """Cancel one managed Slurm array or array task.""" @@ -140,6 +154,21 @@ def _format_selector(selector: JobSelector) -> str: return str(selector) +def _validate_selected_schedulers( + schedulers: Sequence[SchedulerIdentity], + selectors: Sequence[JobSelector], + *, + command: str, +) -> None: + for scheduler in schedulers: + if any( + scheduler == selector if isinstance(selector, SchedulerIdentity) else scheduler.array_job_id == selector + for selector in selectors + ): + continue + raise SlurmParseError(f"{command} returned an unrequested array-task ID") + + def _validate_argument(value: str, *, field_name: str) -> None: if type(value) is not str or not value: raise ValueError(f"{field_name} must not be empty") diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index dc704ebf3..8aaa3c56f 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -9,7 +9,7 @@ import pytest from slurm_test_fakes import FakeCommandResponse, FakeSlurmRunner -from data_designer.slurm.launcher import SlurmCommandClient, SlurmCommandError, SlurmExecutables +from data_designer.slurm.launcher import SlurmCommandClient, SlurmCommandError, SlurmExecutables, SlurmParseError from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -61,6 +61,18 @@ def test_client_deduplicates_explicit_job_selectors(fake_slurm_runner: FakeSlurm assert fake_slurm_runner.calls[-1][-1] == "--jobs=4101,4101_0" +def test_client_rejects_unrequested_scheduler_records(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stdout="9999_0|RUNNING\n")) + + with pytest.raises(SlurmParseError, match="unrequested"): + client.query_queue((4101,)) + + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stdout="9999_0|FAILED|1:0\n")) + with pytest.raises(SlurmParseError, match="unrequested"): + client.query_accounting((SchedulerIdentity(array_job_id=4101, array_task_id=0),)) + + def test_client_rejects_unbounded_or_invalid_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) From ab735626b658166006124530aad57ff84abaf4bd Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:50:40 -0600 Subject: [PATCH 17/23] fix: harden Slurm command boundaries Preserve the caller's PATH as the only ambient lookup input for default Slurm commands while continuing to isolate all other environment variables. Normalize oversized numeric scheduler fields into the launcher parse-error boundary. Part of #868 Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/parsing.py | 32 +++++++++++++++---- .../data_designer/slurm/launcher/runner.py | 7 ++-- .../tests/launcher/test_parsing.py | 20 ++++++++++++ .../tests/launcher/test_runner.py | 9 ++++++ 4 files changed, 60 insertions(+), 8 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index c6a0a8a9e..27c2fda4d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -54,11 +54,14 @@ def parse_submission(output: str) -> SlurmSubmission: """Parse ``sbatch --parsable`` output.""" value = output.strip() job_id, separator, cluster_name = value.partition(";") - if not job_id.isascii() or not job_id.isdecimal() or int(job_id) <= 0: + if not job_id.isascii() or not job_id.isdecimal(): + raise SlurmParseError("sbatch returned an invalid job ID") + array_job_id = _parse_decimal(job_id, message="sbatch returned an invalid job ID") + if array_job_id <= 0: raise SlurmParseError("sbatch returned an invalid job ID") if separator and _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: raise SlurmParseError("sbatch returned an invalid cluster name") - return SlurmSubmission(array_job_id=int(job_id), cluster_name=cluster_name or None) + return SlurmSubmission(array_job_id=array_job_id, cluster_name=cluster_name or None) def parse_queue(output: str) -> tuple[QueueRecord, ...]: @@ -110,7 +113,12 @@ def parse_gpu_counts(output: str) -> tuple[int, ...]: match = _GRES_GPU_PATTERN.fullmatch(gres) if match is None: raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") - line_counts.append(int(match.group("count"))) + line_counts.append( + _parse_decimal( + match.group("count"), + message=f"sinfo line {line_number} contains an invalid GPU resource", + ) + ) if line_counts: counts.append(sum(line_counts)) return tuple(counts) @@ -167,9 +175,10 @@ def _parse_array_identity(value: str, *, command: str, line_number: int) -> Sche match = _ARRAY_ID_PATTERN.fullmatch(value) if match is None: raise SlurmParseError(f"{command} line {line_number} contains an invalid array-task ID") + message = f"{command} line {line_number} contains an invalid array-task ID" return SchedulerIdentity( - array_job_id=int(match.group("job")), - array_task_id=int(match.group("task")), + array_job_id=_parse_decimal(match.group("job"), message=message), + array_task_id=_parse_decimal(match.group("task"), message=message), ) @@ -177,7 +186,18 @@ def _parse_exit_code(value: str, *, line_number: int) -> SlurmExitCode: match = _EXIT_CODE_PATTERN.fullmatch(value) if match is None: raise SlurmParseError(f"sacct line {line_number} contains an invalid exit code") - return SlurmExitCode(status=int(match.group("status")), signal=int(match.group("signal"))) + message = f"sacct line {line_number} contains an invalid exit code" + return SlurmExitCode( + status=_parse_decimal(match.group("status"), message=message), + signal=_parse_decimal(match.group("signal"), message=message), + ) + + +def _parse_decimal(value: str, *, message: str) -> int: + try: + return int(value) + except ValueError as error: + raise SlurmParseError(message) from error def _reject_duplicate( 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 d87781155..2e2f95bf3 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 @@ -6,6 +6,7 @@ from __future__ import annotations import math +import os import subprocess from collections.abc import Mapping, Sequence from types import MappingProxyType @@ -34,7 +35,9 @@ 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 or {}) + explicit_environment = ( + dict(environment) if environment is not None else {"PATH": os.environ.get("PATH", os.defpath)} + ) 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") @@ -45,7 +48,7 @@ def __init__( @property def environment(self) -> Mapping[str, str]: - """Return the explicit environment forwarded to child processes.""" + """Return the allowlisted environment forwarded to child processes.""" return self._environment def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index 16c9dc383..97f480429 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -19,6 +19,7 @@ from data_designer.slurm.state import SchedulerIdentity, SchedulerState GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "slurm" +OVERSIZED_DECIMAL = "9" * 5000 @pytest.mark.parametrize( @@ -115,6 +116,25 @@ def test_scheduler_parsers_reject_malformed_or_ambiguous_rows( parser(output) +@pytest.mark.parametrize( + ("parser", "output", "message"), + ( + (parse_submission, OVERSIZED_DECIMAL, "invalid job ID"), + (parse_queue, f"4101_{OVERSIZED_DECIMAL}|RUNNING", "array-task ID"), + (parse_accounting, f"4101_0|FAILED|{OVERSIZED_DECIMAL}:0", "exit code"), + (parse_gpu_counts, f"gpu:{OVERSIZED_DECIMAL}", "invalid GPU resource"), + ), + ids=("submission-job-id", "queue-task-id", "accounting-exit-code", "gpu-count"), +) +def test_parsers_normalize_oversized_numeric_fields( + parser: Callable[[str], object], + output: str, + message: str, +) -> None: + with pytest.raises(SlurmParseError, match=message): + parser(output) + + @pytest.mark.parametrize( ("output", "expected"), ( diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index 55adc2940..233335bf7 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -60,6 +60,15 @@ def fake_run( } +def test_subprocess_runner_default_environment_forwards_only_search_path(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PATH", "/workspace/slurm/bin:/usr/bin") + monkeypatch.setenv("SECRET", "must-not-leak") + + runner = SubprocessRunner() + + assert runner.environment == {"LC_ALL": "C", "PATH": "/workspace/slurm/bin:/usr/bin"} + + def test_subprocess_runner_environment_is_immutable() -> None: runner = SubprocessRunner() From 57a93777dd38c5c2258a0386db47131e2d383bbb Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 15:24:17 -0600 Subject: [PATCH 18/23] fix: fall back from empty command path Use the platform default search path when the ambient PATH is absent or empty so bare Slurm executables remain resolvable. Part of #868 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/launcher/runner.py | 2 +- .../data-designer-slurm/tests/launcher/test_runner.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) 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 2e2f95bf3..e7c79b4b4 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 @@ -36,7 +36,7 @@ def __init__( 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", os.defpath)} + dict(environment) if environment is not None else {"PATH": os.environ.get("PATH") or os.defpath} ) for name, value in explicit_environment.items(): if type(name) is not str or not name or "=" in name or "\0" in name: diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index 233335bf7..8d51e68ae 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -3,6 +3,7 @@ from __future__ import annotations +import os import subprocess from collections.abc import Mapping, Sequence @@ -69,6 +70,14 @@ def test_subprocess_runner_default_environment_forwards_only_search_path(monkeyp assert runner.environment == {"LC_ALL": "C", "PATH": "/workspace/slurm/bin:/usr/bin"} +def test_subprocess_runner_default_environment_replaces_empty_search_path(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PATH", "") + + runner = SubprocessRunner() + + assert runner.environment == {"LC_ALL": "C", "PATH": os.defpath} + + def test_subprocess_runner_environment_is_immutable() -> None: runner = SubprocessRunner() From ddad8f06937f538cd4f8c8a61853479e09224af7 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 26 Aug 2026 12:59:03 -0600 Subject: [PATCH 19/23] fix: align Slurm launcher semantics --- .../src/data_designer/slurm/launcher/client.py | 2 +- .../src/data_designer/slurm/launcher/parsing.py | 8 ++++---- .../data_designer/slurm/launcher/renderer.py | 1 + .../tests/launcher/test_client.py | 2 +- .../tests/launcher/test_parsing.py | 10 ++++++---- .../tests/launcher/test_renderer.py | 17 ++++++++++++++++- .../golden/rendered/multi_node.sbatch | 1 + .../golden/rendered/single_node.sbatch | 1 + .../tests/slurm_test_fakes/slurm.py | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 5 +++-- .../tests/slurm_test_fakes/test_slurm.py | 4 ++-- 11 files changed, 37 insertions(+), 16 deletions(-) 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 fbb477e06..44c79b368 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 @@ -101,7 +101,7 @@ def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[Accounting "--array", "--allocations", "--parsable2", - "--format=JobIDRaw,State,ExitCode", + "--format=JobID,State,ExitCode", f"--jobs={jobs}", ) ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index 27c2fda4d..43129b15b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -17,7 +17,6 @@ from data_designer.slurm.state import SchedulerIdentity, SchedulerState _ARRAY_ID_PATTERN = re.compile(r"^(?P[1-9][0-9]*)_(?P[0-9]+)$") -_ARRAY_STEP_ID_PATTERN = re.compile(r"^[1-9][0-9]*_[0-9]+\.[^\s|]+$") _JOB_ID_PATTERN = re.compile(r"^[1-9][0-9]*$") _CLUSTER_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") _EXIT_CODE_PATTERN = re.compile(r"^(?P[0-9]+):(?P[0-9]+)$") @@ -38,11 +37,12 @@ "REQUEUED": SchedulerState.REQUEUED, "REQUEUE_FED": SchedulerState.PENDING, "REQUEUE_HOLD": SchedulerState.PENDING, + "RESV_DEL_HOLD": SchedulerState.PENDING, "RESIZING": SchedulerState.RUNNING, "REVOKED": SchedulerState.FAILED, "RUNNING": SchedulerState.RUNNING, "SIGNALING": SchedulerState.RUNNING, - "SPECIAL_EXIT": SchedulerState.FAILED, + "SPECIAL_EXIT": SchedulerState.PENDING, "STAGE_OUT": SchedulerState.RUNNING, "STOPPED": SchedulerState.RUNNING, "SUSPENDED": SchedulerState.RUNNING, @@ -79,14 +79,14 @@ def parse_queue(output: str) -> tuple[QueueRecord, ...]: def parse_accounting(output: str) -> tuple[AccountingRecord, ...]: - """Parse array-task rows from ``sacct --format=JobIDRaw,State,ExitCode``.""" + """Parse array-task rows from ``sacct --format=JobID,State,ExitCode``.""" records: list[AccountingRecord] = [] identities: set[SchedulerIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 3: raise SlurmParseError(f"sacct line {line_number} must contain three fields") - if _JOB_ID_PATTERN.fullmatch(fields[0]) is not None or _ARRAY_STEP_ID_PATTERN.fullmatch(fields[0]) is not None: + if _JOB_ID_PATTERN.fullmatch(fields[0]) is not None: continue scheduler = _parse_array_identity(fields[0], command="sacct", line_number=line_number) _reject_duplicate(scheduler, identities, command="sacct", line_number=line_number) 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 c489ca382..580e31d39 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 @@ -98,6 +98,7 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire ("account", plan.submission.account), ("partition", plan.submission.partition), ("nodes", str(node_count)), + ("cpus-per-task", str(plan.client.authored.cpus)), ("time", plan.submission.time_limit), ("array", array), ] diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 8aaa3c56f..9609ef175 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -46,7 +46,7 @@ def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: "--array", "--allocations", "--parsable2", - "--format=JobIDRaw,State,ExitCode", + "--format=JobID,State,ExitCode", "--jobs=4101_1", ), ] diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index 97f480429..455e41193 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -63,6 +63,9 @@ def test_parse_queue_normalizes_active_array_tasks() -> None: ("NODE_FAIL", SchedulerState.NODE_FAILED), ("PREEMPTED", SchedulerState.PREEMPTED), ("REQUEUED", SchedulerState.REQUEUED), + ("REQUEUE_HOLD", SchedulerState.PENDING), + ("RESV_DEL_HOLD", SchedulerState.PENDING), + ("SPECIAL_EXIT", SchedulerState.PENDING), ("OUT_OF_MEMORY", SchedulerState.OUT_OF_MEMORY), ("A_NEW_STATE", SchedulerState.UNKNOWN), ), @@ -71,10 +74,8 @@ def test_parse_state_normalizes_long_slurm_spellings(raw_state: str, expected: S assert parse_state(raw_state) is expected -def test_parse_accounting_normalizes_terminal_rows_and_ignores_step_rows() -> None: - output = ( - "4101|RUNNING|0:0\n" + (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() + "4101_0.batch|FAILED|1:0\n" - ) +def test_parse_accounting_normalizes_terminal_rows_and_ignores_array_parent() -> None: + output = "4101|RUNNING|0:0\n" + (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() records = parse_accounting(output) @@ -103,6 +104,7 @@ def test_empty_scheduler_output_preserves_absent_evidence_for_reconciliation() - (parse_queue, "4101_0|RUNNING\n4101_0|PENDING\n", "duplicates"), (parse_accounting, "4101_0|FAILED\n", "three fields"), (parse_accounting, "4101_0|FAILED|not-an-exit-code\n", "exit code"), + (parse_accounting, "4101_0.batch|FAILED|1:0\n", "array-task ID"), (parse_accounting, "garbage.step|FAILED|1:0\n", "array-task ID"), (parse_queue, "4101_0|COMPLETED unexpectedly\n", "unexpected whitespace"), ), diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 5012d4635..016cb15b7 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -5,6 +5,7 @@ import subprocess from pathlib import Path +from typing import Literal import pytest @@ -69,6 +70,20 @@ def test_renderer_emits_mem_per_gpu_for_gres_mode(single_node_plan: ResolvedSlur assert "#SBATCH --mem-per-gpu=80G\n" in render_batch_script(plan) +@pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) +def test_renderer_reserves_client_cpus_for_each_gpu_request_mode( + single_node_plan: ResolvedSlurmRunPlan, + gpu_request_mode: Literal["gres", "visible"], +) -> None: + profile = single_node_plan.selected_profile.profile.model_copy(update={"gpu_request_mode": gpu_request_mode}) + client = single_node_plan.client.model_copy( + update={"authored": single_node_plan.client.authored.model_copy(update={"cpus": 17})} + ) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile), "client": client}) + + assert "#SBATCH --cpus-per-task=17\n" in render_batch_script(plan) + + def test_renderer_rejects_mem_per_gpu_without_a_slurm_gpu_request( single_node_plan: ResolvedSlurmRunPlan, ) -> None: @@ -150,5 +165,5 @@ def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) - assert script.count("dd_slurm_run_allocation") == 1 assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script - assert len(script.splitlines()) <= 41 + assert len(script.splitlines()) <= 42 assert script.endswith("\n") 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 405a5ec82..d934fd637 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 @@ -3,6 +3,7 @@ #SBATCH --account=research #SBATCH --partition=batch #SBATCH --nodes=3 +#SBATCH --cpus-per-task=32 #SBATCH --time=03:55:00 #SBATCH --array=0-1%2 #SBATCH --gres=gpu:8 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 b8daa536b..93156136d 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 @@ -3,6 +3,7 @@ #SBATCH --account=research #SBATCH --partition=batch #SBATCH --nodes=1 +#SBATCH --cpus-per-task=32 #SBATCH --time=03:55:00 #SBATCH --array=0 #SBATCH --gres=gpu:8 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 5f29ca68b..2c211eccd 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -20,7 +20,7 @@ "--array", "--allocations", "--parsable2", - "--format=JobIDRaw,State,ExitCode", + "--format=JobID,State,ExitCode", ) 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 d6b4ae84a..5439234e5 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="cc6bb2a035a541b01f422cafc661012a9ecfcb751547b921285e8e4eac19d94a", + expected_fixture_sha256="8ddf07c38a825a7487c487fddfe051f0a1940f63063725b54b32bfe4c03fd9ca", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="2c29159b25f56250cc5e0b502115940f9354c681b1404d10ed599a82fed244c4", + expected_fixture_sha256="17a4c2e16189d22dfdb6885bf76264844ad3168dea0cf94aef70948d5ab2e6b7", ) @@ -62,6 +62,7 @@ def _assert_script_matches_plan( assert f"#SBATCH --account={plan.submission.account}\n" in script assert f"#SBATCH --partition={plan.submission.partition}\n" in script assert f"#SBATCH --nodes={node_count}\n" in script + assert f"#SBATCH --cpus-per-task={plan.client.authored.cpus}\n" in script assert f"#SBATCH --time={plan.submission.time_limit}\n" in script assert f"#SBATCH --array={array}\n" in script assert f"#SBATCH --gres=gpu:{plan.resolved_gpus_per_node}\n" in script diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py index e13b9f586..202143b34 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -18,7 +18,7 @@ "--array", "--allocations", "--parsable2", - "--format=JobIDRaw,State,ExitCode", + "--format=JobID,State,ExitCode", ) @@ -180,7 +180,7 @@ def test_fake_slurm_runner_matches_sbatch_parsable_mode( "command", ( ("squeue", "--noheader"), - ("sacct", "--noheader", "--format=JobIDRaw,State,ExitCode"), + ("sacct", "--noheader", "--format=JobID,State,ExitCode"), ), ) def test_fake_slurm_runner_rejects_underspecified_state_queries( From 709963cdda25fd7e14fe7b34339635a04816e120 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 26 Aug 2026 13:49:16 -0600 Subject: [PATCH 20/23] fix Slurm job observation contracts Support both ordinary jobs and array-task observations while preserving accounting-lag semantics. Bound scheduler numeric fields and diagnostics, and honor unthrottled plan arrays when concurrency is omitted. --- .../data_designer/slurm/launcher/__init__.py | 2 + .../data_designer/slurm/launcher/client.py | 62 +++++++++++++------ .../data_designer/slurm/launcher/models.py | 11 ++-- .../data_designer/slurm/launcher/parsing.py | 50 ++++++++++----- .../data_designer/slurm/launcher/renderer.py | 4 +- .../tests/launcher/test_client.py | 61 ++++++++++++++++-- .../tests/launcher/test_parsing.py | 34 ++++++++-- .../tests/launcher/test_renderer.py | 14 ++++- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- 9 files changed, 192 insertions(+), 50 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py index cc69332b5..9bc45a4a6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py @@ -16,6 +16,7 @@ AccountingRecord, QueueRecord, SlurmExitCode, + SlurmJobIdentity, SlurmSubmission, ) from data_designer.slurm.launcher.renderer import render_batch_script @@ -30,6 +31,7 @@ "SlurmCommandError", "SlurmExecutables", "SlurmExitCode", + "SlurmJobIdentity", "SlurmLauncherError", "SlurmParseError", "SlurmSubmission", 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 44c79b368..fdb62bc83 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 @@ -11,10 +11,11 @@ from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path +from typing import TypeAlias from data_designer.slurm.contracts import Identifier from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmParseError -from data_designer.slurm.launcher.models import AccountingRecord, QueueRecord, SlurmSubmission +from data_designer.slurm.launcher.models import AccountingRecord, QueueRecord, SlurmJobIdentity, SlurmSubmission from data_designer.slurm.launcher.parsing import ( parse_accounting, parse_gpu_counts, @@ -24,8 +25,9 @@ from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner from data_designer.slurm.state import SchedulerIdentity -JobSelector = int | SchedulerIdentity +JobSelector: TypeAlias = SlurmJobIdentity _IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_MAX_SLURM_INTEGER = (1 << 32) - 1 @dataclass(frozen=True, slots=True) @@ -83,12 +85,12 @@ def query_queue(self, selectors: Sequence[JobSelector]) -> tuple[QueueRecord, .. ) ) records = parse_queue(output) - _validate_selected_schedulers( + ignored = _validate_selected_schedulers( tuple(record.scheduler for record in records), requested, command="squeue", ) - return records + return tuple(record for record in records if record.scheduler not in ignored) def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[AccountingRecord, ...]: """Return normalized accounting rows for explicit managed jobs.""" @@ -106,15 +108,15 @@ def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[Accounting ) ) records = parse_accounting(output) - _validate_selected_schedulers( + ignored = _validate_selected_schedulers( tuple(record.scheduler for record in records), requested, command="sacct", ) - return records + return tuple(record for record in records if record.scheduler not in ignored) def cancel(self, selector: JobSelector) -> None: - """Cancel one managed Slurm array or array task.""" + """Cancel one managed Slurm job, array, or array task.""" self._run((self._executables.scancel, _format_selector(selector))) def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, ...]: @@ -148,25 +150,47 @@ def _format_selectors(selectors: Sequence[JobSelector]) -> str: def _format_selector(selector: JobSelector) -> str: if isinstance(selector, SchedulerIdentity): - return f"{selector.array_job_id}_{selector.array_task_id}" - if type(selector) is not int or selector <= 0: - raise ValueError("Slurm job IDs must be positive integers") - return str(selector) + job_id = _format_job_id(selector.array_job_id) + if selector.array_task_id > _MAX_SLURM_INTEGER: + raise ValueError("Slurm array-task IDs must be non-negative 32-bit integers") + return f"{job_id}_{selector.array_task_id}" + return _format_job_id(selector) + + +def _format_job_id(value: object) -> str: + if type(value) is not int or not 0 < value <= _MAX_SLURM_INTEGER: + raise ValueError("Slurm job IDs must be positive 32-bit integers") + return str(value) def _validate_selected_schedulers( - schedulers: Sequence[SchedulerIdentity], + schedulers: Sequence[SlurmJobIdentity], selectors: Sequence[JobSelector], *, command: str, -) -> None: +) -> frozenset[SlurmJobIdentity]: + """Validate result correlation and identify aggregate rows to omit.""" + ignored: set[SlurmJobIdentity] = set() for scheduler in schedulers: - if any( - scheduler == selector if isinstance(selector, SchedulerIdentity) else scheduler.array_job_id == selector - for selector in selectors - ): + explicitly_selected = any(type(selector) is int and selector == scheduler for selector in selectors) + is_array_parent = type(scheduler) is int and any( + isinstance(selector, SchedulerIdentity) and selector.array_job_id == scheduler for selector in selectors + ) + if is_array_parent and not explicitly_selected: + ignored.add(scheduler) + continue + if any(_selector_matches(scheduler, selector) for selector in selectors): continue - raise SlurmParseError(f"{command} returned an unrequested array-task ID") + raise SlurmParseError(f"{command} returned an unrequested job or array-task ID") + return frozenset(ignored) + + +def _selector_matches(scheduler: SlurmJobIdentity, selector: JobSelector) -> bool: + if isinstance(selector, SchedulerIdentity): + return scheduler == selector + if type(scheduler) is int: + return scheduler == selector + return scheduler.array_job_id == selector def _validate_argument(value: str, *, field_name: str) -> None: @@ -179,7 +203,7 @@ def _validate_argument(value: str, *, field_name: str) -> None: def _normalize_bounded_text(value: str, *, limit: int = 512) -> str: sanitized = "".join(" " if unicodedata.category(character).startswith("C") else character for character in value) normalized = " ".join(sanitized.split()) - return normalized if len(normalized) <= limit else f"{normalized[:limit]}..." + return normalized if len(normalized) <= limit else f"{normalized[: limit - 3]}..." def _format_error_detail(error: BaseException) -> str: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py index 7a45ac4e1..c7e3bc68e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -6,16 +6,19 @@ from __future__ import annotations from dataclasses import dataclass +from typing import TypeAlias from data_designer.slurm.contracts import Identifier from data_designer.slurm.state import SchedulerIdentity, SchedulerState +SlurmJobIdentity: TypeAlias = int | SchedulerIdentity + @dataclass(frozen=True, slots=True) class SlurmSubmission: - """Identity assigned by Slurm to one accepted array submission.""" + """Identity assigned by Slurm to one accepted batch submission.""" - array_job_id: int + job_id: int cluster_name: Identifier | None = None @@ -31,7 +34,7 @@ class SlurmExitCode: class QueueRecord: """One normalized active-queue row.""" - scheduler: SchedulerIdentity + scheduler: SlurmJobIdentity state: SchedulerState @@ -39,6 +42,6 @@ class QueueRecord: class AccountingRecord: """One normalized accounting row.""" - scheduler: SchedulerIdentity + scheduler: SlurmJobIdentity state: SchedulerState exit_code: SlurmExitCode diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index 43129b15b..3220bc8fd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -12,6 +12,7 @@ AccountingRecord, QueueRecord, SlurmExitCode, + SlurmJobIdentity, SlurmSubmission, ) from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -21,6 +22,7 @@ _CLUSTER_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") _EXIT_CODE_PATTERN = re.compile(r"^(?P[0-9]+):(?P[0-9]+)$") _GRES_GPU_PATTERN = re.compile(r"^gpu:(?:(?:[^:,()]+):)*(?P[1-9][0-9]*)(?:\([^\r\n]*\))?$") +_MAX_SLURM_INTEGER = (1 << 32) - 1 _STATE_MAP = { "BOOT_FAIL": SchedulerState.FAILED, @@ -56,39 +58,37 @@ def parse_submission(output: str) -> SlurmSubmission: job_id, separator, cluster_name = value.partition(";") if not job_id.isascii() or not job_id.isdecimal(): raise SlurmParseError("sbatch returned an invalid job ID") - array_job_id = _parse_decimal(job_id, message="sbatch returned an invalid job ID") - if array_job_id <= 0: + parsed_job_id = _parse_decimal(job_id, message="sbatch returned an invalid job ID") + if parsed_job_id <= 0: raise SlurmParseError("sbatch returned an invalid job ID") if separator and _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: raise SlurmParseError("sbatch returned an invalid cluster name") - return SlurmSubmission(array_job_id=array_job_id, cluster_name=cluster_name or None) + return SlurmSubmission(job_id=parsed_job_id, cluster_name=cluster_name or None) def parse_queue(output: str) -> tuple[QueueRecord, ...]: """Parse ``squeue --format=%i|%T`` rows.""" records: list[QueueRecord] = [] - identities: set[SchedulerIdentity] = set() + identities: set[SlurmJobIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 2: raise SlurmParseError(f"squeue line {line_number} must contain two fields") - scheduler = _parse_array_identity(fields[0], command="squeue", line_number=line_number) + scheduler = _parse_job_identity(fields[0], command="squeue", line_number=line_number) _reject_duplicate(scheduler, identities, command="squeue", line_number=line_number) records.append(QueueRecord(scheduler=scheduler, state=parse_state(fields[1]))) return tuple(records) def parse_accounting(output: str) -> tuple[AccountingRecord, ...]: - """Parse array-task rows from ``sacct --format=JobID,State,ExitCode``.""" + """Parse job and array-task rows from ``sacct --format=JobID,State,ExitCode``.""" records: list[AccountingRecord] = [] - identities: set[SchedulerIdentity] = set() + identities: set[SlurmJobIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 3: raise SlurmParseError(f"sacct line {line_number} must contain three fields") - if _JOB_ID_PATTERN.fullmatch(fields[0]) is not None: - continue - scheduler = _parse_array_identity(fields[0], command="sacct", line_number=line_number) + scheduler = _parse_job_identity(fields[0], command="sacct", line_number=line_number) _reject_duplicate(scheduler, identities, command="sacct", line_number=line_number) records.append( AccountingRecord( @@ -97,7 +97,12 @@ def parse_accounting(output: str) -> tuple[AccountingRecord, ...]: exit_code=_parse_exit_code(fields[2], line_number=line_number), ) ) - return tuple(records) + array_job_ids = { + record.scheduler.array_job_id for record in records if isinstance(record.scheduler, SchedulerIdentity) + } + return tuple( + record for record in records if not (type(record.scheduler) is int and record.scheduler in array_job_ids) + ) def parse_gpu_counts(output: str) -> tuple[int, ...]: @@ -182,6 +187,16 @@ def _parse_array_identity(value: str, *, command: str, line_number: int) -> Sche ) +def _parse_job_identity(value: str, *, command: str, line_number: int) -> SlurmJobIdentity: + message = f"{command} line {line_number} contains an invalid job or array-task ID" + if _JOB_ID_PATTERN.fullmatch(value) is not None: + return _parse_decimal(value, message=message) + try: + return _parse_array_identity(value, command=command, line_number=line_number) + except SlurmParseError as error: + raise SlurmParseError(message) from error + + def _parse_exit_code(value: str, *, line_number: int) -> SlurmExitCode: match = _EXIT_CODE_PATTERN.fullmatch(value) if match is None: @@ -194,19 +209,24 @@ def _parse_exit_code(value: str, *, line_number: int) -> SlurmExitCode: def _parse_decimal(value: str, *, message: str) -> int: + if len(value) > 10: + raise SlurmParseError(message) try: - return int(value) + parsed = int(value) except ValueError as error: raise SlurmParseError(message) from error + if parsed > _MAX_SLURM_INTEGER: + raise SlurmParseError(message) + return parsed def _reject_duplicate( - scheduler: SchedulerIdentity, - identities: set[SchedulerIdentity], + scheduler: SlurmJobIdentity, + identities: set[SlurmJobIdentity], *, command: str, line_number: int, ) -> None: if scheduler in identities: - raise SlurmParseError(f"{command} line {line_number} duplicates an array-task ID") + raise SlurmParseError(f"{command} line {line_number} duplicates a job or array-task ID") identities.add(scheduler) 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 580e31d39..915fceb02 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 @@ -91,7 +91,9 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire node_count = max(node_indices) + 1 array = "0" if plan.array_tasks.count > 1: - array = f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + array = f"0-{plan.array_tasks.count - 1}" + if plan.array_tasks.max_concurrent is not None: + array = f"{array}%{plan.array_tasks.max_concurrent}" values: list[tuple[str, str | None]] = [ ("job-name", plan.submission.job_name), diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 9609ef175..6118593ca 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -17,9 +17,9 @@ def test_client_submits_and_observes_one_managed_array(fake_slurm_runner: FakeSl client = SlurmCommandClient(fake_slurm_runner) submission = client.submit("/workspace/run.sbatch") - queue = client.query_queue((submission.array_job_id,)) + queue = client.query_queue((submission.job_id,)) - assert submission.array_job_id == 4101 + assert submission.job_id == 4101 assert tuple(record.state for record in queue) == (SchedulerState.PENDING, SchedulerState.RUNNING) assert fake_slurm_runner.calls == [ ("sbatch", "--parsable", "--export=NIL", "/workspace/run.sbatch"), @@ -73,15 +73,56 @@ def test_client_rejects_unrequested_scheduler_records(fake_slurm_runner: FakeSlu client.query_accounting((SchedulerIdentity(array_job_id=4101, array_task_id=0),)) +def test_client_observes_regular_cpu_job() -> None: + runner = FakeSlurmRunner() + runner.script_next("squeue", FakeCommandResponse(stdout="5101|RUNNING\n")) + runner.script_next("sacct", FakeCommandResponse(stdout="5101|COMPLETED|0:0\n")) + client = SlurmCommandClient(runner) + + queue = client.query_queue((5101,)) + accounting = client.query_accounting((5101,)) + + assert queue[0].scheduler == 5101 + assert queue[0].state is SchedulerState.RUNNING + assert accounting[0].scheduler == 5101 + assert accounting[0].state is SchedulerState.COMPLETED + + +def test_client_ignores_array_parent_observation_for_exact_task() -> None: + runner = FakeSlurmRunner() + runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n")) + client = SlurmCommandClient(runner) + + records = client.query_accounting((SchedulerIdentity(array_job_id=4101, array_task_id=0),)) + + assert records == () + + +def test_client_keeps_explicitly_selected_parent_observation() -> None: + runner = FakeSlurmRunner() + runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n")) + client = SlurmCommandClient(runner) + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + + records = client.query_accounting((4101, task)) + + assert len(records) == 1 + assert records[0].scheduler == 4101 + + def test_client_rejects_unbounded_or_invalid_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) with pytest.raises(ValueError, match="at least one"): client.query_queue(()) - with pytest.raises(ValueError, match="positive integers"): + with pytest.raises(ValueError, match="positive 32-bit integers"): client.query_accounting((0,)) - with pytest.raises(ValueError, match="positive integers"): + with pytest.raises(ValueError, match="positive 32-bit integers"): client.cancel(True) + with pytest.raises(ValueError, match="32-bit"): + client.cancel(1 << 32) + with pytest.raises(ValueError, match="array-task IDs"): + client.cancel(SchedulerIdentity(array_job_id=4101, array_task_id=1 << 32)) assert fake_slurm_runner.calls == [] @@ -134,6 +175,18 @@ def test_client_removes_terminal_controls_from_command_failures(fake_slurm_runne assert "\x1b" not in str(error.value) +def test_client_bounds_command_failure_detail(fake_slurm_runner: FakeSlurmRunner) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr="x" * 600, returncode=2)) + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(SlurmCommandError) as error: + client.query_queue((4101,)) + + detail = str(error.value).partition(": ")[2] + assert len(detail) == 512 + assert detail.endswith("...") + + def test_client_normalizes_execution_errors() -> None: client = SlurmCommandClient(_FailingRunner()) diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index 455e41193..a974e22ab 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -33,16 +33,26 @@ def test_parse_submission_accepts_parsable_sbatch_output( ) -> None: submission = parse_submission(output) - assert submission.array_job_id == expected_job_id + assert submission.job_id == expected_job_id assert submission.cluster_name == expected_cluster -@pytest.mark.parametrize("output", ("", "0", "Submitted batch job 4101", "٤١٠١", "4101;", "4101;bad name")) +@pytest.mark.parametrize( + "output", + ("", "0", "Submitted batch job 4101", "٤١٠١", "4101;", "4101;bad name", f"{'0' * 5000}1"), +) def test_parse_submission_rejects_malformed_output(output: str) -> None: with pytest.raises(SlurmParseError, match="invalid"): parse_submission(output) +def test_parse_submission_enforces_slurm_job_id_width() -> None: + assert parse_submission(str((1 << 32) - 1)).job_id == (1 << 32) - 1 + + with pytest.raises(SlurmParseError, match="invalid job ID"): + parse_submission(str(1 << 32)) + + def test_parse_queue_normalizes_active_array_tasks() -> None: records = parse_queue((GOLDEN_DIRECTORY / "squeue_active.txt").read_text()) @@ -52,6 +62,12 @@ def test_parse_queue_normalizes_active_array_tasks() -> None: ) +def test_parse_queue_normalizes_regular_jobs() -> None: + records = parse_queue("5101|RUNNING\n") + + assert records == (QueueRecord(scheduler=5101, state=SchedulerState.RUNNING),) + + @pytest.mark.parametrize( ("raw_state", "expected"), ( @@ -91,6 +107,14 @@ def test_parse_accounting_normalizes_terminal_rows_and_ignores_array_parent() -> assert records[0].exit_code.signal == 125 +def test_parse_accounting_normalizes_regular_jobs() -> None: + records = parse_accounting("5101|COMPLETED|0:0\n") + + assert len(records) == 1 + assert records[0].scheduler == 5101 + assert records[0].state is SchedulerState.COMPLETED + + def test_empty_scheduler_output_preserves_absent_evidence_for_reconciliation() -> None: assert parse_queue("") == () assert parse_accounting("\n") == () @@ -100,12 +124,12 @@ def test_empty_scheduler_output_preserves_absent_evidence_for_reconciliation() - ("parser", "output", "message"), ( (parse_queue, "malformed scheduler output\n", "two fields"), - (parse_queue, "4101|RUNNING\n", "array-task ID"), + (parse_queue, "not-a-job|RUNNING\n", "job or array-task ID"), (parse_queue, "4101_0|RUNNING\n4101_0|PENDING\n", "duplicates"), (parse_accounting, "4101_0|FAILED\n", "three fields"), (parse_accounting, "4101_0|FAILED|not-an-exit-code\n", "exit code"), - (parse_accounting, "4101_0.batch|FAILED|1:0\n", "array-task ID"), - (parse_accounting, "garbage.step|FAILED|1:0\n", "array-task ID"), + (parse_accounting, "4101_0.batch|FAILED|1:0\n", "job or array-task ID"), + (parse_accounting, "garbage.step|FAILED|1:0\n", "job or array-task ID"), (parse_queue, "4101_0|COMPLETED unexpectedly\n", "unexpected whitespace"), ), ) diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 016cb15b7..00616fd19 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -9,7 +9,7 @@ import pytest -from data_designer.slurm.config import SchedulerProfile, injected_profile +from data_designer.slurm.config import ArrayTasksConfig, SchedulerProfile, injected_profile from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.launcher import BatchRenderError, render_batch_script from data_designer.slurm.planning import ResolvedSlurmRunPlan, ResolvedSubmission @@ -84,6 +84,18 @@ def test_renderer_reserves_client_cpus_for_each_gpu_request_mode( assert "#SBATCH --cpus-per-task=17\n" in render_batch_script(plan) +def test_renderer_omits_array_throttle_when_concurrency_is_unset( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + array_tasks = ArrayTasksConfig.model_construct(count=2, max_concurrent=None) + plan = multi_node_plan.model_copy(update={"array_tasks": array_tasks}) + + script = render_batch_script(plan) + + assert "#SBATCH --array=0-1\n" in script + assert "#SBATCH --array=0-1%" not in script + + def test_renderer_rejects_mem_per_gpu_without_a_slurm_gpu_request( single_node_plan: ResolvedSlurmRunPlan, ) -> None: 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 5439234e5..c51af0cac 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 @@ -54,7 +54,9 @@ def _assert_script_matches_plan( *(index for deployment in plan.deployments for index in deployment.node_indices), ) node_count = max(node_indices) + 1 - array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}" + if plan.array_tasks.count > 1 and plan.array_tasks.max_concurrent is not None: + array = f"{array}%{plan.array_tasks.max_concurrent}" plan_path = posixpath.join(posixpath.dirname(plan.authored_config.path), "resolved-plan.json") run_root = posixpath.dirname(plan.authored_config.path) From 87a520f5b8f1d36c9b25ed7faa6bcd3af97cf898 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 27 Aug 2026 11:43:00 -0600 Subject: [PATCH 21/23] refine Slurm launcher boundaries - Internalize implementation APIs and clarify transient types - Harden parsing, rendering, and process error normalization - Extend fakes and executable checksum coverage Part of #868 --- .../data_designer/slurm/launcher/__init__.py | 36 +----- .../data_designer/slurm/launcher/client.py | 101 ++++++++------- .../data_designer/slurm/launcher/errors.py | 8 +- .../data_designer/slurm/launcher/models.py | 28 ++-- .../data_designer/slurm/launcher/parsing.py | 120 +++++++++--------- .../data_designer/slurm/launcher/renderer.py | 40 +++--- .../data_designer/slurm/launcher/runner.py | 2 +- .../tests/launcher/test_client.py | 43 ++++--- .../tests/launcher/test_parsing.py | 53 ++++---- .../tests/launcher/test_renderer.py | 111 +++++++++++----- .../tests/launcher/test_runner.py | 2 +- .../tests/slurm_test_fakes/__init__.py | 2 + .../tests/slurm_test_fakes/slurm.py | 120 +++++++++++++++--- .../tests/slurm_test_fakes/test_slurm.py | 21 ++- scripts/test_slurm_package_install.py | 2 - 15 files changed, 411 insertions(+), 278 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py index 9bc45a4a6..61da221a6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py @@ -1,40 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Structured Slurm submission, observation, and batch rendering.""" +"""Internal Slurm submission, observation, and batch-rendering helpers.""" from __future__ import annotations - -from data_designer.slurm.launcher.client import SlurmCommandClient, SlurmExecutables -from data_designer.slurm.launcher.errors import ( - BatchRenderError, - SlurmCommandError, - SlurmLauncherError, - SlurmParseError, -) -from data_designer.slurm.launcher.models import ( - AccountingRecord, - QueueRecord, - SlurmExitCode, - SlurmJobIdentity, - SlurmSubmission, -) -from data_designer.slurm.launcher.renderer import render_batch_script -from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner - -__all__ = [ - "AccountingRecord", - "BatchRenderError", - "CommandRunner", - "QueueRecord", - "SlurmCommandClient", - "SlurmCommandError", - "SlurmExecutables", - "SlurmExitCode", - "SlurmJobIdentity", - "SlurmLauncherError", - "SlurmParseError", - "SlurmSubmission", - "SubprocessRunner", - "render_batch_script", -] 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 fdb62bc83..3bddd06ef 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 @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Typed argument-vector client for Slurm command-line tools.""" +"""Internal typed argument-vector client for Slurm command-line tools.""" from __future__ import annotations @@ -14,8 +14,13 @@ from typing import TypeAlias from data_designer.slurm.contracts import Identifier -from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmParseError -from data_designer.slurm.launcher.models import AccountingRecord, QueueRecord, SlurmJobIdentity, SlurmSubmission +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmObservedJobIdentity, + SlurmQueueEntry, +) from data_designer.slurm.launcher.parsing import ( parse_accounting, parse_gpu_counts, @@ -25,7 +30,7 @@ from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner from data_designer.slurm.state import SchedulerIdentity -JobSelector: TypeAlias = SlurmJobIdentity +_JobSelector: TypeAlias = int | SchedulerIdentity _IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") _MAX_SLURM_INTEGER = (1 << 32) - 1 @@ -62,7 +67,7 @@ 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) -> SlurmSubmission: + def submit(self, script_path: str | Path) -> SlurmJobSubmissionReceipt: """Submit one rendered batch script and return its assigned job ID.""" path = str(script_path) _validate_argument(path, field_name="batch script path") @@ -71,7 +76,7 @@ def submit(self, script_path: str | Path) -> SlurmSubmission: output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) return parse_submission(output) - def query_queue(self, selectors: Sequence[JobSelector]) -> tuple[QueueRecord, ...]: + def query_queue(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmQueueEntry, ...]: """Return normalized active-queue rows for explicit managed jobs.""" requested = tuple(selectors) jobs = _format_selectors(requested) @@ -84,15 +89,15 @@ def query_queue(self, selectors: Sequence[JobSelector]) -> tuple[QueueRecord, .. f"--jobs={jobs}", ) ) - records = parse_queue(output) - ignored = _validate_selected_schedulers( - tuple(record.scheduler for record in records), + entries = parse_queue(output) + ignored = _validate_observed_job_identities( + tuple(entry.job_identity for entry in entries), requested, command="squeue", ) - return tuple(record for record in records if record.scheduler not in ignored) + return tuple(entry for entry in entries if entry.job_identity not in ignored) - def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[AccountingRecord, ...]: + def query_accounting(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmAccountingEntry, ...]: """Return normalized accounting rows for explicit managed jobs.""" requested = tuple(selectors) jobs = _format_selectors(requested) @@ -107,15 +112,15 @@ def query_accounting(self, selectors: Sequence[JobSelector]) -> tuple[Accounting f"--jobs={jobs}", ) ) - records = parse_accounting(output) - ignored = _validate_selected_schedulers( - tuple(record.scheduler for record in records), + entries = parse_accounting(output) + ignored = _validate_observed_job_identities( + tuple(entry.job_identity for entry in entries), requested, command="sacct", ) - return tuple(record for record in records if record.scheduler not in ignored) + return tuple(entry for entry in entries if entry.job_identity not in ignored) - def cancel(self, selector: JobSelector) -> None: + def cancel(self, selector: _JobSelector) -> None: """Cancel one managed Slurm job, array, or array task.""" self._run((self._executables.scancel, _format_selector(selector))) @@ -134,21 +139,24 @@ def _run(self, command: Sequence[str]) -> str: completed = self._runner.run(command) except (OSError, subprocess.SubprocessError) as error: raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") from error - if completed.returncode: - detail = _normalize_bounded_text(completed.stderr) or "no diagnostic output" - raise SlurmCommandError(f"{command_name} failed with exit code {completed.returncode}: {detail}") - if not isinstance(completed.stdout, str): - raise SlurmCommandError(f"{command_name} did not return text output") - return completed.stdout - - -def _format_selectors(selectors: Sequence[JobSelector]) -> str: + returncode = getattr(completed, "returncode", None) + stdout = getattr(completed, "stdout", None) + stderr = getattr(completed, "stderr", None) + if type(returncode) is not int or not isinstance(stdout, str) or not isinstance(stderr, str): + raise SlurmCommandError(f"{command_name} returned a malformed process result") + if returncode: + detail = _normalize_bounded_text(stderr) or "no diagnostic output" + raise SlurmCommandError(f"{command_name} failed with exit code {returncode}: {detail}") + return stdout + + +def _format_selectors(selectors: Sequence[_JobSelector]) -> str: if not selectors: raise ValueError("at least one managed Slurm job selector is required") return ",".join(dict.fromkeys(_format_selector(selector) for selector in selectors)) -def _format_selector(selector: JobSelector) -> str: +def _format_selector(selector: _JobSelector) -> str: if isinstance(selector, SchedulerIdentity): job_id = _format_job_id(selector.array_job_id) if selector.array_task_id > _MAX_SLURM_INTEGER: @@ -163,36 +171,31 @@ def _format_job_id(value: object) -> str: return str(value) -def _validate_selected_schedulers( - schedulers: Sequence[SlurmJobIdentity], - selectors: Sequence[JobSelector], +def _validate_observed_job_identities( + job_identities: Sequence[SlurmObservedJobIdentity], + selectors: Sequence[_JobSelector], *, command: str, -) -> frozenset[SlurmJobIdentity]: - """Validate result correlation and identify aggregate rows to omit.""" - ignored: set[SlurmJobIdentity] = set() - for scheduler in schedulers: - explicitly_selected = any(type(selector) is int and selector == scheduler for selector in selectors) - is_array_parent = type(scheduler) is int and any( - isinstance(selector, SchedulerIdentity) and selector.array_job_id == scheduler for selector in selectors - ) - if is_array_parent and not explicitly_selected: - ignored.add(scheduler) +) -> frozenset[SlurmObservedJobIdentity]: + """Validate result correlation and return unselected aggregate rows.""" + selected_job_ids = {selector for selector in selectors if type(selector) is int} + selected_array_tasks = {selector for selector in selectors if isinstance(selector, SchedulerIdentity)} + selected_array_job_ids = {selector.array_job_id for selector in selected_array_tasks} + ignored: set[SlurmObservedJobIdentity] = set() + for job_identity in job_identities: + if type(job_identity) is int and job_identity in selected_job_ids: continue - if any(_selector_matches(scheduler, selector) for selector in selectors): + if type(job_identity) is int and job_identity in selected_array_job_ids: + ignored.add(job_identity) continue - raise SlurmParseError(f"{command} returned an unrequested job or array-task ID") + if isinstance(job_identity, SchedulerIdentity) and ( + job_identity in selected_array_tasks or job_identity.array_job_id in selected_job_ids + ): + continue + raise SlurmCommandOutputError(f"{command} returned an unrequested job or array-task ID") return frozenset(ignored) -def _selector_matches(scheduler: SlurmJobIdentity, selector: JobSelector) -> bool: - if isinstance(selector, SchedulerIdentity): - return scheduler == selector - if type(scheduler) is int: - return scheduler == selector - return scheduler.array_job_id == selector - - def _validate_argument(value: str, *, field_name: str) -> None: if type(value) is not str or not value: raise ValueError(f"{field_name} must not be empty") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py index 791ae69b7..4a611152b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Canonical errors for the Slurm launcher boundary.""" +"""Internal normalized errors for the Slurm launcher boundary.""" from __future__ import annotations @@ -14,9 +14,9 @@ class SlurmCommandError(SlurmLauncherError): """A Slurm command could not be executed successfully.""" -class SlurmParseError(SlurmLauncherError, ValueError): - """Slurm returned output that violates the requested format.""" +class SlurmCommandOutputError(SlurmLauncherError, ValueError): + """A Slurm command returned output that violates its requested format.""" -class BatchRenderError(SlurmLauncherError, ValueError): +class SlurmBatchRenderError(SlurmLauncherError, ValueError): """A resolved plan cannot be rendered as a safe batch script.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py index c7e3bc68e..b686cc02f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -8,40 +8,38 @@ from dataclasses import dataclass from typing import TypeAlias -from data_designer.slurm.contracts import Identifier from data_designer.slurm.state import SchedulerIdentity, SchedulerState -SlurmJobIdentity: TypeAlias = int | SchedulerIdentity +SlurmObservedJobIdentity: TypeAlias = int | SchedulerIdentity @dataclass(frozen=True, slots=True) -class SlurmSubmission: - """Identity assigned by Slurm to one accepted batch submission.""" +class SlurmJobSubmissionReceipt: + """Job identity returned for one accepted non-federated submission.""" job_id: int - cluster_name: Identifier | None = None @dataclass(frozen=True, slots=True) -class SlurmExitCode: +class SlurmProcessExitCode: """Slurm's process status and terminating signal pair.""" - status: int - signal: int + exit_status: int + termination_signal: int @dataclass(frozen=True, slots=True) -class QueueRecord: - """One normalized active-queue row.""" +class SlurmQueueEntry: + """One transient normalized active-queue entry.""" - scheduler: SlurmJobIdentity + job_identity: SlurmObservedJobIdentity state: SchedulerState @dataclass(frozen=True, slots=True) -class AccountingRecord: - """One normalized accounting row.""" +class SlurmAccountingEntry: + """One transient normalized accounting entry.""" - scheduler: SlurmJobIdentity + job_identity: SlurmObservedJobIdentity state: SchedulerState - exit_code: SlurmExitCode + process_exit_code: SlurmProcessExitCode diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index 3220bc8fd..2aff70228 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -1,19 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Strict parsers for bounded, machine-readable Slurm output.""" +"""Internal strict parsers for bounded, machine-readable Slurm output.""" from __future__ import annotations import re -from data_designer.slurm.launcher.errors import SlurmParseError +from data_designer.slurm.launcher.errors import SlurmCommandOutputError from data_designer.slurm.launcher.models import ( - AccountingRecord, - QueueRecord, - SlurmExitCode, - SlurmJobIdentity, - SlurmSubmission, + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmObservedJobIdentity, + SlurmProcessExitCode, + SlurmQueueEntry, ) from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -52,56 +52,58 @@ } -def parse_submission(output: str) -> SlurmSubmission: - """Parse ``sbatch --parsable`` output.""" +def parse_submission(output: str) -> SlurmJobSubmissionReceipt: + """Parse non-federated ``sbatch --parsable`` output.""" value = output.strip() job_id, separator, cluster_name = value.partition(";") if not job_id.isascii() or not job_id.isdecimal(): - raise SlurmParseError("sbatch returned an invalid job ID") + raise SlurmCommandOutputError("sbatch returned an invalid job ID") parsed_job_id = _parse_decimal(job_id, message="sbatch returned an invalid job ID") if parsed_job_id <= 0: - raise SlurmParseError("sbatch returned an invalid job ID") - if separator and _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: - raise SlurmParseError("sbatch returned an invalid cluster name") - return SlurmSubmission(job_id=parsed_job_id, cluster_name=cluster_name or None) + raise SlurmCommandOutputError("sbatch returned an invalid job ID") + if separator: + if _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: + raise SlurmCommandOutputError("sbatch returned an invalid cluster name") + raise SlurmCommandOutputError("federated Slurm submissions are not supported") + return SlurmJobSubmissionReceipt(job_id=parsed_job_id) -def parse_queue(output: str) -> tuple[QueueRecord, ...]: +def parse_queue(output: str) -> tuple[SlurmQueueEntry, ...]: """Parse ``squeue --format=%i|%T`` rows.""" - records: list[QueueRecord] = [] - identities: set[SlurmJobIdentity] = set() + entries: list[SlurmQueueEntry] = [] + identities: set[SlurmObservedJobIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 2: - raise SlurmParseError(f"squeue line {line_number} must contain two fields") - scheduler = _parse_job_identity(fields[0], command="squeue", line_number=line_number) - _reject_duplicate(scheduler, identities, command="squeue", line_number=line_number) - records.append(QueueRecord(scheduler=scheduler, state=parse_state(fields[1]))) - return tuple(records) + raise SlurmCommandOutputError(f"squeue line {line_number} must contain two fields") + job_identity = _parse_job_identity(fields[0], command="squeue", line_number=line_number) + _reject_duplicate(job_identity, identities, command="squeue", line_number=line_number) + entries.append(SlurmQueueEntry(job_identity=job_identity, state=parse_state(fields[1]))) + return tuple(entries) -def parse_accounting(output: str) -> tuple[AccountingRecord, ...]: +def parse_accounting(output: str) -> tuple[SlurmAccountingEntry, ...]: """Parse job and array-task rows from ``sacct --format=JobID,State,ExitCode``.""" - records: list[AccountingRecord] = [] - identities: set[SlurmJobIdentity] = set() + entries: list[SlurmAccountingEntry] = [] + identities: set[SlurmObservedJobIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 3: - raise SlurmParseError(f"sacct line {line_number} must contain three fields") - scheduler = _parse_job_identity(fields[0], command="sacct", line_number=line_number) - _reject_duplicate(scheduler, identities, command="sacct", line_number=line_number) - records.append( - AccountingRecord( - scheduler=scheduler, + raise SlurmCommandOutputError(f"sacct line {line_number} must contain three fields") + job_identity = _parse_job_identity(fields[0], command="sacct", line_number=line_number) + _reject_duplicate(job_identity, identities, command="sacct", line_number=line_number) + entries.append( + SlurmAccountingEntry( + job_identity=job_identity, state=parse_state(fields[1]), - exit_code=_parse_exit_code(fields[2], line_number=line_number), + process_exit_code=_parse_exit_code(fields[2], line_number=line_number), ) ) array_job_ids = { - record.scheduler.array_job_id for record in records if isinstance(record.scheduler, SchedulerIdentity) + entry.job_identity.array_job_id for entry in entries if isinstance(entry.job_identity, SchedulerIdentity) } return tuple( - record for record in records if not (type(record.scheduler) is int and record.scheduler in array_job_ids) + entry for entry in entries if not (type(entry.job_identity) is int and entry.job_identity in array_job_ids) ) @@ -117,7 +119,7 @@ def parse_gpu_counts(output: str) -> tuple[int, ...]: continue match = _GRES_GPU_PATTERN.fullmatch(gres) if match is None: - raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") line_counts.append( _parse_decimal( match.group("count"), @@ -137,19 +139,19 @@ def _split_gres_fields(value: str, *, line_number: int) -> tuple[str, ...]: if character == "(": annotation_depth += 1 if annotation_depth > 1: - raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") elif character == ")": annotation_depth -= 1 if annotation_depth < 0: - raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") elif character == "," and annotation_depth == 0: fields.append(value[start:index]) start = index + 1 if annotation_depth: - raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") fields.append(value[start:]) if any(not field for field in fields): - raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") return tuple(fields) @@ -157,14 +159,14 @@ def parse_state(value: str) -> SchedulerState: """Normalize one Slurm long state spelling without guessing unknown states.""" normalized = value.strip().upper().removesuffix("+") if not normalized: - raise SlurmParseError("scheduler state must not be empty") + raise SlurmCommandOutputError("scheduler state must not be empty") if normalized.startswith("CANCELLED BY "): canceller = normalized.removeprefix("CANCELLED BY ") if not canceller.isascii() or not canceller.isdecimal(): - raise SlurmParseError("cancelled scheduler state has an invalid owner") + raise SlurmCommandOutputError("cancelled scheduler state has an invalid owner") normalized = "CANCELLED" elif any(character.isspace() for character in normalized): - raise SlurmParseError("scheduler state contains unexpected whitespace") + raise SlurmCommandOutputError("scheduler state contains unexpected whitespace") return _STATE_MAP.get(normalized, SchedulerState.UNKNOWN) @@ -179,7 +181,7 @@ def _collect_nonempty_lines(output: str) -> tuple[tuple[int, str], ...]: def _parse_array_identity(value: str, *, command: str, line_number: int) -> SchedulerIdentity: match = _ARRAY_ID_PATTERN.fullmatch(value) if match is None: - raise SlurmParseError(f"{command} line {line_number} contains an invalid array-task ID") + raise SlurmCommandOutputError(f"{command} line {line_number} contains an invalid array-task ID") message = f"{command} line {line_number} contains an invalid array-task ID" return SchedulerIdentity( array_job_id=_parse_decimal(match.group("job"), message=message), @@ -187,46 +189,46 @@ def _parse_array_identity(value: str, *, command: str, line_number: int) -> Sche ) -def _parse_job_identity(value: str, *, command: str, line_number: int) -> SlurmJobIdentity: +def _parse_job_identity(value: str, *, command: str, line_number: int) -> SlurmObservedJobIdentity: message = f"{command} line {line_number} contains an invalid job or array-task ID" if _JOB_ID_PATTERN.fullmatch(value) is not None: return _parse_decimal(value, message=message) try: return _parse_array_identity(value, command=command, line_number=line_number) - except SlurmParseError as error: - raise SlurmParseError(message) from error + except SlurmCommandOutputError as error: + raise SlurmCommandOutputError(message) from error -def _parse_exit_code(value: str, *, line_number: int) -> SlurmExitCode: +def _parse_exit_code(value: str, *, line_number: int) -> SlurmProcessExitCode: match = _EXIT_CODE_PATTERN.fullmatch(value) if match is None: - raise SlurmParseError(f"sacct line {line_number} contains an invalid exit code") + raise SlurmCommandOutputError(f"sacct line {line_number} contains an invalid exit code") message = f"sacct line {line_number} contains an invalid exit code" - return SlurmExitCode( - status=_parse_decimal(match.group("status"), message=message), - signal=_parse_decimal(match.group("signal"), message=message), + return SlurmProcessExitCode( + exit_status=_parse_decimal(match.group("status"), message=message), + termination_signal=_parse_decimal(match.group("signal"), message=message), ) def _parse_decimal(value: str, *, message: str) -> int: if len(value) > 10: - raise SlurmParseError(message) + raise SlurmCommandOutputError(message) try: parsed = int(value) except ValueError as error: - raise SlurmParseError(message) from error + raise SlurmCommandOutputError(message) from error if parsed > _MAX_SLURM_INTEGER: - raise SlurmParseError(message) + raise SlurmCommandOutputError(message) return parsed def _reject_duplicate( - scheduler: SlurmJobIdentity, - identities: set[SlurmJobIdentity], + job_identity: SlurmObservedJobIdentity, + identities: set[SlurmObservedJobIdentity], *, command: str, line_number: int, ) -> None: - if scheduler in identities: - raise SlurmParseError(f"{command} line {line_number} duplicates a job or array-task ID") - identities.add(scheduler) + if job_identity in identities: + raise SlurmCommandOutputError(f"{command} line {line_number} duplicates a job or array-task ID") + identities.add(job_identity) 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 915fceb02..931eb977a 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 @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Safe deterministic rendering for thin Slurm batch entrypoints.""" +"""Internal safe rendering for thin deterministic Slurm batch entrypoints.""" from __future__ import annotations @@ -9,7 +9,7 @@ import re from dataclasses import dataclass -from data_designer.slurm.launcher.errors import BatchRenderError +from data_designer.slurm.launcher.errors import SlurmBatchRenderError from data_designer.slurm.planning import ResolvedSlurmRunPlan _DIRECTIVE_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9-]*$") @@ -26,18 +26,18 @@ class _BatchDirective: def render(self) -> str: """Render the directive as one non-executable scheduler line.""" if type(self.name) is not str or _DIRECTIVE_NAME_PATTERN.fullmatch(self.name) is None: - raise BatchRenderError("batch directive name is invalid") + raise SlurmBatchRenderError("batch directive name is invalid") if type(self.value) is not str: - raise BatchRenderError("batch directive value must be text") + raise SlurmBatchRenderError("batch directive value must be text") _reject_control_characters(self.value, field_name=f"--{self.name} value") - value = self.value if _DIRECTIVE_TOKEN_PATTERN.fullmatch(self.value) else _quote_double_value(self.value) + value = self.value if _DIRECTIVE_TOKEN_PATTERN.fullmatch(self.value) else _quote_sbatch_option_value(self.value) return f"#SBATCH --{self.name}={value}" -def render_batch_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int = 1) -> str: +def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: """Render a resolved generation plan as one thin deterministic entrypoint.""" if type(attempt_ordinal) is not int or attempt_ordinal <= 0: - raise BatchRenderError("attempt_ordinal must be a positive integer") + raise SlurmBatchRenderError("attempt_ordinal must be a positive integer") run_root = posixpath.dirname(plan.authored_config.path) plan_path = posixpath.join(run_root, "resolved-plan.json") @@ -50,12 +50,12 @@ def render_batch_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int = 1) set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" -readonly DD_RUNTIME_ARCHIVE={_quote_double_value(plan.runtime_bundle.path)} -readonly DD_RUNTIME_SHA256={_quote_double_value(plan.runtime_bundle.sha256)} -readonly DD_PLAN={_quote_double_value(plan_path)} -readonly DD_PLAN_SHA256={_quote_double_value(plan.compute_sha256())} -readonly DD_RUN_ROOT={_quote_double_value(run_root)} -readonly DD_ATTEMPT_ORDINAL={_quote_double_value(attempt)} +readonly DD_RUNTIME_ARCHIVE={_quote_shell_value(plan.runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={_quote_shell_value(plan.runtime_bundle.sha256)} +readonly DD_PLAN={_quote_shell_value(plan_path)} +readonly DD_PLAN_SHA256={_quote_shell_value(plan.compute_sha256())} +readonly DD_RUN_ROOT={_quote_shell_value(run_root)} +readonly DD_ATTEMPT_ORDINAL={_quote_shell_value(attempt)} verify_sha256() {{ local actual_sha256 @@ -92,6 +92,7 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire array = "0" if plan.array_tasks.count > 1: array = f"0-{plan.array_tasks.count - 1}" + # TODO(#875): Add valid authored-to-rendered coverage once omitted concurrency is supported by the contract. if plan.array_tasks.max_concurrent is not None: array = f"{array}%{plan.array_tasks.max_concurrent}" @@ -108,7 +109,8 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire if profile.gpu_request_mode == "gres": values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) elif profile.scheduler.mem_per_gpu is not None: - raise BatchRenderError("mem_per_gpu requires GRES GPU request mode") + # TODO(#875): Remove this defense once config and plan validation make this state unrepresentable. + raise SlurmBatchRenderError("mem_per_gpu requires GRES GPU request mode") if profile.scheduler.mem_per_gpu is not None: values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) if plan.submission.comment is not None: @@ -116,7 +118,13 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire return tuple(_BatchDirective(name=name, value=value) for name, value in values if value is not None) -def _quote_double_value(value: str) -> str: +def _quote_sbatch_option_value(value: str) -> str: + _reject_control_characters(value, field_name="batch directive value") + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped}"' + + +def _quote_shell_value(value: str) -> str: _reject_control_characters(value, field_name="shell value") escaped = value.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`") return f'"{escaped}"' @@ -124,4 +132,4 @@ def _quote_double_value(value: str) -> str: def _reject_control_characters(value: str, *, field_name: str) -> None: if any(ord(character) < 32 or ord(character) == 127 for character in value): - raise BatchRenderError(f"{field_name} must not contain control characters") + raise SlurmBatchRenderError(f"{field_name} must not contain control characters") 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 e7c79b4b4..3679d3252 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 @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Injectable process execution for Slurm command-line tools.""" +"""Internal injectable process execution for Slurm command-line tools.""" from __future__ import annotations diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 6118593ca..ccc88978a 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -7,9 +7,10 @@ from collections.abc import Sequence import pytest -from slurm_test_fakes import FakeCommandResponse, FakeSlurmRunner +from slurm_test_fakes import FakeCommandResponse, FakeSlurmJob, FakeSlurmRunner -from data_designer.slurm.launcher import SlurmCommandClient, SlurmCommandError, SlurmExecutables, SlurmParseError +from data_designer.slurm.launcher.client import SlurmCommandClient, SlurmExecutables +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -36,7 +37,7 @@ def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: accounting = client.query_accounting((scheduler,)) assert len(accounting) == 1 - assert accounting[0].scheduler == scheduler + assert accounting[0].job_identity == scheduler assert accounting[0].state is SchedulerState.CANCELLED assert fake_slurm_runner.calls[-2:] == [ ("scancel", "4101_1"), @@ -65,26 +66,27 @@ def test_client_rejects_unrequested_scheduler_records(fake_slurm_runner: FakeSlu client = SlurmCommandClient(fake_slurm_runner) fake_slurm_runner.script_next("squeue", FakeCommandResponse(stdout="9999_0|RUNNING\n")) - with pytest.raises(SlurmParseError, match="unrequested"): + with pytest.raises(SlurmCommandOutputError, match="unrequested"): client.query_queue((4101,)) fake_slurm_runner.script_next("sacct", FakeCommandResponse(stdout="9999_0|FAILED|1:0\n")) - with pytest.raises(SlurmParseError, match="unrequested"): + with pytest.raises(SlurmCommandOutputError, match="unrequested"): client.query_accounting((SchedulerIdentity(array_job_id=4101, array_task_id=0),)) def test_client_observes_regular_cpu_job() -> None: - runner = FakeSlurmRunner() - runner.script_next("squeue", FakeCommandResponse(stdout="5101|RUNNING\n")) - runner.script_next("sacct", FakeCommandResponse(stdout="5101|COMPLETED|0:0\n")) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) client = SlurmCommandClient(runner) + submission = client.submit("image-build.sbatch") queue = client.query_queue((5101,)) + runner.set_job_state(5101, queue_state=None, accounting_state="COMPLETED") accounting = client.query_accounting((5101,)) - assert queue[0].scheduler == 5101 - assert queue[0].state is SchedulerState.RUNNING - assert accounting[0].scheduler == 5101 + assert submission.job_id == 5101 + assert queue[0].job_identity == 5101 + assert queue[0].state is SchedulerState.PENDING + assert accounting[0].job_identity == 5101 assert accounting[0].state is SchedulerState.COMPLETED @@ -107,7 +109,7 @@ def test_client_keeps_explicitly_selected_parent_observation() -> None: records = client.query_accounting((4101, task)) assert len(records) == 1 - assert records[0].scheduler == 4101 + assert records[0].job_identity == 4101 def test_client_rejects_unbounded_or_invalid_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: @@ -205,10 +207,11 @@ def test_client_normalizes_command_timeouts() -> None: assert isinstance(error.value.__cause__, subprocess.TimeoutExpired) -def test_client_rejects_non_text_runner_output() -> None: - client = SlurmCommandClient(_NonTextRunner()) +@pytest.mark.parametrize("returncode", (0, 2)) +def test_client_rejects_non_text_runner_output(returncode: int) -> None: + client = SlurmCommandClient(_NonTextRunner(returncode)) - with pytest.raises(SlurmCommandError, match="did not return text output"): + with pytest.raises(SlurmCommandError, match="malformed process result"): client.query_queue((4101,)) @@ -262,7 +265,13 @@ def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: class _NonTextRunner: + def __init__(self, returncode: int) -> None: + self._returncode = returncode + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: - completed = subprocess.CompletedProcess(command, 0, stdout="ok", stderr="") - completed.stdout = b"not text" # type: ignore[assignment] + completed = subprocess.CompletedProcess(command, self._returncode, stdout="ok", stderr="") + if self._returncode: + completed.stderr = b"not text" # type: ignore[assignment] + else: + completed.stdout = b"not text" # type: ignore[assignment] return completed diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index a974e22ab..34e7c5a2f 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -8,7 +8,8 @@ import pytest -from data_designer.slurm.launcher import QueueRecord, SlurmParseError +from data_designer.slurm.launcher.errors import SlurmCommandOutputError +from data_designer.slurm.launcher.models import SlurmQueueEntry from data_designer.slurm.launcher.parsing import ( parse_accounting, parse_gpu_counts, @@ -22,19 +23,15 @@ OVERSIZED_DECIMAL = "9" * 5000 -@pytest.mark.parametrize( - ("output", "expected_job_id", "expected_cluster"), - (("4101\n", 4101, None), ("4101;primary\n", 4101, "primary")), -) -def test_parse_submission_accepts_parsable_sbatch_output( - output: str, - expected_job_id: int, - expected_cluster: str | None, -) -> None: - submission = parse_submission(output) +def test_parse_submission_accepts_non_federated_parsable_sbatch_output() -> None: + submission = parse_submission("4101\n") + + assert submission.job_id == 4101 + - assert submission.job_id == expected_job_id - assert submission.cluster_name == expected_cluster +def test_parse_submission_rejects_federated_receipts() -> None: + with pytest.raises(SlurmCommandOutputError, match="federated"): + parse_submission("4101;primary\n") @pytest.mark.parametrize( @@ -42,14 +39,14 @@ def test_parse_submission_accepts_parsable_sbatch_output( ("", "0", "Submitted batch job 4101", "٤١٠١", "4101;", "4101;bad name", f"{'0' * 5000}1"), ) def test_parse_submission_rejects_malformed_output(output: str) -> None: - with pytest.raises(SlurmParseError, match="invalid"): + with pytest.raises(SlurmCommandOutputError, match="invalid"): parse_submission(output) def test_parse_submission_enforces_slurm_job_id_width() -> None: assert parse_submission(str((1 << 32) - 1)).job_id == (1 << 32) - 1 - with pytest.raises(SlurmParseError, match="invalid job ID"): + with pytest.raises(SlurmCommandOutputError, match="invalid job ID"): parse_submission(str(1 << 32)) @@ -57,15 +54,15 @@ def test_parse_queue_normalizes_active_array_tasks() -> None: records = parse_queue((GOLDEN_DIRECTORY / "squeue_active.txt").read_text()) assert records == ( - _make_queue_record(0, SchedulerState.PENDING), - _make_queue_record(1, SchedulerState.RUNNING), + _make_queue_entry(0, SchedulerState.PENDING), + _make_queue_entry(1, SchedulerState.RUNNING), ) def test_parse_queue_normalizes_regular_jobs() -> None: records = parse_queue("5101|RUNNING\n") - assert records == (QueueRecord(scheduler=5101, state=SchedulerState.RUNNING),) + assert records == (SlurmQueueEntry(job_identity=5101, state=SchedulerState.RUNNING),) @pytest.mark.parametrize( @@ -103,15 +100,15 @@ def test_parse_accounting_normalizes_terminal_rows_and_ignores_array_parent() -> SchedulerState.OUT_OF_MEMORY, SchedulerState.CANCELLED, ) - assert records[0].exit_code.status == 0 - assert records[0].exit_code.signal == 125 + assert records[0].process_exit_code.exit_status == 0 + assert records[0].process_exit_code.termination_signal == 125 def test_parse_accounting_normalizes_regular_jobs() -> None: records = parse_accounting("5101|COMPLETED|0:0\n") assert len(records) == 1 - assert records[0].scheduler == 5101 + assert records[0].job_identity == 5101 assert records[0].state is SchedulerState.COMPLETED @@ -138,7 +135,7 @@ def test_scheduler_parsers_reject_malformed_or_ambiguous_rows( output: str, message: str, ) -> None: - with pytest.raises(SlurmParseError, match=message): + with pytest.raises(SlurmCommandOutputError, match=message): parser(output) @@ -157,7 +154,7 @@ def test_parsers_normalize_oversized_numeric_fields( output: str, message: str, ) -> None: - with pytest.raises(SlurmParseError, match=message): + with pytest.raises(SlurmCommandOutputError, match=message): parser(output) @@ -189,18 +186,18 @@ def test_parse_gpu_counts_normalizes_configured_gres(output: str, expected: tupl ), ) def test_parse_gpu_counts_rejects_malformed_gpu_resources(output: str) -> None: - with pytest.raises(SlurmParseError, match="invalid GPU resource"): + with pytest.raises(SlurmCommandOutputError, match="invalid GPU resource"): parse_gpu_counts(output) @pytest.mark.parametrize("state", ("", "CANCELLED by root")) def test_parse_state_rejects_invalid_spellings(state: str) -> None: - with pytest.raises(SlurmParseError): + with pytest.raises(SlurmCommandOutputError): parse_state(state) -def _make_queue_record(array_task_id: int, state: SchedulerState) -> QueueRecord: - return QueueRecord( - scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=array_task_id), +def _make_queue_entry(array_task_id: int, state: SchedulerState) -> SlurmQueueEntry: + return SlurmQueueEntry( + job_identity=SchedulerIdentity(array_job_id=4101, array_task_id=array_task_id), state=state, ) diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 00616fd19..c80c95f5f 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -3,15 +3,20 @@ from __future__ import annotations +import hashlib +import os +import shlex import subprocess +import tarfile from pathlib import Path from typing import Literal import pytest -from data_designer.slurm.config import ArrayTasksConfig, SchedulerProfile, injected_profile +from data_designer.slurm.config import SchedulerProfile, injected_profile from data_designer.slurm.contracts import ArtifactReference -from data_designer.slurm.launcher import BatchRenderError, render_batch_script +from data_designer.slurm.launcher.errors import SlurmBatchRenderError +from data_designer.slurm.launcher.renderer import render_generation_attempt_script from data_designer.slurm.planning import ResolvedSlurmRunPlan, ResolvedSubmission GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "rendered" @@ -28,7 +33,7 @@ def test_renderer_matches_contract_bound_goldens( ) -> None: plan = request.getfixturevalue(plan_fixture) - assert render_batch_script(plan) == (GOLDEN_DIRECTORY / fixture_name).read_text() + assert render_generation_attempt_script(plan, attempt_ordinal=1) == (GOLDEN_DIRECTORY / fixture_name).read_text() def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fields( @@ -53,7 +58,7 @@ def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fiel } ) - script = render_batch_script(plan) + script = render_generation_attempt_script(plan, attempt_ordinal=1) assert "#SBATCH --gres=" not in script assert "#SBATCH --account=" not in script @@ -67,7 +72,7 @@ def test_renderer_emits_mem_per_gpu_for_gres_mode(single_node_plan: ResolvedSlur ) plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) - assert "#SBATCH --mem-per-gpu=80G\n" in render_batch_script(plan) + assert "#SBATCH --mem-per-gpu=80G\n" in render_generation_attempt_script(plan, attempt_ordinal=1) @pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) @@ -81,19 +86,7 @@ def test_renderer_reserves_client_cpus_for_each_gpu_request_mode( ) plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile), "client": client}) - assert "#SBATCH --cpus-per-task=17\n" in render_batch_script(plan) - - -def test_renderer_omits_array_throttle_when_concurrency_is_unset( - multi_node_plan: ResolvedSlurmRunPlan, -) -> None: - array_tasks = ArrayTasksConfig.model_construct(count=2, max_concurrent=None) - plan = multi_node_plan.model_copy(update={"array_tasks": array_tasks}) - - script = render_batch_script(plan) - - assert "#SBATCH --array=0-1\n" in script - assert "#SBATCH --array=0-1%" not in script + assert "#SBATCH --cpus-per-task=17\n" in render_generation_attempt_script(plan, attempt_ordinal=1) def test_renderer_rejects_mem_per_gpu_without_a_slurm_gpu_request( @@ -107,8 +100,8 @@ def test_renderer_rejects_mem_per_gpu_without_a_slurm_gpu_request( ) plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) - with pytest.raises(BatchRenderError, match="requires GRES"): - render_batch_script(plan) + with pytest.raises(SlurmBatchRenderError, match="requires GRES"): + render_generation_attempt_script(plan, attempt_ordinal=1) def test_renderer_escapes_shell_expansion_in_structured_paths( @@ -123,7 +116,7 @@ def test_renderer_escapes_shell_expansion_in_structured_paths( } ) - script = render_batch_script(plan) + script = render_generation_attempt_script(plan, attempt_ordinal=1) assert ( 'readonly DD_RUNTIME_ARCHIVE="/workspace/runtime/\\$(touch owned)-\\`whoami\\`-\\"bundle\\".tar.gz"' in script @@ -140,12 +133,13 @@ def test_renderer_keeps_user_text_on_one_non_executable_directive( update={"submission": single_node_plan.submission.model_copy(update={"comment": comment})} ) - script = render_batch_script(plan) + script = render_generation_attempt_script(plan, attempt_ordinal=1) comment_lines = [line for line in script.splitlines() if line.startswith("#SBATCH --comment=")] assert len(comment_lines) == 1 - assert "\\$(touch owned)" in comment_lines[0] - assert "\\`whoami\\`" in comment_lines[0] + assert shlex.split(comment_lines[0].removeprefix("#SBATCH ")) == [f"--comment={comment}"] + assert "\\$" not in comment_lines[0] + assert "\\`" not in comment_lines[0] assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 @@ -157,8 +151,8 @@ def test_renderer_rejects_invalid_attempt_ordinals( single_node_plan: ResolvedSlurmRunPlan, attempt_ordinal: object, ) -> None: - with pytest.raises(BatchRenderError, match="positive integer"): - render_batch_script(single_node_plan, attempt_ordinal=attempt_ordinal) # type: ignore[arg-type] + with pytest.raises(SlurmBatchRenderError, match="positive integer"): + render_generation_attempt_script(single_node_plan, attempt_ordinal=attempt_ordinal) # type: ignore[arg-type] def test_renderer_rejects_control_characters_from_unvalidated_plan_copies( @@ -168,14 +162,73 @@ def test_renderer_rejects_control_characters_from_unvalidated_plan_copies( update={"submission": single_node_plan.submission.model_copy(update={"comment": "unsafe\ntext"})} ) - with pytest.raises(BatchRenderError, match="control characters"): - render_batch_script(plan) + with pytest.raises(SlurmBatchRenderError, match="control characters"): + render_generation_attempt_script(plan, attempt_ordinal=1) def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) -> None: - script = render_batch_script(single_node_plan, attempt_ordinal=12) + script = render_generation_attempt_script(single_node_plan, attempt_ordinal=12) assert script.count("dd_slurm_run_allocation") == 1 assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script assert len(script.splitlines()) <= 42 assert script.endswith("\n") + + +def test_rendered_script_verifies_exact_persisted_plan_bytes_before_sourcing_runtime( + single_node_plan: ResolvedSlurmRunPlan, + tmp_path: Path, +) -> None: + run_root = tmp_path / "run" + run_root.mkdir() + captured_plan_path = tmp_path / "captured-plan.json" + entrypoint_path = tmp_path / "entrypoint.sh" + entrypoint_path.write_text( + f'dd_slurm_run_allocation() {{\n cp -- "$1" {shlex.quote(str(captured_plan_path))}\n}}\n' + ) + runtime_archive_path = tmp_path / "runtime.tar.gz" + with tarfile.open(runtime_archive_path, mode="w:gz") as runtime_archive: + runtime_archive.add(entrypoint_path, arcname="entrypoint.sh") + + plan = single_node_plan.model_copy( + update={ + "authored_config": ArtifactReference( + path=str(run_root / "authored-config.json"), + sha256="a" * 64, + ), + "runtime_bundle": ArtifactReference( + path=str(runtime_archive_path), + sha256=hashlib.sha256(runtime_archive_path.read_bytes()).hexdigest(), + ), + } + ) + plan_path = run_root / "resolved-plan.json" + serialized_plan = plan.serialize_json() + serialized_plan_bytes = serialized_plan.encode() + plan_path.write_bytes(serialized_plan_bytes) + script = render_generation_attempt_script(plan, attempt_ordinal=1) + environment = {**os.environ, "SLURM_ARRAY_TASK_ID": "0"} + + valid = subprocess.run( + ("bash",), + input=script, + env=environment, + text=True, + check=False, + ) + + assert valid.returncode == 0 + assert captured_plan_path.read_bytes() == serialized_plan_bytes + + captured_plan_path.unlink() + plan_path.write_bytes(serialized_plan_bytes + b" ") + tampered = subprocess.run( + ("bash",), + input=script, + env=environment, + text=True, + check=False, + ) + + assert tampered.returncode != 0 + assert not captured_plan_path.exists() diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index 8d51e68ae..6bab64a83 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -9,7 +9,7 @@ import pytest -from data_designer.slurm.launcher import SubprocessRunner +from data_designer.slurm.launcher.runner import SubprocessRunner def test_subprocess_runner_uses_argv_and_only_explicit_environment(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py index 0b85d89f3..60d45aa06 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py @@ -12,6 +12,7 @@ from slurm_test_fakes.slurm import ( FakeCommandResponse, FakeSlurmArray, + FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask, ) @@ -25,6 +26,7 @@ "FakeLogicalEndpoint", "FakeServingState", "FakeSlurmArray", + "FakeSlurmJob", "FakeSlurmRunner", "FakeSlurmTask", "FakeVllmBackend", 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 2c211eccd..22da3c51f 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -33,6 +33,20 @@ class FakeCommandResponse: returncode: int = 0 +@dataclass +class FakeSlurmJob: + """Mutable scheduler views for one ordinary batch-job identity.""" + + job_id: int + queue_state: str | None = "PENDING" + accounting_state: str | None = None + exit_code: str = "0:0" + + def __post_init__(self) -> None: + if type(self.job_id) is not int or self.job_id <= 0: + raise ValueError("fake Slurm job IDs must be positive integers") + + @dataclass class FakeSlurmTask: """Mutable scheduler views for one canonical array-task identity.""" @@ -66,16 +80,25 @@ def array_job_id(self) -> int: class FakeSlurmRunner: - """Stateful fake that copies its configured arrays before exposing Slurm commands.""" + """Stateful fake that copies configured jobs before exposing Slurm commands.""" def __init__( self, arrays: Iterable[FakeSlurmArray] = (), *, + jobs: Iterable[FakeSlurmJob] = (), sinfo_responses: Mapping[tuple[str, ...], FakeCommandResponse] | None = None, ) -> None: - self._pending_arrays = deque(copy.deepcopy(tuple(arrays))) + submissions = (*arrays, *jobs) + job_ids = tuple( + submission.array_job_id if isinstance(submission, FakeSlurmArray) else submission.job_id + for submission in submissions + ) + if len(job_ids) != len(set(job_ids)): + raise ValueError("fake Slurm submissions must have unique job IDs") + self._pending_submissions = deque(copy.deepcopy(submissions)) self._submitted_arrays: dict[int, FakeSlurmArray] = {} + self._submitted_jobs: dict[int, FakeSlurmJob] = {} self._scripted_responses: dict[str, deque[FakeCommandResponse]] = {} self._sinfo_responses = dict(sinfo_responses or {}) self.calls: list[tuple[str, ...]] = [] @@ -126,6 +149,20 @@ def set_task_state( task.accounting_state = accounting_state task.exit_code = exit_code + def set_job_state( + self, + job_id: int, + *, + queue_state: str | None, + accounting_state: str | None, + exit_code: str = "0:0", + ) -> None: + """Set the independently observable states for one ordinary job.""" + job = self._submitted_jobs[job_id] + job.queue_state = queue_state + job.accounting_state = accounting_state + job.exit_code = exit_code + def assert_scripts_consumed(self) -> None: """Assert that no scripted command response remains.""" remaining = sum(len(responses) for responses in self._scripted_responses.values()) @@ -147,38 +184,58 @@ def _dispatch(self, command_name: str, argv: tuple[str, ...]) -> FakeCommandResp return handler(argv) def _run_sbatch(self, argv: tuple[str, ...]) -> FakeCommandResponse: - if not self._pending_arrays: + if not self._pending_submissions: return FakeCommandResponse(stderr="no scripted submission\n", returncode=1) - array = self._pending_arrays.popleft() - self._submitted_arrays[array.array_job_id] = array + submission = self._pending_submissions.popleft() + if isinstance(submission, FakeSlurmArray): + job_id = submission.array_job_id + self._submitted_arrays[job_id] = submission + else: + job_id = submission.job_id + self._submitted_jobs[job_id] = submission if "--parsable" in argv[1:]: - return FakeCommandResponse(stdout=f"{array.array_job_id}\n") - return FakeCommandResponse(stdout=f"Submitted batch job {array.array_job_id}\n") + return FakeCommandResponse(stdout=f"{job_id}\n") + return FakeCommandResponse(stdout=f"Submitted batch job {job_id}\n") def _run_squeue(self, argv: tuple[str, ...]) -> FakeCommandResponse: self._require_arguments(argv, _SQUEUE_REQUIRED_ARGUMENTS) rows = [ + f"{job.job_id}|{job.queue_state}" + for job in self._selected_submitted_jobs(argv) + if job.queue_state is not None + ] + [ f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.queue_state}" for task in self._selected_submitted_tasks(argv) if task.queue_state is not None ] - return FakeCommandResponse(stdout="".join(f"{row}\n" for row in rows)) + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in sorted(rows))) def _run_sacct(self, argv: tuple[str, ...]) -> FakeCommandResponse: self._require_arguments(argv, _SACCT_REQUIRED_ARGUMENTS) rows = [ + f"{job.job_id}|{job.accounting_state}|{job.exit_code}" + for job in self._selected_submitted_jobs(argv) + if job.accounting_state is not None + ] + [ (f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.accounting_state}|{task.exit_code}") for task in self._selected_submitted_tasks(argv) if task.accounting_state is not None ] - return FakeCommandResponse(stdout="".join(f"{row}\n" for row in rows)) + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in sorted(rows))) def _run_scancel(self, argv: tuple[str, ...]) -> FakeCommandResponse: targets = tuple(argument for argument in argv[1:] if not argument.startswith("-")) if len(targets) != 1: return FakeCommandResponse(stderr="expected one cancellation target\n", returncode=1) + target = targets[0] + if "_" not in target and target.isdecimal() and int(target) in self._submitted_jobs: + job = self._submitted_jobs[int(target)] + job.queue_state = None + job.accounting_state = "CANCELLED" + job.exit_code = "0:15" + return FakeCommandResponse() try: - tasks = self._tasks_for_target(targets[0]) + tasks = self._tasks_for_target(target) except (KeyError, ValueError): return FakeCommandResponse(stderr="unknown cancellation target\n", returncode=1) for task in tasks: @@ -221,22 +278,12 @@ def _sorted_submitted_tasks(self) -> list[FakeSlurmTask]: ) def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask]: - selectors: list[str] = [] - for index, argument in enumerate(argv[1:]): - if argument in {"-j", "--jobs"}: - try: - selectors.extend(argv[index + 2].split(",")) - except IndexError: - raise AssertionError(f"missing job selector in {argv!r}") from None - elif argument.startswith("--jobs="): - selectors.extend(argument.partition("=")[2].split(",")) + selectors = self._job_selectors(argv) if not selectors: return self._sorted_submitted_tasks() selected: dict[SchedulerIdentity, FakeSlurmTask] = {} for selector in selectors: - if _JOB_SELECTOR_PATTERN.fullmatch(selector) is None: - raise AssertionError(f"malformed job selector {selector!r} in {argv!r}") try: tasks = self._tasks_for_target(selector) except KeyError: @@ -246,3 +293,34 @@ def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask selected.values(), key=lambda task: (task.scheduler.array_job_id, task.scheduler.array_task_id), ) + + def _selected_submitted_jobs(self, argv: tuple[str, ...]) -> list[FakeSlurmJob]: + selectors = self._job_selectors(argv) + if not selectors: + return sorted(self._submitted_jobs.values(), key=lambda job: job.job_id) + return sorted( + ( + self._submitted_jobs[int(selector)] + for selector in selectors + if "_" not in selector and int(selector) in self._submitted_jobs + ), + key=lambda job: job.job_id, + ) + + @staticmethod + def _job_selectors(argv: tuple[str, ...]) -> tuple[str, ...]: + selectors: list[str] = [] + for index, argument in enumerate(argv[1:]): + if argument in {"-j", "--jobs"}: + try: + selectors.extend(argv[index + 2].split(",")) + except IndexError: + raise AssertionError(f"missing job selector in {argv!r}") from None + elif argument.startswith("--jobs="): + selectors.extend(argument.partition("=")[2].split(",")) + if not selectors: + return () + for selector in selectors: + if _JOB_SELECTOR_PATTERN.fullmatch(selector) is None: + raise AssertionError(f"malformed job selector {selector!r} in {argv!r}") + return tuple(selectors) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py index 202143b34..c47b21432 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -9,7 +9,7 @@ import pytest from data_designer.slurm.state import SchedulerIdentity -from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" SQUEUE_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") @@ -141,6 +141,25 @@ def test_fake_slurm_runner_models_cancellation( assert f"4101_{task_id}|CANCELLED|0:15" in accounting +def test_fake_slurm_runner_models_ordinary_job_lifecycle_and_cancellation() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + + assert runner.run(("sbatch", "--parsable", "image-build.sbatch"), check=True).stdout == "5101\n" + assert runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs=5101"), check=True).stdout == "5101|PENDING\n" + + assert runner.run(("scancel", "5101"), check=True).returncode == 0 + assert runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs=5101"), check=True).stdout == "" + assert runner.run(("sacct", *SACCT_ARGUMENTS, "--jobs=5101"), check=True).stdout == "5101|CANCELLED|0:15\n" + + +def test_fake_slurm_runner_rejects_duplicate_job_ids() -> None: + ordinary_job = FakeSlurmJob(job_id=4101) + array = FakeSlurmArray(tasks=(FakeSlurmTask(scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)),)) + + with pytest.raises(ValueError, match="unique job IDs"): + FakeSlurmRunner((array,), jobs=(ordinary_job,)) + + def test_fake_slurm_runner_supports_malformed_output_and_command_failures( fake_slurm_runner: FakeSlurmRunner, ) -> None: diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 7de100df2..f80247355 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -130,7 +130,6 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.contracts import ResumeWorkspace as ContractResumeWorkspace from data_designer.slurm.integration import PlanStateValidator from data_designer.slurm.images.registry import ImageRegistryStore -from data_designer.slurm.launcher import SlurmCommandClient from data_designer.slurm.planning import ArtifactReference as PlanningArtifactReference from data_designer.slurm.planning import RecordRange as PlanningRecordRange from data_designer.slurm.planning import ResumeWorkspace as PlanningResumeWorkspace @@ -140,7 +139,6 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.state import RunManifest assert RunManifest.__name__ == "RunManifest" assert ImageRegistryStore.__name__ == "ImageRegistryStore" -assert "SlurmCommandClient" in str(SlurmCommandClient) assert PlanningArtifactReference is ContractArtifactReference assert PlanningRecordRange is ContractRecordRange assert PlanningResumeWorkspace is ContractResumeWorkspace From ef82bfe32cc41913ade44b190446cdf6cbd94e12 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 27 Aug 2026 14:01:29 -0600 Subject: [PATCH 22/23] fix Slurm selector filtering ownership --- .../src/data_designer/slurm/launcher/parsing.py | 7 +------ .../src/data_designer/slurm/launcher/renderer.py | 5 +---- .../tests/launcher/test_client.py | 16 ++++++++-------- .../tests/launcher/test_parsing.py | 7 ++++--- .../slurm_test_fakes/test_rendered_scripts.py | 2 +- 5 files changed, 15 insertions(+), 22 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index 2aff70228..af663bf1a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -99,12 +99,7 @@ def parse_accounting(output: str) -> tuple[SlurmAccountingEntry, ...]: process_exit_code=_parse_exit_code(fields[2], line_number=line_number), ) ) - array_job_ids = { - entry.job_identity.array_job_id for entry in entries if isinstance(entry.job_identity, SchedulerIdentity) - } - return tuple( - entry for entry in entries if not (type(entry.job_identity) is int and entry.job_identity in array_job_ids) - ) + return tuple(entries) def parse_gpu_counts(output: str) -> tuple[int, ...]: 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 931eb977a..a0dfe02e8 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 @@ -91,10 +91,7 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire node_count = max(node_indices) + 1 array = "0" if plan.array_tasks.count > 1: - array = f"0-{plan.array_tasks.count - 1}" - # TODO(#875): Add valid authored-to-rendered coverage once omitted concurrency is supported by the contract. - if plan.array_tasks.max_concurrent is not None: - array = f"{array}%{plan.array_tasks.max_concurrent}" + array = f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" values: list[tuple[str, str | None]] = [ ("job-name", plan.submission.job_name), diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index ccc88978a..5155ddb02 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -90,26 +90,26 @@ def test_client_observes_regular_cpu_job() -> None: assert accounting[0].state is SchedulerState.COMPLETED -def test_client_ignores_array_parent_observation_for_exact_task() -> None: +def test_client_ignores_only_array_parent_observation_for_exact_task() -> None: runner = FakeSlurmRunner() - runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n")) + runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n4101_0|COMPLETED|0:0\n")) client = SlurmCommandClient(runner) + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) - records = client.query_accounting((SchedulerIdentity(array_job_id=4101, array_task_id=0),)) + records = client.query_accounting((task,)) - assert records == () + assert tuple(record.job_identity for record in records) == (task,) -def test_client_keeps_explicitly_selected_parent_observation() -> None: +def test_client_keeps_explicitly_selected_parent_and_task_observations() -> None: runner = FakeSlurmRunner() - runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n")) + runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n4101_0|COMPLETED|0:0\n")) client = SlurmCommandClient(runner) task = SchedulerIdentity(array_job_id=4101, array_task_id=0) records = client.query_accounting((4101, task)) - assert len(records) == 1 - assert records[0].job_identity == 4101 + assert tuple(record.job_identity for record in records) == (4101, task) def test_client_rejects_unbounded_or_invalid_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index 34e7c5a2f..7509f93c2 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -87,12 +87,13 @@ def test_parse_state_normalizes_long_slurm_spellings(raw_state: str, expected: S assert parse_state(raw_state) is expected -def test_parse_accounting_normalizes_terminal_rows_and_ignores_array_parent() -> None: +def test_parse_accounting_normalizes_parent_and_terminal_array_task_rows() -> None: output = "4101|RUNNING|0:0\n" + (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() records = parse_accounting(output) assert tuple(record.state for record in records) == ( + SchedulerState.RUNNING, SchedulerState.TIMED_OUT, SchedulerState.NODE_FAILED, SchedulerState.PREEMPTED, @@ -100,8 +101,8 @@ def test_parse_accounting_normalizes_terminal_rows_and_ignores_array_parent() -> SchedulerState.OUT_OF_MEMORY, SchedulerState.CANCELLED, ) - assert records[0].process_exit_code.exit_status == 0 - assert records[0].process_exit_code.termination_signal == 125 + assert records[1].process_exit_code.exit_status == 0 + assert records[1].process_exit_code.termination_signal == 125 def test_parse_accounting_normalizes_regular_jobs() -> None: 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 c51af0cac..a8c29e5ef 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 @@ -55,7 +55,7 @@ def _assert_script_matches_plan( ) node_count = max(node_indices) + 1 array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}" - if plan.array_tasks.count > 1 and plan.array_tasks.max_concurrent is not None: + if plan.array_tasks.count > 1: array = f"{array}%{plan.array_tasks.max_concurrent}" plan_path = posixpath.join(posixpath.dirname(plan.authored_config.path), "resolved-plan.json") run_root = posixpath.dirname(plan.authored_config.path) From d5db22e37eeed610d5c812025068b141294ecf35 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 27 Aug 2026 14:48:06 -0600 Subject: [PATCH 23/23] simplify Slurm launcher structure --- .../data_designer/slurm/launcher/client.py | 2 +- .../data_designer/slurm/launcher/models.py | 8 ++--- .../data_designer/slurm/launcher/parsing.py | 30 +++++++++---------- .../data_designer/slurm/launcher/renderer.py | 2 +- 4 files changed, 21 insertions(+), 21 deletions(-) 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 3bddd06ef..3c0d3e858 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 @@ -35,7 +35,7 @@ _MAX_SLURM_INTEGER = (1 << 32) - 1 -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class SlurmExecutables: """Executable paths used for bounded Slurm operations.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py index b686cc02f..7d1168255 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -13,14 +13,14 @@ SlurmObservedJobIdentity: TypeAlias = int | SchedulerIdentity -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class SlurmJobSubmissionReceipt: """Job identity returned for one accepted non-federated submission.""" job_id: int -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class SlurmProcessExitCode: """Slurm's process status and terminating signal pair.""" @@ -28,7 +28,7 @@ class SlurmProcessExitCode: termination_signal: int -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class SlurmQueueEntry: """One transient normalized active-queue entry.""" @@ -36,7 +36,7 @@ class SlurmQueueEntry: state: SchedulerState -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class SlurmAccountingEntry: """One transient normalized accounting entry.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index af663bf1a..adf0e13c2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -126,6 +126,21 @@ def parse_gpu_counts(output: str) -> tuple[int, ...]: return tuple(counts) +def parse_state(value: str) -> SchedulerState: + """Normalize one Slurm long state spelling without guessing unknown states.""" + normalized = value.strip().upper().removesuffix("+") + if not normalized: + raise SlurmCommandOutputError("scheduler state must not be empty") + if normalized.startswith("CANCELLED BY "): + canceller = normalized.removeprefix("CANCELLED BY ") + if not canceller.isascii() or not canceller.isdecimal(): + raise SlurmCommandOutputError("cancelled scheduler state has an invalid owner") + normalized = "CANCELLED" + elif any(character.isspace() for character in normalized): + raise SlurmCommandOutputError("scheduler state contains unexpected whitespace") + return _STATE_MAP.get(normalized, SchedulerState.UNKNOWN) + + def _split_gres_fields(value: str, *, line_number: int) -> tuple[str, ...]: fields: list[str] = [] start = 0 @@ -150,21 +165,6 @@ def _split_gres_fields(value: str, *, line_number: int) -> tuple[str, ...]: return tuple(fields) -def parse_state(value: str) -> SchedulerState: - """Normalize one Slurm long state spelling without guessing unknown states.""" - normalized = value.strip().upper().removesuffix("+") - if not normalized: - raise SlurmCommandOutputError("scheduler state must not be empty") - if normalized.startswith("CANCELLED BY "): - canceller = normalized.removeprefix("CANCELLED BY ") - if not canceller.isascii() or not canceller.isdecimal(): - raise SlurmCommandOutputError("cancelled scheduler state has an invalid owner") - normalized = "CANCELLED" - elif any(character.isspace() for character in normalized): - raise SlurmCommandOutputError("scheduler state contains unexpected whitespace") - return _STATE_MAP.get(normalized, SchedulerState.UNKNOWN) - - def _collect_nonempty_lines(output: str) -> tuple[tuple[int, str], ...]: return tuple( (line_number, line) 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 a0dfe02e8..203cfe14a 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 @@ -16,7 +16,7 @@ _DIRECTIVE_TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/,%+-]*$") -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class _BatchDirective: """One validated ``#SBATCH`` option."""