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..9bc45a4a6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py @@ -0,0 +1,40 @@ +# 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, + SlurmJobIdentity, + SlurmSubmission, +) +from data_designer.slurm.launcher.renderer import render_batch_script +from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner + +__all__ = [ + "AccountingRecord", + "BatchRenderError", + "CommandRunner", + "QueueRecord", + "SlurmCommandClient", + "SlurmCommandError", + "SlurmExecutables", + "SlurmExitCode", + "SlurmJobIdentity", + "SlurmLauncherError", + "SlurmParseError", + "SlurmSubmission", + "SubprocessRunner", + "render_batch_script", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py new file mode 100644 index 000000000..fdb62bc83 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -0,0 +1,212 @@ +# 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 +import unicodedata +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, SlurmJobIdentity, 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: 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) +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") + if path.startswith("-"): + raise ValueError("batch script path must not begin with '-'; prefix relative paths with './'") + output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) + return parse_submission(output) + + def query_queue(self, selectors: Sequence[JobSelector]) -> tuple[QueueRecord, ...]: + """Return normalized active-queue rows for explicit managed jobs.""" + requested = tuple(selectors) + jobs = _format_selectors(requested) + output = self._run( + ( + self._executables.squeue, + "--noheader", + "--array", + "--format=%i|%T", + f"--jobs={jobs}", + ) + ) + records = parse_queue(output) + ignored = _validate_selected_schedulers( + tuple(record.scheduler for record in records), + requested, + command="squeue", + ) + 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.""" + requested = tuple(selectors) + jobs = _format_selectors(requested) + output = self._run( + ( + self._executables.sacct, + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,State,ExitCode", + f"--jobs={jobs}", + ) + ) + records = parse_accounting(output) + ignored = _validate_selected_schedulers( + tuple(record.scheduler for record in records), + requested, + command="sacct", + ) + return tuple(record for record in records if record.scheduler not in ignored) + + def cancel(self, selector: JobSelector) -> None: + """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, ...]: + """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): + 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[SlurmJobIdentity], + selectors: Sequence[JobSelector], + *, + command: str, +) -> frozenset[SlurmJobIdentity]: + """Validate result correlation and identify aggregate rows to omit.""" + ignored: set[SlurmJobIdentity] = set() + for scheduler in schedulers: + explicitly_selected = any(type(selector) is int and selector == scheduler for selector in selectors) + is_array_parent = type(scheduler) is int and any( + isinstance(selector, SchedulerIdentity) and selector.array_job_id == scheduler for selector in selectors + ) + if is_array_parent and not explicitly_selected: + ignored.add(scheduler) + continue + if any(_selector_matches(scheduler, selector) for selector in selectors): + continue + 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: + 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: + 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 - 3]}..." + + +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..c7e3bc68e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -0,0 +1,47 @@ +# 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 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 batch submission.""" + + job_id: int + cluster_name: Identifier | None = None + + +@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: SlurmJobIdentity + state: SchedulerState + + +@dataclass(frozen=True, slots=True) +class AccountingRecord: + """One normalized accounting row.""" + + 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 new file mode 100644 index 000000000..3220bc8fd --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -0,0 +1,232 @@ +# 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, + SlurmJobIdentity, + SlurmSubmission, +) +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + +_ARRAY_ID_PATTERN = re.compile(r"^(?P[1-9][0-9]*)_(?P[0-9]+)$") +_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]*\))?$") +_MAX_SLURM_INTEGER = (1 << 32) - 1 + +_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, + "RESV_DEL_HOLD": SchedulerState.PENDING, + "RESIZING": SchedulerState.RUNNING, + "REVOKED": SchedulerState.FAILED, + "RUNNING": SchedulerState.RUNNING, + "SIGNALING": SchedulerState.RUNNING, + "SPECIAL_EXIT": SchedulerState.PENDING, + "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(): + raise SlurmParseError("sbatch returned an invalid job ID") + parsed_job_id = _parse_decimal(job_id, message="sbatch returned an invalid job ID") + if parsed_job_id <= 0: + raise SlurmParseError("sbatch returned an invalid job ID") + if separator and _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: + raise SlurmParseError("sbatch returned an invalid cluster name") + return SlurmSubmission(job_id=parsed_job_id, cluster_name=cluster_name or None) + + +def parse_queue(output: str) -> tuple[QueueRecord, ...]: + """Parse ``squeue --format=%i|%T`` rows.""" + records: list[QueueRecord] = [] + 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_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 job and array-task rows from ``sacct --format=JobID,State,ExitCode``.""" + records: list[AccountingRecord] = [] + 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") + scheduler = _parse_job_identity(fields[0], command="sacct", line_number=line_number) + _reject_duplicate(scheduler, identities, command="sacct", line_number=line_number) + records.append( + AccountingRecord( + scheduler=scheduler, + state=parse_state(fields[1]), + exit_code=_parse_exit_code(fields[2], line_number=line_number), + ) + ) + array_job_ids = { + record.scheduler.array_job_id for record in records if isinstance(record.scheduler, SchedulerIdentity) + } + 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, ...]: + """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 _split_gres_fields(line, line_number=line_number): + 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( + _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) + + +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 + 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: + 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:]) + if any(not field for field in fields): + raise SlurmParseError(f"sinfo line {line_number} contains an invalid GPU resource") + return tuple(fields) + + +def parse_state(value: str) -> SchedulerState: + """Normalize one Slurm long state spelling without guessing unknown states.""" + normalized = value.strip().upper().removesuffix("+") + if not normalized: + raise 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") + message = f"{command} line {line_number} contains an invalid array-task ID" + return SchedulerIdentity( + array_job_id=_parse_decimal(match.group("job"), message=message), + array_task_id=_parse_decimal(match.group("task"), message=message), + ) + + +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: + raise SlurmParseError(f"sacct line {line_number} contains an invalid exit code") + message = f"sacct line {line_number} contains an invalid exit code" + return SlurmExitCode( + status=_parse_decimal(match.group("status"), message=message), + signal=_parse_decimal(match.group("signal"), message=message), + ) + + +def _parse_decimal(value: str, *, message: str) -> int: + if len(value) > 10: + raise SlurmParseError(message) + try: + 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: SlurmJobIdentity, + identities: set[SlurmJobIdentity], + *, + command: str, + line_number: int, +) -> None: + if scheduler in identities: + raise SlurmParseError(f"{command} line {line_number} duplicates a job or array-task ID") + identities.add(scheduler) 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..915fceb02 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -0,0 +1,127 @@ +# 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 +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" + +readonly DD_RUNTIME_ARCHIVE={_quote_double_value(plan.runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={_quote_double_value(plan.runtime_bundle.sha256)} +readonly DD_PLAN={_quote_double_value(plan_path)} +readonly DD_PLAN_SHA256={_quote_double_value(plan.compute_sha256())} +readonly DD_RUN_ROOT={_quote_double_value(run_root)} +readonly DD_ATTEMPT_ORDINAL={_quote_double_value(attempt)} + +verify_sha256() {{ + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${{actual_sha256%% *}}" == "$1" ]] +}} + +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}" + 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), + ("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), + ] + 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: + 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..e7c79b4b4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py @@ -0,0 +1,66 @@ +# 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 math +import os +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 type(timeout_seconds) not in {int, float} or not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise ValueError("timeout_seconds must be a finite positive number") + explicit_environment = ( + dict(environment) if environment is not None else {"PATH": os.environ.get("PATH") or os.defpath} + ) + 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 = float(timeout_seconds) + + @property + def environment(self) -> Mapping[str, str]: + """Return the allowlisted 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, + encoding="utf-8", + errors="replace", + 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..6118593ca --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -0,0 +1,268 @@ +# 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, SlurmExecutables, SlurmParseError +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.job_id,)) + + 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"), + ("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", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,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_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_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 32-bit integers"): + client.query_accounting((0,)) + 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 == [] + + +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_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) + + 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_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_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()) + + with pytest.raises(SlurmCommandError, match="squeue could not be executed") as error: + client.query_queue((4101,)) + + 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) + + client.submit("/workspace/run; touch injected.sbatch") + + assert fake_slurm_runner.calls[0] == ( + "sbatch", + "--parsable", + "--export=NIL", + "/workspace/run; touch injected.sbatch", + ) + + +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 == [] + + +@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 new file mode 100644 index 000000000..a974e22ab --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -0,0 +1,206 @@ +# 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" +OVERSIZED_DECIMAL = "9" * 5000 + + +@pytest.mark.parametrize( + ("output", "expected_job_id", "expected_cluster"), + (("4101\n", 4101, None), ("4101;primary\n", 4101, "primary")), +) +def test_parse_submission_accepts_parsable_sbatch_output( + output: str, + expected_job_id: int, + expected_cluster: str | None, +) -> None: + submission = parse_submission(output) + + 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", 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()) + + assert records == ( + _make_queue_record(0, SchedulerState.PENDING), + _make_queue_record(1, SchedulerState.RUNNING), + ) + + +def test_parse_queue_normalizes_regular_jobs() -> None: + records = parse_queue("5101|RUNNING\n") + + assert records == (QueueRecord(scheduler=5101, state=SchedulerState.RUNNING),) + + +@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), + ("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), + ), +) +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_array_parent() -> None: + output = "4101|RUNNING|0:0\n" + (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() + + records = parse_accounting(output) + + assert tuple(record.state for record in records) == ( + SchedulerState.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_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") == () + + +@pytest.mark.parametrize( + ("parser", "output", "message"), + ( + (parse_queue, "malformed scheduler output\n", "two fields"), + (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", "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"), + ), +) +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( + ("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"), + ( + ("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", ()), + ), +) +def test_parse_gpu_counts_normalizes_configured_gres(output: str, expected: tuple[int, ...]) -> None: + 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", + "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) + + +@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), + 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..00616fd19 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -0,0 +1,181 @@ +# 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 +from typing import Literal + +import pytest + +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 + +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_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"), + } + ) + 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 --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) + + +@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_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: + 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: + 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_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"): + render_batch_script(plan) + + +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()) <= 42 + 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..8d51e68ae --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +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, + encoding: str, + errors: str, + env: Mapping[str, str], + timeout: float, + ) -> subprocess.CompletedProcess[str]: + observed.update( + command=command, + check=check, + stdin=stdin, + capture_output=capture_output, + text=text, + encoding=encoding, + errors=errors, + 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, + "encoding": "utf-8", + "errors": "replace", + "env": {"LC_ALL": "C", "PATH": "/usr/bin"}, + "timeout": 4.0, + } + + +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_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() + + with pytest.raises(TypeError): + runner.environment["SECRET"] = "value" # type: ignore[index] + + +@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) # type: ignore[arg-type] + + +@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..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,10 +3,12 @@ #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 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" @@ -16,17 +18,24 @@ 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}" 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..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,10 +3,12 @@ #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 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" @@ -16,17 +18,24 @@ 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}" 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/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py index 785af0630..2c211eccd 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -14,8 +14,14 @@ from data_designer.slurm.state import SchedulerIdentity _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") +_SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") +_SACCT_REQUIRED_ARGUMENTS = ( + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,State,ExitCode", +) @dataclass(frozen=True) 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..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 @@ -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="9d0a88e9c6005998755a80694d14b23208e68723fcdb932962054228e81c801e", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="a4542e56b124c2346d0a92ebadc4e89b45d94c9355d7b86a4a1b05180d331a48", + expected_fixture_sha256="1dd8db6daaf5bc97168c7e205741cd22ab727527c0c144f408e6033ed3bf031b", ) @@ -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) @@ -62,6 +64,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 8bfb3de0e..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 @@ -12,8 +12,14 @@ from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" -SQUEUE_ARGUMENTS = ("--noheader", "--format=%i|%T") -SACCT_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") +SQUEUE_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") +SACCT_ARGUMENTS = ( + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,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=JobID,State,ExitCode"), ), ) def test_fake_slurm_runner_rejects_underspecified_state_queries( 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