feat: add Slurm retry and deterministic collection

packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py

@@ -5,23 +5,28 @@
 
 from __future__ import annotations
 
+import os
 import re
 import subprocess
 import unicodedata
 from collections.abc import Sequence
 from dataclasses import dataclass
+from datetime import datetime, timedelta
 from pathlib import Path
 
 from data_designer.slurm.contracts import Identifier
-from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError
+from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError
 from data_designer.slurm.launcher.models import (
     SlurmAccountingEntry,
     SlurmJobSubmissionReceipt,
+    SlurmNamedJobEntry,
     SlurmQueueEntry,
+    SlurmSubmissionMatch,
 )
 from data_designer.slurm.launcher.parsing import (
     parse_accounting,
     parse_gpu_counts,
+    parse_named_jobs,
     parse_queue,
     parse_submission,
 )
@@ -77,11 +82,20 @@ def submit_script(self, script: str) -> SlurmJobSubmissionReceipt:
         """Submit verified batch-script text through standard input."""
         if type(script) is not str or not script or "\0" in script:
             raise ValueError("batch script text must be non-empty UTF-8 text without NUL")
-        output = self._run(
-            (self._executables.sbatch, "--parsable", "--export=NIL"),
-            input_text=script,
-        )
-        return parse_submission(output)
+        try:
+            output = self._run(
+                (self._executables.sbatch, "--parsable", "--export=NIL"),
+                input_text=script,
+            )
+        except SlurmCommandError as error:
+            raise SlurmSubmissionError(
+                str(error),
+                may_have_succeeded=error.command_may_have_completed,
+            ) from error
+        try:
+            return parse_submission(output)
+        except SlurmCommandOutputError as error:
+            raise SlurmSubmissionError(str(error), may_have_succeeded=True) from error
 
     def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]:
         """Return normalized active-queue rows for explicit managed jobs."""
@@ -127,6 +141,47 @@ def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[S
         )
         return tuple(entry for entry in entries if entry.job_identity not in ignored)
 
+    def query_submissions_by_name(
+        self,
+        job_name: Identifier,
+        *,
+        submitted_after: datetime,
+    ) -> tuple[SlurmSubmissionMatch, ...]:
+        """Return current-user allocations matching one exact recovery name."""
+        if type(job_name) is not str or _IDENTIFIER_PATTERN.fullmatch(job_name) is None:
+            raise ValueError("Slurm job name must be a valid identifier")
+        accounting_start = _format_accounting_start(submitted_after)
+        queue_output = self._run(
+            (
+                self._executables.squeue,
+                "--noheader",
+                "--array",
+                "--format=%i|%.128j",
+                "--me",
+                f"--name={job_name}",
+            )
+        )
+        accounting_output = self._run(
+            (
+                self._executables.sacct,
+                "--noheader",
+                "--array",
+                "--allocations",
+                "--parsable2",
+                "--format=JobIDRaw,JobName%128",
+                f"--uid={os.getuid()}",
+                f"--starttime={accounting_start}",
+                f"--name={job_name}",
+            )
+        )
+        entries = (
+            *parse_named_jobs(queue_output, command="squeue"),
+            *parse_named_jobs(accounting_output, command="sacct"),
+        )
+        if any(entry.job_name != job_name for entry in entries):
+            raise SlurmCommandOutputError("scheduler returned a job outside the requested exact name")
+        return _merge_submission_matches(entries)
+
     def cancel(self, selector: SchedulerJobIdentity) -> None:
         """Cancel one managed Slurm job, array, or array task."""
         self._run((self._executables.scancel, _format_selector(selector)))
@@ -146,13 +201,26 @@ def _run(self, command: Sequence[str], *, input_text: str | None = None) -> str:
             completed = (
                 self._runner.run(command) if input_text is None else self._runner.run(command, input_text=input_text)
             )
-        except (OSError, subprocess.SubprocessError) as error:
+        except subprocess.TimeoutExpired as error:
+            raise SlurmCommandError(
+                f"{command_name} could not be executed: {_format_error_detail(error)}",
+                command_may_have_completed=True,
+            ) from error
+        except OSError as error:
             raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") from error
+        except subprocess.SubprocessError as error:
+            raise SlurmCommandError(
+                f"{command_name} could not be executed: {_format_error_detail(error)}",
+                command_may_have_completed=True,
+            ) 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")
+            raise SlurmCommandError(
+                f"{command_name} returned a malformed process result",
+                command_may_have_completed=True,
+            )
         if returncode:
             detail = _normalize_bounded_text(stderr) or "no diagnostic output"
             raise SlurmCommandError(f"{command_name} failed with exit code {returncode}: {detail}")
