From 965291dccd6b2c3613ddf91dc16909fa58ece1a7 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:26:00 -0600 Subject: [PATCH 01/20] 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 5a78c97b7..9d6ea4ade 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 e10478f7d..dd45658f9 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 c7d119389..920db7671 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="faa1dac9b9b0423423c9d06a13d71bee339932ffa45d1e02a8a95012a7934520", + expected_fixture_sha256="f8e485a709ca2b7c5499b7714118987c4afb1ede63ddec96edd04236e6a763bf", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="a4542e56b124c2346d0a92ebadc4e89b45d94c9355d7b86a4a1b05180d331a48", + expected_fixture_sha256="6865fa77f23c3db85898ce3e4176bee5efcf841f98fa74fa0499fc13e8031cf7", ) From de6336a8ad5ce247a38cc755ba0fa146b5c5c9c5 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:35:22 -0600 Subject: [PATCH 02/20] 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 6809a9e74fce2378d8c0eeddcdc4749a079c081e Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:37:05 -0600 Subject: [PATCH 03/20] 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 fc5d366876cb068b14af304d932f140320bec2c7 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:38:07 -0600 Subject: [PATCH 04/20] 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 fec9f8aa08d1d03010ac8ca63ea48372949092b0 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:39:38 -0600 Subject: [PATCH 05/20] 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 932bb46b631f486418af82c2e57705492a383493 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:41:35 -0600 Subject: [PATCH 06/20] 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 6b773948cd49839d4c42dabcda8d49fac58a18be Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:43:53 -0600 Subject: [PATCH 07/20] 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 9d6ea4ade..908e04134 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 dd45658f9..e308b1ea5 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 920db7671..368c06d96 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="f8e485a709ca2b7c5499b7714118987c4afb1ede63ddec96edd04236e6a763bf", + expected_fixture_sha256="06c56a3b335bbf69dfd32eb8ada968839819f379127e1cefebea5b10a18e1a25", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="6865fa77f23c3db85898ce3e4176bee5efcf841f98fa74fa0499fc13e8031cf7", + expected_fixture_sha256="6c26c1c0fd3956ef2299b4dcc7ce9be82b94e246cb72739ecdf8d38b49b64fb7", ) From 6d2ea1c746d204dcd3ed837f2a7f09c167bdf35f Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:45:11 -0600 Subject: [PATCH 08/20] 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 39e10c139ee0be04090db8d2505a3ab47e96c616 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:47:55 -0600 Subject: [PATCH 09/20] 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 ab328aa829a44dd195d98aabe31ea2c429c99165 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:53:09 -0600 Subject: [PATCH 10/20] 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 29de67388c6471d326ad8397dcb24186aba312bc Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:55:44 -0600 Subject: [PATCH 11/20] 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 0ddbb8ec9de4eff78a40c64dbab8f74295df23f1 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 13:57:54 -0600 Subject: [PATCH 12/20] 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 3738c1f1e3bd716deb33afd0c06539d57a1d6918 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:02:38 -0600 Subject: [PATCH 13/20] 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 f6ba85e19..86ad74edb 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -129,6 +129,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.contracts import RecordRange as ContractRecordRange from data_designer.slurm.contracts import ResumeWorkspace as ContractResumeWorkspace from data_designer.slurm.integration import PlanStateValidator +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 @@ -137,6 +138,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.state import ResumeWorkspace as StateResumeWorkspace from data_designer.slurm.state import RunManifest assert RunManifest.__name__ == "RunManifest" +assert "SlurmCommandClient" in str(SlurmCommandClient) assert PlanningArtifactReference is ContractArtifactReference assert PlanningRecordRange is ContractRecordRange assert PlanningResumeWorkspace is ContractResumeWorkspace From 426a753796257764881cb80c18a6859bfe19e165 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:04:55 -0600 Subject: [PATCH 14/20] 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 908e04134..581d120cc 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 e308b1ea5..e7e92a775 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 368c06d96..abf9b1b5a 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="06c56a3b335bbf69dfd32eb8ada968839819f379127e1cefebea5b10a18e1a25", + expected_fixture_sha256="b2d5e4983f2b39e9f9a4a1c6ac15e2cbdf68f87d8fa565645df62db59058dead", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="6c26c1c0fd3956ef2299b4dcc7ce9be82b94e246cb72739ecdf8d38b49b64fb7", + expected_fixture_sha256="dadf878f8815ab35dfd47d723cca7f4248294e04c5f4a478f693407038499001", ) From 4e544e9e3573380c0028492ced8e5c2cb1699732 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:08:55 -0600 Subject: [PATCH 15/20] 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 67c53ee642adbe66405b4b882225c242e229c250 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:11:45 -0600 Subject: [PATCH 16/20] 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 bb5bc8ec209d25afe7807641b7c215845ee18a35 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 14:50:40 -0600 Subject: [PATCH 17/20] 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 ca308e292761251d484d28a87b54b1717686a633 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 25 Aug 2026 15:24:17 -0600 Subject: [PATCH 18/20] 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 7f84dcfc8501ba32a39131d51f241bb0120845dd Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 26 Aug 2026 12:59:03 -0600 Subject: [PATCH 19/20] 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 581d120cc..00f4bd156 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 e7e92a775..9a041d956 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 abf9b1b5a..7595e63b7 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="b2d5e4983f2b39e9f9a4a1c6ac15e2cbdf68f87d8fa565645df62db59058dead", + expected_fixture_sha256="9d0a88e9c6005998755a80694d14b23208e68723fcdb932962054228e81c801e", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="dadf878f8815ab35dfd47d723cca7f4248294e04c5f4a478f693407038499001", + expected_fixture_sha256="1dd8db6daaf5bc97168c7e205741cd22ab727527c0c144f408e6033ed3bf031b", ) @@ -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 7acac4f3a91ea3ac12721f4a13619ce1a3dcbcf5 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 26 Aug 2026 13:49:16 -0600 Subject: [PATCH 20/20] 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 7595e63b7..e4ee256e8 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)