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..61da221a6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal Slurm submission, observation, and batch-rendering helpers.""" + +from __future__ import annotations 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..3c0d3e858 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal 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, SlurmCommandOutputError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmObservedJobIdentity, + SlurmQueueEntry, +) +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 = int | SchedulerIdentity +_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_MAX_SLURM_INTEGER = (1 << 32) - 1 + + +@dataclass(frozen=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) -> SlurmJobSubmissionReceipt: + """Submit one rendered batch script and return its assigned job ID.""" + path = str(script_path) + _validate_argument(path, field_name="batch script path") + if path.startswith("-"): + raise ValueError("batch script path must not begin with '-'; prefix relative paths with './'") + output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) + return parse_submission(output) + + def query_queue(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmQueueEntry, ...]: + """Return normalized active-queue rows for explicit managed jobs.""" + requested = tuple(selectors) + jobs = _format_selectors(requested) + output = self._run( + ( + self._executables.squeue, + "--noheader", + "--array", + "--format=%i|%T", + f"--jobs={jobs}", + ) + ) + entries = parse_queue(output) + ignored = _validate_observed_job_identities( + tuple(entry.job_identity for entry in entries), + requested, + command="squeue", + ) + return tuple(entry for entry in entries if entry.job_identity not in ignored) + + def query_accounting(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmAccountingEntry, ...]: + """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}", + ) + ) + entries = parse_accounting(output) + ignored = _validate_observed_job_identities( + tuple(entry.job_identity for entry in entries), + requested, + command="sacct", + ) + return tuple(entry for entry in entries if entry.job_identity 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 + returncode = getattr(completed, "returncode", None) + stdout = getattr(completed, "stdout", None) + stderr = getattr(completed, "stderr", None) + if type(returncode) is not int or not isinstance(stdout, str) or not isinstance(stderr, str): + raise SlurmCommandError(f"{command_name} returned a malformed process result") + if returncode: + detail = _normalize_bounded_text(stderr) or "no diagnostic output" + raise SlurmCommandError(f"{command_name} failed with exit code {returncode}: {detail}") + return stdout + + +def _format_selectors(selectors: Sequence[_JobSelector]) -> str: + if not selectors: + raise ValueError("at least one managed Slurm job selector is required") + return ",".join(dict.fromkeys(_format_selector(selector) for selector in selectors)) + + +def _format_selector(selector: _JobSelector) -> str: + 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_observed_job_identities( + job_identities: Sequence[SlurmObservedJobIdentity], + selectors: Sequence[_JobSelector], + *, + command: str, +) -> frozenset[SlurmObservedJobIdentity]: + """Validate result correlation and return unselected aggregate rows.""" + selected_job_ids = {selector for selector in selectors if type(selector) is int} + selected_array_tasks = {selector for selector in selectors if isinstance(selector, SchedulerIdentity)} + selected_array_job_ids = {selector.array_job_id for selector in selected_array_tasks} + ignored: set[SlurmObservedJobIdentity] = set() + for job_identity in job_identities: + if type(job_identity) is int and job_identity in selected_job_ids: + continue + if type(job_identity) is int and job_identity in selected_array_job_ids: + ignored.add(job_identity) + continue + if isinstance(job_identity, SchedulerIdentity) and ( + job_identity in selected_array_tasks or job_identity.array_job_id in selected_job_ids + ): + continue + raise SlurmCommandOutputError(f"{command} returned an unrequested job or array-task ID") + return frozenset(ignored) + + +def _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..4a611152b --- /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 + +"""Internal normalized 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 SlurmCommandOutputError(SlurmLauncherError, ValueError): + """A Slurm command returned output that violates its requested format.""" + + +class SlurmBatchRenderError(SlurmLauncherError, ValueError): + """A resolved plan cannot be rendered as a safe batch script.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py new file mode 100644 index 000000000..7d1168255 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -0,0 +1,45 @@ +# 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.state import SchedulerIdentity, SchedulerState + +SlurmObservedJobIdentity: TypeAlias = int | SchedulerIdentity + + +@dataclass(frozen=True) +class SlurmJobSubmissionReceipt: + """Job identity returned for one accepted non-federated submission.""" + + job_id: int + + +@dataclass(frozen=True) +class SlurmProcessExitCode: + """Slurm's process status and terminating signal pair.""" + + exit_status: int + termination_signal: int + + +@dataclass(frozen=True) +class SlurmQueueEntry: + """One transient normalized active-queue entry.""" + + job_identity: SlurmObservedJobIdentity + state: SchedulerState + + +@dataclass(frozen=True) +class SlurmAccountingEntry: + """One transient normalized accounting entry.""" + + job_identity: SlurmObservedJobIdentity + state: SchedulerState + process_exit_code: SlurmProcessExitCode diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py new file mode 100644 index 000000000..adf0e13c2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal strict parsers for bounded, machine-readable Slurm output.""" + +from __future__ import annotations + +import re + +from data_designer.slurm.launcher.errors import SlurmCommandOutputError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmObservedJobIdentity, + SlurmProcessExitCode, + SlurmQueueEntry, +) +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) -> SlurmJobSubmissionReceipt: + """Parse non-federated ``sbatch --parsable`` output.""" + value = output.strip() + job_id, separator, cluster_name = value.partition(";") + if not job_id.isascii() or not job_id.isdecimal(): + raise SlurmCommandOutputError("sbatch returned an invalid job ID") + parsed_job_id = _parse_decimal(job_id, message="sbatch returned an invalid job ID") + if parsed_job_id <= 0: + raise SlurmCommandOutputError("sbatch returned an invalid job ID") + if separator: + if _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: + raise SlurmCommandOutputError("sbatch returned an invalid cluster name") + raise SlurmCommandOutputError("federated Slurm submissions are not supported") + return SlurmJobSubmissionReceipt(job_id=parsed_job_id) + + +def parse_queue(output: str) -> tuple[SlurmQueueEntry, ...]: + """Parse ``squeue --format=%i|%T`` rows.""" + entries: list[SlurmQueueEntry] = [] + identities: set[SlurmObservedJobIdentity] = set() + for line_number, line in _collect_nonempty_lines(output): + fields = line.split("|") + if len(fields) != 2: + raise SlurmCommandOutputError(f"squeue line {line_number} must contain two fields") + job_identity = _parse_job_identity(fields[0], command="squeue", line_number=line_number) + _reject_duplicate(job_identity, identities, command="squeue", line_number=line_number) + entries.append(SlurmQueueEntry(job_identity=job_identity, state=parse_state(fields[1]))) + return tuple(entries) + + +def parse_accounting(output: str) -> tuple[SlurmAccountingEntry, ...]: + """Parse job and array-task rows from ``sacct --format=JobID,State,ExitCode``.""" + entries: list[SlurmAccountingEntry] = [] + identities: set[SlurmObservedJobIdentity] = set() + for line_number, line in _collect_nonempty_lines(output): + fields = line.split("|") + if len(fields) != 3: + raise SlurmCommandOutputError(f"sacct line {line_number} must contain three fields") + job_identity = _parse_job_identity(fields[0], command="sacct", line_number=line_number) + _reject_duplicate(job_identity, identities, command="sacct", line_number=line_number) + entries.append( + SlurmAccountingEntry( + job_identity=job_identity, + state=parse_state(fields[1]), + process_exit_code=_parse_exit_code(fields[2], line_number=line_number), + ) + ) + return tuple(entries) + + +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 SlurmCommandOutputError(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 parse_state(value: str) -> SchedulerState: + """Normalize one Slurm long state spelling without guessing unknown states.""" + normalized = value.strip().upper().removesuffix("+") + if not normalized: + raise SlurmCommandOutputError("scheduler state must not be empty") + if normalized.startswith("CANCELLED BY "): + canceller = normalized.removeprefix("CANCELLED BY ") + if not canceller.isascii() or not canceller.isdecimal(): + raise SlurmCommandOutputError("cancelled scheduler state has an invalid owner") + normalized = "CANCELLED" + elif any(character.isspace() for character in normalized): + raise SlurmCommandOutputError("scheduler state contains unexpected whitespace") + return _STATE_MAP.get(normalized, SchedulerState.UNKNOWN) + + +def _split_gres_fields(value: str, *, line_number: int) -> tuple[str, ...]: + fields: list[str] = [] + start = 0 + annotation_depth = 0 + for index, character in enumerate(value): + if character == "(": + annotation_depth += 1 + if annotation_depth > 1: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + elif character == ")": + annotation_depth -= 1 + if annotation_depth < 0: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + elif character == "," and annotation_depth == 0: + fields.append(value[start:index]) + start = index + 1 + if annotation_depth: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + fields.append(value[start:]) + if any(not field for field in fields): + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + return tuple(fields) + + +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 SlurmCommandOutputError(f"{command} line {line_number} contains an invalid array-task ID") + message = f"{command} line {line_number} contains an invalid array-task ID" + return SchedulerIdentity( + array_job_id=_parse_decimal(match.group("job"), message=message), + array_task_id=_parse_decimal(match.group("task"), message=message), + ) + + +def _parse_job_identity(value: str, *, command: str, line_number: int) -> SlurmObservedJobIdentity: + message = f"{command} line {line_number} contains an invalid job or array-task ID" + if _JOB_ID_PATTERN.fullmatch(value) is not None: + return _parse_decimal(value, message=message) + try: + return _parse_array_identity(value, command=command, line_number=line_number) + except SlurmCommandOutputError as error: + raise SlurmCommandOutputError(message) from error + + +def _parse_exit_code(value: str, *, line_number: int) -> SlurmProcessExitCode: + match = _EXIT_CODE_PATTERN.fullmatch(value) + if match is None: + raise SlurmCommandOutputError(f"sacct line {line_number} contains an invalid exit code") + message = f"sacct line {line_number} contains an invalid exit code" + return SlurmProcessExitCode( + exit_status=_parse_decimal(match.group("status"), message=message), + termination_signal=_parse_decimal(match.group("signal"), message=message), + ) + + +def _parse_decimal(value: str, *, message: str) -> int: + if len(value) > 10: + raise SlurmCommandOutputError(message) + try: + parsed = int(value) + except ValueError as error: + raise SlurmCommandOutputError(message) from error + if parsed > _MAX_SLURM_INTEGER: + raise SlurmCommandOutputError(message) + return parsed + + +def _reject_duplicate( + job_identity: SlurmObservedJobIdentity, + identities: set[SlurmObservedJobIdentity], + *, + command: str, + line_number: int, +) -> None: + if job_identity in identities: + raise SlurmCommandOutputError(f"{command} line {line_number} duplicates a job or array-task ID") + identities.add(job_identity) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py new file mode 100644 index 000000000..203cfe14a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal safe rendering for thin deterministic Slurm batch entrypoints.""" + +from __future__ import annotations + +import posixpath +import re +from dataclasses import dataclass + +from data_designer.slurm.launcher.errors import SlurmBatchRenderError +from data_designer.slurm.planning import ResolvedSlurmRunPlan + +_DIRECTIVE_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9-]*$") +_DIRECTIVE_TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/,%+-]*$") + + +@dataclass(frozen=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 SlurmBatchRenderError("batch directive name is invalid") + if type(self.value) is not str: + raise SlurmBatchRenderError("batch directive value must be text") + _reject_control_characters(self.value, field_name=f"--{self.name} value") + value = self.value if _DIRECTIVE_TOKEN_PATTERN.fullmatch(self.value) else _quote_sbatch_option_value(self.value) + return f"#SBATCH --{self.name}={value}" + + +def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: + """Render a resolved generation plan as one thin deterministic entrypoint.""" + if type(attempt_ordinal) is not int or attempt_ordinal <= 0: + raise SlurmBatchRenderError("attempt_ordinal must be a positive integer") + + run_root = posixpath.dirname(plan.authored_config.path) + plan_path = posixpath.join(run_root, "resolved-plan.json") + 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_shell_value(plan.runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={_quote_shell_value(plan.runtime_bundle.sha256)} +readonly DD_PLAN={_quote_shell_value(plan_path)} +readonly DD_PLAN_SHA256={_quote_shell_value(plan.compute_sha256())} +readonly DD_RUN_ROOT={_quote_shell_value(run_root)} +readonly DD_ATTEMPT_ORDINAL={_quote_shell_value(attempt)} + +verify_sha256() {{ + local actual_sha256 + 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}%{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: + # TODO(#875): Remove this defense once config and plan validation make this state unrepresentable. + raise SlurmBatchRenderError("mem_per_gpu requires GRES GPU request mode") + if profile.scheduler.mem_per_gpu is not None: + values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) + if plan.submission.comment is not None: + 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_sbatch_option_value(value: str) -> str: + _reject_control_characters(value, field_name="batch directive value") + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped}"' + + +def _quote_shell_value(value: str) -> str: + _reject_control_characters(value, field_name="shell value") + escaped = value.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`") + return f'"{escaped}"' + + +def _reject_control_characters(value: str, *, field_name: str) -> None: + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise SlurmBatchRenderError(f"{field_name} must not contain control characters") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py new file mode 100644 index 000000000..3679d3252 --- /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 + +"""Internal 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..5155ddb02 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -0,0 +1,277 @@ +# 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, FakeSlurmJob, FakeSlurmRunner + +from data_designer.slurm.launcher.client import SlurmCommandClient, SlurmExecutables +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + + +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].job_identity == 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(SlurmCommandOutputError, match="unrequested"): + client.query_queue((4101,)) + + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stdout="9999_0|FAILED|1:0\n")) + with pytest.raises(SlurmCommandOutputError, match="unrequested"): + client.query_accounting((SchedulerIdentity(array_job_id=4101, array_task_id=0),)) + + +def test_client_observes_regular_cpu_job() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + client = SlurmCommandClient(runner) + + submission = client.submit("image-build.sbatch") + queue = client.query_queue((5101,)) + runner.set_job_state(5101, queue_state=None, accounting_state="COMPLETED") + accounting = client.query_accounting((5101,)) + + assert submission.job_id == 5101 + assert queue[0].job_identity == 5101 + assert queue[0].state is SchedulerState.PENDING + assert accounting[0].job_identity == 5101 + assert accounting[0].state is SchedulerState.COMPLETED + + +def test_client_ignores_only_array_parent_observation_for_exact_task() -> None: + runner = FakeSlurmRunner() + runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n4101_0|COMPLETED|0:0\n")) + client = SlurmCommandClient(runner) + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + + records = client.query_accounting((task,)) + + assert tuple(record.job_identity for record in records) == (task,) + + +def test_client_keeps_explicitly_selected_parent_and_task_observations() -> None: + runner = FakeSlurmRunner() + runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n4101_0|COMPLETED|0:0\n")) + client = SlurmCommandClient(runner) + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + + records = client.query_accounting((4101, task)) + + assert tuple(record.job_identity for record in records) == (4101, task) + + +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) + + +@pytest.mark.parametrize("returncode", (0, 2)) +def test_client_rejects_non_text_runner_output(returncode: int) -> None: + client = SlurmCommandClient(_NonTextRunner(returncode)) + + with pytest.raises(SlurmCommandError, match="malformed process result"): + 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 __init__(self, returncode: int) -> None: + self._returncode = returncode + + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + completed = subprocess.CompletedProcess(command, self._returncode, stdout="ok", stderr="") + if self._returncode: + completed.stderr = b"not text" # type: ignore[assignment] + else: + completed.stdout = b"not text" # type: ignore[assignment] + return completed diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py new file mode 100644 index 000000000..7509f93c2 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -0,0 +1,204 @@ +# 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.errors import SlurmCommandOutputError +from data_designer.slurm.launcher.models import SlurmQueueEntry +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 + + +def test_parse_submission_accepts_non_federated_parsable_sbatch_output() -> None: + submission = parse_submission("4101\n") + + assert submission.job_id == 4101 + + +def test_parse_submission_rejects_federated_receipts() -> None: + with pytest.raises(SlurmCommandOutputError, match="federated"): + parse_submission("4101;primary\n") + + +@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(SlurmCommandOutputError, match="invalid"): + parse_submission(output) + + +def test_parse_submission_enforces_slurm_job_id_width() -> None: + assert parse_submission(str((1 << 32) - 1)).job_id == (1 << 32) - 1 + + with pytest.raises(SlurmCommandOutputError, 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_entry(0, SchedulerState.PENDING), + _make_queue_entry(1, SchedulerState.RUNNING), + ) + + +def test_parse_queue_normalizes_regular_jobs() -> None: + records = parse_queue("5101|RUNNING\n") + + assert records == (SlurmQueueEntry(job_identity=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_parent_and_terminal_array_task_rows() -> None: + output = "4101|RUNNING|0:0\n" + (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() + + records = parse_accounting(output) + + assert tuple(record.state for record in records) == ( + SchedulerState.RUNNING, + SchedulerState.TIMED_OUT, + SchedulerState.NODE_FAILED, + SchedulerState.PREEMPTED, + SchedulerState.REQUEUED, + SchedulerState.OUT_OF_MEMORY, + SchedulerState.CANCELLED, + ) + assert records[1].process_exit_code.exit_status == 0 + assert records[1].process_exit_code.termination_signal == 125 + + +def test_parse_accounting_normalizes_regular_jobs() -> None: + records = parse_accounting("5101|COMPLETED|0:0\n") + + assert len(records) == 1 + assert records[0].job_identity == 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(SlurmCommandOutputError, 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(SlurmCommandOutputError, 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(SlurmCommandOutputError, match="invalid GPU resource"): + parse_gpu_counts(output) + + +@pytest.mark.parametrize("state", ("", "CANCELLED by root")) +def test_parse_state_rejects_invalid_spellings(state: str) -> None: + with pytest.raises(SlurmCommandOutputError): + parse_state(state) + + +def _make_queue_entry(array_task_id: int, state: SchedulerState) -> SlurmQueueEntry: + return SlurmQueueEntry( + job_identity=SchedulerIdentity(array_job_id=4101, array_task_id=array_task_id), + state=state, + ) diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py new file mode 100644 index 000000000..c80c95f5f --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import os +import shlex +import subprocess +import tarfile +from pathlib import Path +from typing import Literal + +import pytest + +from data_designer.slurm.config import SchedulerProfile, injected_profile +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.launcher.errors import SlurmBatchRenderError +from data_designer.slurm.launcher.renderer import render_generation_attempt_script +from data_designer.slurm.planning import ResolvedSlurmRunPlan, ResolvedSubmission + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "rendered" + + +@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_generation_attempt_script(plan, attempt_ordinal=1) == (GOLDEN_DIRECTORY / fixture_name).read_text() + + +def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fields( + 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_generation_attempt_script(plan, attempt_ordinal=1) + + 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_generation_attempt_script(plan, attempt_ordinal=1) + + +@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_generation_attempt_script(plan, attempt_ordinal=1) + + +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(SlurmBatchRenderError, match="requires GRES"): + render_generation_attempt_script(plan, attempt_ordinal=1) + + +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_generation_attempt_script(plan, attempt_ordinal=1) + + 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_generation_attempt_script(plan, attempt_ordinal=1) + + comment_lines = [line for line in script.splitlines() if line.startswith("#SBATCH --comment=")] + assert len(comment_lines) == 1 + assert shlex.split(comment_lines[0].removeprefix("#SBATCH ")) == [f"--comment={comment}"] + assert "\\$" not in comment_lines[0] + assert "\\`" not in comment_lines[0] + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +@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(SlurmBatchRenderError, match="positive integer"): + render_generation_attempt_script(single_node_plan, attempt_ordinal=attempt_ordinal) # type: ignore[arg-type] + + +def test_renderer_rejects_control_characters_from_unvalidated_plan_copies( + 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(SlurmBatchRenderError, match="control characters"): + render_generation_attempt_script(plan, attempt_ordinal=1) + + +def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) -> None: + script = render_generation_attempt_script(single_node_plan, attempt_ordinal=12) + + assert script.count("dd_slurm_run_allocation") == 1 + assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script + assert len(script.splitlines()) <= 42 + assert script.endswith("\n") + + +def test_rendered_script_verifies_exact_persisted_plan_bytes_before_sourcing_runtime( + single_node_plan: ResolvedSlurmRunPlan, + tmp_path: Path, +) -> None: + run_root = tmp_path / "run" + run_root.mkdir() + captured_plan_path = tmp_path / "captured-plan.json" + entrypoint_path = tmp_path / "entrypoint.sh" + entrypoint_path.write_text( + f'dd_slurm_run_allocation() {{\n cp -- "$1" {shlex.quote(str(captured_plan_path))}\n}}\n' + ) + runtime_archive_path = tmp_path / "runtime.tar.gz" + with tarfile.open(runtime_archive_path, mode="w:gz") as runtime_archive: + runtime_archive.add(entrypoint_path, arcname="entrypoint.sh") + + plan = single_node_plan.model_copy( + update={ + "authored_config": ArtifactReference( + path=str(run_root / "authored-config.json"), + sha256="a" * 64, + ), + "runtime_bundle": ArtifactReference( + path=str(runtime_archive_path), + sha256=hashlib.sha256(runtime_archive_path.read_bytes()).hexdigest(), + ), + } + ) + plan_path = run_root / "resolved-plan.json" + serialized_plan = plan.serialize_json() + serialized_plan_bytes = serialized_plan.encode() + plan_path.write_bytes(serialized_plan_bytes) + script = render_generation_attempt_script(plan, attempt_ordinal=1) + environment = {**os.environ, "SLURM_ARRAY_TASK_ID": "0"} + + valid = subprocess.run( + ("bash",), + input=script, + env=environment, + text=True, + check=False, + ) + + assert valid.returncode == 0 + assert captured_plan_path.read_bytes() == serialized_plan_bytes + + captured_plan_path.unlink() + plan_path.write_bytes(serialized_plan_bytes + b" ") + tampered = subprocess.run( + ("bash",), + input=script, + env=environment, + text=True, + check=False, + ) + + assert tampered.returncode != 0 + assert not captured_plan_path.exists() diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py new file mode 100644 index 000000000..6bab64a83 --- /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.runner 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/__init__.py b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py index 0b85d89f3..60d45aa06 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py @@ -12,6 +12,7 @@ from slurm_test_fakes.slurm import ( FakeCommandResponse, FakeSlurmArray, + FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask, ) @@ -25,6 +26,7 @@ "FakeLogicalEndpoint", "FakeServingState", "FakeSlurmArray", + "FakeSlurmJob", "FakeSlurmRunner", "FakeSlurmTask", "FakeVllmBackend", diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 0cf7a8e85..d934fd637 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -3,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 60751e07c..93156136d 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -3,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..22da3c51f 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) @@ -27,6 +33,20 @@ class FakeCommandResponse: returncode: int = 0 +@dataclass +class FakeSlurmJob: + """Mutable scheduler views for one ordinary batch-job identity.""" + + job_id: int + queue_state: str | None = "PENDING" + accounting_state: str | None = None + exit_code: str = "0:0" + + def __post_init__(self) -> None: + if type(self.job_id) is not int or self.job_id <= 0: + raise ValueError("fake Slurm job IDs must be positive integers") + + @dataclass class FakeSlurmTask: """Mutable scheduler views for one canonical array-task identity.""" @@ -60,16 +80,25 @@ def array_job_id(self) -> int: class FakeSlurmRunner: - """Stateful fake that copies its configured arrays before exposing Slurm commands.""" + """Stateful fake that copies configured jobs before exposing Slurm commands.""" def __init__( self, arrays: Iterable[FakeSlurmArray] = (), *, + jobs: Iterable[FakeSlurmJob] = (), sinfo_responses: Mapping[tuple[str, ...], FakeCommandResponse] | None = None, ) -> None: - self._pending_arrays = deque(copy.deepcopy(tuple(arrays))) + submissions = (*arrays, *jobs) + job_ids = tuple( + submission.array_job_id if isinstance(submission, FakeSlurmArray) else submission.job_id + for submission in submissions + ) + if len(job_ids) != len(set(job_ids)): + raise ValueError("fake Slurm submissions must have unique job IDs") + self._pending_submissions = deque(copy.deepcopy(submissions)) self._submitted_arrays: dict[int, FakeSlurmArray] = {} + self._submitted_jobs: dict[int, FakeSlurmJob] = {} self._scripted_responses: dict[str, deque[FakeCommandResponse]] = {} self._sinfo_responses = dict(sinfo_responses or {}) self.calls: list[tuple[str, ...]] = [] @@ -120,6 +149,20 @@ def set_task_state( task.accounting_state = accounting_state task.exit_code = exit_code + def set_job_state( + self, + job_id: int, + *, + queue_state: str | None, + accounting_state: str | None, + exit_code: str = "0:0", + ) -> None: + """Set the independently observable states for one ordinary job.""" + job = self._submitted_jobs[job_id] + job.queue_state = queue_state + job.accounting_state = accounting_state + job.exit_code = exit_code + def assert_scripts_consumed(self) -> None: """Assert that no scripted command response remains.""" remaining = sum(len(responses) for responses in self._scripted_responses.values()) @@ -141,38 +184,58 @@ def _dispatch(self, command_name: str, argv: tuple[str, ...]) -> FakeCommandResp return handler(argv) def _run_sbatch(self, argv: tuple[str, ...]) -> FakeCommandResponse: - if not self._pending_arrays: + if not self._pending_submissions: return FakeCommandResponse(stderr="no scripted submission\n", returncode=1) - array = self._pending_arrays.popleft() - self._submitted_arrays[array.array_job_id] = array + submission = self._pending_submissions.popleft() + if isinstance(submission, FakeSlurmArray): + job_id = submission.array_job_id + self._submitted_arrays[job_id] = submission + else: + job_id = submission.job_id + self._submitted_jobs[job_id] = submission if "--parsable" in argv[1:]: - return FakeCommandResponse(stdout=f"{array.array_job_id}\n") - return FakeCommandResponse(stdout=f"Submitted batch job {array.array_job_id}\n") + return FakeCommandResponse(stdout=f"{job_id}\n") + return FakeCommandResponse(stdout=f"Submitted batch job {job_id}\n") def _run_squeue(self, argv: tuple[str, ...]) -> FakeCommandResponse: self._require_arguments(argv, _SQUEUE_REQUIRED_ARGUMENTS) rows = [ + f"{job.job_id}|{job.queue_state}" + for job in self._selected_submitted_jobs(argv) + if job.queue_state is not None + ] + [ f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.queue_state}" for task in self._selected_submitted_tasks(argv) if task.queue_state is not None ] - return FakeCommandResponse(stdout="".join(f"{row}\n" for row in rows)) + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in sorted(rows))) def _run_sacct(self, argv: tuple[str, ...]) -> FakeCommandResponse: self._require_arguments(argv, _SACCT_REQUIRED_ARGUMENTS) rows = [ + f"{job.job_id}|{job.accounting_state}|{job.exit_code}" + for job in self._selected_submitted_jobs(argv) + if job.accounting_state is not None + ] + [ (f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.accounting_state}|{task.exit_code}") for task in self._selected_submitted_tasks(argv) if task.accounting_state is not None ] - return FakeCommandResponse(stdout="".join(f"{row}\n" for row in rows)) + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in sorted(rows))) def _run_scancel(self, argv: tuple[str, ...]) -> FakeCommandResponse: targets = tuple(argument for argument in argv[1:] if not argument.startswith("-")) if len(targets) != 1: return FakeCommandResponse(stderr="expected one cancellation target\n", returncode=1) + target = targets[0] + if "_" not in target and target.isdecimal() and int(target) in self._submitted_jobs: + job = self._submitted_jobs[int(target)] + job.queue_state = None + job.accounting_state = "CANCELLED" + job.exit_code = "0:15" + return FakeCommandResponse() try: - tasks = self._tasks_for_target(targets[0]) + tasks = self._tasks_for_target(target) except (KeyError, ValueError): return FakeCommandResponse(stderr="unknown cancellation target\n", returncode=1) for task in tasks: @@ -215,22 +278,12 @@ def _sorted_submitted_tasks(self) -> list[FakeSlurmTask]: ) def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask]: - selectors: list[str] = [] - for index, argument in enumerate(argv[1:]): - if argument in {"-j", "--jobs"}: - try: - selectors.extend(argv[index + 2].split(",")) - except IndexError: - raise AssertionError(f"missing job selector in {argv!r}") from None - elif argument.startswith("--jobs="): - selectors.extend(argument.partition("=")[2].split(",")) + selectors = self._job_selectors(argv) if not selectors: return self._sorted_submitted_tasks() selected: dict[SchedulerIdentity, FakeSlurmTask] = {} for selector in selectors: - if _JOB_SELECTOR_PATTERN.fullmatch(selector) is None: - raise AssertionError(f"malformed job selector {selector!r} in {argv!r}") try: tasks = self._tasks_for_target(selector) except KeyError: @@ -240,3 +293,34 @@ def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask selected.values(), key=lambda task: (task.scheduler.array_job_id, task.scheduler.array_task_id), ) + + def _selected_submitted_jobs(self, argv: tuple[str, ...]) -> list[FakeSlurmJob]: + selectors = self._job_selectors(argv) + if not selectors: + return sorted(self._submitted_jobs.values(), key=lambda job: job.job_id) + return sorted( + ( + self._submitted_jobs[int(selector)] + for selector in selectors + if "_" not in selector and int(selector) in self._submitted_jobs + ), + key=lambda job: job.job_id, + ) + + @staticmethod + def _job_selectors(argv: tuple[str, ...]) -> tuple[str, ...]: + selectors: list[str] = [] + for index, argument in enumerate(argv[1:]): + if argument in {"-j", "--jobs"}: + try: + selectors.extend(argv[index + 2].split(",")) + except IndexError: + raise AssertionError(f"missing job selector in {argv!r}") from None + elif argument.startswith("--jobs="): + selectors.extend(argument.partition("=")[2].split(",")) + if not selectors: + return () + for selector in selectors: + if _JOB_SELECTOR_PATTERN.fullmatch(selector) is None: + raise AssertionError(f"malformed job selector {selector!r} in {argv!r}") + return tuple(selectors) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 1d793f500..a8c29e5ef 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="cdb5204e8ce7a7affda99ea9c121c6b60a16bf42c510c8eb6bc57661969620b9", + expected_fixture_sha256="8ddf07c38a825a7487c487fddfe051f0a1940f63063725b54b32bfe4c03fd9ca", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="c6708cdbc0a03e095062c0642cfd141f066153b958b7ad3dec779afd6414fa34", + expected_fixture_sha256="17a4c2e16189d22dfdb6885bf76264844ad3168dea0cf94aef70948d5ab2e6b7", ) @@ -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: + 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..c47b21432 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -9,11 +9,17 @@ import pytest from data_designer.slurm.state import SchedulerIdentity -from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" -SQUEUE_ARGUMENTS = ("--noheader", "--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: @@ -135,6 +141,25 @@ def test_fake_slurm_runner_models_cancellation( assert f"4101_{task_id}|CANCELLED|0:15" in accounting +def test_fake_slurm_runner_models_ordinary_job_lifecycle_and_cancellation() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + + assert runner.run(("sbatch", "--parsable", "image-build.sbatch"), check=True).stdout == "5101\n" + assert runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs=5101"), check=True).stdout == "5101|PENDING\n" + + assert runner.run(("scancel", "5101"), check=True).returncode == 0 + assert runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs=5101"), check=True).stdout == "" + assert runner.run(("sacct", *SACCT_ARGUMENTS, "--jobs=5101"), check=True).stdout == "5101|CANCELLED|0:15\n" + + +def test_fake_slurm_runner_rejects_duplicate_job_ids() -> None: + ordinary_job = FakeSlurmJob(job_id=4101) + array = FakeSlurmArray(tasks=(FakeSlurmTask(scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)),)) + + with pytest.raises(ValueError, match="unique job IDs"): + FakeSlurmRunner((array,), jobs=(ordinary_job,)) + + def test_fake_slurm_runner_supports_malformed_output_and_command_failures( fake_slurm_runner: FakeSlurmRunner, ) -> None: @@ -174,7 +199,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(