@@ -222,3 +290,26 @@ 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__
+
+
+def _format_accounting_start(value: datetime) -> str:
+    if not isinstance(value, datetime) or value.tzinfo is None or value.utcoffset() is None:
+        raise ValueError("submission lookup timestamp must be timezone-aware")
+    return (value - timedelta(minutes=1)).astimezone().strftime("%Y-%m-%dT%H:%M:%S")
+
+
+def _merge_submission_matches(entries: Sequence[SlurmNamedJobEntry]) -> tuple[SlurmSubmissionMatch, ...]:
+    grouped: dict[tuple[int, str], set[int]] = {}
+    ordinary: set[tuple[int, str]] = set()
+    for entry in entries:
+        key = (entry.job_id, entry.job_name)
+        if entry.array_task_id is None:
+            ordinary.add(key)
+        else:
+            grouped.setdefault(key, set()).add(entry.array_task_id)
+    matches: list[SlurmSubmissionMatch] = []
+    for job_id, job_name in sorted(ordinary | set(grouped)):
+        key = (job_id, job_name)
+        task_ids = tuple(sorted(grouped[key])) if key in grouped else None
+        matches.append(SlurmSubmissionMatch(job_id=job_id, job_name=job_name, array_task_ids=task_ids))
+    return tuple(matches)

packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py

