From 8bdb26e5350f06baa4c0db7845fd343cb285ab9f Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 2 Sep 2026 18:40:49 -0600 Subject: [PATCH 1/4] feat(slurm): add retry and deterministic collection Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/client.py | 38 +- .../slurm/launcher/collection.py | 103 ++ .../data_designer/slurm/launcher/errors.py | 12 + .../data_designer/slurm/launcher/renderer.py | 111 +- .../src/data_designer/slurm/state/__init__.py | 23 + .../data_designer/slurm/state/artifacts.py | 64 +- .../slurm/state/attempt_identity.py | 91 ++ .../data_designer/slurm/state/collection.py | 330 ++++ .../slurm/state/collection_filesystem.py | 357 ++++ .../slurm/state/collection_inputs.py | 81 + .../slurm/state/collection_merge.py | 338 ++++ .../slurm/state/collection_records.py | 100 ++ .../slurm/state/collection_storage.py | 327 ++++ .../slurm/state/collection_validation.py | 111 ++ .../slurm/state/collection_worker.py | 257 +++ .../data_designer/slurm/state/destinations.py | 88 + .../src/data_designer/slurm/state/outputs.py | 34 + .../src/data_designer/slurm/state/retry.py | 440 +++++ .../slurm/state/retry_records.py | 69 + .../slurm/state/retry_storage.py | 210 +++ .../tests/launcher/test_client.py | 28 +- .../tests/launcher/test_collection.py | 112 ++ .../tests/state/test_retry_collection.py | 1454 +++++++++++++++++ scripts/test_slurm_package_install.py | 8 + 24 files changed, 4769 insertions(+), 17 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/attempt_identity.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/retry.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_collection.py create mode 100644 packages/data-designer-slurm/tests/state/test_retry_collection.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py index f21e7c282..e51ff6edc 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -13,7 +13,7 @@ from pathlib import Path from data_designer.slurm.contracts import Identifier -from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError from data_designer.slurm.launcher.models import ( SlurmAccountingEntry, SlurmJobSubmissionReceipt, @@ -77,11 +77,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.""" @@ -146,13 +155,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}") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py new file mode 100644 index 000000000..e8a77f120 --- /dev/null +++ b/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", f"dd-collect-{resolved_plan.run_id}"), + ("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"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py index 4a611152b..6e483d86e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py @@ -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.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index e076f5cb1..fc78ce06f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -7,9 +7,11 @@ import posixpath +from data_designer.slurm.images.records import validate_enroot_mount_path 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.outputs import RetryPlan def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: @@ -60,17 +62,114 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi """ -def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[tuple[str, str | None], ...]: +def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) -> str: + """Render selected failed shards as one immutable retry array submission.""" + run_root = posixpath.dirname(plan.authored_config.path) + plan_path = posixpath.join(run_root, "resolved-plan.json") + if retry.run_id != plan.run_id: + raise SlurmBatchRenderError("retry run identity does not match the resolved plan") + if retry.resolved_plan.path != plan_path or retry.resolved_plan.sha256 != plan.compute_sha256(): + raise SlurmBatchRenderError("retry does not bind the persisted resolved plan") + try: + validate_enroot_mount_path(plan.selected_profile.profile.workspace_root) + except ValueError as error: + raise SlurmBatchRenderError("retry workspace cannot be represented as a safe Enroot mount") from error + planned_by_id = {shard.shard_id: shard for shard in plan.shards} + for retry_shard in retry.planned_shards: + planned = planned_by_id.get(retry_shard.shard_id) + if planned is None or planned.array_task_index != retry_shard.array_task_index: + raise SlurmBatchRenderError("retry shard does not match the resolved plan") + + array_tasks = ",".join(str(shard.array_task_index) for shard in retry.planned_shards) + if plan.array_tasks.max_concurrent is not None: + array_tasks = f"{array_tasks}%{plan.array_tasks.max_concurrent}" + directives = render_batch_directives(_build_generation_directives(plan, array=array_tasks)) + attempt_cases = "\n".join( + f" {shard.array_task_index}) DD_ATTEMPT_ORDINAL={quote_shell_value(f'{shard.attempt_ordinal:04d}')} ;;" + for shard in retry.planned_shards + ) + 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_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_WORKSPACE_ROOT={quote_shell_value(plan.selected_profile.profile.workspace_root)} +readonly DD_RUN_ID={quote_shell_value(plan.run_id)} +readonly DD_CLIENT_IMAGE={quote_shell_value(plan.client.image.path)} +readonly DD_CLIENT_IMAGE_SHA256={quote_shell_value(plan.client.image.sha256)} +readonly DD_EFFECTIVE_RESUME_MODE={quote_shell_value(retry.effective_resume_mode)} + +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}}" +verify_sha256 "${{DD_CLIENT_IMAGE_SHA256}}" "${{DD_CLIENT_IMAGE}}" +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}}" +case "${{DD_ARRAY_TASK_ID}}" in +{attempt_cases} + *) printf '%s\\n' 'array task is absent from retry plan' >&2; exit 64 ;; +esac +readonly DD_ATTEMPT_ORDINAL +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}}" +readonly DD_ATTEMPT_MANIFEST="${{DD_ATTEMPT_DIR}}/attempt.json" +for ((DD_WAIT_COUNT = 0; DD_WAIT_COUNT < 300; DD_WAIT_COUNT++)); do + [[ -f "${{DD_ATTEMPT_MANIFEST}}" ]] && break + sleep 1 +done +if [[ ! -f "${{DD_ATTEMPT_MANIFEST}}" ]]; then + printf '%s\\n' 'retry attempt state was not published before allocation startup' >&2 + exit 70 +fi +if [[ ! ${{SLURM_ARRAY_JOB_ID:-}} =~ ^[1-9][0-9]*$ ]]; then + printf '%s\\n' 'SLURM_ARRAY_JOB_ID must be a positive integer' >&2 + exit 64 +fi +readonly DD_ARRAY_JOB_ID="${{SLURM_ARRAY_JOB_ID}}" +readonly DD_ATTEMPT_ID="attempt-${{DD_ATTEMPT_ORDINAL}}" +enroot start --root --mount "${{DD_WORKSPACE_ROOT}}:${{DD_WORKSPACE_ROOT}}" "${{DD_CLIENT_IMAGE}}" \\ + python -m data_designer.slurm.state.attempt_identity \\ + --workspace-root "${{DD_WORKSPACE_ROOT}}" --run-id "${{DD_RUN_ID}}" \\ + --shard-id "${{DD_SHARD_ID}}" --attempt-id "${{DD_ATTEMPT_ID}}" \\ + --array-job-id "${{DD_ARRAY_JOB_ID}}" --array-task-id "${{DD_ARRAY_TASK_ID}}" +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, + *, + array: str | None = None, +) -> tuple[tuple[str, str | None], ...]: 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}" + resolved_array = array or "0" + if array is None and plan.array_tasks.count > 1: + resolved_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}" + resolved_array = f"{resolved_array}%{plan.array_tasks.max_concurrent}" values: list[tuple[str, str | None]] = [ ("job-name", plan.submission.job_name), @@ -79,7 +178,7 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[tuple[str, ("nodes", str(node_count)), ("cpus-per-task", str(plan.client.authored.cpus)), ("time", plan.submission.time_limit), - ("array", array), + ("array", resolved_array), ] profile = plan.selected_profile.profile if profile.gpu_request_mode == "gres": diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py index 3c587f615..1fc65422c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -26,6 +26,12 @@ StateRecord, StateValue, ) +from data_designer.slurm.state.collection_records import ( + CollectedOutputFile, + CollectionResult, + CollectionState, + CollectionStatus, +) from data_designer.slurm.state.errors import ( SlurmStateError, StateConflictError, @@ -53,6 +59,8 @@ CandidateOutputManifest, CollectionPlan, CollectionShard, + RetryPlan, + RetryShard, ShardWinner, ) from data_designer.slurm.state.readiness import ( @@ -68,6 +76,7 @@ reconcile_attempt_observation, validate_readiness_transition, ) +from data_designer.slurm.state.retry_records import RetryState, RetryStatus from data_designer.slurm.state.scheduler import ( EffectiveAttemptState, SchedulerObservation, @@ -94,11 +103,15 @@ ) if TYPE_CHECKING: + from data_designer.slurm.state.collection import SlurmCollectionCoordinator # noqa: F401 from data_designer.slurm.state.observer import SlurmStateReconciler # noqa: F401 + from data_designer.slurm.state.retry import SlurmRetryCoordinator # noqa: F401 from data_designer.slurm.state.store import SlurmStateWriter # noqa: F401 _LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "SlurmCollectionCoordinator": ("data_designer.slurm.state.collection", "SlurmCollectionCoordinator"), "SlurmStateReconciler": ("data_designer.slurm.state.observer", "SlurmStateReconciler"), + "SlurmRetryCoordinator": ("data_designer.slurm.state.retry", "SlurmRetryCoordinator"), "SlurmStateWriter": ("data_designer.slurm.state.store", "SlurmStateWriter"), } @@ -115,9 +128,13 @@ "MAXIMUM_CANDIDATE_OUTPUT_FILES", "CandidateOutputFile", "CandidateOutputManifest", + "CollectedOutputFile", "compute_candidate_schema_digest", "CollectionPlan", + "CollectionResult", "CollectionShard", + "CollectionState", + "CollectionStatus", "ContractRecord", "ContractValue", "DeploymentReadiness", @@ -131,6 +148,10 @@ "ReadinessState", "ReasonCode", "RecordRange", + "RetryPlan", + "RetryShard", + "RetryState", + "RetryStatus", "RunManifest", "RunStatus", "ResumeWorkspace", @@ -148,6 +169,8 @@ "ShardId", "ShardWinner", "SlurmStateError", + "SlurmCollectionCoordinator", + "SlurmRetryCoordinator", "SlurmStateReconciler", "SlurmStateWriter", "StateConflictError", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py index 3a3b6b466..e60ee8fbe 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py @@ -12,7 +12,7 @@ from contextlib import ExitStack, contextmanager from dataclasses import dataclass from pathlib import Path, PurePosixPath -from typing import Protocol +from typing import BinaryIO, Protocol import data_designer.lazy_heavy_imports as lazy from data_designer.slurm.state.filesystem import ( @@ -101,6 +101,16 @@ def rebind(self) -> None: output_file.validate_lease() +@dataclass(frozen=True, slots=True) +class CandidateArtifactSnapshot: + """Descriptor-free identity snapshot retained during bounded collection.""" + + record_counts: tuple[int, ...] + dataset_schema_digest: str + _dataset_identity: tuple[int, int] + _bindings: tuple[_ArtifactBinding, ...] + + class CandidateArtifactVerifier: """Verify one manifest-bounded candidate and lease its files through publication.""" @@ -124,6 +134,57 @@ def verify(self, candidate: CandidateOutputManifest) -> Iterator[VerifiedCandida _files=tuple(binding for _, _, binding in metadata), ) + def inspect(self, candidate: CandidateOutputManifest) -> CandidateArtifactSnapshot: + """Inspect one candidate and return identities without retaining descriptors.""" + with self.verify(candidate) as verified: + return CandidateArtifactSnapshot( + record_counts=verified.record_counts, + dataset_schema_digest=verified.dataset_schema_digest, + _dataset_identity=_identity(os.fstat(verified._dataset.descriptor)), + _bindings=verified._bindings, + ) + + def rebind(self, candidate: CandidateOutputManifest, expected: CandidateArtifactSnapshot) -> None: + """Reopen a candidate and require the identities captured before collection.""" + with self.verify(candidate) as current: + actual = CandidateArtifactSnapshot( + record_counts=current.record_counts, + dataset_schema_digest=current.dataset_schema_digest, + _dataset_identity=_identity(os.fstat(current._dataset.descriptor)), + _bindings=current._bindings, + ) + if actual != expected: + raise OSError("candidate paths or metadata changed during collection") + + @contextmanager + def open_output( + self, + candidate: CandidateOutputManifest, + output_file: CandidateOutputFile, + ) -> Iterator[BinaryIO]: + """Yield one digest-verified candidate file through a bounded descriptor.""" + if output_file not in candidate.files: + raise ValueError("candidate output file is not declared by the manifest") + dataset_path = Path(candidate.dataset_path) + parts = PurePosixPath(output_file.relative_path).parts + with open_verified_directory(dataset_path, require_private=True) as dataset_descriptor: + with _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) as ( + parent_descriptor, + parent_path, + _, + ): + display_path = parent_path / parts[-1] + with open_verified_regular_file( + parent_descriptor, + parts[-1], + display_path, + expected_size=output_file.byte_size, + expected_sha256=output_file.sha256, + require_private=False, + ) as descriptor: + with os.fdopen(os.dup(descriptor), "rb") as source: + yield source + def compute_candidate_schema_digest(schema: CandidateSchema) -> str: """Compute the version-1 digest for an attempt-local candidate schema.""" @@ -210,6 +271,7 @@ def _is_safe_file(status: os.stat_result) -> bool: __all__ = [ + "CandidateArtifactSnapshot", "CandidateArtifactVerifier", "CandidateSchema", "SerializedCandidateSchema", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/attempt_identity.py b/packages/data-designer-slurm/src/data_designer/slurm/state/attempt_identity.py new file mode 100644 index 000000000..9e87a5f6d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/attempt_identity.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-local binding of a retry task to its persisted attempt.""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +from pathlib import Path + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import AttemptId, Identifier, ShardId, validate_absolute_path +from data_designer.slurm.state.base import SchedulerIdentity +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError +from data_designer.slurm.state.execution import AttemptLifecycleState +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_SHARD_ID_ADAPTER = TypeAdapter(ShardId) +_ATTEMPT_ID_ADAPTER = TypeAdapter(AttemptId) + + +def require_attempt_scheduler_identity( + workspace_root: str | Path, + run_id: Identifier, + shard_id: ShardId, + attempt_id: AttemptId, + scheduler: SchedulerIdentity, +) -> None: + """Require the complete persisted attempt chain to name this allocation.""" + root, normalized_run_id, normalized_shard_id, normalized_attempt_id = _validate_identity( + workspace_root, + run_id, + shard_id, + attempt_id, + ) + storage = StateStorage(root, normalized_run_id) + reader = StateReader(storage, normalized_run_id) + run, plan, shard = reader.load_shard_context(normalized_shard_id) + attempts = reader.load_validated_shard_attempts(run, plan, shard) + attempt = reader.get_attempt(attempts, normalized_attempt_id) + if attempt.scheduler != scheduler: + raise StateConflictError("retry allocation does not match the persisted attempt scheduler identity") + if attempt.state is not AttemptLifecycleState.SUBMITTED: + raise StateConflictError("retry allocation requires an unstarted persisted attempt") + + +def _validate_identity( + workspace_root: str | Path, + run_id: Identifier, + shard_id: ShardId, + attempt_id: AttemptId, +) -> tuple[Path, Identifier, ShardId, AttemptId]: + try: + root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + normalized_shard_id = _SHARD_ID_ADAPTER.validate_python(shard_id, strict=True) + normalized_attempt_id = _ATTEMPT_ID_ADAPTER.validate_python(attempt_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid retry allocation identity") from error + return Path(root), normalized_run_id, normalized_shard_id, normalized_attempt_id + + +def main(argv: Sequence[str] | None = None) -> int: + """Validate one allocation identity from explicit scheduler arguments.""" + parser = argparse.ArgumentParser() + parser.add_argument("--workspace-root", required=True) + parser.add_argument("--run-id", required=True) + parser.add_argument("--shard-id", required=True) + parser.add_argument("--attempt-id", required=True) + parser.add_argument("--array-job-id", required=True, type=int) + parser.add_argument("--array-task-id", required=True, type=int) + arguments = parser.parse_args(argv) + require_attempt_scheduler_identity( + arguments.workspace_root, + arguments.run_id, + arguments.shard_id, + arguments.attempt_id, + SchedulerIdentity(array_job_id=arguments.array_job_id, array_task_id=arguments.array_task_id), + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = ["main", "require_attempt_scheduler_identity"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py new file mode 100644 index 000000000..1ac68d282 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -0,0 +1,330 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process preparation and reconciliation of CPU collection jobs.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from pathlib import Path +from typing import Protocol + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier, validate_absolute_path +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.collection import render_collection_script +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError +from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt +from data_designer.slurm.state.collection_filesystem import ( + derive_collection_staging_directory, + prepare_collection_destination, + remove_collection_stage, +) +from data_designer.slurm.state.collection_inputs import CollectionInputResolver +from data_designer.slurm.state.collection_records import CollectionResult, CollectionState, CollectionStatus +from data_designer.slurm.state.collection_storage import CollectionStorage +from data_designer.slurm.state.collection_validation import ( + derive_collection_state, + validate_collection_result, + validate_collection_status_transition, +) +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError, StateCorruptionError +from data_designer.slurm.state.observation import SchedulerObservationClient, SchedulerObservationCollector +from data_designer.slurm.state.outputs import CollectionPlan +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.scheduler import SchedulerState +from data_designer.slurm.state.storage import StateStorage + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class CollectionScheduler(SchedulerObservationClient, Protocol): + """Scheduler operations required by collection submission and refresh.""" + + def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + """Submit one rendered CPU collection job.""" + ... + + +class SlurmCollectionCoordinator: + """Persist, submit, and refresh winner-driven collection jobs.""" + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + scheduler: CollectionScheduler | None = None, + ) -> None: + root, normalized_run_id = _validate_location(workspace_root, run_id) + self._scheduler = scheduler if scheduler is not None else SlurmCommandClient() + self._state = StateStorage(root, normalized_run_id) + self._reader = StateReader(self._state, normalized_run_id) + self._collections = CollectionStorage(self._state) + self._inputs = CollectionInputResolver(self._state, self._reader) + self._destinations = CollectionDestinationResolver() + self._collector = SchedulerObservationCollector(self._scheduler) + self._run_id = normalized_run_id + + def submit( + self, + *, + destination: str | Path | None = None, + submitted_at: datetime | None = None, + ) -> CollectionStatus: + """Validate winners and submit a CPU-only collection job.""" + timestamp = datetime.now(timezone.utc) if submitted_at is None else submitted_at + try: + with self._collections.acquire_lock(): + self._collections.discard_incomplete_tail() + current = self._get_current_status() + if current is not None: + current_plan = self._load_bound_plan(current) + if current.state is CollectionState.PREPARED: + raise StateConflictError("previous collection submission has an ambiguous scheduler outcome") + if current.state is not CollectionState.FAILED: + self._validate_existing_destination(current, destination) + return current + resolved_destination = self._destinations.validate_persisted( + self._reader.load_resolved_plan(), + current_plan, + ) + remove_collection_stage( + Path(current_plan.host_destination), + current.staging_directory, + Path(resolved_destination.mount.source), + ) + run, resolved_plan, _ = self._reader.load_context() + resolved_destination = self._destinations.resolve(resolved_plan, destination) + collection_plan = CollectionPlan( + schema_version=1, + collection_id=self._collections.get_next_collection_id(), + run_id=run.run_id, + created_at=timestamp, + resolved_plan=run.resolved_plan, + planned_shards=self._inputs.get_winner_shards(), + host_destination=resolved_destination.host_path, + container_destination=resolved_destination.container_path, + num_partitions=resolved_plan.output.partitions, + ) + self._inputs.resolve(collection_plan) + prepare_collection_destination( + Path(collection_plan.host_destination), + Path(resolved_destination.mount.source), + ) + self._collections.ensure_collection(collection_plan.collection_id) + self._collections.publish_plan(collection_plan) + prepared = CollectionStatus( + schema_version=1, + collection_id=collection_plan.collection_id, + run_id=run.run_id, + collection_plan=self._collections.get_plan_reference(collection_plan), + staging_directory=derive_collection_staging_directory(collection_plan), + revision=1, + updated_at=timestamp, + state=CollectionState.PREPARED, + ) + self._collections.publish_status(prepared) + script = render_collection_script(resolved_plan, collection_plan, resolved_destination) + return self._submit_prepared(collection_plan, prepared, script, timestamp) + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot submit collection for run {self._run_id!r}") from error + + def refresh( + self, + *, + collection_id: Identifier | None = None, + observed_at: datetime | None = None, + ) -> CollectionStatus: + """Reconcile one persisted collection from scheduler and publication evidence.""" + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + try: + with self._collections.acquire_lock(): + selected_id = self._get_selected_collection_id(collection_id) + previous = self._collections.read_status(selected_id) + plan = self._load_bound_plan(previous) + resolved_plan, _ = self._inputs.resolve(plan) + destination = self._destinations.validate_persisted(resolved_plan, plan) + if previous.state is CollectionState.SUCCEEDED: + self._load_valid_result(plan, previous) + return previous + recovered = self._load_optional_result(plan) + if recovered is not None: + return self._publish_succeeded(plan, previous, recovered, timestamp) + if previous.state is CollectionState.FAILED: + remove_collection_stage( + Path(plan.host_destination), + previous.staging_directory, + Path(destination.mount.source), + ) + return previous + if previous.scheduler is None: + raise StateConflictError("collection submission has an ambiguous scheduler outcome") + observations = self._collector.collect( + (previous.scheduler,), + observed_at=timestamp, + previous={previous.scheduler: previous.scheduler_observation}, + ) + observation = observations[0] + state = derive_collection_state(observation.state) + if observation.state is SchedulerState.COMPLETED: + state = CollectionState.FAILED + current = _updated_status( + previous, + revision=previous.revision + 1, + updated_at=timestamp, + state=state, + scheduler_observation=observation, + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + if current.state is CollectionState.FAILED: + remove_collection_stage( + Path(plan.host_destination), + current.staging_directory, + Path(destination.mount.source), + ) + return current + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot refresh collection for run {self._run_id!r}") from error + + def _submit_prepared( + self, + plan: CollectionPlan, + prepared: CollectionStatus, + script: str, + submitted_at: datetime, + ) -> CollectionStatus: + try: + receipt = self._scheduler.submit_script(script) + except SlurmSubmissionError as error: + if not error.may_have_succeeded: + failed = _updated_status( + prepared, + revision=2, + updated_at=submitted_at, + state=CollectionState.FAILED, + ) + validate_collection_status_transition(prepared, failed) + self._collections.replace_status(failed) + raise SlurmStateError(f"cannot submit collection {plan.collection_id!r}") from error + except SlurmLauncherError as error: + raise SlurmStateError(f"cannot submit collection {plan.collection_id!r}") from error + submitted = _updated_status( + prepared, + revision=2, + updated_at=submitted_at, + state=CollectionState.SUBMITTED, + scheduler=receipt.job_id, + ) + validate_collection_status_transition(prepared, submitted) + self._collections.replace_status(submitted) + return submitted + + def _get_current_status(self) -> CollectionStatus | None: + collection_ids = self._collections.list_collection_ids() + return None if not collection_ids else self._collections.read_status(collection_ids[-1]) + + def _validate_existing_destination( + self, + status: CollectionStatus, + requested_destination: str | Path | None, + ) -> None: + plan = self._load_bound_plan(status) + resolved_plan = self._reader.load_resolved_plan() + resolved = self._destinations.validate_persisted(resolved_plan, plan) + requested = self._destinations.resolve(resolved_plan, requested_destination) + if requested != resolved: + raise StateCorruptionError("persisted collection destination does not match the requested destination") + if plan.host_destination != resolved.host_path or plan.container_destination != resolved.container_path: + raise StateCorruptionError("persisted collection destination does not match the resolved plan") + if status.state is CollectionState.SUCCEEDED: + self._load_valid_result(plan, status) + + def _get_selected_collection_id(self, collection_id: Identifier | None) -> Identifier: + collection_ids = self._collections.list_collection_ids() + if not collection_ids: + raise StateConflictError("run has no persisted collection") + selected = collection_ids[-1] if collection_id is None else collection_id + if selected not in collection_ids: + raise StateConflictError("requested collection is not persisted for this run") + return selected + + def _load_optional_result(self, plan: CollectionPlan) -> CollectionResult | None: + try: + return self._load_valid_result(plan) + except FileNotFoundError: + return None + + def _load_valid_result( + self, + plan: CollectionPlan, + status: CollectionStatus | None = None, + ) -> CollectionResult: + self._inputs.resolve(plan) + result = self._collections.read_result(plan) + resolved_plan = self._reader.load_resolved_plan() + validated = validate_collection_result( + plan, + result, + expected_records=resolved_plan.invocation.authored.num_records, + output_format=resolved_plan.output.format, + ) + if status is not None and status.result != self._collections.get_result_reference(plan, validated): + raise StateCorruptionError("collection status does not bind its published result") + self._collections.verify_result_files( + plan, + validated, + Path(plan.host_destination), + verify_digests=False, + ) + return validated + + def _publish_succeeded( + self, + plan: CollectionPlan, + previous: CollectionStatus, + result: CollectionResult, + timestamp: datetime, + ) -> CollectionStatus: + current = _updated_status( + previous, + revision=previous.revision + 1, + updated_at=timestamp, + state=CollectionState.SUCCEEDED, + result=self._collections.get_result_reference(plan, result), + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + return current + + def _load_bound_plan(self, status: CollectionStatus) -> CollectionPlan: + plan = self._collections.read_plan(status.collection_id) + if status.collection_plan != self._collections.get_plan_reference(plan): + raise StateCorruptionError("collection status does not bind its persisted collection plan") + if status.staging_directory != derive_collection_staging_directory(plan): + raise StateCorruptionError("collection status does not bind its exact staging directory") + return plan + + +def _validate_location(workspace_root: str | Path, run_id: Identifier) -> tuple[Path, Identifier]: + try: + root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid persisted collection location") from error + return Path(root), normalized_run_id + + +def _updated_status(previous: CollectionStatus, **updates: object) -> CollectionStatus: + payload = previous.model_dump(mode="python") + payload.update(updates) + return CollectionStatus.model_validate(payload) + + +__all__ = ["CollectionScheduler", "SlurmCollectionCoordinator"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py new file mode 100644 index 000000000..72330c1c0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py @@ -0,0 +1,357 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Same-parent staging and atomic no-overwrite collection publication.""" + +from __future__ import annotations + +import ctypes +import errno +import os +import stat +import sys +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path, PurePosixPath + +from data_designer.slurm.contracts import is_path_below +from data_designer.slurm.filesystem import PRIVATE_DIRECTORY_MODE, open_verified_directory +from data_designer.slurm.state.errors import StateConflictError +from data_designer.slurm.state.filesystem import open_verified_child_directory, open_verified_regular_file +from data_designer.slurm.state.outputs import CollectionPlan + +_RENAME_NOREPLACE = 1 +_RENAME_EXCL = 0x00000004 + + +@dataclass(frozen=True, slots=True) +class StagedFile: + """Expected immutable bytes in a collection staging directory.""" + + name: str + sha256: str + byte_size: int + + +@dataclass(slots=True) +class StagedCollection: + """Private sibling directory that becomes visible only after publication.""" + + path: Path + destination: Path + _parent_descriptor: int + _parent_identity: tuple[int, int] + _stage_identity: tuple[int, int] + _published: bool = False + + def publish(self, expected_files: tuple[StagedFile, ...]) -> None: + """Atomically rename the complete stage while refusing any collision.""" + if self._published: + return + self._rebind() + _verify_stage_files(self, expected_files) + _rename_without_overwrite( + self._parent_descriptor, + self.path.name, + self._parent_descriptor, + self.destination.name, + ) + self._published = True + try: + self._rebind_published_destination() + os.fsync(self._parent_descriptor) + self._rebind_published_destination() + except OSError: + self._restore_stage() + raise + + def _rebind(self) -> None: + self._rebind_parent() + stage = os.stat(self.path.name, dir_fd=self._parent_descriptor, follow_symlinks=False) + if not stat.S_ISDIR(stage.st_mode) or _identity(stage) != self._stage_identity: + raise OSError(f"collection stage {self.path} changed") + + def _rebind_parent(self) -> None: + opened = os.fstat(self._parent_descriptor) + current = self.destination.parent.lstat() + if ( + not stat.S_ISDIR(current.st_mode) + or _identity(opened) != self._parent_identity + or _identity(current) != self._parent_identity + ): + raise OSError(f"collection destination parent {self.destination.parent} changed") + + def _rebind_published_destination(self) -> None: + self._rebind_parent() + opened_view = os.stat(self.destination.name, dir_fd=self._parent_descriptor, follow_symlinks=False) + path_view = self.destination.lstat() + if ( + not stat.S_ISDIR(opened_view.st_mode) + or not stat.S_ISDIR(path_view.st_mode) + or _identity(opened_view) != self._stage_identity + or _identity(path_view) != self._stage_identity + ): + raise OSError(f"published collection destination {self.destination} changed") + + def _restore_stage(self) -> None: + published = os.stat(self.destination.name, dir_fd=self._parent_descriptor, follow_symlinks=False) + if not stat.S_ISDIR(published.st_mode) or _identity(published) != self._stage_identity: + raise OSError(f"published collection destination {self.destination} changed before rollback") + _rename_without_overwrite( + self._parent_descriptor, + self.destination.name, + self._parent_descriptor, + self.path.name, + ) + self._published = False + os.fsync(self._parent_descriptor) + + +def derive_collection_staging_directory(plan: CollectionPlan) -> str: + """Derive one collision-resistant, persisted stage identity from the plan.""" + return f".dd-collection-{plan.compute_sha256()[:32]}.tmp" + + +def remove_collection_stage(destination: Path, staging_directory: str, authorized_root: Path) -> None: + """Remove only the exact persisted stage for one terminal collection.""" + try: + with _open_authorized_parent(destination, authorized_root) as parent_descriptor: + _require_restrictive_parent(parent_descriptor) + _remove_existing_stage(parent_descriptor, destination.parent, staging_directory) + except _MissingDestinationParent: + return + + +def prepare_collection_destination(destination: Path, authorized_root: Path) -> None: + """Create and validate the destination parent without creating the dataset.""" + with _open_authorized_parent(destination, authorized_root, create_missing=True) as parent_descriptor: + _require_restrictive_parent(parent_descriptor) + _require_absent(parent_descriptor, destination.name, destination) + + +@contextmanager +def stage_collection( + destination: Path, + staging_directory: str, + authorized_root: Path, +) -> Iterator[StagedCollection]: + """Yield a private sibling stage and remove it unless atomically published.""" + if destination.parent == destination: + raise StateConflictError("collection destination cannot be the filesystem root") + with _open_authorized_parent(destination, authorized_root) as parent_descriptor: + parent_status = _require_restrictive_parent(parent_descriptor) + _require_absent(parent_descriptor, destination.name, destination) + _remove_existing_stage(parent_descriptor, destination.parent, staging_directory) + stage_name = _create_stage_directory(parent_descriptor, staging_directory) + stage_path = destination.parent / stage_name + staged = StagedCollection( + path=stage_path, + destination=destination, + _parent_descriptor=parent_descriptor, + _parent_identity=_identity(parent_status), + _stage_identity=_identity(os.stat(stage_name, dir_fd=parent_descriptor, follow_symlinks=False)), + ) + try: + yield staged + finally: + if not staged._published: + _remove_stage(staged) + + +class _MissingDestinationParent(FileNotFoundError): + """A child beneath an existing authorized root has not been created yet.""" + + +@contextmanager +def _open_authorized_parent( + destination: Path, + authorized_root: Path, + *, + create_missing: bool = False, +) -> Iterator[int]: + destination_text = destination.as_posix() + root_text = authorized_root.as_posix() + if destination_text == root_text or not is_path_below(destination_text, root_text): + raise StateConflictError("collection destination must be below its authorized mount root") + relative_parent = PurePosixPath(destination.parent.relative_to(authorized_root)).parts + with ExitStack() as resources: + descriptor = resources.enter_context(open_verified_directory(authorized_root, resource_name="collection")) + current_path = authorized_root + for part in relative_parent: + current_path /= part + if create_missing: + _ensure_private_child_directory(descriptor, part) + try: + descriptor = resources.enter_context( + open_verified_child_directory( + descriptor, + part, + current_path, + require_private=False, + ) + ) + except FileNotFoundError as error: + raise _MissingDestinationParent(current_path) from error + yield descriptor + + +def _ensure_private_child_directory(parent_descriptor: int, name: str) -> None: + try: + os.mkdir(name, PRIVATE_DIRECTORY_MODE, dir_fd=parent_descriptor) + except FileExistsError: + return + os.fsync(parent_descriptor) + + +def _require_restrictive_parent(parent_descriptor: int) -> os.stat_result: + status = os.fstat(parent_descriptor) + if status.st_mode & 0o022: + raise StateConflictError("collection destination parent must not be group- or world-writable") + return status + + +def _require_absent(parent_descriptor: int, name: str, display_path: Path) -> None: + try: + os.stat(name, dir_fd=parent_descriptor, follow_symlinks=False) + except FileNotFoundError: + return + raise StateConflictError(f"collection destination {display_path} already exists") + + +def _create_stage_directory(parent_descriptor: int, name: str) -> str: + os.mkdir(name, PRIVATE_DIRECTORY_MODE, dir_fd=parent_descriptor) + os.fsync(parent_descriptor) + return name + + +def _remove_existing_stage(parent_descriptor: int, parent_path: Path, name: str) -> None: + try: + status = os.stat(name, dir_fd=parent_descriptor, follow_symlinks=False) + except FileNotFoundError: + return + if not stat.S_ISDIR(status.st_mode) or status.st_mode & 0o077: + raise OSError(f"persisted collection stage {parent_path / name} is not a private directory") + _delete_stage_directory(parent_descriptor, name, parent_path / name, _identity(status)) + + +def _remove_stage(staged: StagedCollection) -> None: + try: + current = os.stat(staged.path.name, dir_fd=staged._parent_descriptor, follow_symlinks=False) + except FileNotFoundError: + return + if not stat.S_ISDIR(current.st_mode) or _identity(current) != staged._stage_identity: + raise OSError(f"collection stage {staged.path} changed before cleanup") + _delete_stage_directory( + staged._parent_descriptor, + staged.path.name, + staged.path, + staged._stage_identity, + ) + + +def _delete_stage_directory( + parent_descriptor: int, + name: str, + display_path: Path, + expected_identity: tuple[int, int], +) -> None: + with open_verified_child_directory(parent_descriptor, name, display_path) as descriptor: + if _identity(os.fstat(descriptor)) != expected_identity: + raise OSError(f"collection stage {display_path} changed before cleanup") + _clear_directory(descriptor, display_path) + current = os.stat(name, dir_fd=parent_descriptor, follow_symlinks=False) + if _identity(current) != expected_identity: + raise OSError(f"collection stage {display_path} changed during cleanup") + os.rmdir(name, dir_fd=parent_descriptor) + os.fsync(parent_descriptor) + + +def _clear_directory(directory_descriptor: int, display_path: Path) -> None: + for name in os.listdir(directory_descriptor): + status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + child_path = display_path / name + if stat.S_ISDIR(status.st_mode): + with open_verified_child_directory(directory_descriptor, name, child_path) as child_descriptor: + if _identity(os.fstat(child_descriptor)) != _identity(status): + raise OSError(f"collection staging directory {child_path} changed before cleanup") + _clear_directory(child_descriptor, child_path) + current = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if _identity(current) != _identity(status): + raise OSError(f"collection staging directory {child_path} changed during cleanup") + os.rmdir(name, dir_fd=directory_descriptor) + else: + os.unlink(name, dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + + +def _verify_stage_files(staged: StagedCollection, expected_files: tuple[StagedFile, ...]) -> None: + expected_names = tuple(file.name for file in expected_files) + if len(expected_names) != len(set(expected_names)): + raise OSError("collection stage expectation contains duplicate paths") + with open_verified_child_directory( + staged._parent_descriptor, + staged.path.name, + staged.path, + ) as stage_descriptor: + if set(os.listdir(stage_descriptor)) != set(expected_names): + raise OSError("collection stage inventory changed before publication") + for expected in expected_files: + with open_verified_regular_file( + stage_descriptor, + expected.name, + staged.path / expected.name, + expected_size=expected.byte_size, + expected_sha256=expected.sha256, + require_private=False, + ): + pass + + +def _rename_without_overwrite( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, +) -> None: + library = ctypes.CDLL(None, use_errno=True) + source = os.fsencode(source_name) + destination = os.fsencode(destination_name) + if sys.platform.startswith("linux") and hasattr(library, "renameat2"): + result = library.renameat2( + source_directory, + ctypes.c_char_p(source), + destination_directory, + ctypes.c_char_p(destination), + _RENAME_NOREPLACE, + ) + elif sys.platform == "darwin" and hasattr(library, "renameatx_np"): + result = library.renameatx_np( + source_directory, + ctypes.c_char_p(source), + destination_directory, + ctypes.c_char_p(destination), + _RENAME_EXCL, + ) + else: + raise OSError(errno.ENOTSUP, "atomic no-overwrite directory rename is unavailable") + if result == 0: + return + error_number = ctypes.get_errno() + if error_number in {errno.EEXIST, errno.ENOTEMPTY}: + raise StateConflictError(f"collection destination {destination_name!r} already exists") + raise OSError(error_number, os.strerror(error_number), destination_name) + + +def _identity(status: os.stat_result) -> tuple[int, int]: + return status.st_dev, status.st_ino + + +__all__ = [ + "StagedCollection", + "StagedFile", + "derive_collection_staging_directory", + "prepare_collection_destination", + "remove_collection_stage", + "stage_collection", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py new file mode 100644 index 000000000..764cf1ac3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Winner-only resolution of deterministic collection inputs.""" + +from __future__ import annotations + +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.collection_validation import validate_collection_inputs +from data_designer.slurm.state.errors import StateCorruptionError, StateNotFoundError +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.outputs import CandidateOutputManifest, CollectionPlan, CollectionShard, ShardWinner +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import StateContractError, validate_collection_plan + + +class CollectionInputResolver: + """Resolve planned winner chains without scanning attempt output trees.""" + + def __init__(self, storage: StateStorage, reader: StateReader) -> None: + self._storage = storage + self._reader = reader + self._finalizer = WinnerFinalizer(storage, reader) + + def get_winner_shards(self) -> tuple[CollectionShard, ...]: + """Return one canonical winner reference for every ordered run shard.""" + run, plan, shards = self._reader.load_context() + planned: list[CollectionShard] = [] + for shard in shards: + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self._finalizer.load_optional_winner(run, plan, shard, attempts) + if winner is None: + raise StateNotFoundError(f"shard {shard.shard_id!r} has no winner") + planned.append( + CollectionShard( + shard_id=shard.shard_id, + winner_manifest=ArtifactReference( + path=self._storage.get_winner_path(shard.shard_id).as_posix(), + sha256=winner.compute_sha256(), + ), + ) + ) + return tuple(planned) + + def resolve( + self, collection_plan: CollectionPlan + ) -> tuple[ResolvedSlurmRunPlan, tuple[CandidateOutputManifest, ...]]: + """Validate the complete winner chain and return ordered candidate manifests.""" + run, plan, shards = self._reader.load_context() + winners: list[ShardWinner] = [] + candidates: list[CandidateOutputManifest] = [] + for planned_collection_shard, shard in zip(collection_plan.planned_shards, shards, strict=True): + if planned_collection_shard.shard_id != shard.shard_id: + raise StateCorruptionError("collection plan does not preserve planned shard order") + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self._finalizer.load_optional_winner(run, plan, shard, attempts) + if winner is None: + raise StateNotFoundError(f"shard {shard.shard_id!r} has no winner") + expected_reference = ArtifactReference( + path=self._storage.get_winner_path(shard.shard_id).as_posix(), + sha256=winner.compute_sha256(), + ) + if planned_collection_shard.winner_manifest != expected_reference: + raise StateCorruptionError(f"collection winner changed for shard {shard.shard_id!r}") + attempt = self._reader.get_attempt(attempts, winner.attempt_id) + result = self._reader.load_optional_attempt_result(plan, shard, attempt) + if result is None: + raise StateCorruptionError(f"winning attempt {winner.attempt_id!r} has no result records") + winners.append(winner) + candidates.append(result[1]) + try: + validate_collection_plan(run, collection_plan, shards, tuple(winners)) + validate_collection_inputs(plan, collection_plan, tuple(candidates)) + except StateContractError as error: + raise StateCorruptionError("collection inputs violate persisted run intent") from error + return plan, tuple(candidates) + + +__all__ = ["CollectionInputResolver"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py new file mode 100644 index 000000000..d7ec2803a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py @@ -0,0 +1,338 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bounded-memory merging of validated Parquet shard winners.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Generator +from contextlib import AbstractContextManager, contextmanager +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import BinaryIO, Protocol, TextIO + +import data_designer.lazy_heavy_imports as lazy +from data_designer.slurm.filesystem import get_file_facts +from data_designer.slurm.state.artifacts import CandidateArtifactSnapshot, CandidateArtifactVerifier +from data_designer.slurm.state.collection_filesystem import StagedCollection, StagedFile +from data_designer.slurm.state.collection_records import CollectedOutputFile, CollectionResult +from data_designer.slurm.state.filesystem import open_verified_directory, publish_immutable_text, sync_directory +from data_designer.slurm.state.outputs import CandidateOutputManifest, CollectionPlan + +_BATCH_SIZE = 65_536 +_RESULT_FILENAME = "collection-result.json" +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 + + +class CollectionPartitionWriter(Protocol): + """Bounded writer for one deterministic output partition.""" + + def write(self, batch: object) -> None: + """Append one Arrow record batch.""" + ... + + def finish(self, relative_path: str, record_count: int) -> CollectedOutputFile: + """Seal and describe the exact file descriptor that received the records.""" + ... + + +class _ArrowBatchWriter(Protocol): + def write_batch(self, batch: object) -> None: ... + + def close(self) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _PartitionLayout: + stage_path: Path + output_format: str + schema: object + record_counts: tuple[int, ...] + + +class CollectionMerger: + """Merge ordered winner files while retaining only bounded data and descriptors.""" + + def __init__( + self, + output_format: str, + *, + verifier: CandidateArtifactVerifier | None = None, + completed_at: datetime | None = None, + ) -> None: + self._output_format = output_format + self._verifier = verifier if verifier is not None else CandidateArtifactVerifier() + self._completed_at = completed_at + + def merge( + self, + collection_plan: CollectionPlan, + candidates: tuple[CandidateOutputManifest, ...], + staged: StagedCollection, + ) -> CollectionResult: + """Write deterministic partitions, rebind inputs, and publish the stage.""" + snapshots = tuple(self._inspect_candidate(candidate) for candidate in candidates) + batches = self._iter_batches(candidates) + try: + first_batch = next(batches, None) + if first_batch is None: + raise OSError("collection inputs contain no records") + layout = _PartitionLayout( + stage_path=staged.path, + output_format=self._output_format, + schema=first_batch.schema.remove_metadata(), + record_counts=_partition_record_counts( + sum(candidate.actual_records for candidate in candidates), + collection_plan.num_partitions, + ), + ) + files = self._write_partitions(layout, first_batch, batches) + finally: + batches.close() + completion_time = self._completed_at if self._completed_at is not None else _utc_now() + result = CollectionResult( + schema_version=1, + collection_id=collection_plan.collection_id, + run_id=collection_plan.run_id, + completed_at=completion_time, + collection_plan_sha256=collection_plan.compute_sha256(), + actual_records=sum(output.record_count for output in files), + files=files, + ) + result_bytes = self._write_result(staged.path, result) + for candidate, snapshot in zip(candidates, snapshots, strict=True): + self._verifier.rebind(candidate, snapshot) + staged.publish(_stage_expectations(files, result, result_bytes)) + return result + + def _inspect_candidate(self, candidate: CandidateOutputManifest) -> CandidateArtifactSnapshot: + snapshot = self._verifier.inspect(candidate) + if snapshot.record_counts != tuple(output.record_count for output in candidate.files): + raise OSError("candidate Parquet row counts changed before collection") + if snapshot.dataset_schema_digest != candidate.dataset_schema_digest: + raise OSError("candidate Parquet schema changed before collection") + return snapshot + + def _iter_batches(self, candidates: tuple[CandidateOutputManifest, ...]) -> Generator[object, None, None]: + for candidate in candidates: + for output_file in candidate.files: + with self._verifier.open_output(candidate, output_file) as source: + parquet_file = lazy.pq.ParquetFile(source) + yield from parquet_file.iter_batches(batch_size=_BATCH_SIZE) + + def _write_partitions( + self, + layout: _PartitionLayout, + first_batch: object, + remaining_batches: Generator[object, None, None], + ) -> tuple[CollectedOutputFile, ...]: + cursor = _BatchCursor(first_batch, remaining_batches) + outputs: list[CollectedOutputFile] = [] + for partition_index, record_count in enumerate(layout.record_counts): + suffix = "jsonl" if layout.output_format == "jsonl" else layout.output_format + relative_path = f"part-{partition_index:05d}.{suffix}" + output_path = layout.stage_path / relative_path + with _open_partition_writer(output_path, layout.output_format, layout.schema) as writer: + cursor.write_records(writer, record_count) + outputs.append(writer.finish(relative_path, record_count)) + if cursor.has_remaining_records(): + raise OSError("collection inputs contain more rows than declared") + return tuple(outputs) + + @staticmethod + def _write_result(stage_path: Path, result: CollectionResult) -> bytes: + serialized = result.serialize_json() + with open_verified_directory(stage_path, require_private=True) as descriptor: + publish_immutable_text( + descriptor, + _RESULT_FILENAME, + serialized, + stage_path / _RESULT_FILENAME, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + sync_directory(descriptor) + return serialized.encode("utf-8") + + +class _BatchCursor: + def __init__(self, first_batch: object, remaining: Generator[object, None, None]) -> None: + self._batch = first_batch + self._remaining = remaining + self._offset = 0 + + def write_records(self, writer: CollectionPartitionWriter, record_count: int) -> None: + remaining = record_count + while remaining: + available = self._batch.num_rows - self._offset + if available == 0: + self._advance() + continue + size = min(remaining, available) + writer.write(self._batch.slice(self._offset, size)) + self._offset += size + remaining -= size + + def has_remaining_records(self) -> bool: + if self._offset < self._batch.num_rows: + return True + return next(self._remaining, None) is not None + + def _advance(self) -> None: + next_batch = next(self._remaining, None) + if next_batch is None: + raise OSError("collection inputs contain fewer rows than declared") + self._batch = next_batch + self._offset = 0 + + +class _BoundOutput: + def __init__(self, path: Path, output: BinaryIO | TextIO) -> None: + self.path = path + self._output = output + + def describe(self, relative_path: str, record_count: int) -> CollectedOutputFile: + _sync_output(self._output) + return _describe_open_output(self._output.fileno(), self.path, relative_path, record_count) + + +class _ParquetWriter: + def __init__(self, writer: _ArrowBatchWriter, output: _BoundOutput) -> None: + self._writer = writer + self._output = output + self._closed = False + + def write(self, batch: object) -> None: + self._writer.write_batch(batch.replace_schema_metadata(None)) + + def finish(self, relative_path: str, record_count: int) -> CollectedOutputFile: + self.close() + return self._output.describe(relative_path, record_count) + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._writer.close() + + +class _TextWriter: + def __init__(self, output: TextIO, output_format: str, bound_output: _BoundOutput) -> None: + self._output = output + self._bound_output = bound_output + self._format = output_format + self._is_first = True + + def write(self, batch: object) -> None: + frame = batch.to_pandas() + if self._format == "csv": + frame.to_csv(self._output, header=self._is_first, index=False) + else: + content = frame.to_json(orient="records", lines=True, force_ascii=False, date_format="iso") + if content: + self._output.write(content) + self._is_first = False + + def finish(self, relative_path: str, record_count: int) -> CollectedOutputFile: + return self._bound_output.describe(relative_path, record_count) + + +def _open_partition_writer( + output_path: Path, + output_format: str, + schema: object, +) -> AbstractContextManager[CollectionPartitionWriter]: + if output_format == "parquet": + return _open_parquet_writer(output_path, schema) + if output_format in {"csv", "jsonl"}: + return _open_text_writer(output_path, output_format) + raise ValueError(f"unsupported collection output format {output_format!r}") + + +@contextmanager +def _open_parquet_writer(output_path: Path, schema: object) -> Generator[CollectionPartitionWriter, None, None]: + with output_path.open("x+b") as output: + os.fchmod(output.fileno(), 0o600) + partition_writer = _ParquetWriter(lazy.pq.ParquetWriter(output, schema), _BoundOutput(output_path, output)) + try: + yield partition_writer + finally: + partition_writer.close() + _sync_output(output) + + +@contextmanager +def _open_text_writer(output_path: Path, output_format: str) -> Generator[CollectionPartitionWriter, None, None]: + with output_path.open("x+", encoding="utf-8", newline="") as output: + os.fchmod(output.fileno(), 0o600) + try: + yield _TextWriter(output, output_format, _BoundOutput(output_path, output)) + finally: + _sync_output(output) + + +def _sync_output(output: BinaryIO | TextIO) -> None: + output.flush() + os.fsync(output.fileno()) + + +def _partition_record_counts(record_count: int, partition_count: int) -> tuple[int, ...]: + floor_count = record_count // partition_count + return tuple( + record_count - floor_count * (partition_count - 1) if index == partition_count - 1 else floor_count + for index in range(partition_count) + ) + + +def _describe_open_output( + descriptor: int, + output_path: Path, + relative_path: str, + record_count: int, +) -> CollectedOutputFile: + before = os.fstat(descriptor) + _require_safe_output(before, output_path) + digest = hashlib.sha256() + offset = 0 + while block := os.pread(descriptor, 1024 * 1024, offset): + digest.update(block) + offset += len(block) + after = os.fstat(descriptor) + path_status = output_path.lstat() + if get_file_facts(before) != get_file_facts(after) or get_file_facts(after) != get_file_facts(path_status): + raise OSError(f"collection output {output_path} changed while it was being described") + _require_safe_output(after, output_path) + return CollectedOutputFile( + relative_path=relative_path, + sha256=digest.hexdigest(), + byte_size=after.st_size, + record_count=record_count, + modified_at_ns=after.st_mtime_ns, + changed_at_ns=after.st_ctime_ns, + ) + + +def _require_safe_output(status: os.stat_result, output_path: Path) -> None: + if not stat.S_ISREG(status.st_mode) or status.st_nlink != 1 or status.st_mode & 0o077: + raise OSError(f"collection output {output_path} is not a private single-link regular file") + + +def _stage_expectations( + files: tuple[CollectedOutputFile, ...], + result: CollectionResult, + result_bytes: bytes, +) -> tuple[StagedFile, ...]: + return ( + *(StagedFile(file.relative_path, file.sha256, file.byte_size) for file in files), + StagedFile(_RESULT_FILENAME, result.compute_sha256(), len(result_bytes)), + ) + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +__all__ = ["CollectionMerger", "CollectionPartitionWriter"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py new file mode 100644 index 000000000..3d6e52107 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Persisted collection lifecycle and output records.""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum +from typing import Annotated + +from pydantic import Field, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ArtifactReference, Identifier, Sha256Digest, validate_relative_path +from data_designer.slurm.state.base import SchedulerJobIdentity, StateRecord, StateValue, validate_utc_timestamp +from data_designer.slurm.state.scheduler import SchedulerObservation + + +class CollectionState(str, Enum): + """Persisted lifecycle for one CPU collection job.""" + + PREPARED = "prepared" + SUBMITTED = "submitted" + PENDING = "pending" + RUNNING = "running" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class CollectedOutputFile(StateValue): + """One deterministic file in a published collected dataset.""" + + relative_path: str + sha256: Sha256Digest + byte_size: NonNegativeInt + record_count: NonNegativeInt + modified_at_ns: NonNegativeInt + changed_at_ns: NonNegativeInt + + _relative_path_is_safe = field_validator("relative_path")(validate_relative_path) + + +class CollectionResult(StateRecord): + """Immutable proof of a completely staged collection output.""" + + collection_id: Identifier + run_id: Identifier + completed_at: datetime + collection_plan_sha256: Sha256Digest + actual_records: NonNegativeInt + files: tuple[CollectedOutputFile, ...] = Field(min_length=1) + + _completed_at_is_utc = field_validator("completed_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_files(self) -> CollectionResult: + paths = tuple(output.relative_path for output in self.files) + if len(paths) != len(set(paths)): + raise ValueError("collected output paths must be unique") + if sum(output.record_count for output in self.files) != self.actual_records: + raise ValueError("collected output row counts must equal actual_records") + return self + + +class CollectionStatus(StateRecord): + """Atomically replaced scheduler and publication state for one collection.""" + + collection_id: Identifier + run_id: Identifier + collection_plan: ArtifactReference + staging_directory: Annotated[ + str, + StringConstraints(pattern=r"^\.dd-collection-[0-9a-f]{32}\.tmp$"), + ] + revision: PositiveInt + updated_at: datetime + state: CollectionState + scheduler: SchedulerJobIdentity | None = None + scheduler_observation: SchedulerObservation | None = None + result: ArtifactReference | None = None + + _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_evidence(self) -> CollectionStatus: + if self.state is CollectionState.PREPARED: + if self.scheduler is not None or self.scheduler_observation is not None or self.result is not None: + raise ValueError("prepared collection cannot contain scheduler or result evidence") + return self + if self.state is not CollectionState.FAILED and self.scheduler is None: + raise ValueError("submitted collection states require a scheduler identity") + if self.scheduler is None and self.scheduler_observation is not None: + raise ValueError("collection observation requires a scheduler identity") + if self.scheduler_observation is not None and self.scheduler_observation.scheduler != self.scheduler: + raise ValueError("collection scheduler observation identity does not match") + if (self.state is CollectionState.SUCCEEDED) != (self.result is not None): + raise ValueError("collection result is required exactly for succeeded state") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py new file mode 100644 index 000000000..0740f5345 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Descriptor-bound persistence for collection plans and lifecycle.""" + +from __future__ import annotations + +import os +import re +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.filesystem import get_file_facts +from data_designer.slurm.state.collection_records import CollectedOutputFile, CollectionResult, CollectionStatus +from data_designer.slurm.state.filesystem import ( + acquire_file_lock, + ensure_private_child_directory, + is_state_temporary_name, + open_verified_child_directory, + open_verified_directory, + open_verified_regular_file, + publish_immutable_text, + replace_text, +) +from data_designer.slurm.state.outputs import CollectionPlan +from data_designer.slurm.state.storage import StateStorage + +_COLLECTIONS_DIRECTORY = "collections" +_COLLECTION_LOCK = "collection.lock" +_PLAN_FILENAME = "plan.json" +_STATUS_FILENAME = "status.json" +_RESULT_FILENAME = "collection-result.json" +_COLLECTION_PATTERN = re.compile(r"^collection-[0-9]{4,}$") +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 + + +class CollectionStorage: + """Persist collection state without adding collection concerns to run storage.""" + + def __init__(self, state_storage: StateStorage) -> None: + self._state = state_storage + self.collections_root = state_storage.run_root / _COLLECTIONS_DIRECTORY + + @contextmanager + def acquire_lock(self) -> Iterator[None]: + """Serialize collection preparation, refresh, and worker publication.""" + with self._state.open_run_directory() as run_descriptor: + with acquire_file_lock( + run_descriptor, + _COLLECTION_LOCK, + self._state.run_root / _COLLECTION_LOCK, + ): + yield + + def get_next_collection_id(self) -> Identifier: + """Return the next monotonic identity after validating existing directories.""" + names = self.list_collection_ids() + return f"collection-{len(names) + 1:04d}" + + def discard_incomplete_tail(self) -> None: + """Discard one trailing collection journal that predates submission.""" + try: + with self._open_collections_directory() as collections_descriptor: + collection_ids = _validated_collection_ids(tuple(os.listdir(collections_descriptor))) + if not collection_ids: + return + collection_id = collection_ids[-1] + with open_verified_child_directory( + collections_descriptor, + collection_id, + self.get_collection_root(collection_id), + ) as collection_descriptor: + if _record_exists(collection_descriptor, _STATUS_FILENAME): + return + _discard_prepared_files( + collection_descriptor, + self.get_collection_root(collection_id), + ) + os.rmdir(collection_id, dir_fd=collections_descriptor) + os.fsync(collections_descriptor) + except FileNotFoundError: + return + + def list_collection_ids(self) -> tuple[Identifier, ...]: + """List only a complete monotonic set of managed collection directories.""" + try: + with self._open_collections_directory() as descriptor: + names = tuple(os.listdir(descriptor)) + except FileNotFoundError: + return () + return _validated_collection_ids(names) + + def ensure_collection(self, collection_id: Identifier) -> None: + """Create one private collection state directory.""" + with self._state.open_run_directory() as run_descriptor: + ensure_private_child_directory(run_descriptor, _COLLECTIONS_DIRECTORY, self.collections_root) + with open_verified_child_directory( + run_descriptor, + _COLLECTIONS_DIRECTORY, + self.collections_root, + ) as collections_descriptor: + ensure_private_child_directory( + collections_descriptor, + collection_id, + self.get_collection_root(collection_id), + ) + + def get_collection_root(self, collection_id: Identifier) -> Path: + return self.collections_root / collection_id + + def get_plan_path(self, collection_id: Identifier) -> Path: + return self.get_collection_root(collection_id) / _PLAN_FILENAME + + def get_status_path(self, collection_id: Identifier) -> Path: + return self.get_collection_root(collection_id) / _STATUS_FILENAME + + def get_result_path(self, plan: CollectionPlan) -> Path: + return Path(plan.host_destination) / _RESULT_FILENAME + + def get_result_reference(self, plan: CollectionPlan, result: CollectionResult) -> ArtifactReference: + """Return the canonical host-view reference for a published result.""" + return ArtifactReference(path=self.get_result_path(plan).as_posix(), sha256=result.compute_sha256()) + + def publish_plan(self, plan: CollectionPlan) -> None: + self._require_run_id(plan.run_id) + with self._open_collection_directory(plan.collection_id) as descriptor: + publish_immutable_text( + descriptor, + _PLAN_FILENAME, + plan.serialize_json(), + self.get_plan_path(plan.collection_id), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def read_plan(self, collection_id: Identifier) -> CollectionPlan: + with self._open_collection_directory(collection_id) as descriptor: + plan = self._state.read_record( + descriptor, + _PLAN_FILENAME, + self.get_plan_path(collection_id), + CollectionPlan, + ) + if plan.collection_id != collection_id or plan.run_id != self._state.run_id: + raise OSError("collection plan identity does not match its persisted location") + return plan + + def get_plan_reference(self, plan: CollectionPlan) -> ArtifactReference: + """Return the canonical persisted reference for an immutable collection plan.""" + return ArtifactReference(path=self.get_plan_path(plan.collection_id).as_posix(), sha256=plan.compute_sha256()) + + def publish_status(self, status: CollectionStatus) -> None: + self._require_run_id(status.run_id) + with self._open_collection_directory(status.collection_id) as descriptor: + publish_immutable_text( + descriptor, + _STATUS_FILENAME, + status.serialize_json(), + self.get_status_path(status.collection_id), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def replace_status(self, status: CollectionStatus) -> None: + self._require_run_id(status.run_id) + with self._open_collection_directory(status.collection_id) as descriptor: + replace_text( + descriptor, + _STATUS_FILENAME, + status.serialize_json(), + self.get_status_path(status.collection_id), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def read_status(self, collection_id: Identifier) -> CollectionStatus: + with self._open_collection_directory(collection_id) as descriptor: + status = self._state.read_record( + descriptor, + _STATUS_FILENAME, + self.get_status_path(collection_id), + CollectionStatus, + ) + if status.collection_id != collection_id or status.run_id != self._state.run_id: + raise OSError("collection status identity does not match its persisted location") + return status + + def read_result(self, plan: CollectionPlan) -> CollectionResult: + return self.read_result_from(plan, Path(plan.host_destination)) + + def read_result_from(self, plan: CollectionPlan, destination: Path) -> CollectionResult: + """Read a result through an explicitly selected host or container view.""" + if destination.as_posix() not in {plan.host_destination, plan.container_destination}: + raise OSError("collection result destination does not match its immutable plan") + with open_verified_directory(destination, require_private=True) as descriptor: + return self._state.read_record( + descriptor, + _RESULT_FILENAME, + destination / _RESULT_FILENAME, + CollectionResult, + ) + + def verify_result_files( + self, + plan: CollectionPlan, + result: CollectionResult, + destination: Path, + *, + verify_digests: bool = True, + ) -> None: + """Verify exact inventory with bounded metadata or full output digests.""" + if destination.as_posix() not in {plan.host_destination, plan.container_destination}: + raise OSError("collection result destination does not match its immutable plan") + expected_names = tuple(output.relative_path for output in result.files) + (_RESULT_FILENAME,) + if any("/" in name for name in expected_names): + raise OSError("collection result inventory must contain only direct child files") + with open_verified_directory(destination, require_private=True) as descriptor: + if set(os.listdir(descriptor)) != set(expected_names): + raise OSError("published collection inventory does not match its result manifest") + for output in result.files: + if verify_digests: + with open_verified_regular_file( + descriptor, + output.relative_path, + destination / output.relative_path, + expected_size=output.byte_size, + expected_sha256=output.sha256, + require_private=False, + ): + pass + else: + _verify_output_metadata(descriptor, output, destination) + result_bytes = result.serialize_json().encode("utf-8") + with open_verified_regular_file( + descriptor, + _RESULT_FILENAME, + destination / _RESULT_FILENAME, + expected_size=len(result_bytes), + expected_sha256=result.compute_sha256(), + ): + pass + + @contextmanager + def _open_collections_directory(self) -> Iterator[int]: + with self._state.open_run_directory() as run_descriptor: + with open_verified_child_directory( + run_descriptor, + _COLLECTIONS_DIRECTORY, + self.collections_root, + ) as collections_descriptor: + yield collections_descriptor + + @contextmanager + def _open_collection_directory(self, collection_id: Identifier) -> Iterator[int]: + with self._open_collections_directory() as collections_descriptor: + with open_verified_child_directory( + collections_descriptor, + collection_id, + self.get_collection_root(collection_id), + ) as descriptor: + yield descriptor + + def _require_run_id(self, run_id: Identifier) -> None: + if run_id != self._state.run_id: + raise OSError("collection record run identity does not match storage") + + +def _verify_output_metadata( + directory_descriptor: int, + output: CollectedOutputFile, + destination: Path, +) -> None: + before = os.stat(output.relative_path, dir_fd=directory_descriptor, follow_symlinks=False) + if ( + not stat.S_ISREG(before.st_mode) + or before.st_nlink != 1 + or before.st_mode & 0o022 + or before.st_size != output.byte_size + or before.st_mtime_ns != output.modified_at_ns + or before.st_ctime_ns != output.changed_at_ns + ): + raise OSError(f"collected output {destination / output.relative_path} is not a safe regular file") + descriptor = os.open( + output.relative_path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + try: + after = os.fstat(descriptor) + rebound = os.stat(output.relative_path, dir_fd=directory_descriptor, follow_symlinks=False) + finally: + os.close(descriptor) + if get_file_facts(before) != get_file_facts(after) or get_file_facts(after) != get_file_facts(rebound): + raise OSError(f"collected output {destination / output.relative_path} changed during validation") + + +def _validated_collection_ids(names: tuple[str, ...]) -> tuple[Identifier, ...]: + if any(_COLLECTION_PATTERN.fullmatch(name) is None for name in names): + raise OSError("collection state contains an unowned directory") + ordered = tuple(sorted(names, key=lambda name: int(name.rsplit("-", maxsplit=1)[1]))) + expected = tuple(f"collection-{index:04d}" for index in range(1, len(ordered) + 1)) + if ordered != expected: + raise OSError("collection state identities are not a complete monotonic sequence") + return ordered + + +def _record_exists(directory_descriptor: int, name: str) -> bool: + try: + os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + except FileNotFoundError: + return False + return True + + +def _discard_prepared_files(directory_descriptor: int, display_path: Path) -> None: + names = tuple(os.listdir(directory_descriptor)) + if any(name != _PLAN_FILENAME and not is_state_temporary_name(name) for name in names): + raise OSError(f"incomplete collection journal {display_path} contains an unowned entry") + for name in names: + status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not stat.S_ISREG(status.st_mode) or status.st_mode & 0o077: + raise OSError(f"incomplete collection journal entry {display_path / name} is unsafe") + os.unlink(name, dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + + +__all__ = ["CollectionStorage"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py new file mode 100644 index 000000000..f79357060 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cross-record validation for collection lifecycle and inputs.""" + +from __future__ import annotations + +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.collection_records import CollectionResult, CollectionState, CollectionStatus +from data_designer.slurm.state.outputs import CandidateOutputManifest, CollectionPlan +from data_designer.slurm.state.scheduler import SchedulerState +from data_designer.slurm.state.validation import StateContractError + +_TERMINAL_COLLECTION_STATES = frozenset({CollectionState.SUCCEEDED, CollectionState.FAILED}) + + +def validate_collection_inputs( + resolved_plan: ResolvedSlurmRunPlan, + collection_plan: CollectionPlan, + candidates: tuple[CandidateOutputManifest, ...], +) -> tuple[CandidateOutputManifest, ...]: + """Require one compatible, complete candidate for every planned shard.""" + _require(len(candidates) == len(collection_plan.planned_shards), "collection candidate set is incomplete") + expected_shards = tuple(shard.shard_id for shard in collection_plan.planned_shards) + actual_shards = tuple(candidate.shard_id for candidate in candidates) + _require(actual_shards == expected_shards, "collection candidates do not match ordered planned shards") + _require(all(candidate.winner_eligible for candidate in candidates), "collection candidate is not complete") + + schema_digests = {candidate.dataset_schema_digest for candidate in candidates} + provenance_digests = {candidate.provenance_digest for candidate in candidates} + _require(len(schema_digests) == 1, "collection candidates have incompatible schemas") + _require(len(provenance_digests) == 1, "collection candidates have incompatible provenance") + + expected_records = resolved_plan.invocation.authored.num_records + actual_records = sum(candidate.actual_records for candidate in candidates) + _require(actual_records == expected_records, "collection candidate rows do not match requested run rows") + return candidates + + +def validate_collection_status_transition( + previous: CollectionStatus, + current: CollectionStatus, +) -> CollectionStatus: + """Validate immutable identity, monotonic revisions, and terminal evidence.""" + _require(current.collection_id == previous.collection_id, "collection identity cannot change") + _require(current.run_id == previous.run_id, "collection run identity cannot change") + _require(current.collection_plan == previous.collection_plan, "collection plan identity cannot change") + _require(current.staging_directory == previous.staging_directory, "collection staging identity cannot change") + _require(previous.state not in _TERMINAL_COLLECTION_STATES, "terminal collection status is immutable") + _require(current.revision == previous.revision + 1, "collection status revision must increase by one") + _require(current.updated_at >= previous.updated_at, "collection status timestamp cannot move backward") + if previous.scheduler is not None: + _require(current.scheduler == previous.scheduler, "collection scheduler identity cannot change") + if previous.scheduler_observation is not None and current.scheduler_observation is not None: + _require( + current.scheduler_observation.observed_at >= previous.scheduler_observation.observed_at, + "collection scheduler observation cannot move backward", + ) + return current + + +def validate_collection_result( + collection_plan: CollectionPlan, + result: CollectionResult, + *, + expected_records: int, + output_format: str, +) -> CollectionResult: + """Bind a collected result to its immutable plan and exact row count.""" + _require(result.collection_id == collection_plan.collection_id, "collection result identity does not match") + _require(result.run_id == collection_plan.run_id, "collection result run identity does not match") + _require( + result.collection_plan_sha256 == collection_plan.compute_sha256(), + "collection result does not bind the collection plan", + ) + _require(result.completed_at >= collection_plan.created_at, "collection completion precedes plan creation") + _require(result.actual_records == expected_records, "collected result has the wrong row count") + _require(len(result.files) == collection_plan.num_partitions, "collected result has the wrong partition count") + suffix = "jsonl" if output_format == "jsonl" else output_format + expected_paths = tuple(f"part-{index:05d}.{suffix}" for index in range(collection_plan.num_partitions)) + _require( + tuple(output.relative_path for output in result.files) == expected_paths, + "collected result files do not match deterministic partition intent", + ) + return result + + +def derive_collection_state(observation_state: SchedulerState) -> CollectionState: + """Map normalized scheduler evidence to collection lifecycle state.""" + if observation_state is SchedulerState.PENDING: + return CollectionState.PENDING + if observation_state is SchedulerState.RUNNING: + return CollectionState.RUNNING + if observation_state is SchedulerState.ACCOUNTING_LAG: + return CollectionState.ACCOUNTING_LAG + if observation_state is SchedulerState.UNKNOWN: + return CollectionState.UNKNOWN + return CollectionState.FAILED + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise StateContractError(message) + + +__all__ = [ + "derive_collection_state", + "validate_collection_inputs", + "validate_collection_result", + "validate_collection_status_transition", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py new file mode 100644 index 000000000..d675e80a6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py @@ -0,0 +1,257 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-only worker for one persisted collection plan.""" + +from __future__ import annotations + +import argparse +import os +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from pathlib import Path + +from pydantic import TypeAdapter, ValidationError + +import data_designer.lazy_heavy_imports as lazy +from data_designer.slurm.contracts import Identifier, validate_absolute_path +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.collection_filesystem import derive_collection_staging_directory, stage_collection +from data_designer.slurm.state.collection_inputs import CollectionInputResolver +from data_designer.slurm.state.collection_merge import CollectionMerger +from data_designer.slurm.state.collection_records import CollectionResult, CollectionState, CollectionStatus +from data_designer.slurm.state.collection_storage import CollectionStorage +from data_designer.slurm.state.collection_validation import ( + validate_collection_result, + validate_collection_status_transition, +) +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError, StateCorruptionError +from data_designer.slurm.state.outputs import CandidateOutputManifest, CollectionPlan +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class SlurmCollectionWorker: + """Execute bulk collection only inside its recorded CPU Slurm job.""" + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + collection_id: Identifier, + *, + environment: Mapping[str, str] | None = None, + ) -> None: + root, normalized_run_id, normalized_collection_id = _validate_location(workspace_root, run_id, collection_id) + self._state = StateStorage(root, normalized_run_id) + self._collections = CollectionStorage(self._state) + self._reader = StateReader(self._state, normalized_run_id) + self._inputs = CollectionInputResolver(self._state, self._reader) + self._destinations = CollectionDestinationResolver() + self._environment = dict(os.environ if environment is None else environment) + self._run_id = normalized_run_id + self._collection_id = normalized_collection_id + + def run(self, *, completed_at: datetime | None = None) -> CollectionResult: + """Validate, merge, and atomically publish one collection output.""" + started_at = _utc_now() if completed_at is None else completed_at + try: + return self._run_locked(started_at, completed_at) + except (StateConflictError, SlurmStateError): + raise + except (OSError, ValueError, lazy.pa.ArrowException) as error: + raise SlurmStateError(f"collection {self._collection_id!r} failed") from error + + def _run_locked(self, started_at: datetime, completed_at: datetime | None) -> CollectionResult: + with self._collections.acquire_lock(): + status = self._collections.read_status(self._collection_id) + plan = self._load_bound_plan(status) + self._validate_identity(plan.run_id, status.run_id) + self._require_scheduler_job(status) + resolved_plan, candidates = self._inputs.resolve(plan) + self._destinations.validate_persisted(resolved_plan, plan) + if status.state is CollectionState.SUCCEEDED: + return self._load_valid_result(plan, status) + recovered = self._load_optional_result(plan) + if recovered is not None: + return self._publish_success_status(plan, status, recovered, recovered.completed_at) + running = self._advance_status(status, CollectionState.RUNNING, started_at) + return self._execute_collection(plan, running, resolved_plan, candidates, completed_at) + + def _execute_collection( + self, + plan: CollectionPlan, + running: CollectionStatus, + resolved_plan: ResolvedSlurmRunPlan, + candidates: tuple[CandidateOutputManifest, ...], + completed_at: datetime | None, + ) -> CollectionResult: + try: + destination = self._destinations.validate_persisted(resolved_plan, plan) + with stage_collection( + Path(plan.container_destination), + running.staging_directory, + Path(destination.mount.target), + ) as staged: + merger = CollectionMerger(resolved_plan.output.format, completed_at=completed_at) + result = merger.merge( + plan, + candidates, + staged, + ) + except (SlurmStateError, OSError, ValueError, lazy.pa.ArrowException) as error: + return self._recover_or_fail(plan, running, completed_at, error) + return self._publish_success_status(plan, running, result, result.completed_at) + + def _recover_or_fail( + self, + plan: CollectionPlan, + running: CollectionStatus, + completed_at: datetime | None, + error: Exception, + ) -> CollectionResult: + try: + recovered = self._load_optional_result(plan) + except (SlurmStateError, OSError, ValueError) as recovery_error: + self._advance_status(running, CollectionState.FAILED, self._completion_time(completed_at)) + raise recovery_error from error + if recovered is not None: + return self._publish_success_status(plan, running, recovered, recovered.completed_at) + self._advance_status(running, CollectionState.FAILED, self._completion_time(completed_at)) + raise error + + def _load_optional_result(self, plan: CollectionPlan) -> CollectionResult | None: + try: + return self._load_valid_result(plan) + except FileNotFoundError: + return None + + def _load_valid_result( + self, + plan: CollectionPlan, + status: CollectionStatus | None = None, + ) -> CollectionResult: + result = self._collections.read_result_from(plan, Path(plan.container_destination)) + resolved_plan = self._reader.load_resolved_plan() + validated = validate_collection_result( + plan, + result, + expected_records=resolved_plan.invocation.authored.num_records, + output_format=resolved_plan.output.format, + ) + if status is not None and status.result != self._collections.get_result_reference(plan, validated): + raise StateCorruptionError("collection status does not bind its published result") + self._collections.verify_result_files(plan, validated, Path(plan.container_destination)) + return validated + + def _publish_success_status( + self, + plan: CollectionPlan, + previous: CollectionStatus, + result: CollectionResult, + updated_at: datetime, + ) -> CollectionResult: + current = CollectionStatus( + schema_version=1, + collection_id=previous.collection_id, + run_id=previous.run_id, + collection_plan=previous.collection_plan, + staging_directory=previous.staging_directory, + revision=previous.revision + 1, + updated_at=updated_at, + state=CollectionState.SUCCEEDED, + scheduler=previous.scheduler, + scheduler_observation=previous.scheduler_observation, + result=self._collections.get_result_reference(plan, result), + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + return result + + def _advance_status( + self, + previous: CollectionStatus, + state: CollectionState, + updated_at: datetime, + ) -> CollectionStatus: + current = _updated_status( + previous, + revision=previous.revision + 1, + updated_at=updated_at, + state=state, + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + return current + + def _require_scheduler_job(self, status: CollectionStatus) -> None: + scheduler = status.scheduler + if type(scheduler) is not int: + raise StateConflictError("collection worker requires an ordinary Slurm job identity") + observed_job_id = self._environment.get("SLURM_JOB_ID") + if observed_job_id != str(scheduler): + raise StateConflictError("collection worker must run inside its recorded Slurm job") + + def _validate_identity(self, plan_run_id: Identifier, status_run_id: Identifier) -> None: + if plan_run_id != self._run_id or status_run_id != self._run_id: + raise StateConflictError("collection records do not match the requested run") + + def _load_bound_plan(self, status: CollectionStatus) -> CollectionPlan: + plan = self._collections.read_plan(status.collection_id) + if status.collection_plan != self._collections.get_plan_reference(plan): + raise StateCorruptionError("collection status does not bind its persisted collection plan") + if status.staging_directory != derive_collection_staging_directory(plan): + raise StateCorruptionError("collection status does not bind its exact staging directory") + return plan + + def _completion_time(self, completed_at: datetime | None) -> datetime: + return _utc_now() if completed_at is None else completed_at + + +def _updated_status(previous: CollectionStatus, **updates: object) -> CollectionStatus: + payload = previous.model_dump(mode="python") + payload.update(updates) + return CollectionStatus.model_validate(payload) + + +def _validate_location( + workspace_root: str | Path, + run_id: Identifier, + collection_id: Identifier, +) -> tuple[Path, Identifier, Identifier]: + try: + root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + normalized_collection_id = _IDENTIFIER_ADAPTER.validate_python(collection_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid persisted collection worker location") from error + return Path(root), normalized_run_id, normalized_collection_id + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def main(argv: Sequence[str] | None = None) -> int: + """Run one allocation-local collection from explicit persisted identities.""" + parser = argparse.ArgumentParser() + parser.add_argument("--workspace-root", required=True) + parser.add_argument("--run-id", required=True) + parser.add_argument("--collection-id", required=True) + arguments = parser.parse_args(argv) + SlurmCollectionWorker( + arguments.workspace_root, + arguments.run_id, + arguments.collection_id, + ).run() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = ["SlurmCollectionWorker", "main"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py new file mode 100644 index 000000000..872eaf2e3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Profile-authorized host and container collection destinations.""" + +from __future__ import annotations + +import posixpath +from dataclasses import dataclass +from pathlib import Path + +from data_designer.slurm.config import ContainerMount +from data_designer.slurm.contracts import is_path_below, validate_absolute_path +from data_designer.slurm.images.records import validate_enroot_mount_path +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.errors import StateConflictError +from data_designer.slurm.state.outputs import CollectionPlan + + +@dataclass(frozen=True, slots=True) +class CollectionDestination: + """One normalized output directory and its authorized writable mount.""" + + host_path: str + container_path: str + mount: ContainerMount + + +class CollectionDestinationResolver: + """Resolve the pinned output root through the selected writable mount map.""" + + def resolve( + self, + plan: ResolvedSlurmRunPlan, + requested_destination: str | Path | None = None, + ) -> CollectionDestination: + """Return the exact plan output root in host and container namespaces.""" + raw_destination = plan.output.root if requested_destination is None else Path(requested_destination).as_posix() + try: + host_path = validate_absolute_path(raw_destination) + except ValueError as error: + raise StateConflictError("collection destination must be a normalized absolute path") from error + if host_path != plan.output.root: + raise StateConflictError("collection destination must match the pinned resolved output root") + + writable = tuple( + mount + for mount in plan.container_mounts + if not mount.read_only and (host_path == mount.source or is_path_below(host_path, mount.source)) + ) + if not writable: + raise StateConflictError("collection destination is not covered by a profile-authorized writable mount") + longest = max(len(mount.source) for mount in writable) + matches = tuple(mount for mount in writable if len(mount.source) == longest) + if len(matches) != 1: + raise StateConflictError("collection destination has an ambiguous writable mount mapping") + mount = matches[0] + if host_path == mount.source: + raise StateConflictError("collection destination must be below its writable mount source") + try: + validate_enroot_mount_path(plan.selected_profile.profile.workspace_root) + validate_enroot_mount_path(mount.source) + validate_enroot_mount_path(mount.target) + except ValueError as error: + raise StateConflictError("collection paths cannot be represented as safe Enroot mounts") from error + relative = posixpath.relpath(host_path, mount.source) + container_path = mount.target if relative == "." else posixpath.join(mount.target, relative) + return CollectionDestination(host_path, validate_absolute_path(container_path), mount) + + def validate_persisted( + self, + resolved_plan: ResolvedSlurmRunPlan, + collection_plan: CollectionPlan, + ) -> CollectionDestination: + """Reauthorize persisted collection intent against its pinned run plan.""" + destination = self.resolve(resolved_plan) + if collection_plan.run_id != resolved_plan.run_id: + raise StateConflictError("collection run identity does not match the resolved plan") + if collection_plan.host_destination != destination.host_path: + raise StateConflictError("collection host destination no longer matches resolved intent") + if collection_plan.container_destination != destination.container_path: + raise StateConflictError("collection container destination no longer matches resolved intent") + if collection_plan.num_partitions != resolved_plan.output.partitions: + raise StateConflictError("collection partition count no longer matches resolved intent") + return destination + + +__all__ = ["CollectionDestination", "CollectionDestinationResolver"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py index d332f7769..11759739a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -115,6 +115,40 @@ class CollectionShard(StateValue): winner_manifest: ArtifactReference +class RetryShard(StateValue): + """One failed planned shard selected for a new attempt.""" + + shard_id: ShardId + attempt_id: AttemptId + attempt_ordinal: PositiveInt + array_task_index: NonNegativeInt + + +class RetryPlan(StateRecord): + """Immutable failed-shard selection used for one retry submission.""" + + retry_id: Identifier + run_id: Identifier + created_at: datetime + resolved_plan: ArtifactReference + planned_shards: tuple[RetryShard, ...] = Field(min_length=1) + effective_resume_mode: Literal["never", "always"] + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_shards(self) -> RetryPlan: + shard_ids = tuple(shard.shard_id for shard in self.planned_shards) + task_indices = tuple(shard.array_task_index for shard in self.planned_shards) + if len(shard_ids) != len(set(shard_ids)): + raise ValueError("retry shard IDs must be unique") + if len(task_indices) != len(set(task_indices)): + raise ValueError("retry array-task indices must be unique") + if task_indices != tuple(sorted(task_indices)): + raise ValueError("retry shards must be ordered by array-task index") + return self + + class CollectionPlan(StateRecord): """Immutable inputs and destinations for deterministic collection.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py new file mode 100644 index 000000000..79a43b84f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process failed-shard retry orchestration.""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from contextlib import ExitStack, contextmanager +from datetime import datetime, timezone +from pathlib import Path +from typing import Literal, Protocol + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier, ShardId, validate_absolute_path +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError +from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt +from data_designer.slurm.launcher.renderer import render_generation_retry_script +from data_designer.slurm.state.base import SchedulerIdentity +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError, StateCorruptionError +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.observation import SchedulerObservationClient +from data_designer.slurm.state.observer import SlurmStateReconciler +from data_designer.slurm.state.outputs import RetryPlan, RetryShard +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.retry_records import RetryState, RetryStatus, validate_retry_status_transition +from data_designer.slurm.state.retry_storage import RetryStorage +from data_designer.slurm.state.scheduler import EffectiveAttemptState +from data_designer.slurm.state.status import RunStatus, ShardStatus +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_transition, + validate_shard_attempt_set, +) + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class RetryScheduler(SchedulerObservationClient, Protocol): + """Scheduler operations required by fresh-process retry.""" + + def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + """Submit one rendered retry array.""" + ... + + +class SlurmRetryCoordinator: + """Select retryable shards and durably submit their next attempts.""" + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + scheduler: RetryScheduler | None = None, + ) -> None: + root, normalized_run_id = _validate_location(workspace_root, run_id) + self._scheduler = scheduler if scheduler is not None else SlurmCommandClient() + self._state = StateStorage(root, normalized_run_id) + self._reader = StateReader(self._state, normalized_run_id) + self._retries = RetryStorage(self._state) + self._finalizer = WinnerFinalizer(self._state, self._reader) + self._reconciler = SlurmStateReconciler(root, normalized_run_id, self._scheduler) + self._run_id = normalized_run_id + + def retry( + self, + *, + shard_ids: Sequence[ShardId] | None = None, + effective_resume_mode: Literal["never", "always"], + observed_at: datetime | None = None, + ) -> tuple[AttemptManifest, ...]: + """Refresh state, submit exactly the retryable selection, and publish attempts.""" + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + try: + if effective_resume_mode not in {"never", "always"}: + raise StateConflictError("effective resume mode must be 'never' or 'always'") + with self._retries.acquire_lock(): + self._retries.discard_incomplete_tail() + self._settle_submitted_retry(timestamp) + status = self._reconciler.refresh(observed_at=timestamp) + active = self._load_active_retry(status, shard_ids, effective_resume_mode) + if active is not None: + return active + selected = _select_retryable_shards(status, shard_ids) + plan = self._build_retry_plan(status, selected, effective_resume_mode, timestamp) + script = render_generation_retry_script(self._reader.load_resolved_plan(status.run), plan) + with self._acquire_selection_locks(selected): + self._require_fresh_selection(status, selected) + self._persist_prepared_retry(plan, timestamp) + receipt = self._submit(plan, script, timestamp) + attempts, superseded = self._publish_attempts(plan, receipt.job_id, timestamp) + self._settle_retry(plan, receipt.job_id, timestamp, superseded=superseded) + return attempts + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot retry persisted run {self._run_id!r}") from error + + def _settle_submitted_retry(self, updated_at: datetime) -> None: + retry_ids = self._retries.list_retry_ids() + if not retry_ids: + return + latest_id = retry_ids[-1] + status = self._retries.read_status(latest_id) + plan = self._load_bound_plan(status) + if status.state is RetryState.PREPARED: + raise StateConflictError("previous retry submission has an ambiguous scheduler outcome") + if status.state is not RetryState.SUBMITTED: + return + assert status.array_job_id is not None + shard_ids = tuple(shard.shard_id for shard in plan.planned_shards) + run_status = self._reconciler.refresh(observed_at=updated_at) + selected = tuple(_get_shard_status(run_status, shard_id) for shard_id in shard_ids) + with self._acquire_selection_locks(selected): + _, superseded = self._publish_attempts(plan, status.array_job_id, plan.created_at) + self._settle_retry(plan, status.array_job_id, updated_at, superseded=superseded) + + def _build_retry_plan( + self, + status: RunStatus, + selected: tuple[ShardStatus, ...], + effective_resume_mode: Literal["never", "always"], + created_at: datetime, + ) -> RetryPlan: + plan = self._reader.load_resolved_plan(status.run) + requested_resume = plan.invocation.authored.resume + if requested_resume != "if_possible" and effective_resume_mode != requested_resume: + raise StateConflictError("effective resume mode does not match the pinned resolved plan") + return RetryPlan( + schema_version=1, + retry_id=self._retries.get_next_retry_id(), + run_id=status.run.run_id, + created_at=created_at, + resolved_plan=status.run.resolved_plan, + planned_shards=tuple( + RetryShard( + shard_id=shard_status.shard.shard_id, + attempt_id=f"attempt-{len(shard_status.attempts) + 1:04d}", + attempt_ordinal=len(shard_status.attempts) + 1, + array_task_index=plan.shards[shard_status.shard.shard_index].array_task_index, + ) + for shard_status in selected + ), + effective_resume_mode=effective_resume_mode, + ) + + def _load_active_retry( + self, + status: RunStatus, + requested_shard_ids: Sequence[ShardId] | None, + effective_resume_mode: Literal["never", "always"], + ) -> tuple[AttemptManifest, ...] | None: + retry_ids = self._retries.list_retry_ids() + if not retry_ids: + return None + retry_status = self._retries.read_status(retry_ids[-1]) + retry_plan = self._load_bound_plan(retry_status) + if retry_status.state is not RetryState.COMPLETED: + return None + planned_ids = tuple(shard.shard_id for shard in retry_plan.planned_shards) + if requested_shard_ids is None and any( + _is_retryable(shard) and shard.shard.shard_id not in planned_ids for shard in status.shards + ): + return None + if retry_plan.effective_resume_mode != effective_resume_mode or not self._matches_requested_shards( + requested_shard_ids, planned_ids + ): + return None + return self._get_active_attempts(status, retry_plan) + + @staticmethod + def _matches_requested_shards( + requested_shard_ids: Sequence[ShardId] | None, + planned_ids: tuple[ShardId, ...], + ) -> bool: + if requested_shard_ids is None: + return True + requested = tuple(requested_shard_ids) + return len(requested) == len(set(requested)) and set(requested) == set(planned_ids) + + @staticmethod + def _get_active_attempts(status: RunStatus, retry_plan: RetryPlan) -> tuple[AttemptManifest, ...] | None: + statuses_by_shard = {shard.shard.shard_id: shard for shard in status.shards} + attempts: list[AttemptManifest] = [] + for planned in retry_plan.planned_shards: + shard_status = statuses_by_shard.get(planned.shard_id) + if shard_status is None or not shard_status.attempts: + return None + matching = next( + (item for item in shard_status.attempts if item.attempt.attempt_id == planned.attempt_id), + None, + ) + if shard_status.winner is not None: + if ( + shard_status.winner.attempt_id != planned.attempt_id + or matching is None + or matching.effective_state is not EffectiveAttemptState.SUCCEEDED + ): + return None + attempts.append(matching.attempt) + continue + latest = shard_status.attempts[-1] + if latest.attempt.attempt_id != planned.attempt_id or latest.effective_state not in { + EffectiveAttemptState.PENDING, + EffectiveAttemptState.RUNNING, + EffectiveAttemptState.ACCOUNTING_LAG, + }: + return None + attempts.append(latest.attempt) + return tuple(attempts) + + @contextmanager + def _acquire_selection_locks(self, selected: tuple[ShardStatus, ...]) -> Iterator[None]: + with ExitStack() as resources: + for shard_status in selected: + resources.enter_context(self._state.acquire_resume_and_shard_locks(shard_status.shard.shard_id)) + yield + + def _require_fresh_selection(self, status: RunStatus, selected: tuple[ShardStatus, ...]) -> None: + for shard_status in selected: + run, plan, shard = self._reader.load_shard_context(shard_status.shard.shard_id) + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self._finalizer.load_optional_winner(run, plan, shard, attempts) + expected_attempts = tuple(attempt_status.attempt for attempt_status in shard_status.attempts) + expected_observations = tuple(attempt_status.scheduler for attempt_status in shard_status.attempts) + current_observations = tuple( + self._reader.load_optional_scheduler_observation(attempt) for attempt in attempts + ) + if run != status.run or shard != shard_status.shard or attempts != expected_attempts: + raise StateConflictError("persisted shard changed after retry reconciliation; retry again") + if current_observations != expected_observations: + raise StateConflictError("scheduler evidence changed after retry reconciliation; retry again") + if winner is not None: + raise StateConflictError(f"shard {shard.shard_id!r} already has an immutable winner") + + def _persist_prepared_retry(self, plan: RetryPlan, timestamp: datetime) -> None: + self._retries.ensure_retry(plan.retry_id) + self._retries.publish_plan(plan) + self._retries.publish_status( + RetryStatus( + schema_version=1, + retry_id=plan.retry_id, + run_id=plan.run_id, + retry_plan=self._retries.get_plan_reference(plan), + revision=1, + updated_at=timestamp, + state=RetryState.PREPARED, + ) + ) + + def _submit(self, plan: RetryPlan, script: str, timestamp: datetime) -> SlurmJobSubmissionReceipt: + try: + receipt = self._scheduler.submit_script(script) + except SlurmSubmissionError as error: + if not error.may_have_succeeded: + self._fail_retry(plan, timestamp) + raise SlurmStateError(f"cannot submit retry {plan.retry_id!r}") from error + except SlurmLauncherError as error: + raise SlurmStateError(f"cannot submit retry {plan.retry_id!r}") from error + plan_reference = self._retries.get_plan_reference(plan) + submitted = RetryStatus( + schema_version=1, + retry_id=plan.retry_id, + run_id=plan.run_id, + retry_plan=plan_reference, + revision=2, + updated_at=timestamp, + state=RetryState.SUBMITTED, + array_job_id=receipt.job_id, + ) + validate_retry_status_transition(self._retries.read_status(plan.retry_id), submitted) + self._retries.replace_status(submitted) + return receipt + + def _publish_attempts( + self, + retry_plan: RetryPlan, + array_job_id: int, + timestamp: datetime, + ) -> tuple[tuple[AttemptManifest, ...], bool]: + run = self._reader.load_run() + plan = self._reader.load_resolved_plan(run) + if retry_plan.run_id != run.run_id or retry_plan.resolved_plan != run.resolved_plan: + raise StateConflictError("retry plan does not bind the current persisted run") + prepared: list[tuple[AttemptManifest, bool]] = [] + superseded = False + for selected in retry_plan.planned_shards: + shard = self._reader.load_shard_context(selected.shard_id)[2] + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self._finalizer.load_optional_winner(run, plan, shard, attempts) + attempt = AttemptManifest( + schema_version=1, + run_id=run.run_id, + shard_id=selected.shard_id, + attempt_id=selected.attempt_id, + attempt_ordinal=selected.attempt_ordinal, + resolved_plan=run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity( + array_job_id=array_job_id, + array_task_id=selected.array_task_index, + ), + created_at=timestamp, + updated_at=timestamp, + ) + existing = next((item for item in attempts if item.attempt_id == selected.attempt_id), None) + if existing is not None: + try: + validate_attempt_transition(attempt, existing) + except StateContractError as error: + raise StateConflictError( + f"retry attempt {selected.attempt_id!r} contains incompatible state" + ) from error + if winner is not None and winner.attempt_id != selected.attempt_id: + superseded = True + continue + prepared.append((existing, False)) + continue + if winner is not None: + superseded = True + continue + self._finalizer.require_no_winner(run, plan, shard, attempts) + if selected.attempt_ordinal != len(attempts) + 1: + raise StateConflictError("retry attempt ordinal is no longer next for its shard") + self._reader.validate_attempt_against_plan(run, plan, shard, attempt) + validate_shard_attempt_set(run, shard, attempts + (attempt,)) + prepared.append((attempt, True)) + published: list[AttemptManifest] = [] + for attempt, requires_publication in prepared: + if requires_publication: + self._state.publish_attempt(attempt) + published.append(attempt) + return tuple(published), superseded + + def _settle_retry( + self, + plan: RetryPlan, + array_job_id: int, + timestamp: datetime, + *, + superseded: bool, + ) -> None: + if superseded: + self._fail_retry(plan, timestamp, array_job_id=array_job_id) + else: + self._complete_retry(plan, array_job_id, timestamp) + + def _fail_retry(self, plan: RetryPlan, timestamp: datetime, *, array_job_id: int | None = None) -> None: + previous = self._retries.read_status(plan.retry_id) + failed = RetryStatus( + schema_version=1, + retry_id=plan.retry_id, + run_id=plan.run_id, + retry_plan=self._retries.get_plan_reference(plan), + revision=previous.revision + 1, + updated_at=timestamp, + state=RetryState.FAILED, + array_job_id=array_job_id, + ) + validate_retry_status_transition(previous, failed) + self._retries.replace_status(failed) + + def _complete_retry(self, plan: RetryPlan, array_job_id: int, timestamp: datetime) -> None: + previous = self._retries.read_status(plan.retry_id) + if previous.state is RetryState.COMPLETED: + return + completed = RetryStatus( + schema_version=1, + retry_id=plan.retry_id, + run_id=plan.run_id, + retry_plan=self._retries.get_plan_reference(plan), + revision=previous.revision + 1, + updated_at=timestamp, + state=RetryState.COMPLETED, + array_job_id=array_job_id, + ) + validate_retry_status_transition(previous, completed) + self._retries.replace_status(completed) + + def _load_bound_plan(self, status: RetryStatus) -> RetryPlan: + plan = self._retries.read_plan(status.retry_id) + if status.retry_plan != self._retries.get_plan_reference(plan): + raise StateCorruptionError("retry status does not bind its persisted retry plan") + return plan + + +def _select_retryable_shards( + status: RunStatus, + shard_ids: Sequence[ShardId] | None, +) -> tuple[ShardStatus, ...]: + requested = None if shard_ids is None else tuple(shard_ids) + if requested is not None and len(requested) != len(set(requested)): + raise StateConflictError("explicit retry shard IDs must be unique") + known = {shard.shard.shard_id for shard in status.shards} + if requested is not None and not set(requested).issubset(known): + raise StateConflictError("explicit retry selection contains an unknown shard") + selected = tuple( + shard + for shard in status.shards + if (requested is None or shard.shard.shard_id in requested) and _is_retryable(shard) + ) + expected_count = len( + tuple(shard for shard in status.shards if requested is None or shard.shard.shard_id in requested) + ) + if requested is not None and len(selected) != expected_count: + raise StateConflictError("explicit retry selection includes a sealed or nonterminal shard") + if not selected: + raise StateConflictError("run has no retryable shards") + return selected + + +def _is_retryable(shard: ShardStatus) -> bool: + return ( + shard.winner is None + and bool(shard.attempts) + and shard.attempts[-1].effective_state in {EffectiveAttemptState.FAILED, EffectiveAttemptState.UNKNOWN} + ) + + +def _get_shard_status(status: RunStatus, shard_id: ShardId) -> ShardStatus: + shard_status = next((shard for shard in status.shards if shard.shard.shard_id == shard_id), None) + if shard_status is None: + raise StateCorruptionError(f"retry plan references unknown shard {shard_id!r}") + return shard_status + + +def _validate_location(workspace_root: str | Path, run_id: Identifier) -> tuple[Path, Identifier]: + try: + root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid persisted retry location") from error + return Path(root), normalized_run_id + + +__all__ = ["RetryScheduler", "SlurmRetryCoordinator"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py new file mode 100644 index 000000000..961f54bba --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Persisted retry submission lifecycle.""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from pydantic import PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.state.base import StateRecord, validate_utc_timestamp + + +class RetryState(str, Enum): + """Durable state of one failed-shard retry request.""" + + PREPARED = "prepared" + SUBMITTED = "submitted" + COMPLETED = "completed" + FAILED = "failed" + + +class RetryStatus(StateRecord): + """Atomically replaced submission progress for one retry plan.""" + + retry_id: Identifier + run_id: Identifier + retry_plan: ArtifactReference + revision: PositiveInt + updated_at: datetime + state: RetryState + array_job_id: PositiveInt | None = None + + _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_scheduler_identity(self) -> RetryStatus: + if self.state in {RetryState.SUBMITTED, RetryState.COMPLETED} and self.array_job_id is None: + raise ValueError("submitted retry state requires an array job identity") + if self.state is RetryState.PREPARED and self.array_job_id is not None: + raise ValueError("prepared retry state cannot contain an array job identity") + return self + + +def validate_retry_status_transition(previous: RetryStatus, current: RetryStatus) -> RetryStatus: + """Require immutable retry identity and one-way submission progress.""" + if previous.retry_id != current.retry_id or previous.run_id != current.run_id: + raise ValueError("retry status identity cannot change") + if previous.retry_plan != current.retry_plan: + raise ValueError("retry status plan identity cannot change") + if current.revision != previous.revision + 1: + raise ValueError("retry status revision must increase by one") + if current.updated_at < previous.updated_at: + raise ValueError("retry status timestamp cannot move backward") + allowed = { + RetryState.PREPARED: {RetryState.SUBMITTED, RetryState.FAILED}, + RetryState.SUBMITTED: {RetryState.COMPLETED, RetryState.FAILED}, + RetryState.COMPLETED: set(), + RetryState.FAILED: set(), + } + if current.state not in allowed[previous.state]: + raise ValueError(f"retry status cannot move from {previous.state.value!r} to {current.state.value!r}") + return current + + +__all__ = ["RetryState", "RetryStatus", "validate_retry_status_transition"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py new file mode 100644 index 000000000..489af26b0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Descriptor-bound persistence for retry plans and submission progress.""" + +from __future__ import annotations + +import os +import re +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.state.filesystem import ( + acquire_file_lock, + ensure_private_child_directory, + is_state_temporary_name, + open_verified_child_directory, + publish_immutable_text, + replace_text, +) +from data_designer.slurm.state.outputs import RetryPlan +from data_designer.slurm.state.retry_records import RetryStatus +from data_designer.slurm.state.storage import StateStorage + +_RETRIES_DIRECTORY = "retries" +_RETRY_LOCK = "retry.lock" +_PLAN_FILENAME = "plan.json" +_STATUS_FILENAME = "status.json" +_RETRY_PATTERN = re.compile(r"^retry-[0-9]{4,}$") +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 + + +class RetryStorage: + """Persist retry journals separately from run and attempt storage.""" + + def __init__(self, state_storage: StateStorage) -> None: + self._state = state_storage + self.retries_root = state_storage.run_root / _RETRIES_DIRECTORY + + @contextmanager + def acquire_lock(self) -> Iterator[None]: + """Serialize retry selection and submission for one run.""" + with self._state.open_run_directory() as run_descriptor: + with acquire_file_lock(run_descriptor, _RETRY_LOCK, self._state.run_root / _RETRY_LOCK): + yield + + def get_next_retry_id(self) -> Identifier: + """Return the next monotonic retry identity.""" + return f"retry-{len(self.list_retry_ids()) + 1:04d}" + + def discard_incomplete_tail(self) -> None: + """Discard one trailing journal that cannot have reached submission.""" + try: + with self._open_retries_directory() as retries_descriptor: + retry_ids = _validated_retry_ids(tuple(os.listdir(retries_descriptor))) + if not retry_ids: + return + retry_id = retry_ids[-1] + with open_verified_child_directory( + retries_descriptor, + retry_id, + self.get_retry_root(retry_id), + ) as retry_descriptor: + if _record_exists(retry_descriptor, _STATUS_FILENAME): + return + _discard_prepared_files(retry_descriptor, self.get_retry_root(retry_id)) + os.rmdir(retry_id, dir_fd=retries_descriptor) + os.fsync(retries_descriptor) + except FileNotFoundError: + return + + def list_retry_ids(self) -> tuple[Identifier, ...]: + """List a complete monotonic set of managed retry directories.""" + try: + with self._open_retries_directory() as descriptor: + names = tuple(os.listdir(descriptor)) + except FileNotFoundError: + return () + return _validated_retry_ids(names) + + def ensure_retry(self, retry_id: Identifier) -> None: + """Create one private retry journal directory.""" + with self._state.open_run_directory() as run_descriptor: + ensure_private_child_directory(run_descriptor, _RETRIES_DIRECTORY, self.retries_root) + with open_verified_child_directory(run_descriptor, _RETRIES_DIRECTORY, self.retries_root) as descriptor: + ensure_private_child_directory(descriptor, retry_id, self.get_retry_root(retry_id)) + + def get_retry_root(self, retry_id: Identifier) -> Path: + return self.retries_root / retry_id + + def get_plan_path(self, retry_id: Identifier) -> Path: + """Return the canonical immutable retry-plan path.""" + return self.get_retry_root(retry_id) / _PLAN_FILENAME + + def get_plan_reference(self, plan: RetryPlan) -> ArtifactReference: + """Return the exact path and digest bound by retry status.""" + return ArtifactReference(path=self.get_plan_path(plan.retry_id).as_posix(), sha256=plan.compute_sha256()) + + def publish_plan(self, plan: RetryPlan) -> None: + self._require_run_id(plan.run_id) + with self._open_retry_directory(plan.retry_id) as descriptor: + publish_immutable_text( + descriptor, + _PLAN_FILENAME, + plan.serialize_json(), + self.get_plan_path(plan.retry_id), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def read_plan(self, retry_id: Identifier) -> RetryPlan: + with self._open_retry_directory(retry_id) as descriptor: + plan = self._state.read_record( + descriptor, + _PLAN_FILENAME, + self.get_plan_path(retry_id), + RetryPlan, + ) + if plan.retry_id != retry_id or plan.run_id != self._state.run_id: + raise OSError("retry plan identity does not match its persisted location") + return plan + + def publish_status(self, status: RetryStatus) -> None: + self._require_run_id(status.run_id) + with self._open_retry_directory(status.retry_id) as descriptor: + publish_immutable_text( + descriptor, + _STATUS_FILENAME, + status.serialize_json(), + self.get_retry_root(status.retry_id) / _STATUS_FILENAME, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def replace_status(self, status: RetryStatus) -> None: + self._require_run_id(status.run_id) + with self._open_retry_directory(status.retry_id) as descriptor: + replace_text( + descriptor, + _STATUS_FILENAME, + status.serialize_json(), + self.get_retry_root(status.retry_id) / _STATUS_FILENAME, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def read_status(self, retry_id: Identifier) -> RetryStatus: + with self._open_retry_directory(retry_id) as descriptor: + status = self._state.read_record( + descriptor, + _STATUS_FILENAME, + self.get_retry_root(retry_id) / _STATUS_FILENAME, + RetryStatus, + ) + if status.retry_id != retry_id or status.run_id != self._state.run_id: + raise OSError("retry status identity does not match its persisted location") + return status + + @contextmanager + def _open_retries_directory(self) -> Iterator[int]: + with self._state.open_run_directory() as run_descriptor: + with open_verified_child_directory(run_descriptor, _RETRIES_DIRECTORY, self.retries_root) as descriptor: + yield descriptor + + @contextmanager + def _open_retry_directory(self, retry_id: Identifier) -> Iterator[int]: + with self._open_retries_directory() as retries_descriptor: + with open_verified_child_directory( + retries_descriptor, + retry_id, + self.get_retry_root(retry_id), + ) as descriptor: + yield descriptor + + def _require_run_id(self, run_id: Identifier) -> None: + if run_id != self._state.run_id: + raise OSError("retry record run identity does not match storage") + + +def _validated_retry_ids(names: tuple[str, ...]) -> tuple[Identifier, ...]: + if any(_RETRY_PATTERN.fullmatch(name) is None for name in names): + raise OSError("retry state contains an unowned directory") + ordered = tuple(sorted(names, key=lambda name: int(name.rsplit("-", maxsplit=1)[1]))) + expected = tuple(f"retry-{index:04d}" for index in range(1, len(ordered) + 1)) + if ordered != expected: + raise OSError("retry identities are not a complete monotonic sequence") + return ordered + + +def _record_exists(directory_descriptor: int, name: str) -> bool: + try: + os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + except FileNotFoundError: + return False + return True + + +def _discard_prepared_files(directory_descriptor: int, display_path: Path) -> None: + names = tuple(os.listdir(directory_descriptor)) + if any(name != _PLAN_FILENAME and not is_state_temporary_name(name) for name in names): + raise OSError(f"incomplete retry journal {display_path} contains an unowned entry") + for name in names: + status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not stat.S_ISREG(status.st_mode) or status.st_mode & 0o077: + raise OSError(f"incomplete retry journal entry {display_path / name} is unsafe") + os.unlink(name, dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + + +__all__ = ["RetryStorage"] diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 183680d34..e2cee1699 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -10,7 +10,7 @@ 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.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -40,6 +40,23 @@ def test_client_submits_verified_script_text_through_standard_input() -> None: assert runner.inputs == [script] +def test_script_submission_classifies_scheduler_rejection_as_definite() -> None: + runner = FakeSlurmRunner() + runner.script_next("sbatch", FakeCommandResponse(stderr="rejected", returncode=2)) + + with pytest.raises(SlurmSubmissionError, match="rejected") as error: + SlurmCommandClient(runner).submit_script("#!/bin/sh\n") + + assert not error.value.may_have_succeeded + + +def test_script_submission_classifies_timeout_as_ambiguous() -> None: + with pytest.raises(SlurmSubmissionError, match="timed out") as error: + SlurmCommandClient(_TimeoutRunner()).submit_script("#!/bin/sh\n") + + assert error.value.may_have_succeeded + + def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) client.submit("run.sbatch") @@ -282,7 +299,14 @@ def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: class _TimeoutRunner: - def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + def run( + self, + command: Sequence[str], + *, + check: bool = False, + input_text: str | None = None, + ) -> subprocess.CompletedProcess[str]: + del check, input_text raise subprocess.TimeoutExpired(command, 30.0) diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py new file mode 100644 index 000000000..d8df08bed --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import subprocess +from datetime import datetime, timezone + +from data_designer.slurm.config import ContainerMount +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.launcher.collection import render_collection_script +from data_designer.slurm.launcher.renderer import render_generation_retry_script +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import CollectionPlan, CollectionShard, RetryPlan, RetryShard +from data_designer.slurm.state.destinations import CollectionDestinationResolver + + +def test_collection_renderer_uses_authorized_mounts_and_no_gpu_directives( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + workspace_mount = ContainerMount(source="/workspace", target="/workspace", read_only=False) + output_mount = ContainerMount( + source="/workspace/primary/runs/run-001", + target="/exports", + read_only=False, + ) + mounts = (workspace_mount, output_mount) + profile = multi_node_plan.selected_profile.profile.model_copy(update={"container_mounts": list(mounts)}) + selection = multi_node_plan.selected_profile.model_copy( + update={ + "profile": profile, + "profile_sha256": compute_canonical_json_sha256(profile.model_dump(mode="json")), + } + ) + plan = ResolvedSlurmRunPlan.model_validate_json( + json.dumps( + multi_node_plan.model_copy(update={"container_mounts": mounts, "selected_profile": selection}).model_dump( + mode="json" + ) + ) + ) + destination = CollectionDestinationResolver().resolve(plan) + collection = CollectionPlan( + schema_version=1, + collection_id="collection-0001", + run_id=plan.run_id, + created_at=datetime(2026, 9, 2, tzinfo=timezone.utc), + resolved_plan=ArtifactReference( + path="/workspace/primary/runs/run-001/resolved-plan.json", + sha256=plan.compute_sha256(), + ), + planned_shards=( + CollectionShard( + shard_id="shard-00000", + winner_manifest=ArtifactReference( + path="/workspace/primary/runs/run-001/shards/shard-00000/winner.json", + sha256="a" * 64, + ), + ), + ), + host_destination=destination.host_path, + container_destination=destination.container_path, + num_partitions=plan.output.partitions, + ) + + script = render_collection_script(plan, collection, destination) + + assert 'readonly DD_STATE_MOUNT="/workspace/primary:/workspace/primary"' in script + assert 'readonly DD_OUTPUT_MOUNT="/workspace/primary/runs/run-001:/exports"' in script + assert ( + 'readonly DD_COLLECTION_PLAN="/workspace/primary/runs/run-001/collections/collection-0001/plan.json"' in script + ) + assert "data_designer.slurm.state.collection_worker" in script + assert "#SBATCH --gres=" not in script + assert "#SBATCH --gpus=" not in script + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=multi_node_plan.run_id, + created_at=datetime(2026, 9, 2, tzinfo=timezone.utc), + resolved_plan=ArtifactReference( + path="/workspace/primary/runs/run-001/resolved-plan.json", + sha256=multi_node_plan.compute_sha256(), + ), + planned_shards=( + RetryShard( + shard_id="shard-00001", + attempt_id="attempt-0002", + attempt_ordinal=2, + array_task_index=1, + ), + ), + effective_resume_mode="never", + ) + + script = render_generation_retry_script(multi_node_plan, retry) + + assert "#SBATCH --array=1%2" in script + assert 'DD_ATTEMPT_ORDINAL="0002"' in script + assert 'readonly DD_ATTEMPT_MANIFEST="${DD_ATTEMPT_DIR}/attempt.json"' in script + assert script.index("DD_ATTEMPT_MANIFEST") < script.index("DD_RUNTIME_DIR") + assert "data_designer.slurm.state.attempt_identity" in script + assert '--array-job-id "${DD_ARRAY_JOB_ID}" --array-task-id "${DD_ARRAY_TASK_ID}"' in script + assert script.index("data_designer.slurm.state.attempt_identity") < script.index("DD_RUNTIME_DIR") + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py new file mode 100644 index 000000000..eff457973 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -0,0 +1,1454 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import csv +import hashlib +import json +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import cast + +import pytest +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask + +import data_designer.lazy_heavy_imports as lazy +import data_designer.slurm.state.collection_filesystem as collection_filesystem +import data_designer.slurm.state.collection_merge as collection_merge +import data_designer.slurm.state.collection_storage as collection_storage_module +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmSubmissionError +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + CollectionState, + RetryState, + RunManifest, + SchedulerIdentity, + ShardManifest, + SlurmCollectionCoordinator, + SlurmRetryCoordinator, + SlurmStateError, + SlurmStateReconciler, + SlurmStateWriter, + StateConflictError, + StateCorruptionError, + StateNotFoundError, + compute_candidate_schema_digest, +) +from data_designer.slurm.state.attempt_identity import require_attempt_scheduler_identity +from data_designer.slurm.state.collection_filesystem import derive_collection_staging_directory +from data_designer.slurm.state.collection_storage import CollectionStorage +from data_designer.slurm.state.collection_worker import SlurmCollectionWorker +from data_designer.slurm.state.retry_storage import RetryStorage +from data_designer.slurm.state.storage import StateStorage + + +@dataclass(frozen=True, slots=True) +class _RunCase: + workspace: Path + plan: ResolvedSlurmRunPlan + run: RunManifest + shards: tuple[ShardManifest, ...] + writer: SlurmStateWriter + created_at: datetime + + +def test_retry_refreshes_failed_shard_and_publishes_exact_next_attempt( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + attempts = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + assert len(attempts) == 1 + assert attempts[0].attempt_id == "attempt-0002" + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=0) + require_attempt_scheduler_identity( + case.workspace, + case.plan.run_id, + attempts[0].shard_id, + attempts[0].attempt_id, + SchedulerIdentity(array_job_id=4201, array_task_id=0), + ) + with pytest.raises(StateConflictError, match="scheduler identity"): + require_attempt_scheduler_identity( + case.workspace, + case.plan.run_id, + attempts[0].shard_id, + attempts[0].attempt_id, + SchedulerIdentity(array_job_id=9999, array_task_id=0), + ) + assert case.writer.load_attempts(case.shards[0].shard_id) == (attempt, attempts[0]) + assert "#SBATCH --array=0" in cast(str, runner.inputs[-1]) + assert 'DD_ATTEMPT_ORDINAL="0002"' in cast(str, runner.inputs[-1]) + assert ( + coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + == attempts + ) + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 + + +def test_retry_rejects_a_nonterminal_explicit_shard( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + case.writer.create_attempt( + _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + ) + + with pytest.raises(StateConflictError, match="sealed or nonterminal"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + +def test_retry_rejects_a_shard_with_an_immutable_winner( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + _publish_all_winners(case) + + with pytest.raises(StateConflictError, match="sealed or nonterminal"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=7), + ) + + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 1 + + +def test_retry_accepts_unknown_after_bounded_accounting_lag( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state=None) + SlurmStateReconciler(case.workspace, case.plan.run_id, scheduler).refresh( + observed_at=case.created_at + timedelta(minutes=3) + ) + + attempts = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=9), + ) + + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=0) + + +def test_retry_submits_only_the_failed_sparse_array_task( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + initial_tasks = tuple( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=shard.shard_index)) for shard in case.shards + ) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=initial_tasks), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=1)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + for shard in case.shards: + case.writer.create_attempt( + _submitted_attempt( + case, + shard, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=shard.shard_index), + ) + ) + runner.set_task_state(initial_tasks[0].scheduler, queue_state="RUNNING", accounting_state=None) + runner.set_task_state(initial_tasks[1].scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + + attempts = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + assert tuple(attempt.shard_id for attempt in attempts) == ("shard-00001",) + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=1) + assert "#SBATCH --array=1%2" in cast(str, runner.inputs[-1]) + assert 'case "${DD_ARRAY_TASK_ID}"' in cast(str, runner.inputs[-1]) + + +def test_retry_ambiguous_submission_remains_prepared_and_blocks_duplicate( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + submissions = 0 + + def ambiguous_submit(script: str) -> object: + nonlocal submissions + del script + submissions += 1 + raise SlurmSubmissionError("sbatch could not be executed: command timed out", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", ambiguous_submit) + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + + with pytest.raises(SlurmStateError, match="cannot submit retry"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=5)) + + retry_storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert retry_storage.read_status("retry-0001").state is RetryState.PREPARED + with pytest.raises(StateConflictError, match="ambiguous scheduler outcome"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=6)) + assert submissions == 1 + + +def test_retry_definite_submission_failure_settles_and_can_be_retried( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + runner.script_next("sbatch", FakeCommandResponse(stderr="submission rejected", returncode=2)) + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + + with pytest.raises(SlurmStateError, match="cannot submit retry"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=5)) + + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("retry-0001").state is RetryState.FAILED + attempts = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=0) + assert storage.read_status("retry-0002").state is RetryState.COMPLETED + + +def test_concurrent_retry_requests_converge_on_one_array( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + + def retry() -> tuple[AttemptManifest, ...]: + return SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + results = tuple(executor.map(lambda _: retry(), range(2))) + + assert results[0] == results[1] + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 + + +def test_default_retry_does_not_hide_failures_outside_an_active_explicit_subset( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + initial_tasks = tuple( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=index)) for index in range(2) + ) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=initial_tasks), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4301, array_task_id=1)),)), + ) + ) + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + for shard, task in zip(case.shards, initial_tasks, strict=True): + case.writer.create_attempt(_submitted_attempt(case, shard, scheduler=task.scheduler)) + runner.set_task_state(task.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + + first = coordinator.retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + second = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + + assert tuple(attempt.shard_id for attempt in first) == (case.shards[0].shard_id,) + assert tuple(attempt.shard_id for attempt in second) == (case.shards[1].shard_id,) + assert second[0].scheduler == SchedulerIdentity(array_job_id=4301, array_task_id=1) + + +def test_retry_discards_trailing_journal_interrupted_before_submission( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + original_publish = RetryStorage.publish_status + + def interrupt_before_status(self: RetryStorage, status: object) -> None: + del self, status + raise OSError("injected journal interruption") + + monkeypatch.setattr(RetryStorage, "publish_status", interrupt_before_status) + with pytest.raises(SlurmStateError, match="cannot retry persisted run"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=5)) + monkeypatch.setattr(RetryStorage, "publish_status", original_publish) + + attempts = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + + assert attempts[0].attempt_id == "attempt-0002" + assert RetryStorage(StateStorage(case.workspace, case.plan.run_id)).list_retry_ids() == ("retry-0001",) + + +def test_retry_recovers_evolved_attempt_and_its_exact_winner( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + initial = FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)) + retried = FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)) + runner = FakeSlurmRunner(arrays=(FakeSlurmArray(tasks=(initial,)), FakeSlurmArray(tasks=(retried,)))) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + first_attempt = _submitted_attempt(case, case.shards[0], scheduler=initial.scheduler) + case.writer.create_attempt(first_attempt) + runner.set_task_state(initial.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + original_complete = SlurmRetryCoordinator._complete_retry + + def interrupt_after_attempts(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("injected completion interruption") + + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", interrupt_after_attempts) + with pytest.raises(SlurmStateError, match="cannot retry persisted run"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=5)) + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", original_complete) + retry_attempt = case.writer.load_attempts(case.shards[0].shard_id)[-1] + running = _copy_attempt( + retry_attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=6), + ) + case.writer.update_attempt(running) + with case.writer.acquire_dataset_workspace(case.shards[0].shard_id, running.attempt_id, "never") as dataset_path: + _publish_candidate(case, case.shards[0], running, dataset_path) + succeeded = case.writer.load_attempts(case.shards[0].shard_id)[-1] + case.writer.finalize_winner( + case.shards[0].shard_id, + succeeded.attempt_id, + published_at=case.created_at + timedelta(minutes=10), + ) + runner.set_task_state(retried.scheduler, queue_state=None, accounting_state="COMPLETED") + + recovered = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=11), + ) + + assert recovered == (succeeded,) + assert ( + RetryStorage(StateStorage(case.workspace, case.plan.run_id)).read_status("retry-0001").state + is RetryState.COMPLETED + ) + + +def test_retry_settles_submitted_journal_before_serving_a_disjoint_selection( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + initial_tasks = tuple( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=index)) for index in range(2) + ) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=initial_tasks), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4301, array_task_id=1)),)), + ) + ) + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + for shard, task in zip(case.shards, initial_tasks, strict=True): + case.writer.create_attempt(_submitted_attempt(case, shard, scheduler=task.scheduler)) + runner.set_task_state(task.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + original_complete = SlurmRetryCoordinator._complete_retry + + def interrupt_after_attempts(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("injected completion interruption") + + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", interrupt_after_attempts) + with pytest.raises(SlurmStateError, match="cannot retry persisted run"): + coordinator.retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", original_complete) + + attempts = coordinator.retry( + shard_ids=(case.shards[1].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + + assert tuple(attempt.shard_id for attempt in attempts) == (case.shards[1].shard_id,) + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4301, array_task_id=1) + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("retry-0001").state is RetryState.COMPLETED + assert storage.read_status("retry-0002").state is RetryState.COMPLETED + + +def test_retry_does_not_return_recovered_attempt_for_a_different_resume_mode( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + initial_tasks = tuple( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=index)) for index in range(2) + ) + retried = FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)) + runner = FakeSlurmRunner(arrays=(FakeSlurmArray(tasks=initial_tasks), FakeSlurmArray(tasks=(retried,)))) + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + for shard, task in zip(case.shards, initial_tasks, strict=True): + case.writer.create_attempt(_submitted_attempt(case, shard, scheduler=task.scheduler)) + runner.set_task_state(task.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + original_complete = SlurmRetryCoordinator._complete_retry + + def interrupt_after_attempts(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("injected completion interruption") + + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", interrupt_after_attempts) + with pytest.raises(SlurmStateError, match="cannot retry persisted run"): + coordinator.retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", original_complete) + + with pytest.raises(StateConflictError, match="sealed or nonterminal"): + coordinator.retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="always", + observed_at=case.created_at + timedelta(minutes=6), + ) + + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("retry-0001").state is RetryState.COMPLETED + assert storage.list_retry_ids() == ("retry-0001",) + + +def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + + assert submitted.state is CollectionState.SUBMITTED + assert submitted.scheduler == 5101 + script = cast(str, runner.inputs[-1]) + assert "data_designer.slurm.state.collection_worker" in script + assert "--gpus" not in script + assert "--gres" not in script + stale_stage = Path(case.plan.output.root).parent / submitted.staging_directory + stale_stage.mkdir(mode=0o700) + (stale_stage / "partial").write_text("incomplete") + unrelated_stage = Path(case.plan.output.root).parent / f".dd-collection-{'f' * 32}.tmp" + unrelated_stage.mkdir(mode=0o700) + (unrelated_stage / "active").write_text("preserve") + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + destination = Path(case.plan.output.root) + assert result.actual_records == case.plan.invocation.authored.num_records + assert len(result.files) == case.plan.output.partitions + assert lazy.pq.read_table(destination / result.files[0].relative_path).column("record_id").to_pylist() == list( + range(case.plan.invocation.authored.num_records) + ) + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.SUCCEEDED + assert coordinator.submit() == persisted + assert not stale_stage.exists() + assert (unrelated_stage / "active").read_text() == "preserve" + + +def test_collection_requires_every_planned_winner_before_submission( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + + with pytest.raises(StateNotFoundError, match="has no winner"): + SlurmCollectionCoordinator(case.workspace, case.plan.run_id).submit( + submitted_at=case.created_at + timedelta(minutes=10) + ) + + +def test_collection_prepares_missing_authorized_destination_parents( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + workspace_root = Path(multi_node_plan.selected_profile.profile.workspace_root) + output = multi_node_plan.output.model_copy( + update={"root": (workspace_root / "new" / "nested" / "output").as_posix()} + ) + plan = ResolvedSlurmRunPlan.model_validate_json( + json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) + ) + case = _initialize_run(tmp_path, authored_run, plan) + destination = Path(case.plan.output.root) + assert not destination.parent.exists() + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + assert destination.parent.is_dir() + assert not destination.exists() + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + assert result.actual_records == case.plan.invocation.authored.num_records + assert destination.is_dir() + + +def test_concurrent_collection_submission_converges_on_one_job( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + + def submit() -> object: + return SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + with ThreadPoolExecutor(max_workers=2) as executor: + results = tuple(executor.map(lambda _: submit(), range(2))) + + assert results[0] == results[1] + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 1 + + +def test_collection_ambiguous_submission_remains_prepared_and_blocks_duplicate( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + scheduler = SlurmCommandClient(FakeSlurmRunner()) + submissions = 0 + + def ambiguous_submit(script: str) -> object: + nonlocal submissions + del script + submissions += 1 + raise SlurmSubmissionError("sbatch could not be executed: command timed out", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", ambiguous_submit) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler) + + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("collection-0001").state is CollectionState.PREPARED + with pytest.raises(StateConflictError, match="ambiguous scheduler outcome"): + coordinator.submit(submitted_at=case.created_at + timedelta(minutes=11)) + assert submissions == 1 + + +def test_collection_definite_submission_failure_settles_and_can_be_retried( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + runner.script_next("sbatch", FakeCommandResponse(stderr="submission rejected", returncode=2)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("collection-0001").state is CollectionState.FAILED + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=11)) + assert submitted.collection_id == "collection-0002" + assert submitted.state is CollectionState.SUBMITTED + + +def test_collection_discards_trailing_journal_interrupted_before_submission( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + original_publish = CollectionStorage.publish_status + + def interrupt_before_status(self: CollectionStorage, status: object) -> None: + del self, status + raise OSError("injected journal interruption") + + monkeypatch.setattr(CollectionStorage, "publish_status", interrupt_before_status) + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + monkeypatch.setattr(CollectionStorage, "publish_status", original_publish) + + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=11)) + + assert submitted.collection_id == "collection-0001" + assert CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).list_collection_ids() == ( + "collection-0001", + ) + + +def test_collection_failure_removes_staging_without_publishing_partial_output( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + def fail_after_writes(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("injected collection failure") + + monkeypatch.setattr( + "data_designer.slurm.state.collection_merge.CollectionMerger._write_result", + fail_after_writes, + ) + with pytest.raises(SlurmStateError, match="collection .* failed"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + destination = Path(case.plan.output.root) + assert not destination.exists() + assert not tuple(destination.parent.glob(".dd-*.tmp")) + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.FAILED + + +def test_collection_rejects_output_replacement_before_descriptor_digest( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + original_describe = collection_merge._BoundOutput.describe + injected = False + + def replace_output_before_describe( + output: collection_merge._BoundOutput, + relative_path: str, + record_count: int, + ) -> object: + nonlocal injected + if not injected: + injected = True + output.path.rename(output.path.with_suffix(f"{output.path.suffix}.written")) + output.path.write_bytes(b"replacement") + output.path.chmod(0o600) + return original_describe(output, relative_path, record_count) + + monkeypatch.setattr(collection_merge._BoundOutput, "describe", replace_output_before_describe) + + with pytest.raises(SlurmStateError, match="collection .* failed"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + destination = Path(case.plan.output.root) + assert not destination.exists() + assert not (destination.parent / submitted.staging_directory).exists() + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.FAILED + + +def test_collection_refresh_cleans_only_its_exact_stage_after_oom( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + parent = Path(case.plan.output.root).parent + exact_stage = parent / submitted.staging_directory + exact_stage.mkdir(mode=0o700) + (exact_stage / "partial.parquet").write_text("incomplete") + unrelated_stage = parent / f".dd-collection-{'e' * 32}.tmp" + unrelated_stage.mkdir(mode=0o700) + (unrelated_stage / "active").write_text("preserve") + runner.set_job_state(5101, queue_state=None, accounting_state="OUT_OF_MEMORY", exit_code="0:125") + + refreshed = coordinator.refresh(observed_at=case.created_at + timedelta(minutes=11)) + + assert refreshed.state is CollectionState.FAILED + assert not exact_stage.exists() + assert (unrelated_stage / "active").read_text() == "preserve" + assert not Path(case.plan.output.root).exists() + + +def test_collection_worker_reauthorizes_persisted_partition_intent( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + plan = storage.read_plan(submitted.collection_id) + tampered = plan.model_copy(update={"num_partitions": plan.num_partitions + 1}) + storage.get_plan_path(plan.collection_id).write_text(tampered.serialize_json()) + storage.replace_status( + submitted.model_copy( + update={ + "collection_plan": storage.get_plan_reference(tampered), + "staging_directory": derive_collection_staging_directory(tampered), + } + ) + ) + + with pytest.raises(StateConflictError, match="partition count"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert not Path(case.plan.output.root).exists() + + +def test_succeeded_collection_rejects_modified_partition_bytes( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + output_path = Path(case.plan.output.root) / result.files[0].relative_path + output_path.write_bytes(output_path.read_bytes() + b"tampered") + + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit() + + +def test_succeeded_collection_rejects_same_size_partition_mutation( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + output_path = Path(case.plan.output.root) / result.files[0].relative_path + content = bytearray(output_path.read_bytes()) + content[-1] ^= 1 + output_path.write_bytes(content) + + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit() + + +def test_succeeded_collection_status_check_does_not_read_partition_payloads( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + original_open = collection_storage_module.open_verified_regular_file + + def reject_partition_reads(*args: object, **kwargs: object) -> object: + if args[1] != "collection-result.json": + raise AssertionError("login-host validation attempted to read collected partition bytes") + return original_open(*args, **kwargs) + + monkeypatch.setattr(collection_storage_module, "open_verified_regular_file", reject_partition_reads) + + assert coordinator.submit().state is CollectionState.SUCCEEDED + + +def test_succeeded_collection_rejects_status_result_digest_mismatch( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + succeeded = storage.read_status(submitted.collection_id) + assert succeeded.result is not None + storage.replace_status( + succeeded.model_copy( + update={ + "result": succeeded.result.model_copy(update={"sha256": "f" * 64}), + } + ) + ) + + with pytest.raises(StateCorruptionError, match="does not bind its published result"): + coordinator.submit() + + +def test_collection_refuses_destination_collision_before_submission( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + destination = Path(case.plan.output.root) + destination.mkdir(mode=0o700) + marker = destination / "existing.txt" + marker.write_text("preserve") + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + + with pytest.raises(StateConflictError, match="already exists"): + SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + assert marker.read_text() == "preserve" + assert not any(Path(call[0]).name == "sbatch" for call in runner.calls) + + +def test_collection_refuses_atomic_publication_collision_without_replacement( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + destination = Path(case.plan.output.root) + original_rename = collection_filesystem._rename_without_overwrite + + def collide( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, + ) -> None: + destination.mkdir(mode=0o700) + (destination / "existing.txt").write_text("preserve") + original_rename(source_directory, source_name, destination_directory, destination_name) + + monkeypatch.setattr(collection_filesystem, "_rename_without_overwrite", collide) + + with pytest.raises(StateConflictError, match="already exists"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert (destination / "existing.txt").read_text() == "preserve" + assert not (destination.parent / submitted.staging_directory).exists() + + +def test_collection_detects_destination_parent_replacement_before_success( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace_root = Path(multi_node_plan.selected_profile.profile.workspace_root) + output = multi_node_plan.output.model_copy(update={"root": (workspace_root / "exports" / "output").as_posix()}) + plan = ResolvedSlurmRunPlan.model_validate_json( + json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) + ) + case = _initialize_run(tmp_path, authored_run, plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + destination = Path(case.plan.output.root) + original_parent = destination.parent + moved_parent = original_parent.with_name("moved-exports") + original_rename = collection_filesystem._rename_without_overwrite + injected = False + + def replace_parent( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, + ) -> None: + nonlocal injected + if not injected: + injected = True + original_parent.rename(moved_parent) + original_parent.mkdir(mode=0o700) + original_rename(source_directory, source_name, destination_directory, destination_name) + + monkeypatch.setattr(collection_filesystem, "_rename_without_overwrite", replace_parent) + + with pytest.raises(SlurmStateError, match="collection .* failed"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert not destination.exists() + assert not (moved_parent / destination.name).exists() + assert not (moved_parent / submitted.staging_directory).exists() + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.FAILED + + +def test_collection_recovers_when_publication_completed_before_reported_error( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + original_rename = collection_filesystem._rename_without_overwrite + + def publish_then_fail( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, + ) -> None: + original_rename(source_directory, source_name, destination_directory, destination_name) + raise OSError("injected post-rename failure") + + monkeypatch.setattr(collection_filesystem, "_rename_without_overwrite", publish_then_fail) + + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert result.actual_records == case.plan.invocation.authored.num_records + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.SUCCEEDED + assert Path(case.plan.output.root).is_dir() + + +def test_collection_worker_rejects_login_node_execution( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + with pytest.raises(StateConflictError, match="inside its recorded Slurm job"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + +def test_collection_worker_rejects_login_node_recovery_before_reading_outputs( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + storage.replace_status(submitted) + + def reject_output_reads(*args: object, **kwargs: object) -> None: + del args, kwargs + raise AssertionError("login-host worker attempted to read collected outputs") + + monkeypatch.setattr(CollectionStorage, "verify_result_files", reject_output_reads) + worker = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={}, + ) + + with pytest.raises(StateConflictError, match="inside its recorded Slurm job"): + worker.run(completed_at=case.created_at + timedelta(minutes=12)) + + assert storage.read_status(submitted.collection_id) == submitted + + +def test_collection_worker_validates_location_before_constructing_state(tmp_path: Path) -> None: + with pytest.raises(SlurmStateError, match="invalid persisted collection worker location"): + SlurmCollectionWorker("relative/workspace", "run-001", "collection-0001") + with pytest.raises(SlurmStateError, match="invalid persisted collection worker location"): + SlurmCollectionWorker(tmp_path, "../run", "collection-0001") + + +@pytest.mark.parametrize("output_format", ["csv", "jsonl"]) +def test_collection_exports_non_parquet_formats_in_deterministic_partitions( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + output_format: str, +) -> None: + output = multi_node_plan.output.model_copy(update={"format": output_format, "partitions": 3}) + plan = ResolvedSlurmRunPlan.model_validate_json( + json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) + ) + case = _initialize_run(tmp_path, authored_run, plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert len(result.files) == 3 + requested_records = case.plan.invocation.authored.num_records + floor_count = requested_records // 3 + assert tuple(file.record_count for file in result.files) == ( + floor_count, + floor_count, + requested_records - 2 * floor_count, + ) + values: list[int] = [] + for output_file in result.files: + path = Path(case.plan.output.root) / output_file.relative_path + if output_format == "csv": + with path.open(newline="") as source: + values.extend(int(row["record_id"]) for row in csv.DictReader(source)) + else: + values.extend(int(json.loads(line)["record_id"]) for line in path.read_text().splitlines()) + assert values == list(range(requested_records)) + + +def _initialize_run( + tmp_path: Path, + authored_config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, +) -> _RunCase: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o700) + relocated_plan = _relocate_plan(plan, workspace) + created_at = datetime(2026, 9, 1, 12, tzinfo=timezone.utc) + run_root = workspace / "runs" / relocated_plan.run_id + run = RunManifest( + schema_version=1, + run_id=relocated_plan.run_id, + created_at=created_at, + authored_config=relocated_plan.authored_config, + resolved_plan=ArtifactReference( + path=(run_root / "resolved-plan.json").as_posix(), + sha256=relocated_plan.compute_sha256(), + ), + shard_count=len(relocated_plan.shards), + ) + shards = tuple( + ShardManifest( + schema_version=1, + run_id=run.run_id, + shard_id=planned.shard_id, + shard_index=planned.shard_index, + record_range=planned.record_range, + input_partition=planned.input_partition, + resume_workspace=planned.resume_workspace, + created_at=created_at, + ) + for planned in relocated_plan.shards + ) + writer = SlurmStateWriter(workspace, run.run_id) + writer.initialize_run(authored_config, relocated_plan, run, shards) + return _RunCase(workspace, relocated_plan, run, shards, writer, created_at) + + +def _relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurmRunPlan: + previous_workspace = plan.selected_profile.profile.workspace_root + payload = cast( + dict[str, object], + json.loads(plan.serialize_json().replace(previous_workspace, workspace.as_posix())), + ) + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected_profile["profile"]) + profile_mounts = cast(list[dict[str, object]], profile_payload["container_mounts"]) + resolved_mounts = cast(list[dict[str, object]], payload["container_mounts"]) + for mount in (*profile_mounts, *resolved_mounts): + mount["source"] = workspace.parent.as_posix() + mount["target"] = workspace.parent.as_posix() + profile = SlurmProfile.model_validate(selected_profile["profile"]) + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def _submitted_attempt( + case: _RunCase, + shard: ShardManifest, + *, + scheduler: SchedulerIdentity, +) -> AttemptManifest: + return AttemptManifest( + schema_version=1, + run_id=case.run.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=case.run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED, + scheduler=scheduler, + created_at=case.created_at + timedelta(minutes=1), + updated_at=case.created_at + timedelta(minutes=1), + ) + + +def _publish_all_winners(case: _RunCase) -> None: + for shard in case.shards: + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=shard.shard_index) + submitted = _submitted_attempt(case, shard, scheduler=scheduler) + case.writer.create_attempt(submitted) + running = _copy_attempt( + submitted, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=2), + ) + case.writer.update_attempt(running) + with case.writer.acquire_dataset_workspace(shard.shard_id, submitted.attempt_id, "never") as dataset_path: + _publish_candidate(case, shard, running, dataset_path) + case.writer.finalize_winner( + shard.shard_id, + running.attempt_id, + published_at=case.created_at + timedelta(minutes=6), + ) + + +def _publish_candidate(case: _RunCase, shard: ShardManifest, running: AttemptManifest, dataset_path: Path) -> None: + output_path = dataset_path / "part-00000.parquet" + values = range(shard.record_range.start_index, shard.record_range.end_index_exclusive) + table = lazy.pa.table({"record_id": values}) + lazy.pq.write_table(table, output_path) + output_path.chmod(0o644) + content = output_path.read_bytes() + candidate_path = ( + case.writer.run_root / "shards" / shard.shard_id / "attempts" / running.attempt_id / "output-manifest.json" + ) + candidate = CandidateOutputManifest( + schema_version=1, + run_id=case.run.run_id, + shard_id=shard.shard_id, + attempt_id=running.attempt_id, + attempt_ordinal=running.attempt_ordinal, + created_at=running.updated_at + timedelta(minutes=1), + dataset_path=dataset_path.as_posix(), + requested_records=shard.record_range.record_count, + actual_records=shard.record_range.record_count, + outcome=CandidateOutcome.COMPLETE, + files=( + CandidateOutputFile( + relative_path=output_path.name, + sha256=hashlib.sha256(content).hexdigest(), + byte_size=len(content), + record_count=shard.record_range.record_count, + ), + ), + dataset_schema_digest=compute_candidate_schema_digest(table.schema), + provenance_digest=case.plan.compute_sha256(), + ) + candidate_reference = ArtifactReference(path=candidate_path.as_posix(), sha256=candidate.compute_sha256()) + result = ClientResult( + schema_version=1, + run_id=case.run.run_id, + shard_id=shard.shard_id, + attempt_id=running.attempt_id, + completed_at=running.updated_at + timedelta(minutes=2), + requested_records=shard.record_range.record_count, + actual_records=shard.record_range.record_count, + outcome=ClientOutcome.COMPLETE, + dataset_path=dataset_path.as_posix(), + early_shutdown=False, + requested_resume_mode=case.plan.invocation.authored.resume, + effective_resume_mode="never", + candidate_output_manifest=candidate_reference, + ) + case.writer.publish_attempt_result(result, candidate) + completed = _copy_attempt( + running, + state=AttemptLifecycleState.SUCCEEDED, + terminal_classification=AttemptTerminalClassification.SUCCEEDED, + candidate_output=candidate_reference, + updated_at=running.updated_at + timedelta(minutes=3), + ) + case.writer.update_attempt(completed) + + +def _copy_attempt(attempt: AttemptManifest, **updates: object) -> AttemptManifest: + payload = attempt.model_dump(mode="python") + payload.update(updates) + return AttemptManifest.model_validate(payload) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 8375421d3..da8da21a0 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -137,14 +137,22 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.state import RecordRange as StateRecordRange from data_designer.slurm.state import ResumeWorkspace as StateResumeWorkspace from data_designer.slurm.state import ( + CollectionResult, + RetryPlan, RunManifest, RunStatus, SchedulerObservationCollector, + SlurmCollectionCoordinator, + SlurmRetryCoordinator, SlurmStateReconciler, ) +assert CollectionResult.__name__ == "CollectionResult" +assert RetryPlan.__name__ == "RetryPlan" assert RunManifest.__name__ == "RunManifest" assert RunStatus.__name__ == "RunStatus" assert SchedulerObservationCollector.__name__ == "SchedulerObservationCollector" +assert SlurmCollectionCoordinator.__name__ == "SlurmCollectionCoordinator" +assert SlurmRetryCoordinator.__name__ == "SlurmRetryCoordinator" assert SlurmStateReconciler.__name__ == "SlurmStateReconciler" assert ImageRegistryStore.__name__ == "ImageRegistryStore" assert PlanningArtifactReference is ContractArtifactReference From 5387820963b836c488aba3f0863593751a3a24ae Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 2 Sep 2026 19:23:22 -0600 Subject: [PATCH 2/4] fix(slurm): recover ambiguous submissions Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/client.py | 69 +++++ .../slurm/launcher/collection.py | 2 +- .../data_designer/slurm/launcher/models.py | 18 ++ .../data_designer/slurm/launcher/parsing.py | 21 ++ .../data_designer/slurm/launcher/renderer.py | 7 +- .../data_designer/slurm/state/collection.py | 63 ++++- .../slurm/state/collection_records.py | 14 +- .../slurm/state/collection_worker.py | 19 +- .../src/data_designer/slurm/state/outputs.py | 10 + .../src/data_designer/slurm/state/retry.py | 53 +++- .../slurm/state/retry_records.py | 13 +- .../slurm/state/submission_recovery.py | 74 +++++ .../tests/launcher/test_client.py | 46 ++++ .../tests/launcher/test_collection.py | 4 + .../tests/state/test_retry_collection.py | 258 +++++++++++++++++- .../tests/state/test_submission_recovery.py | 113 ++++++++ 16 files changed, 759 insertions(+), 25 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py create mode 100644 packages/data-designer-slurm/tests/state/test_submission_recovery.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py index e51ff6edc..83bbe1644 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -5,11 +5,13 @@ 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 @@ -17,11 +19,14 @@ 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, ) @@ -136,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))) @@ -244,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) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py index e8a77f120..e75552506 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py @@ -36,7 +36,7 @@ def render_collection_script( collection_plan_path = posixpath.join(collection_root, "plan.json") directives = render_batch_directives( ( - ("job-name", f"dd-collect-{resolved_plan.run_id}"), + ("job-name", collection_plan.submission_job_name), ("account", resolved_plan.submission.account), ("partition", resolved_plan.submission.partition), ("nodes", "1"), diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py index 649460dbf..64111099e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index 213ce09ca..32766d782 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -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] = [] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index fc78ce06f..e8e3bf202 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -83,7 +83,9 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) array_tasks = ",".join(str(shard.array_task_index) for shard in retry.planned_shards) if plan.array_tasks.max_concurrent is not None: array_tasks = f"{array_tasks}%{plan.array_tasks.max_concurrent}" - directives = render_batch_directives(_build_generation_directives(plan, array=array_tasks)) + directives = render_batch_directives( + _build_generation_directives(plan, array=array_tasks, job_name=retry.submission_job_name) + ) attempt_cases = "\n".join( f" {shard.array_task_index}) DD_ATTEMPT_ORDINAL={quote_shell_value(f'{shard.attempt_ordinal:04d}')} ;;" for shard in retry.planned_shards @@ -159,6 +161,7 @@ def _build_generation_directives( plan: ResolvedSlurmRunPlan, *, array: str | None = None, + job_name: str | None = None, ) -> tuple[tuple[str, str | None], ...]: node_indices = ( plan.client.host_node_index, @@ -172,7 +175,7 @@ def _build_generation_directives( resolved_array = f"{resolved_array}%{plan.array_tasks.max_concurrent}" values: list[tuple[str, str | None]] = [ - ("job-name", plan.submission.job_name), + ("job-name", plan.submission.job_name if job_name is None else job_name), ("account", plan.submission.account), ("partition", plan.submission.partition), ("nodes", str(node_count)), diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py index 1ac68d282..396ccdcf5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -15,7 +15,7 @@ from data_designer.slurm.launcher.client import SlurmCommandClient from data_designer.slurm.launcher.collection import render_collection_script from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError -from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt +from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt, SlurmSubmissionMatch from data_designer.slurm.state.collection_filesystem import ( derive_collection_staging_directory, prepare_collection_destination, @@ -36,6 +36,11 @@ from data_designer.slurm.state.reader import StateReader from data_designer.slurm.state.scheduler import SchedulerState from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.submission_recovery import ( + SUBMISSION_VISIBILITY_WINDOW, + PreparedSubmission, + resolve_prepared_submission, +) _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) @@ -47,6 +52,15 @@ def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: """Submit one rendered CPU collection job.""" ... + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + """Return allocations matching one exact collection submission name.""" + ... + class SlurmCollectionCoordinator: """Persist, submit, and refresh winner-driven collection jobs.""" @@ -82,7 +96,7 @@ def submit( if current is not None: current_plan = self._load_bound_plan(current) if current.state is CollectionState.PREPARED: - raise StateConflictError("previous collection submission has an ambiguous scheduler outcome") + current = self._reconcile_prepared_collection(current_plan, current, timestamp) if current.state is not CollectionState.FAILED: self._validate_existing_destination(current, destination) return current @@ -124,6 +138,7 @@ def submit( revision=1, updated_at=timestamp, state=CollectionState.PREPARED, + reconciliation_deadline=timestamp + SUBMISSION_VISIBILITY_WINDOW, ) self._collections.publish_status(prepared) script = render_collection_script(resolved_plan, collection_plan, resolved_destination) @@ -161,8 +176,16 @@ def refresh( Path(destination.mount.source), ) return previous - if previous.scheduler is None: - raise StateConflictError("collection submission has an ambiguous scheduler outcome") + if previous.state is CollectionState.PREPARED: + previous = self._reconcile_prepared_collection(plan, previous, timestamp) + if previous.state is CollectionState.FAILED: + remove_collection_stage( + Path(plan.host_destination), + previous.staging_directory, + Path(destination.mount.source), + ) + return previous + assert previous.scheduler is not None observations = self._collector.collect( (previous.scheduler,), observed_at=timestamp, @@ -209,6 +232,7 @@ def _submit_prepared( revision=2, updated_at=submitted_at, state=CollectionState.FAILED, + reconciliation_deadline=None, ) validate_collection_status_transition(prepared, failed) self._collections.replace_status(failed) @@ -221,11 +245,42 @@ def _submit_prepared( updated_at=submitted_at, state=CollectionState.SUBMITTED, scheduler=receipt.job_id, + reconciliation_deadline=None, ) validate_collection_status_transition(prepared, submitted) self._collections.replace_status(submitted) return submitted + def _reconcile_prepared_collection( + self, + plan: CollectionPlan, + prepared: CollectionStatus, + observed_at: datetime, + ) -> CollectionStatus: + assert prepared.reconciliation_deadline is not None + job_id = resolve_prepared_submission( + self._scheduler, + PreparedSubmission( + job_name=plan.submission_job_name, + submitted_after=plan.created_at, + reconciliation_deadline=prepared.reconciliation_deadline, + expected_array_task_ids=None, + ), + observed_at=observed_at, + ) + state = CollectionState.SUBMITTED if job_id is not None else CollectionState.FAILED + current = _updated_status( + prepared, + revision=prepared.revision + 1, + updated_at=observed_at, + state=state, + scheduler=job_id, + reconciliation_deadline=None, + ) + validate_collection_status_transition(prepared, current) + self._collections.replace_status(current) + return current + def _get_current_status(self) -> CollectionStatus | None: collection_ids = self._collections.list_collection_ids() return None if not collection_ids else self._collections.read_status(collection_ids[-1]) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py index 3d6e52107..b8774661f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py @@ -12,7 +12,13 @@ from pydantic import Field, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator from data_designer.slurm.contracts import ArtifactReference, Identifier, Sha256Digest, validate_relative_path -from data_designer.slurm.state.base import SchedulerJobIdentity, StateRecord, StateValue, validate_utc_timestamp +from data_designer.slurm.state.base import ( + SchedulerJobIdentity, + StateRecord, + StateValue, + validate_optional_utc_timestamp, + validate_utc_timestamp, +) from data_designer.slurm.state.scheduler import SchedulerObservation @@ -80,15 +86,21 @@ class CollectionStatus(StateRecord): scheduler: SchedulerJobIdentity | None = None scheduler_observation: SchedulerObservation | None = None result: ArtifactReference | None = None + reconciliation_deadline: datetime | None = None _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + _reconciliation_deadline_is_utc = field_validator("reconciliation_deadline")(validate_optional_utc_timestamp) @model_validator(mode="after") def validate_evidence(self) -> CollectionStatus: if self.state is CollectionState.PREPARED: if self.scheduler is not None or self.scheduler_observation is not None or self.result is not None: raise ValueError("prepared collection cannot contain scheduler or result evidence") + if self.reconciliation_deadline is None or self.reconciliation_deadline <= self.updated_at: + raise ValueError("prepared collection requires a future reconciliation deadline") return self + if self.reconciliation_deadline is not None: + raise ValueError("settled collection cannot contain a reconciliation deadline") if self.state is not CollectionState.FAILED and self.scheduler is None: raise ValueError("submitted collection states require a scheduler identity") if self.scheduler is None and self.scheduler_observation is not None: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py index d675e80a6..2a43828cf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py @@ -10,6 +10,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime, timezone from pathlib import Path +from time import sleep from pydantic import TypeAdapter, ValidationError @@ -32,6 +33,7 @@ from data_designer.slurm.state.storage import StateStorage _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_SCHEDULER_BINDING_WAIT_ATTEMPTS = 300 class SlurmCollectionWorker: @@ -57,8 +59,9 @@ def __init__( def run(self, *, completed_at: datetime | None = None) -> CollectionResult: """Validate, merge, and atomically publish one collection output.""" - started_at = _utc_now() if completed_at is None else completed_at try: + self._wait_for_scheduler_binding() + started_at = _utc_now() if completed_at is None else completed_at return self._run_locked(started_at, completed_at) except (StateConflictError, SlurmStateError): raise @@ -195,6 +198,20 @@ def _require_scheduler_job(self, status: CollectionStatus) -> None: if observed_job_id != str(scheduler): raise StateConflictError("collection worker must run inside its recorded Slurm job") + def _wait_for_scheduler_binding(self) -> None: + if self._environment.get("SLURM_JOB_ID") is None: + raise StateConflictError("collection worker must run inside its recorded Slurm job") + for attempt in range(_SCHEDULER_BINDING_WAIT_ATTEMPTS): + status = self._collections.read_status(self._collection_id) + if status.scheduler is not None: + self._require_scheduler_job(status) + return + if status.state is not CollectionState.PREPARED: + raise StateConflictError("collection worker requires an ordinary Slurm job identity") + if attempt + 1 < _SCHEDULER_BINDING_WAIT_ATTEMPTS: + sleep(1) + raise StateConflictError("collection scheduler identity was not published before allocation startup") + def _validate_identity(self, plan_run_id: Identifier, status_run_id: Identifier) -> None: if plan_run_id != self._run_id or status_run_id != self._run_id: raise StateConflictError("collection records do not match the requested run") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py index 11759739a..13aa11887 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -136,6 +136,11 @@ class RetryPlan(StateRecord): _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + @property + def submission_job_name(self) -> Identifier: + """Return the immutable scheduler lookup key for this retry.""" + return f"dd-retry-{self.compute_sha256()[:32]}" + @model_validator(mode="after") def validate_shards(self) -> RetryPlan: shard_ids = tuple(shard.shard_id for shard in self.planned_shards) @@ -165,6 +170,11 @@ class CollectionPlan(StateRecord): _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) _destinations_are_safe = field_validator("host_destination", "container_destination")(validate_absolute_path) + @property + def submission_job_name(self) -> Identifier: + """Return the immutable scheduler lookup key for this collection.""" + return f"dd-collect-{self.compute_sha256()[:32]}" + @model_validator(mode="after") def validate_shards(self) -> CollectionPlan: shard_ids = [shard.shard_id for shard in self.planned_shards] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py index 79a43b84f..607ab3fc9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -16,7 +16,7 @@ from data_designer.slurm.contracts import Identifier, ShardId, validate_absolute_path from data_designer.slurm.launcher.client import SlurmCommandClient from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError -from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt +from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt, SlurmSubmissionMatch from data_designer.slurm.launcher.renderer import render_generation_retry_script from data_designer.slurm.state.base import SchedulerIdentity from data_designer.slurm.state.errors import SlurmStateError, StateConflictError, StateCorruptionError @@ -31,6 +31,11 @@ from data_designer.slurm.state.scheduler import EffectiveAttemptState from data_designer.slurm.state.status import RunStatus, ShardStatus from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.submission_recovery import ( + SUBMISSION_VISIBILITY_WINDOW, + PreparedSubmission, + resolve_prepared_submission, +) from data_designer.slurm.state.validation import ( StateContractError, validate_attempt_transition, @@ -47,6 +52,15 @@ def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: """Submit one rendered retry array.""" ... + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + """Return allocations matching one exact retry submission name.""" + ... + class SlurmRetryCoordinator: """Select retryable shards and durably submit their next attempts.""" @@ -80,7 +94,7 @@ def retry( raise StateConflictError("effective resume mode must be 'never' or 'always'") with self._retries.acquire_lock(): self._retries.discard_incomplete_tail() - self._settle_submitted_retry(timestamp) + self._settle_pending_retry(timestamp) status = self._reconciler.refresh(observed_at=timestamp) active = self._load_active_retry(status, shard_ids, effective_resume_mode) if active is not None: @@ -100,7 +114,7 @@ def retry( except (OSError, ValidationError, ValueError) as error: raise SlurmStateError(f"cannot retry persisted run {self._run_id!r}") from error - def _settle_submitted_retry(self, updated_at: datetime) -> None: + def _settle_pending_retry(self, updated_at: datetime) -> None: retry_ids = self._retries.list_retry_ids() if not retry_ids: return @@ -108,7 +122,7 @@ def _settle_submitted_retry(self, updated_at: datetime) -> None: status = self._retries.read_status(latest_id) plan = self._load_bound_plan(status) if status.state is RetryState.PREPARED: - raise StateConflictError("previous retry submission has an ambiguous scheduler outcome") + status = self._reconcile_prepared_retry(plan, status, updated_at) if status.state is not RetryState.SUBMITTED: return assert status.array_job_id is not None @@ -119,6 +133,28 @@ def _settle_submitted_retry(self, updated_at: datetime) -> None: _, superseded = self._publish_attempts(plan, status.array_job_id, plan.created_at) self._settle_retry(plan, status.array_job_id, updated_at, superseded=superseded) + def _reconcile_prepared_retry( + self, + plan: RetryPlan, + status: RetryStatus, + updated_at: datetime, + ) -> RetryStatus: + assert status.reconciliation_deadline is not None + job_id = resolve_prepared_submission( + self._scheduler, + PreparedSubmission( + job_name=plan.submission_job_name, + submitted_after=plan.created_at, + reconciliation_deadline=status.reconciliation_deadline, + expected_array_task_ids=tuple(shard.array_task_index for shard in plan.planned_shards), + ), + observed_at=updated_at, + ) + if job_id is not None: + return self._publish_submitted_status(plan, job_id, updated_at) + self._fail_retry(plan, updated_at) + return self._retries.read_status(plan.retry_id) + def _build_retry_plan( self, status: RunStatus, @@ -249,6 +285,7 @@ def _persist_prepared_retry(self, plan: RetryPlan, timestamp: datetime) -> None: revision=1, updated_at=timestamp, state=RetryState.PREPARED, + reconciliation_deadline=timestamp + SUBMISSION_VISIBILITY_WINDOW, ) ) @@ -261,6 +298,10 @@ def _submit(self, plan: RetryPlan, script: str, timestamp: datetime) -> SlurmJob raise SlurmStateError(f"cannot submit retry {plan.retry_id!r}") from error except SlurmLauncherError as error: raise SlurmStateError(f"cannot submit retry {plan.retry_id!r}") from error + self._publish_submitted_status(plan, receipt.job_id, timestamp) + return receipt + + def _publish_submitted_status(self, plan: RetryPlan, array_job_id: int, timestamp: datetime) -> RetryStatus: plan_reference = self._retries.get_plan_reference(plan) submitted = RetryStatus( schema_version=1, @@ -270,11 +311,11 @@ def _submit(self, plan: RetryPlan, script: str, timestamp: datetime) -> SlurmJob revision=2, updated_at=timestamp, state=RetryState.SUBMITTED, - array_job_id=receipt.job_id, + array_job_id=array_job_id, ) validate_retry_status_transition(self._retries.read_status(plan.retry_id), submitted) self._retries.replace_status(submitted) - return receipt + return submitted def _publish_attempts( self, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py index 961f54bba..d5045c20c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py @@ -11,7 +11,7 @@ from pydantic import PositiveInt, field_validator, model_validator from data_designer.slurm.contracts import ArtifactReference, Identifier -from data_designer.slurm.state.base import StateRecord, validate_utc_timestamp +from data_designer.slurm.state.base import StateRecord, validate_optional_utc_timestamp, validate_utc_timestamp class RetryState(str, Enum): @@ -33,15 +33,22 @@ class RetryStatus(StateRecord): updated_at: datetime state: RetryState array_job_id: PositiveInt | None = None + reconciliation_deadline: datetime | None = None _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + _reconciliation_deadline_is_utc = field_validator("reconciliation_deadline")(validate_optional_utc_timestamp) @model_validator(mode="after") def validate_scheduler_identity(self) -> RetryStatus: if self.state in {RetryState.SUBMITTED, RetryState.COMPLETED} and self.array_job_id is None: raise ValueError("submitted retry state requires an array job identity") - if self.state is RetryState.PREPARED and self.array_job_id is not None: - raise ValueError("prepared retry state cannot contain an array job identity") + if self.state is RetryState.PREPARED: + if self.array_job_id is not None: + raise ValueError("prepared retry state cannot contain an array job identity") + if self.reconciliation_deadline is None or self.reconciliation_deadline <= self.updated_at: + raise ValueError("prepared retry state requires a future reconciliation deadline") + elif self.reconciliation_deadline is not None: + raise ValueError("settled retry state cannot contain a reconciliation deadline") return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py b/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py new file mode 100644 index 000000000..f7ea62853 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared recovery policy for ambiguous Slurm submission receipts.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Protocol + +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.launcher.errors import SlurmLauncherError +from data_designer.slurm.launcher.models import SlurmSubmissionMatch +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError + +SUBMISSION_VISIBILITY_WINDOW = timedelta(minutes=5) + + +class SubmissionLookup(Protocol): + """Scheduler lookup needed to recover one immutable submission plan.""" + + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + """Return allocations matching one exact submission name.""" + ... + + +@dataclass(frozen=True) +class PreparedSubmission: + """Immutable scheduler correlation facts for one prepared operation.""" + + job_name: Identifier + submitted_after: datetime + reconciliation_deadline: datetime + expected_array_task_ids: tuple[int, ...] | None + + +def resolve_prepared_submission( + scheduler: SubmissionLookup, + prepared: PreparedSubmission, + *, + observed_at: datetime, +) -> int | None: + """Return one recovered job ID, or ``None`` after definitive bounded absence.""" + try: + matches = scheduler.query_submissions_by_name( + prepared.job_name, + submitted_after=prepared.submitted_after, + ) + except SlurmLauncherError as error: + raise SlurmStateError("cannot reconcile ambiguous Slurm submission") from error + if len(matches) > 1: + raise StateConflictError("multiple scheduler jobs match the prepared submission") + if matches: + match = matches[0] + if match.array_task_ids != prepared.expected_array_task_ids: + raise StateConflictError("scheduler job shape does not match the prepared submission") + return match.job_id + if observed_at <= prepared.reconciliation_deadline: + raise StateConflictError("prepared submission is still being reconciled") + return None + + +__all__ = [ + "SUBMISSION_VISIBILITY_WINDOW", + "PreparedSubmission", + "SubmissionLookup", + "resolve_prepared_submission", +] diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index e2cee1699..a3ef7949f 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -5,6 +5,7 @@ import subprocess from collections.abc import Sequence +from datetime import datetime, timezone import pytest from slurm_test_fakes import FakeCommandResponse, FakeSlurmJob, FakeSlurmRunner @@ -82,6 +83,51 @@ def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: ] +def test_client_finds_one_exact_named_array_across_queue_and_accounting() -> None: + runner = FakeSlurmRunner() + job_name = f"dd-retry-{'a' * 32}" + runner.script_next("squeue", FakeCommandResponse(stdout=f"4201_0|{job_name}\n")) + runner.script_next("sacct", FakeCommandResponse(stdout=f"4201_0|{job_name}\n4201_1|{job_name}\n")) + submitted_after = datetime(2026, 9, 2, 18, tzinfo=timezone.utc) + + matches = SlurmCommandClient(runner).query_submissions_by_name(job_name, submitted_after=submitted_after) + + assert len(matches) == 1 + assert matches[0].job_id == 4201 + assert matches[0].array_task_ids == (0, 1) + assert runner.calls[0] == ( + "squeue", + "--noheader", + "--array", + "--format=%i|%.128j", + "--me", + f"--name={job_name}", + ) + assert runner.calls[1][0:6] == ( + "sacct", + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobIDRaw,JobName%128", + ) + assert runner.calls[1][6].startswith("--uid=") + assert runner.calls[1][7].startswith("--starttime=") + assert runner.calls[1][8] == f"--name={job_name}" + + +def test_client_rejects_a_named_lookup_result_with_a_different_name() -> None: + runner = FakeSlurmRunner() + runner.script_next("squeue", FakeCommandResponse(stdout="4201|unrelated\n")) + runner.script_next("sacct", FakeCommandResponse()) + + with pytest.raises(SlurmCommandOutputError, match="outside the requested exact name"): + SlurmCommandClient(runner).query_submissions_by_name( + f"dd-retry-{'a' * 32}", + submitted_after=datetime(2026, 9, 2, 18, tzinfo=timezone.utc), + ) + + def test_client_deduplicates_explicit_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) client.submit("run.sbatch") diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index d8df08bed..a937d6e9a 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -66,6 +66,8 @@ def test_collection_renderer_uses_authorized_mounts_and_no_gpu_directives( script = render_collection_script(plan, collection, destination) + assert f"#SBATCH --job-name={collection.submission_job_name}" in script + assert f"dd-collect-{plan.run_id}" not in script assert 'readonly DD_STATE_MOUNT="/workspace/primary:/workspace/primary"' in script assert 'readonly DD_OUTPUT_MOUNT="/workspace/primary/runs/run-001:/exports"' in script assert ( @@ -102,6 +104,8 @@ def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( script = render_generation_retry_script(multi_node_plan, retry) + assert f"#SBATCH --job-name={retry.submission_job_name}" in script + assert multi_node_plan.submission.job_name not in script assert "#SBATCH --array=1%2" in script assert 'DD_ATTEMPT_ORDINAL="0002"' in script assert 'readonly DD_ATTEMPT_MANIFEST="${DD_ATTEMPT_DIR}/attempt.json"' in script diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index eff457973..d7a4b0d42 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -10,6 +10,7 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path +from threading import Event from typing import cast import pytest @@ -19,6 +20,7 @@ import data_designer.slurm.state.collection_filesystem as collection_filesystem import data_designer.slurm.state.collection_merge as collection_merge import data_designer.slurm.state.collection_storage as collection_storage_module +import data_designer.slurm.state.collection_worker as collection_worker_module from data_designer.slurm.client import ClientOutcome, ClientResult from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 @@ -234,7 +236,7 @@ def test_retry_submits_only_the_failed_sparse_array_task( assert 'case "${DD_ARRAY_TASK_ID}"' in cast(str, runner.inputs[-1]) -def test_retry_ambiguous_submission_remains_prepared_and_blocks_duplicate( +def test_retry_ambiguous_submission_waits_for_scheduler_visibility_without_a_duplicate( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, single_node_plan: ResolvedSlurmRunPlan, @@ -265,11 +267,136 @@ def ambiguous_submit(script: str) -> object: retry_storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) assert retry_storage.read_status("retry-0001").state is RetryState.PREPARED - with pytest.raises(StateConflictError, match="ambiguous scheduler outcome"): - coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=6)) + runner.script_next("squeue", FakeCommandResponse()) + runner.script_next("sacct", FakeCommandResponse()) + with pytest.raises(StateConflictError, match="still being reconciled"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) assert submissions == 1 +def test_retry_recovers_an_accepted_submission_after_the_receipt_is_lost( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + initial = FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)) + retried = FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)) + runner = FakeSlurmRunner(arrays=(FakeSlurmArray(tasks=(initial,)), FakeSlurmArray(tasks=(retried,)))) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + first_attempt = _submitted_attempt(case, case.shards[0], scheduler=initial.scheduler) + case.writer.create_attempt(first_attempt) + runner.set_task_state(initial.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + original_submit = scheduler.submit_script + + def accept_then_lose_receipt(script: str) -> object: + original_submit(script) + raise SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", accept_then_lose_receipt) + with pytest.raises(SlurmStateError, match="cannot submit retry"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + monkeypatch.setattr(scheduler, "submit_script", original_submit) + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + retry_plan = storage.read_plan("retry-0001") + runner.script_next( + "squeue", + FakeCommandResponse(stdout=f"4201_0|{retry_plan.submission_job_name}\n"), + ) + runner.script_next("sacct", FakeCommandResponse()) + + recovered = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + + assert recovered[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=0) + assert case.writer.load_attempts(case.shards[0].shard_id) == (first_attempt, recovered[0]) + assert storage.read_status("retry-0001").state is RetryState.COMPLETED + assert ( + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=7), + ) + == recovered + ) + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 + + +@pytest.mark.parametrize( + ("accepted_before_receipt_loss", "replacement_job_id", "submission_count"), + [ + pytest.param(False, 4201, 2, id="unaccepted"), + pytest.param(True, 4301, 3, id="accepted-but-still-invisible"), + ], +) +def test_retry_replaces_or_fences_an_ambiguous_submission_after_its_deadline( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, + accepted_before_receipt_loss: bool, + replacement_job_id: int, + submission_count: int, +) -> None: + initial = FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)) + hidden = FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)) + replacement = FakeSlurmTask(SchedulerIdentity(array_job_id=replacement_job_id, array_task_id=0)) + retry_arrays = (FakeSlurmArray(tasks=(hidden,)),) if accepted_before_receipt_loss else () + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(initial,)), *retry_arrays, FakeSlurmArray(tasks=(replacement,))) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + first_attempt = _submitted_attempt(case, case.shards[0], scheduler=initial.scheduler) + case.writer.create_attempt(first_attempt) + runner.set_task_state(initial.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + original_submit = scheduler.submit_script + + def lose_submission_receipt(script: str) -> object: + if accepted_before_receipt_loss: + original_submit(script) + raise SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", lose_submission_receipt) + with pytest.raises(SlurmStateError, match="cannot submit retry"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + monkeypatch.setattr(scheduler, "submit_script", original_submit) + runner.script_next("squeue", FakeCommandResponse()) + runner.script_next("sacct", FakeCommandResponse()) + + attempts = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=11), + ) + + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("retry-0001").state is RetryState.FAILED + assert storage.read_status("retry-0002").state is RetryState.COMPLETED + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=replacement_job_id, array_task_id=0) + with pytest.raises(StateConflictError, match="scheduler identity"): + require_attempt_scheduler_identity( + case.workspace, + case.plan.run_id, + attempts[0].shard_id, + attempts[0].attempt_id, + SchedulerIdentity(array_job_id=4201 if accepted_before_receipt_loss else 4301, array_task_id=0), + ) + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == submission_count + + def test_retry_definite_submission_failure_settles_and_can_be_retried( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -678,7 +805,7 @@ def submit() -> object: assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 1 -def test_collection_ambiguous_submission_remains_prepared_and_blocks_duplicate( +def test_collection_ambiguous_submission_waits_for_scheduler_visibility_without_a_duplicate( tmp_path: Path, authored_run: DataDesignerSlurmConfig, multi_node_plan: ResolvedSlurmRunPlan, @@ -686,7 +813,8 @@ def test_collection_ambiguous_submission_remains_prepared_and_blocks_duplicate( ) -> None: case = _initialize_run(tmp_path, authored_run, multi_node_plan) _publish_all_winners(case) - scheduler = SlurmCommandClient(FakeSlurmRunner()) + runner = FakeSlurmRunner() + scheduler = SlurmCommandClient(runner) submissions = 0 def ambiguous_submit(script: str) -> object: @@ -703,11 +831,127 @@ def ambiguous_submit(script: str) -> object: storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) assert storage.read_status("collection-0001").state is CollectionState.PREPARED - with pytest.raises(StateConflictError, match="ambiguous scheduler outcome"): - coordinator.submit(submitted_at=case.created_at + timedelta(minutes=11)) + runner.script_next("squeue", FakeCommandResponse()) + runner.script_next("sacct", FakeCommandResponse()) + with pytest.raises(StateConflictError, match="still being reconciled"): + SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=11) + ) assert submissions == 1 +def test_collection_recovers_an_accepted_submission_after_the_receipt_is_lost( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + scheduler = SlurmCommandClient(runner) + original_submit = scheduler.submit_script + + def accept_then_lose_receipt(script: str) -> object: + original_submit(script) + raise SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", accept_then_lose_receipt) + with pytest.raises(SlurmStateError, match="cannot submit collection"): + SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=10) + ) + monkeypatch.setattr(scheduler, "submit_script", original_submit) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + collection_plan = storage.read_plan("collection-0001") + runner.script_next( + "squeue", + FakeCommandResponse(stdout=f"5101|{collection_plan.submission_job_name}\n"), + ) + runner.script_next("sacct", FakeCommandResponse()) + waiting_for_binding = Event() + binding_published = Event() + + def wait_for_binding(seconds: float) -> None: + assert seconds == 1 + waiting_for_binding.set() + assert binding_published.wait(timeout=2) + + monkeypatch.setattr(collection_worker_module, "sleep", wait_for_binding) + + with ThreadPoolExecutor(max_workers=1) as executor: + worker_result = executor.submit( + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + "collection-0001", + environment={"SLURM_JOB_ID": "5101"}, + ).run, + completed_at=case.created_at + timedelta(minutes=12), + ) + assert waiting_for_binding.wait(timeout=2) + recovered = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).refresh( + observed_at=case.created_at + timedelta(minutes=11) + ) + binding_published.set() + result = worker_result.result(timeout=5) + + assert recovered.collection_id == "collection-0001" + assert recovered.state is CollectionState.PENDING + assert recovered.scheduler == 5101 + assert result.actual_records == case.plan.invocation.authored.num_records + settled = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=13) + ) + assert settled.state is CollectionState.SUCCEEDED + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 1 + + +def test_collection_fences_an_invisible_accepted_submission_before_replacement_writes( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101), FakeSlurmJob(5102))) + scheduler = SlurmCommandClient(runner) + original_submit = scheduler.submit_script + + def accept_then_lose_receipt(script: str) -> object: + original_submit(script) + raise SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", accept_then_lose_receipt) + with pytest.raises(SlurmStateError, match="cannot submit collection"): + SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=10) + ) + monkeypatch.setattr(scheduler, "submit_script", original_submit) + runner.script_next("squeue", FakeCommandResponse()) + runner.script_next("sacct", FakeCommandResponse()) + + submitted = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=16) + ) + + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("collection-0001").state is CollectionState.FAILED + assert submitted.collection_id == "collection-0002" + assert submitted.state is CollectionState.SUBMITTED + assert submitted.scheduler == 5102 + with pytest.raises(StateConflictError, match="ordinary Slurm job identity"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + "collection-0001", + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=17)) + assert not Path(case.plan.output.root).exists() + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 + + def test_collection_definite_submission_failure_settles_and_can_be_retried( tmp_path: Path, authored_run: DataDesignerSlurmConfig, diff --git a/packages/data-designer-slurm/tests/state/test_submission_recovery.py b/packages/data-designer-slurm/tests/state/test_submission_recovery.py new file mode 100644 index 000000000..795a3994e --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_submission_recovery.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone + +import pytest + +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.launcher.models import SlurmSubmissionMatch +from data_designer.slurm.state import StateConflictError +from data_designer.slurm.state.submission_recovery import PreparedSubmission, resolve_prepared_submission + + +@dataclass(frozen=True) +class _SubmissionLookup: + matches: tuple[SlurmSubmissionMatch, ...] + + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + assert job_name == "dd-retry-0123456789abcdef0123456789abcdef" + assert submitted_after == datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + return self.matches + + +def test_prepared_submission_recovery_rejects_multiple_exact_matches() -> None: + submitted_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + + with pytest.raises(StateConflictError, match="multiple scheduler jobs"): + resolve_prepared_submission( + _SubmissionLookup( + ( + SlurmSubmissionMatch( + job_id=4201, + job_name="dd-retry-0123456789abcdef0123456789abcdef", + array_task_ids=(0,), + ), + SlurmSubmissionMatch( + job_id=4301, + job_name="dd-retry-0123456789abcdef0123456789abcdef", + array_task_ids=(0,), + ), + ) + ), + PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=submitted_at + timedelta(minutes=5), + expected_array_task_ids=(0,), + ), + observed_at=submitted_at + timedelta(minutes=1), + ) + + +def test_prepared_submission_recovery_returns_definitive_absence_only_after_deadline() -> None: + submitted_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + deadline = submitted_at + timedelta(minutes=5) + lookup = _SubmissionLookup(()) + + with pytest.raises(StateConflictError, match="still being reconciled"): + resolve_prepared_submission( + lookup, + PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=deadline, + expected_array_task_ids=(0,), + ), + observed_at=deadline, + ) + assert ( + resolve_prepared_submission( + lookup, + PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=deadline, + expected_array_task_ids=(0,), + ), + observed_at=deadline + timedelta(microseconds=1), + ) + is None + ) + + +@pytest.mark.parametrize("actual_shape", [None, (0,), (0, 2)]) +def test_prepared_submission_recovery_rejects_the_wrong_scheduler_shape( + actual_shape: tuple[int, ...] | None, +) -> None: + submitted_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + match = SlurmSubmissionMatch( + job_id=4201, + job_name="dd-retry-0123456789abcdef0123456789abcdef", + array_task_ids=actual_shape, + ) + + with pytest.raises(StateConflictError, match="shape"): + resolve_prepared_submission( + _SubmissionLookup((match,)), + PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=submitted_at + timedelta(minutes=5), + expected_array_task_ids=(0, 1), + ), + observed_at=submitted_at + timedelta(minutes=1), + ) From c78de4666e60933fc817f400032752cf2c6bcda1 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 3 Sep 2026 08:03:46 -0600 Subject: [PATCH 3/4] fix(slurm): preserve collection snapshots Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/state/artifacts.py | 90 +++++++++++++++---- .../tests/state/test_store.py | 16 ++++ 2 files changed, 90 insertions(+), 16 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py index e60ee8fbe..74b8131dd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py @@ -72,6 +72,7 @@ def rebind(self, dataset_descriptor: int, dataset_path: Path) -> None: with _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) as ( parent_descriptor, parent_path, + _, ): current = os.stat(parts[-1], dir_fd=parent_descriptor, follow_symlinks=False) if not _is_safe_file(current) or _file_facts(current) != self.file_facts: @@ -83,6 +84,13 @@ def validate_lease(self) -> None: raise OSError(f"candidate file {self.relative_path!r} changed during finalization") +@dataclass(frozen=True, slots=True) +class _ArtifactSnapshot: + relative_path: str + directory_identities: tuple[tuple[int, int], ...] + file_facts: tuple[int, int, int, int, int] + + @dataclass(frozen=True, slots=True) class VerifiedCandidateArtifacts: """Bounded live candidate leases plus metadata derived from artifact bytes.""" @@ -108,7 +116,7 @@ class CandidateArtifactSnapshot: record_counts: tuple[int, ...] dataset_schema_digest: str _dataset_identity: tuple[int, int] - _bindings: tuple[_ArtifactBinding, ...] + _artifacts: tuple[_ArtifactSnapshot, ...] class CandidateArtifactVerifier: @@ -136,23 +144,23 @@ def verify(self, candidate: CandidateOutputManifest) -> Iterator[VerifiedCandida def inspect(self, candidate: CandidateOutputManifest) -> CandidateArtifactSnapshot: """Inspect one candidate and return identities without retaining descriptors.""" - with self.verify(candidate) as verified: + dataset_path = Path(candidate.dataset_path) + with open_verified_directory(dataset_path, require_private=True) as dataset_descriptor: + record_counts, schema_digest, artifacts = _inspect_candidate_files( + dataset_descriptor, + dataset_path, + candidate.files, + ) return CandidateArtifactSnapshot( - record_counts=verified.record_counts, - dataset_schema_digest=verified.dataset_schema_digest, - _dataset_identity=_identity(os.fstat(verified._dataset.descriptor)), - _bindings=verified._bindings, + record_counts=record_counts, + dataset_schema_digest=schema_digest, + _dataset_identity=_identity(os.fstat(dataset_descriptor)), + _artifacts=artifacts, ) def rebind(self, candidate: CandidateOutputManifest, expected: CandidateArtifactSnapshot) -> None: """Reopen a candidate and require the identities captured before collection.""" - with self.verify(candidate) as current: - actual = CandidateArtifactSnapshot( - record_counts=current.record_counts, - dataset_schema_digest=current.dataset_schema_digest, - _dataset_identity=_identity(os.fstat(current._dataset.descriptor)), - _bindings=current._bindings, - ) + actual = self.inspect(candidate) if actual != expected: raise OSError("candidate paths or metadata changed during collection") @@ -199,7 +207,7 @@ def _open_output_file( output_file: CandidateOutputFile, ) -> tuple[int, str, _FileBinding]: parts = PurePosixPath(output_file.relative_path).parts - parent_descriptor, parent_path = resources.enter_context( + parent_descriptor, parent_path, _ = resources.enter_context( _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) ) name = parts[-1] @@ -223,14 +231,63 @@ def _open_output_file( return record_count, schema_digest, binding +def _inspect_candidate_files( + dataset_descriptor: int, + dataset_path: Path, + output_files: tuple[CandidateOutputFile, ...], +) -> tuple[tuple[int, ...], str, tuple[_ArtifactSnapshot, ...]]: + metadata = tuple( + _inspect_output_file(dataset_descriptor, dataset_path, output_file) for output_file in output_files + ) + schema_digests = tuple(schema_digest for _, schema_digest, _ in metadata) + if not schema_digests or any(digest != schema_digests[0] for digest in schema_digests[1:]): + raise OSError("candidate Parquet files do not share one dataset schema") + return ( + tuple(record_count for record_count, _, _ in metadata), + schema_digests[0], + tuple(artifact for _, _, artifact in metadata), + ) + + +def _inspect_output_file( + dataset_descriptor: int, + dataset_path: Path, + output_file: CandidateOutputFile, +) -> tuple[int, str, _ArtifactSnapshot]: + parts = PurePosixPath(output_file.relative_path).parts + with _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) as ( + parent_descriptor, + parent_path, + directory_identities, + ): + name = parts[-1] + display_path = parent_path / name + with open_verified_regular_file( + parent_descriptor, + name, + display_path, + expected_size=output_file.byte_size, + expected_sha256=output_file.sha256, + require_private=False, + ) as descriptor: + record_count, schema_digest = _read_parquet_metadata(descriptor, display_path) + artifact = _ArtifactSnapshot( + relative_path=output_file.relative_path, + directory_identities=directory_identities, + file_facts=_file_facts(os.fstat(descriptor)), + ) + return record_count, schema_digest, artifact + + @contextmanager def _open_parent_directory( dataset_descriptor: int, dataset_path: Path, parts: tuple[str, ...], -) -> Iterator[tuple[int, Path]]: +) -> Iterator[tuple[int, Path, tuple[tuple[int, int], ...]]]: parent_descriptor = os.dup(dataset_descriptor) parent_path = dataset_path + directory_identities: list[tuple[int, int]] = [] try: for part in parts: child_path = parent_path / part @@ -244,7 +301,8 @@ def _open_parent_directory( os.close(parent_descriptor) parent_descriptor = next_descriptor parent_path = child_path - yield parent_descriptor, parent_path + directory_identities.append(_identity(os.fstat(parent_descriptor))) + yield parent_descriptor, parent_path, tuple(directory_identities) finally: os.close(parent_descriptor) diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 8557b113e..91a0bb998 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -1427,6 +1427,22 @@ def track_open_file( assert active_files == 0 assert maximum_active_files == len(files) + maximum_active_files = 0 + verifier = state_artifacts.CandidateArtifactVerifier() + snapshot = verifier.inspect(candidate) + verifier.rebind(candidate, snapshot) + + assert active_files == 0 + assert maximum_active_files == 1 + + replacement = dataset_path / "replacement.parquet" + replacement.write_bytes((dataset_path / files[0].relative_path).read_bytes()) + replacement.chmod(0o644) + os.replace(replacement, dataset_path / files[0].relative_path) + + with pytest.raises(OSError, match="changed during collection"): + verifier.rebind(candidate, snapshot) + def test_candidate_schema_digest_ignores_arrow_metadata() -> None: schema = lazy.pa.schema([("record_id", lazy.pa.int64())]) From 4ddab96b8f9b0dfbc24c24450214b4e9370fa62f Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 3 Sep 2026 09:50:14 -0600 Subject: [PATCH 4/4] fix(slurm): bound partial array recovery Signed-off-by: Nabin Mulepati --- .../slurm/state/submission_recovery.py | 17 ++++++++-- .../tests/state/test_submission_recovery.py | 33 ++++++++++++++++++- 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py b/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py index f7ea62853..99a5a2d04 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py @@ -58,14 +58,25 @@ def resolve_prepared_submission( raise StateConflictError("multiple scheduler jobs match the prepared submission") if matches: match = matches[0] - if match.array_task_ids != prepared.expected_array_task_ids: - raise StateConflictError("scheduler job shape does not match the prepared submission") - return match.job_id + if match.array_task_ids == prepared.expected_array_task_ids: + return match.job_id + if _is_partial_array_view(match.array_task_ids, prepared.expected_array_task_ids): + if observed_at <= prepared.reconciliation_deadline: + raise StateConflictError("prepared submission is still being reconciled") + return None + raise StateConflictError("scheduler job shape does not match the prepared submission") if observed_at <= prepared.reconciliation_deadline: raise StateConflictError("prepared submission is still being reconciled") return None +def _is_partial_array_view( + observed: tuple[int, ...] | None, + expected: tuple[int, ...] | None, +) -> bool: + return observed is not None and expected is not None and set(observed) < set(expected) + + __all__ = [ "SUBMISSION_VISIBILITY_WINDOW", "PreparedSubmission", diff --git a/packages/data-designer-slurm/tests/state/test_submission_recovery.py b/packages/data-designer-slurm/tests/state/test_submission_recovery.py index 795a3994e..5e29db0d0 100644 --- a/packages/data-designer-slurm/tests/state/test_submission_recovery.py +++ b/packages/data-designer-slurm/tests/state/test_submission_recovery.py @@ -89,7 +89,7 @@ def test_prepared_submission_recovery_returns_definitive_absence_only_after_dead ) -@pytest.mark.parametrize("actual_shape", [None, (0,), (0, 2)]) +@pytest.mark.parametrize("actual_shape", [None, (0, 2)]) def test_prepared_submission_recovery_rejects_the_wrong_scheduler_shape( actual_shape: tuple[int, ...] | None, ) -> None: @@ -111,3 +111,34 @@ def test_prepared_submission_recovery_rejects_the_wrong_scheduler_shape( ), observed_at=submitted_at + timedelta(minutes=1), ) + + +def test_prepared_submission_recovery_bounds_a_partial_array_view() -> None: + submitted_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + deadline = submitted_at + timedelta(minutes=5) + lookup = _SubmissionLookup( + ( + SlurmSubmissionMatch( + job_id=4201, + job_name="dd-retry-0123456789abcdef0123456789abcdef", + array_task_ids=(0,), + ), + ) + ) + prepared = PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=deadline, + expected_array_task_ids=(0, 1), + ) + + with pytest.raises(StateConflictError, match="still being reconciled"): + resolve_prepared_submission(lookup, prepared, observed_at=deadline) + assert ( + resolve_prepared_submission( + lookup, + prepared, + observed_at=deadline + timedelta(microseconds=1), + ) + is None + )