@@ -0,0 +1,103 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Safe rendering for zero-GPU CPU collection jobs."""
+
+from __future__ import annotations
+
+import posixpath
+from pathlib import PurePosixPath
+
+from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives
+from data_designer.slurm.launcher.errors import SlurmBatchRenderError
+from data_designer.slurm.planning import ResolvedSlurmRunPlan
+from data_designer.slurm.state.destinations import CollectionDestination
+from data_designer.slurm.state.outputs import CollectionPlan
+
+
+def render_collection_script(
+    resolved_plan: ResolvedSlurmRunPlan,
+    collection_plan: CollectionPlan,
+    destination: CollectionDestination,
+) -> str:
+    """Render one CPU-only job that invokes the allocation-gated collection worker."""
+    if collection_plan.run_id != resolved_plan.run_id:
+        raise SlurmBatchRenderError("collection run identity does not match the resolved plan")
+    if collection_plan.host_destination != destination.host_path:
+        raise SlurmBatchRenderError("collection host destination does not match its resolved mount")
+    if collection_plan.container_destination != destination.container_path:
+        raise SlurmBatchRenderError("collection container destination does not match its resolved mount")
+
+    collection_root = posixpath.join(
+        posixpath.dirname(resolved_plan.authored_config.path),
+        "collections",
+        collection_plan.collection_id,
+    )
+    collection_plan_path = posixpath.join(collection_root, "plan.json")
+    directives = render_batch_directives(
+        (
+            ("job-name", collection_plan.submission_job_name),
+            ("account", resolved_plan.submission.account),
+            ("partition", resolved_plan.submission.partition),
+            ("nodes", "1"),
+            ("ntasks", "1"),
+            ("cpus-per-task", str(resolved_plan.client.authored.cpus)),
+            ("time", resolved_plan.submission.time_limit),
+            ("chdir", collection_root),
+            ("output", f"{collection_root}/slurm-%j.out"),
+            ("error", f"{collection_root}/slurm-%j.err"),
+        )
+    )
+    workspace_root = resolved_plan.selected_profile.profile.workspace_root
+    state_mount = f"{workspace_root}:{workspace_root}"
+    output_mount = f"{destination.mount.source}:{destination.mount.target}"
+    mount_arguments = _render_mount_arguments(
+        ("DD_STATE_MOUNT", workspace_root, workspace_root),
+        ("DD_OUTPUT_MOUNT", destination.mount.source, destination.mount.target),
+    )
+    return f"""#!/usr/bin/env bash
+{directives}
+set -Eeuo pipefail
+export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
+
+readonly DD_CLIENT_IMAGE={quote_shell_value(resolved_plan.client.image.path)}
+readonly DD_CLIENT_IMAGE_SHA256={quote_shell_value(resolved_plan.client.image.sha256)}
+readonly DD_COLLECTION_PLAN={quote_shell_value(collection_plan_path)}
+readonly DD_COLLECTION_PLAN_SHA256={quote_shell_value(collection_plan.compute_sha256())}
+readonly DD_WORKSPACE_ROOT={quote_shell_value(resolved_plan.selected_profile.profile.workspace_root)}
+readonly DD_RUN_ID={quote_shell_value(resolved_plan.run_id)}
+readonly DD_COLLECTION_ID={quote_shell_value(collection_plan.collection_id)}
+readonly DD_STATE_MOUNT={quote_shell_value(state_mount)}
+readonly DD_OUTPUT_MOUNT={quote_shell_value(output_mount)}
+
+verify_sha256() {{
+    local actual_sha256
+    actual_sha256="$(sha256sum < "$2")"
+    [[ "${{actual_sha256%% *}}" == "$1" ]]
+}}
+
+verify_sha256 "${{DD_CLIENT_IMAGE_SHA256}}" "${{DD_CLIENT_IMAGE}}"
+verify_sha256 "${{DD_COLLECTION_PLAN_SHA256}}" "${{DD_COLLECTION_PLAN}}"
+DD_ENROOT_MOUNTS=({mount_arguments})
+readonly DD_ENROOT_MOUNTS
+exec enroot start --root "${{DD_ENROOT_MOUNTS[@]}}" "${{DD_CLIENT_IMAGE}}" \
+    python -m data_designer.slurm.state.collection_worker \
+    --workspace-root "${{DD_WORKSPACE_ROOT}}" --run-id "${{DD_RUN_ID}}" --collection-id "${{DD_COLLECTION_ID}}"
+"""
+
+
+def _render_mount_arguments(*mounts: tuple[str, str, str]) -> str:
+    unique: dict[str, tuple[str, str, str]] = {}
+    targets: dict[str, str] = {}
+    for variable, source, target in mounts:
+        mount = f"{source}:{target}"
+        existing_source = targets.get(target)
+        if existing_source is not None and existing_source != source:
+            raise SlurmBatchRenderError("collection state and output mounts cannot share a target")
+        targets[target] = source
+        unique.setdefault(mount, (variable, source, target))
+    ordered = sorted(unique.values(), key=lambda item: len(PurePosixPath(item[2]).parts))
+    return " ".join(f'--mount "${{{variable}}}"' for variable, _, _ in ordered)
+
+
+__all__ = ["render_collection_script"]

packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py

@@ -13,6 +13,18 @@ class SlurmLauncherError(RuntimeError):
 class SlurmCommandError(SlurmLauncherError):
     """A Slurm command could not be executed successfully."""
 
+    def __init__(self, message: str, *, command_may_have_completed: bool = False) -> None:
+        super().__init__(message)
+        self.command_may_have_completed = command_may_have_completed
+
+
+class SlurmSubmissionError(SlurmLauncherError):
+    """An sbatch submission failed with an explicit ambiguity classification."""
+
+    def __init__(self, message: str, *, may_have_succeeded: bool) -> None:
+        super().__init__(message)
+        self.may_have_succeeded = may_have_succeeded
+
 
 class SlurmCommandOutputError(SlurmLauncherError, ValueError):
     """A Slurm command returned output that violates its requested format."""

packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py

@@ -40,3 +40,21 @@ class SlurmAccountingEntry:
     job_identity: SchedulerJobIdentity
     state: SchedulerState
     process_exit_code: SlurmProcessExitCode
+
+
+@dataclass(frozen=True)
+class SlurmNamedJobEntry:
+    """One scheduler allocation found through an exact job-name lookup."""
+
+    job_id: int
+    array_task_id: int | None
+    job_name: str
+
+
+@dataclass(frozen=True)
+class SlurmSubmissionMatch:
+    """One named scheduler allocation with its ordinary or array shape."""
+
+    job_id: int
+    job_name: str
+    array_task_ids: tuple[int, ...] | None

packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py

@@ -6,11 +6,13 @@
 from __future__ import annotations
 
 import re
+from typing import Literal
 
 from data_designer.slurm.launcher.errors import SlurmCommandOutputError
 from data_designer.slurm.launcher.models import (
     SlurmAccountingEntry,
     SlurmJobSubmissionReceipt,
+    SlurmNamedJobEntry,
     SlurmProcessExitCode,
     SlurmQueueEntry,
 )
@@ -101,6 +103,25 @@ def parse_accounting(output: str) -> tuple[SlurmAccountingEntry, ...]:
     return tuple(entries)
 
 
+def parse_named_jobs(output: str, *, command: Literal["sacct", "squeue"]) -> tuple[SlurmNamedJobEntry, ...]:
+    """Parse scheduler allocations returned for an exact job-name lookup."""
+    entries: list[SlurmNamedJobEntry] = []
+    identities: set[tuple[int, int | None, str]] = set()
+    for line_number, line in _collect_nonempty_lines(output):
+        fields = tuple(field.strip() for field in line.split("|"))
+        if len(fields) != 2 or not fields[1]:
+            raise SlurmCommandOutputError(f"{command} line {line_number} must contain a job ID and name")
+        identity = _parse_job_identity(fields[0], command=command, line_number=line_number)
+        job_id = identity.array_job_id if isinstance(identity, SchedulerIdentity) else identity
+        array_task_id = identity.array_task_id if isinstance(identity, SchedulerIdentity) else None
+        key = (job_id, array_task_id, fields[1])
+        if key in identities:
+            continue
+        identities.add(key)
+        entries.append(SlurmNamedJobEntry(job_id=job_id, array_task_id=array_task_id, job_name=fields[1]))
+    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] = []