From 05d3c72a7791296afa414a10897d1d46394090ef Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 2 Sep 2026 16:26:12 -0600 Subject: [PATCH 1/5] feat: reconcile persisted Slurm state Persist normalized scheduler observations and compose fresh-process run, shard, attempt, readiness, generation, and winner status. Preserve bounded accounting lag and immutable terminal evidence for status and benchmark refresh consumers. Part of #869 Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/client.py | 23 +- .../data_designer/slurm/launcher/models.py | 9 +- .../data_designer/slurm/launcher/parsing.py | 13 +- .../src/data_designer/slurm/state/__init__.py | 27 + .../src/data_designer/slurm/state/base.py | 5 + .../data_designer/slurm/state/observation.py | 202 +++++++ .../src/data_designer/slurm/state/observer.py | 283 ++++++++++ .../src/data_designer/slurm/state/reader.py | 64 ++- .../slurm/state/reconciliation.py | 15 +- .../data_designer/slurm/state/scheduler.py | 22 +- .../src/data_designer/slurm/state/status.py | 208 ++++++++ .../src/data_designer/slurm/state/storage.py | 54 ++ .../data_designer/slurm/state/validation.py | 8 +- .../tests/state/test_observer.py | 501 ++++++++++++++++++ scripts/test_slurm_package_install.py | 10 +- 15 files changed, 1394 insertions(+), 50 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/observation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/observer.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/status.py create mode 100644 packages/data-designer-slurm/tests/state/test_observer.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 2d822a7f1..e45e08048 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -11,14 +11,12 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from typing import TypeAlias from data_designer.slurm.contracts import Identifier from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError from data_designer.slurm.launcher.models import ( SlurmAccountingEntry, SlurmJobSubmissionReceipt, - SlurmObservedJobIdentity, SlurmQueueEntry, ) from data_designer.slurm.launcher.parsing import ( @@ -28,9 +26,8 @@ parse_submission, ) from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner -from data_designer.slurm.state import SchedulerIdentity +from data_designer.slurm.state import SchedulerIdentity, SchedulerJobIdentity -_JobSelector: TypeAlias = int | SchedulerIdentity _IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") _ENVIRONMENT_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") _MAX_SLURM_INTEGER = (1 << 32) - 1 @@ -108,7 +105,7 @@ def submit_script( ) return parse_submission(output) - def query_queue(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmQueueEntry, ...]: + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: """Return normalized active-queue rows for explicit managed jobs.""" requested = tuple(selectors) jobs = _format_selectors(requested) @@ -129,7 +126,7 @@ def query_queue(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmQueueEntr ) return tuple(entry for entry in entries if entry.job_identity not in ignored) - def query_accounting(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmAccountingEntry, ...]: + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: """Return normalized accounting rows for explicit managed jobs.""" requested = tuple(selectors) jobs = _format_selectors(requested) @@ -152,7 +149,7 @@ def query_accounting(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmAcco ) return tuple(entry for entry in entries if entry.job_identity not in ignored) - def cancel(self, selector: _JobSelector) -> None: + def cancel(self, selector: SchedulerJobIdentity) -> None: """Cancel one managed Slurm job, array, or array task.""" self._run((self._executables.scancel, _format_selector(selector))) @@ -212,13 +209,13 @@ def _format_export_environment(environment: Mapping[str, str] | None) -> tuple[s return f"--export={','.join(names)}", environment -def _format_selectors(selectors: Sequence[_JobSelector]) -> str: +def _format_selectors(selectors: Sequence[SchedulerJobIdentity]) -> str: if not selectors: raise ValueError("at least one managed Slurm job selector is required") return ",".join(dict.fromkeys(_format_selector(selector) for selector in selectors)) -def _format_selector(selector: _JobSelector) -> str: +def _format_selector(selector: SchedulerJobIdentity) -> str: if isinstance(selector, SchedulerIdentity): job_id = _format_job_id(selector.array_job_id) if selector.array_task_id > _MAX_SLURM_INTEGER: @@ -234,16 +231,16 @@ def _format_job_id(value: object) -> str: def _validate_observed_job_identities( - job_identities: Sequence[SlurmObservedJobIdentity], - selectors: Sequence[_JobSelector], + job_identities: Sequence[SchedulerJobIdentity], + selectors: Sequence[SchedulerJobIdentity], *, command: str, -) -> frozenset[SlurmObservedJobIdentity]: +) -> frozenset[SchedulerJobIdentity]: """Validate result correlation and return unselected aggregate rows.""" selected_job_ids = {selector for selector in selectors if type(selector) is int} selected_array_tasks = {selector for selector in selectors if isinstance(selector, SchedulerIdentity)} selected_array_job_ids = {selector.array_job_id for selector in selected_array_tasks} - ignored: set[SlurmObservedJobIdentity] = set() + ignored: set[SchedulerJobIdentity] = set() for job_identity in job_identities: if type(job_identity) is int and job_identity in selected_job_ids: continue 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 7d1168255..649460dbf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -6,11 +6,8 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TypeAlias -from data_designer.slurm.state import SchedulerIdentity, SchedulerState - -SlurmObservedJobIdentity: TypeAlias = int | SchedulerIdentity +from data_designer.slurm.state import SchedulerJobIdentity, SchedulerState @dataclass(frozen=True) @@ -32,7 +29,7 @@ class SlurmProcessExitCode: class SlurmQueueEntry: """One transient normalized active-queue entry.""" - job_identity: SlurmObservedJobIdentity + job_identity: SchedulerJobIdentity state: SchedulerState @@ -40,6 +37,6 @@ class SlurmQueueEntry: class SlurmAccountingEntry: """One transient normalized accounting entry.""" - job_identity: SlurmObservedJobIdentity + job_identity: SchedulerJobIdentity state: SchedulerState process_exit_code: SlurmProcessExitCode diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index adf0e13c2..213ce09ca 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 @@ -11,11 +11,10 @@ from data_designer.slurm.launcher.models import ( SlurmAccountingEntry, SlurmJobSubmissionReceipt, - SlurmObservedJobIdentity, SlurmProcessExitCode, SlurmQueueEntry, ) -from data_designer.slurm.state import SchedulerIdentity, SchedulerState +from data_designer.slurm.state import SchedulerIdentity, SchedulerJobIdentity, SchedulerState _ARRAY_ID_PATTERN = re.compile(r"^(?P[1-9][0-9]*)_(?P[0-9]+)$") _JOB_ID_PATTERN = re.compile(r"^[1-9][0-9]*$") @@ -71,7 +70,7 @@ def parse_submission(output: str) -> SlurmJobSubmissionReceipt: def parse_queue(output: str) -> tuple[SlurmQueueEntry, ...]: """Parse ``squeue --format=%i|%T`` rows.""" entries: list[SlurmQueueEntry] = [] - identities: set[SlurmObservedJobIdentity] = set() + identities: set[SchedulerJobIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 2: @@ -85,7 +84,7 @@ def parse_queue(output: str) -> tuple[SlurmQueueEntry, ...]: def parse_accounting(output: str) -> tuple[SlurmAccountingEntry, ...]: """Parse job and array-task rows from ``sacct --format=JobID,State,ExitCode``.""" entries: list[SlurmAccountingEntry] = [] - identities: set[SlurmObservedJobIdentity] = set() + identities: set[SchedulerJobIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 3: @@ -184,7 +183,7 @@ def _parse_array_identity(value: str, *, command: str, line_number: int) -> Sche ) -def _parse_job_identity(value: str, *, command: str, line_number: int) -> SlurmObservedJobIdentity: +def _parse_job_identity(value: str, *, command: str, line_number: int) -> SchedulerJobIdentity: message = f"{command} line {line_number} contains an invalid job or array-task ID" if _JOB_ID_PATTERN.fullmatch(value) is not None: return _parse_decimal(value, message=message) @@ -218,8 +217,8 @@ def _parse_decimal(value: str, *, message: str) -> int: def _reject_duplicate( - job_identity: SlurmObservedJobIdentity, - identities: set[SlurmObservedJobIdentity], + job_identity: SchedulerJobIdentity, + identities: set[SchedulerJobIdentity], *, command: str, line_number: int, 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 ce69166e3..3c587f615 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 @@ -22,6 +22,7 @@ from data_designer.slurm.state.artifacts import compute_candidate_schema_digest from data_designer.slurm.state.base import ( SchedulerIdentity, + SchedulerJobIdentity, StateRecord, StateValue, ) @@ -38,6 +39,12 @@ RunManifest, ShardManifest, ) +from data_designer.slurm.state.observation import ( + SchedulerAccountingRecord, + SchedulerObservationClient, + SchedulerObservationCollector, + SchedulerQueueRecord, +) from data_designer.slurm.state.outputs import ( CANDIDATE_OUTPUT_FORMAT, MAXIMUM_CANDIDATE_OUTPUT_FILES, @@ -66,6 +73,13 @@ SchedulerObservation, SchedulerState, ) +from data_designer.slurm.state.status import ( + AttemptStatus, + EffectiveRunState, + GenerationState, + RunStatus, + ShardStatus, +) from data_designer.slurm.state.validation import ( StateContractError, validate_attempt_manifest, @@ -80,9 +94,11 @@ ) if TYPE_CHECKING: + from data_designer.slurm.state.observer import SlurmStateReconciler # noqa: F401 from data_designer.slurm.state.store import SlurmStateWriter # noqa: F401 _LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "SlurmStateReconciler": ("data_designer.slurm.state.observer", "SlurmStateReconciler"), "SlurmStateWriter": ("data_designer.slurm.state.store", "SlurmStateWriter"), } @@ -92,6 +108,7 @@ "AttemptManifest", "AttemptId", "AttemptReadiness", + "AttemptStatus", "AttemptTerminalClassification", "CandidateOutcome", "CANDIDATE_OUTPUT_FORMAT", @@ -105,23 +122,33 @@ "ContractValue", "DeploymentReadiness", "EffectiveAttemptState", + "EffectiveRunState", "EndpointPublicationState", "Identifier", + "GenerationState", "ProbeEvidence", "ProbeOutcome", "ReadinessState", "ReasonCode", "RecordRange", "RunManifest", + "RunStatus", "ResumeWorkspace", "SchedulerIdentity", + "SchedulerJobIdentity", + "SchedulerAccountingRecord", + "SchedulerObservationClient", + "SchedulerObservationCollector", + "SchedulerQueueRecord", "SchedulerObservation", "SchedulerState", "Sha256Digest", "ShardManifest", + "ShardStatus", "ShardId", "ShardWinner", "SlurmStateError", + "SlurmStateReconciler", "SlurmStateWriter", "StateConflictError", "StateContractError", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py index e9825bd95..dc3ffcf19 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -4,6 +4,7 @@ from __future__ import annotations from datetime import datetime, timedelta +from typing import TypeAlias from pydantic import NonNegativeInt, PositiveInt @@ -44,10 +45,14 @@ class SchedulerIdentity(StateValue): array_task_id: NonNegativeInt +SchedulerJobIdentity: TypeAlias = PositiveInt | SchedulerIdentity + + __all__ = [ "ArtifactReference", "Identifier", "SchedulerIdentity", + "SchedulerJobIdentity", "Sha256Digest", "StateRecord", "StateValue", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py new file mode 100644 index 000000000..3bf039715 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize active and accounting evidence into scheduler observations.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta +from typing import Protocol + +from data_designer.slurm.state.base import SchedulerJobIdentity, validate_utc_timestamp +from data_designer.slurm.state.errors import SlurmStateError +from data_designer.slurm.state.scheduler import ( + SchedulerObservation, + SchedulerState, + is_scheduler_terminal_state, +) +from data_designer.slurm.state.validation import StateContractError, validate_scheduler_observation_transition + +_ACCOUNTING_LAG_WINDOW = timedelta(minutes=5) + + +class SchedulerQueueRecord(Protocol): + """Normalized active-queue record consumed by reconciliation.""" + + job_identity: SchedulerJobIdentity + state: SchedulerState + + +class SchedulerAccountingRecord(Protocol): + """Normalized accounting record consumed by reconciliation.""" + + job_identity: SchedulerJobIdentity + state: SchedulerState + + +class SchedulerObservationClient(Protocol): + """Query normalized active and accounting scheduler records.""" + + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SchedulerQueueRecord, ...]: + """Return active queue records for the requested identities.""" + ... + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SchedulerAccountingRecord, ...]: + """Return accounting records for the requested identities.""" + ... + + +class SchedulerObservationCollector: + """Apply terminal-accounting precedence and bounded lag semantics.""" + + def __init__(self, client: SchedulerObservationClient) -> None: + self._client = client + + def collect( + self, + selectors: Sequence[SchedulerJobIdentity], + *, + observed_at: datetime, + previous: Mapping[SchedulerJobIdentity, SchedulerObservation | None] | None = None, + ) -> tuple[SchedulerObservation, ...]: + """Return one deterministic observation for every requested identity.""" + validate_utc_timestamp(observed_at) + requested = tuple(dict.fromkeys(selectors)) + if not requested: + return () + prior = {} if previous is None else previous + queue, accounting = self._query_scheduler(requested) + queue_by_identity = self._index_records(queue, requested, source="active queue") + accounting_by_identity = self._index_records(accounting, requested, source="accounting") + return tuple( + self._resolve_observation( + identity, + observed_at, + queue_by_identity.get(identity), + accounting_by_identity.get(identity), + prior.get(identity), + ) + for identity in requested + ) + + def _query_scheduler( + self, + selectors: tuple[SchedulerJobIdentity, ...], + ) -> tuple[tuple[SchedulerQueueRecord, ...], tuple[SchedulerAccountingRecord, ...]]: + try: + return self._client.query_queue(selectors), self._client.query_accounting(selectors) + except (OSError, RuntimeError, ValueError) as error: + raise SlurmStateError("cannot query normalized scheduler observations") from error + + @staticmethod + def _index_records( + records: Sequence[SchedulerQueueRecord | SchedulerAccountingRecord], + requested: tuple[SchedulerJobIdentity, ...], + *, + source: str, + ) -> dict[SchedulerJobIdentity, SchedulerState]: + expected = set(requested) + indexed: dict[SchedulerJobIdentity, SchedulerState] = {} + for record in records: + identity = record.job_identity + if identity not in expected: + raise SlurmStateError(f"{source} returned an unrequested scheduler identity") + if identity in indexed: + raise SlurmStateError(f"{source} returned a duplicate scheduler identity") + if not isinstance(record.state, SchedulerState): + raise SlurmStateError(f"{source} returned an invalid normalized scheduler state") + indexed[identity] = record.state + return indexed + + @staticmethod + def _resolve_observation( + identity: SchedulerJobIdentity, + observed_at: datetime, + queue_state: SchedulerState | None, + accounting_state: SchedulerState | None, + previous: SchedulerObservation | None, + ) -> SchedulerObservation: + state = _select_observed_state(queue_state, accounting_state) + if ( + previous is not None + and is_scheduler_terminal_state(previous.state) + and (accounting_state is None or not is_scheduler_terminal_state(accounting_state)) + ): + state = previous.state + observation = ( + _resolve_missing_observation(identity, observed_at, previous) + if state is None + else SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=state, + ) + ) + if previous is not None: + try: + validate_scheduler_observation_transition(previous, observation) + except StateContractError as error: + raise SlurmStateError("scheduler observation violates persisted chronology") from error + return observation + + +def _select_observed_state( + queue_state: SchedulerState | None, + accounting_state: SchedulerState | None, +) -> SchedulerState | None: + if accounting_state is not None and is_scheduler_terminal_state(accounting_state): + return accounting_state + if queue_state is not None: + return queue_state + return accounting_state + + +def _resolve_missing_observation( + identity: SchedulerJobIdentity, + observed_at: datetime, + previous: SchedulerObservation | None, +) -> SchedulerObservation: + if previous is not None and previous.state is SchedulerState.ACCOUNTING_LAG: + deadline = previous.reconciliation_deadline + if deadline is None: + raise SlurmStateError("persisted accounting lag has no reconciliation deadline") + if observed_at > deadline: + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=SchedulerState.UNKNOWN, + ) + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=SchedulerState.ACCOUNTING_LAG, + reconciliation_deadline=deadline, + ) + if previous is not None and ( + previous.state is SchedulerState.UNKNOWN or is_scheduler_terminal_state(previous.state) + ): + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=previous.state, + ) + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=SchedulerState.ACCOUNTING_LAG, + reconciliation_deadline=observed_at + _ACCOUNTING_LAG_WINDOW, + ) + + +__all__ = [ + "SchedulerAccountingRecord", + "SchedulerObservationClient", + "SchedulerObservationCollector", + "SchedulerQueueRecord", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py new file mode 100644 index 000000000..eeef347c5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process reconciliation of one persisted Slurm run.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier, ShardId, validate_absolute_path +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.base import SchedulerIdentity, SchedulerJobIdentity +from data_designer.slurm.state.errors import ( + SlurmStateError, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.observation import SchedulerObservationClient, SchedulerObservationCollector +from data_designer.slurm.state.outputs import ShardWinner +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.reconciliation import reconcile_attempt_observation +from data_designer.slurm.state.scheduler import EffectiveAttemptState, SchedulerObservation +from data_designer.slurm.state.status import ( + AttemptStatus, + RunStatus, + ShardStatus, + derive_generation_state, + derive_run_state, + derive_shard_state, +) +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import StateContractError, validate_scheduler_observation_transition + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +@dataclass(frozen=True, slots=True) +class _ShardSnapshot: + run: RunManifest + plan: ResolvedSlurmRunPlan + shard: ShardManifest + attempts: tuple[AttemptManifest, ...] + + +@dataclass(frozen=True, slots=True) +class _ObservationBatch: + previous: dict[SchedulerIdentity, SchedulerObservation | None] + current: dict[SchedulerJobIdentity, SchedulerObservation] + observed_at: datetime + + +class SlurmStateReconciler: + """Refresh persisted run status from normalized scheduler observations. + + Each refresh reconstructs state from the compute-visible workspace. No + controller memory participates in status, wait, or benchmark refreshes. + + Args: + workspace_root: Selected compute-visible workspace root. + run_id: Stable application-owned run identity. + scheduler: Client returning normalized active and accounting records. + """ + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + scheduler: SchedulerObservationClient, + ) -> None: + normalized_root, normalized_run_id = _validate_location(workspace_root, run_id) + self._storage = StateStorage(normalized_root, normalized_run_id) + self._reader = StateReader(self._storage, normalized_run_id) + self._finalizer = WinnerFinalizer(self._storage, self._reader) + self._collector = SchedulerObservationCollector(scheduler) + self._run_id = normalized_run_id + + @property + def run_root(self) -> Path: + """Return the workspace-derived root for this run.""" + return self._storage.run_root + + def refresh(self, *, observed_at: datetime | None = None) -> RunStatus: + """Persist current scheduler evidence and return complete run status. + + Raises: + SlurmStateError: If scheduler evidence cannot be queried or state + cannot be reconstructed safely. + """ + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + _validate_observed_at(timestamp) + run, plan, shards = self._reader.load_context() + if timestamp < run.created_at: + raise SlurmStateError("observation timestamp cannot precede run creation") + attempts_by_shard = self._reader.load_validated_attempts(run, plan, shards) + previous = self._load_previous_observations(attempts_by_shard) + selectors = tuple(previous.keys()) + current = self._collector.collect(selectors, observed_at=timestamp, previous=previous) + batch = _ObservationBatch( + previous=previous, + current={observation.scheduler: observation for observation in current}, + observed_at=timestamp, + ) + shard_statuses = tuple( + self._refresh_shard( + _ShardSnapshot(run, plan, shard, attempts_by_shard[shard.shard_id]), + batch, + ) + for shard in shards + ) + return self._compose_run_status(run, timestamp, shard_statuses) + + def _compose_run_status( + self, + run: RunManifest, + observed_at: datetime, + shards: tuple[ShardStatus, ...], + ) -> RunStatus: + try: + return RunStatus( + run=run, + observed_at=observed_at, + shards=shards, + effective_state=derive_run_state(shards), + ) + except ValidationError as error: + raise StateCorruptionError(f"cannot reconcile run {self._run_id!r}") from error + + def _load_previous_observations( + self, + attempts_by_shard: dict[ShardId, tuple[AttemptManifest, ...]], + ) -> dict[SchedulerIdentity, SchedulerObservation | None]: + previous: dict[SchedulerIdentity, SchedulerObservation | None] = {} + for attempts in attempts_by_shard.values(): + for attempt in attempts: + if attempt.scheduler is not None: + previous[attempt.scheduler] = self._reader.load_optional_scheduler_observation(attempt) + return previous + + def _refresh_shard( + self, + expected: _ShardSnapshot, + batch: _ObservationBatch, + ) -> ShardStatus: + try: + with self._storage.acquire_shard_lock(expected.shard.shard_id): + current_run, current_plan, current_shard = self._reader.load_shard_context(expected.shard.shard_id) + attempts = self._reader.load_validated_shard_attempts(current_run, current_plan, current_shard) + self._require_unchanged_context( + expected, + _ShardSnapshot(current_run, current_plan, current_shard, attempts), + ) + winner = self._finalizer.load_optional_winner( + expected.run, + expected.plan, + expected.shard, + attempts, + ) + statuses = tuple( + self._build_attempt_status( + expected, + batch, + attempt, + winner, + ) + for attempt in attempts + ) + self._validate_winner_scheduler_consistency(winner, statuses) + return ShardStatus( + shard=expected.shard, + attempts=statuses, + winner=winner, + effective_state=derive_shard_state(statuses, winner), + ) + except (StateConflictError, StateCorruptionError, StateNotFoundError): + raise + except (OSError, StateContractError, ValidationError) as error: + raise StateCorruptionError(f"cannot reconcile shard {expected.shard.shard_id!r}") from error + + def _build_attempt_status( + self, + snapshot: _ShardSnapshot, + batch: _ObservationBatch, + attempt: AttemptManifest, + winner: ShardWinner | None, + ) -> AttemptStatus: + readiness = self._reader.load_optional_readiness(snapshot.plan, attempt) + result = self._reader.load_optional_attempt_result(snapshot.plan, snapshot.shard, attempt) + scheduler = None + if attempt.scheduler is not None: + scheduler = batch.current[attempt.scheduler] + self._persist_observation(attempt, batch.previous[attempt.scheduler], scheduler) + effective_state = reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=batch.observed_at, + ) + else: + if attempt.state is not AttemptLifecycleState.CREATED: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has no scheduler identity") + effective_state = EffectiveAttemptState.PENDING + if attempt.candidate_output is not None and result is None: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} is missing its generation results") + client_result, candidate = (None, None) if result is None else result + is_winner = winner is not None and winner.attempt_id == attempt.attempt_id + generation_state = derive_generation_state( + effective_state, + has_candidate=candidate is not None, + is_winner=is_winner, + ) + return AttemptStatus( + attempt=attempt, + readiness=readiness, + scheduler=scheduler, + client_result=client_result, + candidate_output=candidate, + effective_state=effective_state, + generation_state=generation_state, + is_winner=is_winner, + ) + + def _persist_observation( + self, + attempt: AttemptManifest, + expected_previous: SchedulerObservation | None, + current: SchedulerObservation, + ) -> None: + persisted = self._reader.load_optional_scheduler_observation(attempt) + if persisted != expected_previous: + raise StateConflictError("scheduler evidence changed during reconciliation; refresh again") + if persisted == current: + self._storage.sync_attempt_directory(attempt.shard_id, attempt.attempt_id) + return + if persisted is None: + self._storage.publish_scheduler_observation(attempt.shard_id, attempt.attempt_id, current) + return + validate_scheduler_observation_transition(persisted, current) + self._storage.replace_scheduler_observation(attempt.shard_id, attempt.attempt_id, current) + + @staticmethod + def _require_unchanged_context( + expected: _ShardSnapshot, + current: _ShardSnapshot, + ) -> None: + if current != expected: + raise StateConflictError("persisted state changed during reconciliation; refresh again") + + @staticmethod + def _validate_winner_scheduler_consistency( + winner: ShardWinner | None, + statuses: tuple[AttemptStatus, ...], + ) -> None: + if winner is None: + return + winning = next(status for status in statuses if status.attempt.attempt_id == winner.attempt_id) + if winning.effective_state is not EffectiveAttemptState.SUCCEEDED: + raise StateCorruptionError("persisted winner conflicts with terminal scheduler evidence") + + +def _validate_location(workspace_root: str | Path, run_id: Identifier) -> tuple[Path, Identifier]: + try: + normalized_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 run location") from error + return Path(normalized_root), normalized_run_id + + +def _validate_observed_at(observed_at: datetime) -> None: + if not isinstance(observed_at, datetime) or observed_at.tzinfo is None or observed_at.utcoffset() is None: + raise SlurmStateError("observation timestamp must be timezone-aware UTC") + if observed_at.utcoffset().total_seconds() != 0: + raise SlurmStateError("observation timestamp must be timezone-aware UTC") + + +__all__ = ["SlurmStateReconciler"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py index 67e23c2a9..a20828838 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py @@ -5,13 +5,16 @@ from __future__ import annotations +from data_designer.slurm.client import ClientResult from data_designer.slurm.config import DataDesignerSlurmConfig from data_designer.slurm.contracts import AttemptId, Identifier, ShardId from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.state.errors import StateCorruptionError, StateNotFoundError from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.outputs import CandidateOutputManifest from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError from data_designer.slurm.state.readiness import AttemptReadiness +from data_designer.slurm.state.scheduler import SchedulerObservation from data_designer.slurm.state.storage import StateStorage from data_designer.slurm.state.validation import ( StateContractError, @@ -139,18 +142,65 @@ def load_readiness(self, shard_id: ShardId, attempt_id: AttemptId) -> AttemptRea context = self.load_shard_context(shard_id) _, plan, _ = context attempt = self.get_attempt(self.load_attempts(shard_id, context), attempt_id) + readiness = self.load_optional_readiness(plan, attempt) + if readiness is None: + raise StateNotFoundError(f"attempt {attempt_id!r} has no readiness snapshot") + return readiness + + def load_optional_readiness( + self, + plan: ResolvedSlurmRunPlan, + attempt: AttemptManifest, + ) -> AttemptReadiness | None: + """Load validated readiness when the runtime has published it.""" try: - readiness = self._storage.read_readiness(shard_id, attempt_id) + readiness = self._storage.read_readiness(attempt.shard_id, attempt.attempt_id) PersistedPlanStateValidator(plan).validate_readiness_snapshot(attempt, readiness) return readiness - except FileNotFoundError as error: - raise StateNotFoundError(f"attempt {attempt_id!r} has no readiness snapshot") from error - except StateCorruptionError: - raise + except FileNotFoundError: + return None except PlanStateContractError as error: - raise StateCorruptionError(f"attempt {attempt_id!r} has invalid persisted readiness") from error + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has invalid persisted readiness") from error + except OSError as error: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has unreadable readiness") from error + + def load_optional_scheduler_observation(self, attempt: AttemptManifest) -> SchedulerObservation | None: + """Load and identity-check the most recent scheduler observation.""" + try: + observation = self._storage.read_scheduler_observation(attempt.shard_id, attempt.attempt_id) + except FileNotFoundError: + return None except OSError as error: - raise StateCorruptionError(f"attempt {attempt_id!r} has unreadable readiness") from error + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has unreadable scheduler evidence") from error + if attempt.scheduler is None or observation.scheduler != attempt.scheduler: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has mismatched scheduler evidence") + if observation.observed_at < attempt.created_at: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has scheduler evidence before its creation") + return observation + + def load_optional_attempt_result( + self, + plan: ResolvedSlurmRunPlan, + shard: ShardManifest, + attempt: AttemptManifest, + ) -> tuple[ClientResult, CandidateOutputManifest] | None: + """Load and validate a complete producer result pair when present.""" + try: + client_result, candidate = self._storage.read_finalization_records(shard.shard_id, attempt.attempt_id) + except FileNotFoundError: + return None + except OSError as error: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has unreadable generation results") from error + try: + PersistedPlanStateValidator(plan).validate_attempt_result( + plan.shards[shard.shard_index], + attempt, + client_result, + candidate, + ) + except PlanStateContractError as error: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has invalid generation results") from error + return client_result, candidate def load_validated_attempts( self, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index 63c0d2a1a..a210353f4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -15,6 +15,7 @@ EffectiveAttemptState, SchedulerObservation, SchedulerState, + is_scheduler_failure_state, ) from data_designer.slurm.state.validation import StateContractError @@ -59,18 +60,6 @@ EndpointPublicationState.FAILED: frozenset({EndpointPublicationState.FAILED}), } -_SCHEDULER_FAILURE_STATES = frozenset( - { - SchedulerState.FAILED, - SchedulerState.CANCELLED, - SchedulerState.TIMED_OUT, - SchedulerState.NODE_FAILED, - SchedulerState.PREEMPTED, - SchedulerState.REQUEUED, - SchedulerState.OUT_OF_MEMORY, - } -) - def validate_readiness_transition( previous: AttemptReadiness, @@ -165,7 +154,7 @@ def reconcile_attempt_observation( if readiness is not None: _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") - if scheduler.state in _SCHEDULER_FAILURE_STATES: + if is_scheduler_failure_state(scheduler.state): return EffectiveAttemptState.FAILED if attempt.state is AttemptLifecycleState.FAILED: return EffectiveAttemptState.FAILED diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py index 3245c75fa..e06e11cc2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py @@ -9,7 +9,7 @@ from pydantic import field_validator, model_validator from data_designer.slurm.state.base import ( - SchedulerIdentity, + SchedulerJobIdentity, StateRecord, validate_optional_utc_timestamp, validate_utc_timestamp, @@ -34,7 +34,7 @@ class SchedulerState(str, Enum): class SchedulerObservation(StateRecord): """Normalized scheduler observation used for deterministic reconciliation.""" - scheduler: SchedulerIdentity + scheduler: SchedulerJobIdentity observed_at: datetime state: SchedulerState reconciliation_deadline: datetime | None = None @@ -61,3 +61,21 @@ class EffectiveAttemptState(str, Enum): FAILED = "failed" ACCOUNTING_LAG = "accounting_lag" UNKNOWN = "unknown" + + +def is_scheduler_failure_state(state: SchedulerState) -> bool: + """Return whether a scheduler state is terminal failure evidence.""" + return state in { + SchedulerState.FAILED, + SchedulerState.CANCELLED, + SchedulerState.TIMED_OUT, + SchedulerState.NODE_FAILED, + SchedulerState.PREEMPTED, + SchedulerState.REQUEUED, + SchedulerState.OUT_OF_MEMORY, + } + + +def is_scheduler_terminal_state(state: SchedulerState) -> bool: + """Return whether a scheduler state is terminal accounting evidence.""" + return state is SchedulerState.COMPLETED or is_scheduler_failure_state(state) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/status.py b/packages/data-designer-slurm/src/data_designer/slurm/state/status.py new file mode 100644 index 000000000..581126a0b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/status.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process status values derived from persisted and scheduler evidence.""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from pydantic import Field, field_validator, model_validator + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.state.base import StateValue, validate_utc_timestamp +from data_designer.slurm.state.execution import AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner +from data_designer.slurm.state.readiness import AttemptReadiness +from data_designer.slurm.state.scheduler import EffectiveAttemptState, SchedulerObservation + + +class GenerationState(str, Enum): + """Effective progress of one attempt's dataset generation.""" + + NOT_STARTED = "not_started" + ACTIVE = "active" + CANDIDATE_READY = "candidate_ready" + WON = "won" + FAILED = "failed" + UNKNOWN = "unknown" + + +class EffectiveRunState(str, Enum): + """Aggregated state of all planned shards in one run.""" + + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" + + +class AttemptStatus(StateValue): + """Validated persisted and observed evidence for one attempt.""" + + attempt: AttemptManifest + readiness: AttemptReadiness | None + scheduler: SchedulerObservation | None + client_result: ClientResult | None + candidate_output: CandidateOutputManifest | None + effective_state: EffectiveAttemptState + generation_state: GenerationState + is_winner: bool = False + + @model_validator(mode="after") + def validate_evidence(self) -> AttemptStatus: + if self.attempt.scheduler is None: + if self.scheduler is not None: + raise ValueError("created attempts cannot have scheduler evidence") + elif self.scheduler is None or self.scheduler.scheduler != self.attempt.scheduler: + raise ValueError("attempt status requires matching scheduler evidence") + if self.readiness is not None and ( + self.readiness.run_id, + self.readiness.shard_id, + self.readiness.attempt_id, + ) != (self.attempt.run_id, self.attempt.shard_id, self.attempt.attempt_id): + raise ValueError("attempt status readiness identity does not match") + if (self.client_result is None) != (self.candidate_output is None): + raise ValueError("attempt status requires a complete generation-result pair") + if self.client_result is not None and self.candidate_output is not None: + expected = (self.attempt.run_id, self.attempt.shard_id, self.attempt.attempt_id) + if ( + self.client_result.run_id, + self.client_result.shard_id, + self.client_result.attempt_id, + ) != expected: + raise ValueError("client result identity does not match the attempt") + if ( + self.candidate_output.run_id, + self.candidate_output.shard_id, + self.candidate_output.attempt_id, + ) != expected: + raise ValueError("candidate output identity does not match the attempt") + expected_generation = derive_generation_state( + self.effective_state, + has_candidate=self.candidate_output is not None, + is_winner=self.is_winner, + ) + if self.generation_state is not expected_generation: + raise ValueError("generation state does not match the attempt evidence") + return self + + +class ShardStatus(StateValue): + """Effective state and attempt history for one planned shard.""" + + shard: ShardManifest + attempts: tuple[AttemptStatus, ...] + winner: ShardWinner | None + effective_state: EffectiveAttemptState + + @model_validator(mode="after") + def validate_status(self) -> ShardStatus: + if any(status.attempt.shard_id != self.shard.shard_id for status in self.attempts): + raise ValueError("shard status contains an attempt for another shard") + ordinals = tuple(status.attempt.attempt_ordinal for status in self.attempts) + if ordinals != tuple(range(1, len(self.attempts) + 1)): + raise ValueError("shard status attempts must be in complete ordinal order") + winning_attempts = tuple(status for status in self.attempts if status.is_winner) + if self.winner is None: + if winning_attempts: + raise ValueError("shard status marks a winner without a winner record") + elif ( + self.winner.shard_id != self.shard.shard_id + or len(winning_attempts) != 1 + or winning_attempts[0].attempt.attempt_id != self.winner.attempt_id + ): + raise ValueError("shard winner does not match its observed attempt") + if self.effective_state is not derive_shard_state(self.attempts, self.winner): + raise ValueError("effective shard state does not match its evidence") + return self + + +class RunStatus(StateValue): + """Fresh-process status for every planned shard in one run.""" + + run: RunManifest + observed_at: datetime + shards: tuple[ShardStatus, ...] = Field(min_length=1) + effective_state: EffectiveRunState + + _observed_at_is_utc = field_validator("observed_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_status(self) -> RunStatus: + if self.observed_at < self.run.created_at: + raise ValueError("run observation cannot precede run creation") + if len(self.shards) != self.run.shard_count: + raise ValueError("run status must include every planned shard") + if tuple(status.shard.shard_index for status in self.shards) != tuple(range(self.run.shard_count)): + raise ValueError("run status shards must be in planned order") + if any(status.shard.run_id != self.run.run_id for status in self.shards): + raise ValueError("run status contains a shard for another run") + if self.effective_state is not derive_run_state(self.shards): + raise ValueError("effective run state does not match its shard evidence") + return self + + +def derive_generation_state( + effective_state: EffectiveAttemptState, + *, + has_candidate: bool, + is_winner: bool, +) -> GenerationState: + """Derive generation progress without treating readiness as success.""" + if effective_state is EffectiveAttemptState.FAILED: + return GenerationState.FAILED + if effective_state is EffectiveAttemptState.UNKNOWN: + return GenerationState.UNKNOWN + if is_winner: + return GenerationState.WON + if has_candidate: + return GenerationState.CANDIDATE_READY + if effective_state is EffectiveAttemptState.RUNNING: + return GenerationState.ACTIVE + return GenerationState.NOT_STARTED + + +def derive_shard_state( + attempts: tuple[AttemptStatus, ...], + winner: ShardWinner | None, +) -> EffectiveAttemptState: + """Derive one shard's effective state from its newest attempt and winner.""" + if winner is not None: + winning = next((status for status in attempts if status.attempt.attempt_id == winner.attempt_id), None) + if winning is not None and winning.effective_state is EffectiveAttemptState.SUCCEEDED: + return EffectiveAttemptState.SUCCEEDED + if not attempts: + return EffectiveAttemptState.PENDING + return attempts[-1].effective_state + + +def derive_run_state(shards: tuple[ShardStatus, ...]) -> EffectiveRunState: + """Aggregate shard states without declaring a partially active run terminal.""" + states = tuple(shard.effective_state for shard in shards) + if all(state is EffectiveAttemptState.SUCCEEDED for state in states): + return EffectiveRunState.SUCCEEDED + if any(state is EffectiveAttemptState.RUNNING for state in states): + return EffectiveRunState.RUNNING + if any(state is EffectiveAttemptState.ACCOUNTING_LAG for state in states): + return EffectiveRunState.ACCOUNTING_LAG + if any(state is EffectiveAttemptState.PENDING for state in states): + return EffectiveRunState.PENDING + if any(state is EffectiveAttemptState.UNKNOWN for state in states): + return EffectiveRunState.UNKNOWN + return EffectiveRunState.FAILED + + +__all__ = [ + "AttemptStatus", + "EffectiveRunState", + "GenerationState", + "RunStatus", + "ShardStatus", + "derive_generation_state", + "derive_run_state", + "derive_shard_state", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index f940c3238..4f6f60694 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -34,6 +34,7 @@ ) from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner from data_designer.slurm.state.readiness import AttemptReadiness +from data_designer.slurm.state.scheduler import SchedulerObservation _RUN_FILENAME = "run.json" _AUTHORED_CONFIG_FILENAME = "authored-config.json" @@ -44,6 +45,7 @@ _SHARD_LOCK_FILENAME = "shard.lock" _ATTEMPT_FILENAME = "attempt.json" _READINESS_FILENAME = "readiness.json" +_SCHEDULER_OBSERVATION_FILENAME = "scheduler.json" _CLIENT_RESULT_FILENAME = "client-result.json" _CANDIDATE_OUTPUT_FILENAME = "output-manifest.json" _WINNER_FILENAME = "winner.json" @@ -115,6 +117,9 @@ def get_attempt_path(self, shard_id: str, attempt_id: str) -> Path: def get_readiness_path(self, shard_id: str, attempt_id: str) -> Path: return self.get_attempt_path(shard_id, attempt_id) / _READINESS_FILENAME + def get_scheduler_observation_path(self, shard_id: str, attempt_id: str) -> Path: + return self.get_attempt_path(shard_id, attempt_id) / _SCHEDULER_OBSERVATION_FILENAME + def get_winner_path(self, shard_id: str) -> Path: return self.get_shard_path(shard_id) / _WINNER_FILENAME @@ -310,6 +315,55 @@ def replace_readiness(self, readiness: AttemptReadiness) -> None: with self.open_attempt_directory(readiness.shard_id, readiness.attempt_id) as attempt_descriptor: self._replace_record(attempt_descriptor, _READINESS_FILENAME, readiness) + def read_scheduler_observation( + self, + shard_id: ShardId, + attempt_id: AttemptId, + ) -> SchedulerObservation: + """Read one attempt's latest reconciled scheduler evidence.""" + path = self.get_scheduler_observation_path(shard_id, attempt_id) + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + return self.read_record( + attempt_descriptor, + _SCHEDULER_OBSERVATION_FILENAME, + path, + SchedulerObservation, + ) + + def publish_scheduler_observation( + self, + shard_id: ShardId, + attempt_id: AttemptId, + observation: SchedulerObservation, + ) -> None: + """Publish the first scheduler observation for one attempt.""" + path = self.get_scheduler_observation_path(shard_id, attempt_id) + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + publish_immutable_text( + attempt_descriptor, + _SCHEDULER_OBSERVATION_FILENAME, + observation.serialize_json(), + path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def replace_scheduler_observation( + self, + shard_id: ShardId, + attempt_id: AttemptId, + observation: SchedulerObservation, + ) -> None: + """Atomically replace one attempt's scheduler observation.""" + path = self.get_scheduler_observation_path(shard_id, attempt_id) + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + replace_text( + attempt_descriptor, + _SCHEDULER_OBSERVATION_FILENAME, + observation.serialize_json(), + path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + def sync_attempt_directory(self, shard_id: ShardId, attempt_id: AttemptId) -> None: with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: sync_directory(attempt_descriptor) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py index f990b65e8..8c337b89b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -15,7 +15,11 @@ CollectionPlan, ShardWinner, ) -from data_designer.slurm.state.scheduler import SchedulerObservation, SchedulerState +from data_designer.slurm.state.scheduler import ( + SchedulerObservation, + SchedulerState, + is_scheduler_terminal_state, +) _ATTEMPT_STATE_ORDER = { AttemptLifecycleState.CREATED: 0, @@ -183,6 +187,8 @@ def validate_scheduler_observation_transition( """Validate scheduler identity, chronology, and a fixed accounting-lag deadline.""" _require(current.scheduler == previous.scheduler, "scheduler identity cannot change between observations") _require(current.observed_at >= previous.observed_at, "scheduler observed_at cannot move backward") + if is_scheduler_terminal_state(previous.state): + _require(current.state is previous.state, "terminal scheduler evidence cannot change") if previous.state is SchedulerState.ACCOUNTING_LAG: deadline = previous.reconciliation_deadline diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py new file mode 100644 index 000000000..bbae29feb --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -0,0 +1,501 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import stat +from collections.abc import Sequence +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 FakeSlurmRunner + +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.models import SlurmAccountingEntry, SlurmProcessExitCode, SlurmQueueEntry +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + EffectiveAttemptState, + EffectiveRunState, + GenerationState, + RunManifest, + SchedulerIdentity, + SchedulerJobIdentity, + SchedulerObservation, + SchedulerObservationCollector, + SchedulerState, + ShardManifest, + ShardWinner, + SlurmStateError, + SlurmStateReconciler, + SlurmStateWriter, + StateConflictError, + StateCorruptionError, +) + + +@dataclass(frozen=True, slots=True) +class _ReconciliationCase: + workspace: Path + plan: ResolvedSlurmRunPlan + run: RunManifest + shard: ShardManifest + attempt: AttemptManifest + writer: SlurmStateWriter + created_at: datetime + + +@dataclass(frozen=True, slots=True) +class _StaticSchedulerClient: + queue: tuple[SlurmQueueEntry, ...] + accounting: tuple[SlurmAccountingEntry, ...] + + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + del selectors + return self.queue + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + del selectors + return self.accounting + + +def test_collector_prefers_terminal_accounting_for_array_and_collection_jobs() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + queue = ( + SlurmQueueEntry(job_identity=task, state=SchedulerState.RUNNING), + SlurmQueueEntry(job_identity=5101, state=SchedulerState.RUNNING), + ) + accounting = ( + _accounting(task, SchedulerState.NODE_FAILED), + _accounting(5101, SchedulerState.COMPLETED), + ) + observed_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + + observations = SchedulerObservationCollector(_StaticSchedulerClient(queue, accounting)).collect( + (task, 5101), + observed_at=observed_at, + ) + + assert tuple(observation.state for observation in observations) == ( + SchedulerState.NODE_FAILED, + SchedulerState.COMPLETED, + ) + assert tuple(observation.scheduler for observation in observations) == (task, 5101) + + +def test_fresh_process_refresh_persists_one_fixed_accounting_lag_deadline( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state(scheduler, queue_state=None, accounting_state=None) + first_time = case.created_at + timedelta(minutes=3) + + first = SlurmStateReconciler(case.workspace, case.plan.run_id, SlurmCommandClient(fake_slurm_runner)).refresh( + observed_at=first_time + ) + second = SlurmStateReconciler(case.workspace, case.plan.run_id, SlurmCommandClient(fake_slurm_runner)).refresh( + observed_at=first_time + timedelta(minutes=1) + ) + expired = SlurmStateReconciler(case.workspace, case.plan.run_id, SlurmCommandClient(fake_slurm_runner)).refresh( + observed_at=first_time + timedelta(minutes=6) + ) + + first_scheduler = first.shards[0].attempts[0].scheduler + second_scheduler = second.shards[0].attempts[0].scheduler + assert first.effective_state is EffectiveRunState.ACCOUNTING_LAG + assert second.effective_state is EffectiveRunState.ACCOUNTING_LAG + assert expired.effective_state is EffectiveRunState.UNKNOWN + assert first_scheduler is not None and second_scheduler is not None + assert first_scheduler.reconciliation_deadline == second_scheduler.reconciliation_deadline + scheduler_path = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/scheduler.json" + assert stat.S_IMODE(scheduler_path.stat().st_mode) == 0o600 + + +def test_refresh_uses_terminal_accounting_over_stale_active_queue_state( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state="RUNNING", + accounting_state="FAILED", + exit_code="1:0", + ) + + status = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=case.created_at + timedelta(minutes=3)) + + attempt = status.shards[0].attempts[0] + assert attempt.scheduler is not None and attempt.scheduler.state is SchedulerState.FAILED + assert attempt.effective_state is EffectiveAttemptState.FAILED + assert attempt.generation_state is GenerationState.FAILED + assert status.effective_state is EffectiveRunState.FAILED + + +def test_refresh_rejects_winner_that_conflicts_with_terminal_scheduler_evidence( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + completed, winner = _publish_winner_state(case) + scheduler = cast(SchedulerIdentity, completed.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state="RUNNING", + accounting_state="NODE_FAIL", + exit_code="1:0", + ) + + with pytest.raises(StateCorruptionError, match="winner conflicts"): + SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=winner.published_at + timedelta(minutes=1)) + + +def test_refresh_reports_a_validated_winner_as_succeeded( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + completed, winner = _publish_winner_state(case) + scheduler = cast(SchedulerIdentity, completed.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="COMPLETED", + exit_code="0:0", + ) + + status = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=winner.published_at + timedelta(minutes=1)) + + attempt = status.shards[0].attempts[0] + assert attempt.client_result is not None + assert attempt.candidate_output is not None + assert attempt.effective_state is EffectiveAttemptState.SUCCEEDED + assert attempt.generation_state is GenerationState.WON + assert status.shards[0].winner == winner + assert status.effective_state is EffectiveRunState.SUCCEEDED + + +def test_refresh_rejects_a_concurrent_attempt_change_instead_of_guessing_status( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + observed_at = case.created_at + timedelta(minutes=4) + + class MutatingClient: + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + del selectors + return (SlurmQueueEntry(job_identity=scheduler, state=SchedulerState.RUNNING),) + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + del selectors + case.writer.update_attempt( + _copy_attempt(case.attempt, state=AttemptLifecycleState.RUNNING, updated_at=observed_at) + ) + return () + + with pytest.raises(StateConflictError, match="changed during reconciliation"): + SlurmStateReconciler(case.workspace, case.plan.run_id, MutatingClient()).refresh(observed_at=observed_at) + + +def test_terminal_observation_remains_authoritative_during_later_accounting_gap() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + first = SchedulerObservationCollector( + _StaticSchedulerClient((), (_accounting(task, SchedulerState.COMPLETED),)) + ).collect((task,), observed_at=first_time)[0] + later = SchedulerObservationCollector( + _StaticSchedulerClient((SlurmQueueEntry(job_identity=task, state=SchedulerState.RUNNING),), ()) + ).collect((task,), observed_at=first_time + timedelta(minutes=1), previous={task: first})[0] + + assert later.state is SchedulerState.COMPLETED + + +def test_collector_rejects_conflicting_terminal_accounting_evidence() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + previous = SchedulerObservationCollector( + _StaticSchedulerClient((), (_accounting(task, SchedulerState.COMPLETED),)) + ).collect((task,), observed_at=first_time)[0] + + with pytest.raises(SlurmStateError, match="violates persisted chronology"): + SchedulerObservationCollector(_StaticSchedulerClient((), (_accounting(task, SchedulerState.FAILED),))).collect( + (task,), observed_at=first_time + timedelta(minutes=1), previous={task: previous} + ) + + +def test_refresh_normalizes_scheduler_query_failures( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + + class FailingClient: + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + del selectors + raise RuntimeError("scheduler unavailable") + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + del selectors + return () + + with pytest.raises(SlurmStateError, match="cannot query normalized scheduler observations"): + SlurmStateReconciler(case.workspace, case.plan.run_id, FailingClient()).refresh( + observed_at=case.created_at + timedelta(minutes=3) + ) + + +def test_refresh_keeps_an_unsubmitted_attempt_pending_without_querying_slurm( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan, submitted=False) + + class UnexpectedClient: + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + raise AssertionError(f"unexpected queue query for {selectors!r}") + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + raise AssertionError(f"unexpected accounting query for {selectors!r}") + + status = SlurmStateReconciler(case.workspace, case.plan.run_id, UnexpectedClient()).refresh( + observed_at=case.created_at + timedelta(minutes=3) + ) + + attempt = status.shards[0].attempts[0] + assert attempt.scheduler is None + assert attempt.effective_state is EffectiveAttemptState.PENDING + assert attempt.generation_state is GenerationState.NOT_STARTED + assert status.effective_state is EffectiveRunState.PENDING + + +def test_fresh_process_refresh_rejects_mismatched_persisted_scheduler_evidence( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state(scheduler, queue_state="RUNNING", accounting_state=None) + observed_at = case.created_at + timedelta(minutes=3) + SlurmStateReconciler(case.workspace, case.plan.run_id, SlurmCommandClient(fake_slurm_runner)).refresh( + observed_at=observed_at + ) + scheduler_path = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/scheduler.json" + mismatched = SchedulerObservation( + schema_version=1, + scheduler=SchedulerIdentity(array_job_id=9999, array_task_id=0), + observed_at=observed_at, + state=SchedulerState.RUNNING, + ) + scheduler_path.write_text(mismatched.serialize_json()) + + with pytest.raises(StateCorruptionError, match="mismatched scheduler evidence"): + SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=observed_at + timedelta(minutes=1)) + + +def _accounting(identity: SchedulerJobIdentity, state: SchedulerState) -> SlurmAccountingEntry: + return SlurmAccountingEntry( + job_identity=identity, + state=state, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=0), + ) + + +def _initialized_case( + tmp_path: Path, + authored_config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + *, + submitted: bool = True, +) -> _ReconciliationCase: + workspace = tmp_path / "workspace" + workspace.mkdir() + 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=1, + ) + planned_shard = relocated_plan.shards[0] + shard = ShardManifest( + schema_version=1, + run_id=relocated_plan.run_id, + shard_id=planned_shard.shard_id, + shard_index=planned_shard.shard_index, + record_range=planned_shard.record_range, + input_partition=planned_shard.input_partition, + resume_workspace=planned_shard.resume_workspace, + created_at=created_at, + ) + writer = SlurmStateWriter(workspace, relocated_plan.run_id) + writer.initialize_run(authored_config, relocated_plan, run, (shard,)) + attempt = AttemptManifest( + schema_version=1, + run_id=relocated_plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED if submitted else AttemptLifecycleState.CREATED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0) if submitted else None, + created_at=created_at + timedelta(minutes=1), + updated_at=created_at + timedelta(minutes=2), + ) + writer.create_attempt(attempt) + return _ReconciliationCase(workspace, relocated_plan, run, shard, attempt, 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 = SlurmProfile.model_validate_json(json.dumps(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 _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, ShardWinner]: + running = _copy_attempt( + case.attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=3), + ) + case.writer.update_attempt(running) + candidate_path = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/output-manifest.json" + dataset_path = candidate_path.parent / "dataset" + requested = case.plan.shards[0].requested_records + candidate = CandidateOutputManifest( + schema_version=1, + run_id=case.plan.run_id, + shard_id=running.shard_id, + attempt_id=running.attempt_id, + attempt_ordinal=running.attempt_ordinal, + created_at=case.created_at + timedelta(minutes=4), + dataset_path=dataset_path.as_posix(), + requested_records=requested, + actual_records=requested, + outcome=CandidateOutcome.COMPLETE, + files=( + CandidateOutputFile( + relative_path="part-00000.parquet", + sha256="a" * 64, + byte_size=1, + record_count=requested, + ), + ), + dataset_schema_digest="b" * 64, + 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.plan.run_id, + shard_id=running.shard_id, + attempt_id=running.attempt_id, + completed_at=case.created_at + timedelta(minutes=5), + requested_records=requested, + actual_records=requested, + 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=case.created_at + timedelta(minutes=6), + ) + case.writer.update_attempt(completed) + winner = ShardWinner( + schema_version=1, + run_id=case.plan.run_id, + shard_id=completed.shard_id, + attempt_id=completed.attempt_id, + attempt_ordinal=completed.attempt_ordinal, + candidate_manifest=candidate_reference, + published_at=case.created_at + timedelta(minutes=7), + ) + winner_path = case.writer.run_root / "shards/shard-00000/winner.json" + winner_path.write_text(winner.serialize_json()) + winner_path.chmod(0o600) + return completed, winner + + +def _copy_attempt(attempt: AttemptManifest, **updates: object) -> AttemptManifest: + payload = attempt.model_dump(mode="json") + payload.update(updates) + return AttemptManifest.model_validate_json(json.dumps(payload, default=_json_value)) + + +def _json_value(value: object) -> object: + if isinstance(value, datetime): + return value.isoformat() + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + raise TypeError(f"unsupported test value: {type(value).__name__}") diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 3ff48707c..17e535c3b 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -136,8 +136,16 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.state import ArtifactReference as StateArtifactReference from data_designer.slurm.state import RecordRange as StateRecordRange from data_designer.slurm.state import ResumeWorkspace as StateResumeWorkspace -from data_designer.slurm.state import RunManifest +from data_designer.slurm.state import ( + RunManifest, + RunStatus, + SchedulerObservationCollector, + SlurmStateReconciler, +) assert RunManifest.__name__ == "RunManifest" +assert RunStatus.__name__ == "RunStatus" +assert SchedulerObservationCollector.__name__ == "SchedulerObservationCollector" +assert SlurmStateReconciler.__name__ == "SlurmStateReconciler" assert ImageRegistryStore.__name__ == "ImageRegistryStore" assert PlanningArtifactReference is ContractArtifactReference assert PlanningRecordRange is ContractRecordRange From 09a42b6aa3d5a653b81ab8c44058087431e415da Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 8 Sep 2026 09:21:29 -0600 Subject: [PATCH 2/5] fix(slurm): harden state reconciliation Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/client.py | 32 ++++++-- .../data_designer/slurm/state/observation.py | 22 ++++- .../slurm/state/reconciliation.py | 2 +- .../data_designer/slurm/state/scheduler.py | 2 - .../src/data_designer/slurm/state/status.py | 3 +- .../tests/launcher/test_client.py | 33 ++++++++ .../tests/state/test_observer.py | 81 ++++++++++++++++++- .../tests/state/test_validation.py | 2 + 8 files changed, 160 insertions(+), 17 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py index e45e08048..81ee4d2b3 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 @@ -31,6 +31,12 @@ _IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") _ENVIRONMENT_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") _MAX_SLURM_INTEGER = (1 << 32) - 1 +_UNKNOWN_QUEUE_JOB_DETAILS = frozenset( + { + "Invalid job id specified", + "slurm_load_jobs error: Invalid job id specified", + } +) @dataclass(frozen=True) @@ -109,15 +115,20 @@ def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQ """Return normalized active-queue rows for explicit managed jobs.""" requested = tuple(selectors) jobs = _format_selectors(requested) - output = self._run( - ( - self._executables.squeue, - "--noheader", - "--array", - "--format=%i|%T", - f"--jobs={jobs}", + try: + output = self._run( + ( + self._executables.squeue, + "--noheader", + "--array", + "--format=%i|%T", + f"--jobs={jobs}", + ) ) - ) + except SlurmCommandError as error: + if _is_unknown_queue_job_error(error): + return () + raise entries = parse_queue(output) ignored = _validate_observed_job_identities( tuple(entry.job_identity for entry in entries), @@ -272,3 +283,8 @@ 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 _is_unknown_queue_job_error(error: SlurmCommandError) -> bool: + message = str(error) + return any(message.endswith(f": {detail}") for detail in _UNKNOWN_QUEUE_JOB_DETAILS) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py index 3bf039715..54165a193 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py @@ -7,7 +7,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime, timedelta -from typing import Protocol +from typing import Callable, Protocol, TypeVar from data_designer.slurm.state.base import SchedulerJobIdentity, validate_utc_timestamp from data_designer.slurm.state.errors import SlurmStateError @@ -19,6 +19,8 @@ from data_designer.slurm.state.validation import StateContractError, validate_scheduler_observation_transition _ACCOUNTING_LAG_WINDOW = timedelta(minutes=5) +_SchedulerRecordT = TypeVar("_SchedulerRecordT", bound=object) +_SchedulerQueryError = OSError | RuntimeError | ValueError class SchedulerQueueRecord(Protocol): @@ -84,10 +86,12 @@ def _query_scheduler( self, selectors: tuple[SchedulerJobIdentity, ...], ) -> tuple[tuple[SchedulerQueueRecord, ...], tuple[SchedulerAccountingRecord, ...]]: - try: - return self._client.query_queue(selectors), self._client.query_accounting(selectors) - except (OSError, RuntimeError, ValueError) as error: + queue, queue_error = _capture_scheduler_query(self._client.query_queue, selectors) + accounting, accounting_error = _capture_scheduler_query(self._client.query_accounting, selectors) + error = queue_error if queue_error is not None else accounting_error + if error is not None: raise SlurmStateError("cannot query normalized scheduler observations") from error + return queue, accounting @staticmethod def _index_records( @@ -194,6 +198,16 @@ def _resolve_missing_observation( ) +def _capture_scheduler_query( + query: Callable[[Sequence[SchedulerJobIdentity]], tuple[_SchedulerRecordT, ...]], + selectors: tuple[SchedulerJobIdentity, ...], +) -> tuple[tuple[_SchedulerRecordT, ...], _SchedulerQueryError | None]: + try: + return query(selectors), None + except (OSError, RuntimeError, ValueError) as error: + return (), error + + __all__ = [ "SchedulerAccountingRecord", "SchedulerObservationClient", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index a210353f4..8b61117d6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -173,7 +173,7 @@ def reconcile_attempt_observation( return EffectiveAttemptState.UNKNOWN if readiness is not None and readiness.state is ReadinessState.FAILED: return EffectiveAttemptState.FAILED - if scheduler.state is SchedulerState.PENDING: + if scheduler.state in {SchedulerState.PENDING, SchedulerState.PREEMPTED, SchedulerState.REQUEUED}: return EffectiveAttemptState.PENDING if scheduler.state is SchedulerState.RUNNING: return EffectiveAttemptState.RUNNING diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py index e06e11cc2..111513060 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py @@ -70,8 +70,6 @@ def is_scheduler_failure_state(state: SchedulerState) -> bool: SchedulerState.CANCELLED, SchedulerState.TIMED_OUT, SchedulerState.NODE_FAILED, - SchedulerState.PREEMPTED, - SchedulerState.REQUEUED, SchedulerState.OUT_OF_MEMORY, } diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/status.py b/packages/data-designer-slurm/src/data_designer/slurm/state/status.py index 581126a0b..2d8aeaa5a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/status.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/status.py @@ -177,7 +177,8 @@ def derive_shard_state( return EffectiveAttemptState.SUCCEEDED if not attempts: return EffectiveAttemptState.PENDING - return attempts[-1].effective_state + latest = attempts[-1].effective_state + return EffectiveAttemptState.UNKNOWN if latest is EffectiveAttemptState.SUCCEEDED else latest def derive_run_state(shards: tuple[ShardStatus, ...]) -> EffectiveRunState: diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 55b8258ff..52951158a 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -208,6 +208,39 @@ def test_client_normalizes_command_failures(fake_slurm_runner: FakeSlurmRunner) client.query_accounting((4101,)) +@pytest.mark.parametrize( + "detail", + ( + "Invalid job id specified", + "slurm_load_jobs error: Invalid job id specified", + ), +) +def test_client_treats_unknown_queue_job_as_an_absent_row( + fake_slurm_runner: FakeSlurmRunner, + detail: str, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=f"{detail}\n", returncode=1)) + + assert SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) == () + + +@pytest.mark.parametrize( + "detail", + ( + "Invalid job id specified for partition", + "queue unavailable", + ), +) +def test_client_preserves_non_missing_queue_failures( + fake_slurm_runner: FakeSlurmRunner, + detail: str, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=f"{detail}\n", returncode=1)) + + with pytest.raises(SlurmCommandError, match=detail): + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + def test_client_removes_terminal_controls_from_command_failures(fake_slurm_runner: FakeSlurmRunner) -> None: fake_slurm_runner.script_next( "squeue", diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py index bbae29feb..48a416548 100644 --- a/packages/data-designer-slurm/tests/state/test_observer.py +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -12,7 +12,7 @@ from typing import cast import pytest -from slurm_test_fakes import FakeSlurmRunner +from slurm_test_fakes import FakeCommandResponse, FakeSlurmRunner from data_designer.slurm.client import ClientOutcome, ClientResult from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile @@ -95,6 +95,48 @@ def test_collector_prefers_terminal_accounting_for_array_and_collection_jobs() - assert tuple(observation.scheduler for observation in observations) == (task, 5101) +@pytest.mark.parametrize("transient_state", (SchedulerState.PREEMPTED, SchedulerState.REQUEUED)) +def test_collector_allows_transient_scheduler_states_to_advance( + transient_state: SchedulerState, +) -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + previous = SchedulerObservation( + schema_version=1, + scheduler=task, + observed_at=first_time, + state=transient_state, + ) + + pending = SchedulerObservationCollector( + _StaticSchedulerClient((SlurmQueueEntry(job_identity=task, state=SchedulerState.PENDING),), ()) + ).collect((task,), observed_at=first_time + timedelta(minutes=1), previous={task: previous})[0] + completed = SchedulerObservationCollector( + _StaticSchedulerClient((), (_accounting(task, SchedulerState.COMPLETED),)) + ).collect((task,), observed_at=first_time + timedelta(minutes=2), previous={task: pending})[0] + + assert pending.state is SchedulerState.PENDING + assert completed.state is SchedulerState.COMPLETED + + +def test_collector_reads_accounting_when_queue_no_longer_knows_job() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + runner = FakeSlurmRunner() + runner.script_next( + "squeue", + FakeCommandResponse(stderr="slurm_load_jobs error: Invalid job id specified\n", returncode=1), + ) + runner.script_next("sacct", FakeCommandResponse(stdout="4101_0|COMPLETED|0:0\n")) + + observation = SchedulerObservationCollector(SlurmCommandClient(runner)).collect( + (task,), + observed_at=datetime(2026, 9, 2, 12, tzinfo=timezone.utc), + )[0] + + assert observation.state is SchedulerState.COMPLETED + assert tuple(call[0] for call in runner.calls) == ("squeue", "sacct") + + def test_fresh_process_refresh_persists_one_fixed_accounting_lag_deadline( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -214,6 +256,38 @@ def test_refresh_reports_a_validated_winner_as_succeeded( assert status.effective_state is EffectiveRunState.SUCCEEDED +def test_refresh_does_not_succeed_a_winnerless_persisted_candidate( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + completed, winner = _publish_winner_state(case) + (case.writer.run_root / "shards/shard-00000/winner.json").unlink() + scheduler = cast(SchedulerIdentity, completed.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="COMPLETED", + exit_code="0:0", + ) + + status = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=winner.published_at + timedelta(minutes=1)) + + attempt = status.shards[0].attempts[0] + assert attempt.effective_state is EffectiveAttemptState.SUCCEEDED + assert attempt.generation_state is GenerationState.CANDIDATE_READY + assert status.shards[0].winner is None + assert status.shards[0].effective_state is EffectiveAttemptState.UNKNOWN + assert status.effective_state is EffectiveRunState.UNKNOWN + + def test_refresh_rejects_a_concurrent_attempt_change_instead_of_guessing_status( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -271,14 +345,17 @@ def test_refresh_normalizes_scheduler_query_failures( single_node_plan: ResolvedSlurmRunPlan, ) -> None: case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + calls: list[str] = [] class FailingClient: def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: del selectors + calls.append("queue") raise RuntimeError("scheduler unavailable") def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: del selectors + calls.append("accounting") return () with pytest.raises(SlurmStateError, match="cannot query normalized scheduler observations"): @@ -286,6 +363,8 @@ def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[S observed_at=case.created_at + timedelta(minutes=3) ) + assert calls == ["queue", "accounting"] + def test_refresh_keeps_an_unsubmitted_attempt_pending_without_querying_slurm( tmp_path: Path, diff --git a/packages/data-designer-slurm/tests/state/test_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py index ff0493b19..8817b4339 100644 --- a/packages/data-designer-slurm/tests/state/test_validation.py +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -569,6 +569,8 @@ def test_reconciliation_covers_nonterminal_and_fallback_states() -> None: cases = ( (attempt, ready, SchedulerState.PENDING, EffectiveAttemptState.PENDING), + (attempt, ready, SchedulerState.PREEMPTED, EffectiveAttemptState.PENDING), + (attempt, ready, SchedulerState.REQUEUED, EffectiveAttemptState.PENDING), (attempt, ready, SchedulerState.RUNNING, EffectiveAttemptState.RUNNING), (attempt, ready, SchedulerState.COMPLETED, EffectiveAttemptState.FAILED), (attempt, pending, SchedulerState.UNKNOWN, EffectiveAttemptState.UNKNOWN), From 6fedd31d41aff321687a108006a41e25354d4009 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 8 Sep 2026 12:16:00 -0600 Subject: [PATCH 3/5] fix(slurm): bound preemption reconciliation Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/state/observation.py | 39 ++++- .../slurm/state/reconciliation.py | 7 +- .../data_designer/slurm/state/scheduler.py | 8 +- .../data_designer/slurm/state/validation.py | 6 + .../tests/state/test_observer.py | 141 +++++++++++++++++- 5 files changed, 185 insertions(+), 16 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py index 54165a193..e8da67c8d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py @@ -19,6 +19,7 @@ from data_designer.slurm.state.validation import StateContractError, validate_scheduler_observation_transition _ACCOUNTING_LAG_WINDOW = timedelta(minutes=5) +_PREEMPTION_REQUEUE_WINDOW = timedelta(minutes=5) _SchedulerRecordT = TypeVar("_SchedulerRecordT", bound=object) _SchedulerQueryError = OSError | RuntimeError | ValueError @@ -128,16 +129,17 @@ def _resolve_observation( and (accounting_state is None or not is_scheduler_terminal_state(accounting_state)) ): state = previous.state - observation = ( - _resolve_missing_observation(identity, observed_at, previous) - if state is None - else SchedulerObservation( + if state is None: + observation = _resolve_missing_observation(identity, observed_at, previous) + elif state is SchedulerState.PREEMPTED and queue_state is None: + observation = _resolve_preemption_observation(identity, observed_at, previous) + else: + observation = SchedulerObservation( schema_version=1, scheduler=identity, observed_at=observed_at, state=state, ) - ) if previous is not None: try: validate_scheduler_observation_transition(previous, observation) @@ -157,11 +159,38 @@ def _select_observed_state( return accounting_state +def _resolve_preemption_observation( + identity: SchedulerJobIdentity, + observed_at: datetime, + previous: SchedulerObservation | None, +) -> SchedulerObservation: + deadline = ( + previous.reconciliation_deadline + if previous is not None + and previous.state is SchedulerState.PREEMPTED + and previous.reconciliation_deadline is not None + else observed_at + _PREEMPTION_REQUEUE_WINDOW + ) + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=SchedulerState.FAILED if observed_at > deadline else SchedulerState.PREEMPTED, + reconciliation_deadline=None if observed_at > deadline else deadline, + ) + + def _resolve_missing_observation( identity: SchedulerJobIdentity, observed_at: datetime, previous: SchedulerObservation | None, ) -> SchedulerObservation: + if ( + previous is not None + and previous.state is SchedulerState.PREEMPTED + and previous.reconciliation_deadline is not None + ): + return _resolve_preemption_observation(identity, observed_at, previous) if previous is not None and previous.state is SchedulerState.ACCOUNTING_LAG: deadline = previous.reconciliation_deadline if deadline is None: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index 8b61117d6..473150f43 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -173,7 +173,12 @@ def reconcile_attempt_observation( return EffectiveAttemptState.UNKNOWN if readiness is not None and readiness.state is ReadinessState.FAILED: return EffectiveAttemptState.FAILED - if scheduler.state in {SchedulerState.PENDING, SchedulerState.PREEMPTED, SchedulerState.REQUEUED}: + if scheduler.state is SchedulerState.PREEMPTED: + deadline = scheduler.reconciliation_deadline + if deadline is not None and current_time > deadline: + return EffectiveAttemptState.FAILED + return EffectiveAttemptState.PENDING + if scheduler.state in {SchedulerState.PENDING, SchedulerState.REQUEUED}: return EffectiveAttemptState.PENDING if scheduler.state is SchedulerState.RUNNING: return EffectiveAttemptState.RUNNING diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py index 111513060..6961a1295 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py @@ -47,10 +47,10 @@ def validate_reconciliation_deadline(self) -> SchedulerObservation: if self.state is SchedulerState.ACCOUNTING_LAG: if self.reconciliation_deadline is None: raise ValueError("accounting lag requires a reconciliation deadline") - if self.reconciliation_deadline < self.observed_at: - raise ValueError("reconciliation deadline must not precede the observation") - elif self.reconciliation_deadline is not None: - raise ValueError("only accounting lag may have a reconciliation deadline") + elif self.state is not SchedulerState.PREEMPTED and self.reconciliation_deadline is not None: + raise ValueError("only accounting lag or preemption may have a reconciliation deadline") + if self.reconciliation_deadline is not None and self.reconciliation_deadline < self.observed_at: + raise ValueError("reconciliation deadline must not precede the observation") return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py index 8c337b89b..1e20fc7c4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -203,6 +203,12 @@ def validate_scheduler_observation_transition( current.observed_at > deadline, "accounting lag cannot become unknown before its reconciliation deadline expires", ) + if previous.state is SchedulerState.PREEMPTED and previous.reconciliation_deadline is not None: + if current.state is SchedulerState.PREEMPTED and current.reconciliation_deadline is not None: + _require( + current.reconciliation_deadline == previous.reconciliation_deadline, + "preemption reconciliation deadline cannot change", + ) return current diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py index 48a416548..0a49de6be 100644 --- a/packages/data-designer-slurm/tests/state/test_observer.py +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -101,12 +101,12 @@ def test_collector_allows_transient_scheduler_states_to_advance( ) -> None: task = SchedulerIdentity(array_job_id=4101, array_task_id=0) first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) - previous = SchedulerObservation( - schema_version=1, - scheduler=task, - observed_at=first_time, - state=transient_state, - ) + previous = SchedulerObservationCollector( + _StaticSchedulerClient( + (SlurmQueueEntry(job_identity=task, state=transient_state),), + (_accounting(task, transient_state),), + ) + ).collect((task,), observed_at=first_time)[0] pending = SchedulerObservationCollector( _StaticSchedulerClient((SlurmQueueEntry(job_identity=task, state=SchedulerState.PENDING),), ()) @@ -115,10 +115,33 @@ def test_collector_allows_transient_scheduler_states_to_advance( _StaticSchedulerClient((), (_accounting(task, SchedulerState.COMPLETED),)) ).collect((task,), observed_at=first_time + timedelta(minutes=2), previous={task: pending})[0] + assert previous.state is transient_state assert pending.state is SchedulerState.PENDING assert completed.state is SchedulerState.COMPLETED +def test_collector_keeps_active_queue_preemption_nonterminal() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + client = _StaticSchedulerClient( + (SlurmQueueEntry(job_identity=task, state=SchedulerState.PREEMPTED),), + (), + ) + + first = SchedulerObservationCollector(client).collect((task,), observed_at=first_time)[0] + second = SchedulerObservationCollector(client).collect( + (task,), observed_at=first_time + timedelta(minutes=1), previous={task: first} + )[0] + later = SchedulerObservationCollector(client).collect( + (task,), observed_at=first_time + timedelta(minutes=6), previous={task: second} + )[0] + + assert first.reconciliation_deadline is None + assert second.reconciliation_deadline is None + assert later.state is SchedulerState.PREEMPTED + assert later.reconciliation_deadline is None + + def test_collector_reads_accounting_when_queue_no_longer_knows_job() -> None: task = SchedulerIdentity(array_job_id=4101, array_task_id=0) runner = FakeSlurmRunner() @@ -199,6 +222,112 @@ def test_refresh_uses_terminal_accounting_over_stale_active_queue_state( assert status.effective_state is EffectiveRunState.FAILED +def test_refresh_marks_unrequeued_preemption_as_failure_after_fixed_deadline( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="PREEMPTED", + exit_code="0:15", + ) + first_time = case.created_at + timedelta(minutes=3) + + first = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time) + fake_slurm_runner.set_task_state(scheduler, queue_state=None, accounting_state=None) + second = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time + timedelta(minutes=1)) + expired = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time + timedelta(minutes=6)) + + first_attempt = first.shards[0].attempts[0] + second_attempt = second.shards[0].attempts[0] + attempt = expired.shards[0].attempts[0] + assert first_attempt.effective_state is EffectiveAttemptState.PENDING + assert second_attempt.effective_state is EffectiveAttemptState.PENDING + assert first_attempt.scheduler is not None and second_attempt.scheduler is not None + assert first_attempt.scheduler.reconciliation_deadline == second_attempt.scheduler.reconciliation_deadline + assert attempt.scheduler is not None and attempt.scheduler.state is SchedulerState.FAILED + assert attempt.effective_state is EffectiveAttemptState.FAILED + assert attempt.generation_state is GenerationState.FAILED + assert expired.effective_state is EffectiveRunState.FAILED + + +def test_refresh_allows_accounting_preemption_to_reappear_in_queue( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="PREEMPTED", + exit_code="0:15", + ) + first_time = case.created_at + timedelta(minutes=3) + first = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time) + + fake_slurm_runner.set_task_state( + scheduler, + queue_state="PREEMPTED", + accounting_state="PREEMPTED", + exit_code="0:15", + ) + still_preempted = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time + timedelta(seconds=30)) + + fake_slurm_runner.set_task_state( + scheduler, + queue_state="PENDING", + accounting_state="PREEMPTED", + exit_code="0:15", + ) + requeued = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time + timedelta(minutes=1)) + + first_attempt = first.shards[0].attempts[0] + still_preempted_attempt = still_preempted.shards[0].attempts[0] + requeued_attempt = requeued.shards[0].attempts[0] + assert first_attempt.scheduler is not None and still_preempted_attempt.scheduler is not None + assert first_attempt.scheduler.state is SchedulerState.PREEMPTED + assert first_attempt.effective_state is EffectiveAttemptState.PENDING + assert still_preempted_attempt.scheduler.reconciliation_deadline is None + assert requeued_attempt.scheduler is not None + assert requeued_attempt.scheduler.state is SchedulerState.PENDING + assert requeued_attempt.effective_state is EffectiveAttemptState.PENDING + assert requeued.effective_state is EffectiveRunState.PENDING + + def test_refresh_rejects_winner_that_conflicts_with_terminal_scheduler_evidence( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, From ce1fe44bd3ae05cec973c4958375bc568e2076c8 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 10 Sep 2026 08:04:25 -0600 Subject: [PATCH 4/5] fix(slurm): bound public preemption status Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/services/wiring.py | 106 +++++++++--------- .../tests/services/test_wiring.py | 39 ++++++- 2 files changed, 88 insertions(+), 57 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index b4d7dbd69..df3ab4709 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -65,15 +65,15 @@ from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, + AttemptStatus, AttemptTerminalClassification, EffectiveAttemptState, - SchedulerObservation, SchedulerState, SlurmStateError, + SlurmStateReconciler, SlurmStateWriter, StateConflictError, StateNotFoundError, - reconcile_attempt_observation, ) RunIdFactory = Callable[[], str] @@ -548,63 +548,57 @@ def _reconcile_attempts( ) if not active: return - identities = tuple(attempt.scheduler for attempt in active if attempt.scheduler is not None) - try: - queue = {entry.job_identity: entry.state for entry in self._launcher.query_queue(identities)} - missing = tuple(identity for identity in identities if identity not in queue) - accounting = ( - {entry.job_identity: entry.state for entry in self._launcher.query_accounting(missing)} - if missing - else {} - ) - except SlurmLauncherError: - return + observed_at = max(self._clock(), *(attempt.updated_at for attempt in active)) for attempt in active: - scheduler = attempt.scheduler - assert scheduler is not None - scheduler_state = queue.get(scheduler, accounting.get(scheduler)) - if scheduler_state is None: - continue readiness = _load_optional(lambda: writer.load_readiness(attempt.shard_id, attempt.attempt_id)) - observed_at = max( - self._clock(), - attempt.updated_at, - readiness.updated_at if readiness is not None else attempt.updated_at, - ) - observation = SchedulerObservation( - schema_version=1, - scheduler=scheduler, - observed_at=observed_at, - state=scheduler_state, - ) - effective = reconcile_attempt_observation( - attempt, - readiness, - observation, - current_time=observed_at, + if readiness is not None: + observed_at = max(observed_at, readiness.updated_at) + try: + reconciled = SlurmStateReconciler( + self._profile.profile.workspace_root, + writer.load_run().run_id, + self._launcher, + ).refresh(observed_at=observed_at) + except SlurmStateError as error: + if isinstance(error.__cause__, SlurmLauncherError): + return + raise + active_identities = {(attempt.shard_id, attempt.attempt_id) for attempt in active} + for shard in reconciled.shards: + for status in shard.attempts: + if (status.attempt.shard_id, status.attempt.attempt_id) in active_identities: + self._update_reconciled_attempt(writer, status, observed_at) + + @staticmethod + def _update_reconciled_attempt( + writer: SlurmStateWriter, + status: AttemptStatus, + observed_at: datetime, + ) -> None: + attempt = status.attempt + update: dict[str, object] = {"updated_at": observed_at} + if status.effective_state is EffectiveAttemptState.PENDING and attempt.state is AttemptLifecycleState.SUBMITTED: + update["state"] = AttemptLifecycleState.PENDING + elif status.effective_state is EffectiveAttemptState.RUNNING and attempt.state in { + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + }: + update["state"] = AttemptLifecycleState.RUNNING + elif status.effective_state is EffectiveAttemptState.FAILED: + scheduler_state = SchedulerState.UNKNOWN if status.scheduler is None else status.scheduler.state + update.update( + state=AttemptLifecycleState.FAILED, + terminal_classification=_FAILURE_CLASSIFICATIONS.get( + scheduler_state, + AttemptTerminalClassification.UNKNOWN, + ), ) - update: dict[str, object] = {"updated_at": observed_at} - if effective is EffectiveAttemptState.PENDING and attempt.state is AttemptLifecycleState.SUBMITTED: - update["state"] = AttemptLifecycleState.PENDING - elif effective is EffectiveAttemptState.RUNNING and attempt.state in { - AttemptLifecycleState.SUBMITTED, - AttemptLifecycleState.PENDING, - }: - update["state"] = AttemptLifecycleState.RUNNING - elif effective is EffectiveAttemptState.FAILED: - update.update( - state=AttemptLifecycleState.FAILED, - terminal_classification=_FAILURE_CLASSIFICATIONS.get( - scheduler_state, - AttemptTerminalClassification.UNKNOWN, - ), - ) - else: - continue - try: - writer.update_attempt(attempt.model_copy(update=update)) - except StateConflictError: - continue + else: + return + try: + writer.update_attempt(attempt.model_copy(update=update)) + except StateConflictError: + return def cancel(self, run_id: Identifier) -> SlurmRunCancellation: status = self.status(run_id) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 8e5c1976a..03dbd24df 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -5,7 +5,7 @@ import hashlib from collections.abc import Mapping -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path import pytest @@ -399,6 +399,43 @@ def test_status_reconciles_cancelled_scheduler_attempt( assert attempt.terminal_classification is AttemptTerminalClassification.CANCELLED +def test_status_expires_unrequeued_preemption_and_cancel_skips_terminal_job( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher() + current_time = [datetime(2026, 9, 8, tzinfo=timezone.utc)] + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: current_time[0], + package_version="0.9.2", + ) + result = service.execute(authored_run_single, source_root=tmp_path) + scheduler = SchedulerIdentity(array_job_id=42, array_task_id=0) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=scheduler, + state=SchedulerState.PREEMPTED, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=0), + ), + ) + + pending = service.status(result.run_id) + current_time[0] += timedelta(minutes=5, seconds=1) + failed = service.status(result.run_id) + cancellation = service.cancel(result.run_id) + + assert pending.shards[0].attempts[0].attempt.state is AttemptLifecycleState.PENDING + assert failed.shards[0].attempts[0].attempt.state is AttemptLifecycleState.FAILED + assert cancellation.job_ids == () + assert launcher.cancellations == [] + + def test_auto_gpu_resolution_rejects_mixed_node_shapes( tmp_path: Path, profile_catalog: SlurmProfileCatalog, From b5bbb6dbaedb5bc863c627b30fe874b070f157b5 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 10 Sep 2026 09:41:45 -0600 Subject: [PATCH 5/5] feat: add Slurm retry and deterministic collection (#915) * feat(slurm): add retry and deterministic collection Signed-off-by: Nabin Mulepati * fix(slurm): recover ambiguous submissions Signed-off-by: Nabin Mulepati * fix(slurm): preserve collection snapshots Signed-off-by: Nabin Mulepati * fix(slurm): bind retries and collection ownership Signed-off-by: Nabin Mulepati * fix(slurm): preserve scheduler path on retry Render retry allocations with the same profile-configured Slurm command path as initial generation attempts so cluster-local srun and scontrol binaries remain discoverable. Signed-off-by: Nabin Mulepati * fix(slurm): isolate attempt observation clocks Apply persisted timestamp floors per scheduler identity so a clock-ahead sibling cannot expire another attempt's preemption-requeue window. Persist each active attempt with its own reconciled observation time. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/client/execution.py | 43 +- .../src/data_designer/slurm/client/worker.py | 19 +- .../data_designer/slurm/launcher/client.py | 109 +- .../slurm/launcher/collection.py | 103 + .../data_designer/slurm/launcher/errors.py | 12 + .../data_designer/slurm/launcher/models.py | 18 + .../data_designer/slurm/launcher/parsing.py | 21 + .../data_designer/slurm/launcher/renderer.py | 130 +- .../data_designer/slurm/runtime/bootstrap.py | 23 +- .../data_designer/slurm/runtime/context.py | 41 +- .../data_designer/slurm/runtime/controller.py | 17 +- .../data_designer/slurm/runtime/entrypoint.py | 43 +- .../data_designer/slurm/runtime/entrypoint.sh | 30 +- .../src/data_designer/slurm/runtime/models.py | 25 +- .../src/data_designer/slurm/runtime/steps.py | 25 +- .../data_designer/slurm/services/wiring.py | 18 +- .../src/data_designer/slurm/state/__init__.py | 23 + .../data_designer/slurm/state/artifacts.py | 128 +- .../slurm/state/attempt_identity.py | 91 + .../data_designer/slurm/state/collection.py | 385 ++++ .../slurm/state/collection_filesystem.py | 357 ++++ .../slurm/state/collection_inputs.py | 81 + .../slurm/state/collection_merge.py | 338 ++++ .../slurm/state/collection_records.py | 112 ++ .../slurm/state/collection_storage.py | 339 ++++ .../slurm/state/collection_validation.py | 113 ++ .../slurm/state/collection_worker.py | 292 +++ .../data_designer/slurm/state/destinations.py | 92 + .../data_designer/slurm/state/observation.py | 20 +- .../src/data_designer/slurm/state/observer.py | 33 +- .../src/data_designer/slurm/state/outputs.py | 44 + .../src/data_designer/slurm/state/retry.py | 481 +++++ .../slurm/state/retry_records.py | 76 + .../slurm/state/retry_storage.py | 210 ++ .../src/data_designer/slurm/state/store.py | 20 +- .../slurm/state/submission_recovery.py | 85 + .../tests/client/test_worker.py | 74 + .../tests/launcher/test_client.py | 74 +- .../tests/launcher/test_collection.py | 205 ++ .../tests/runtime/conftest.py | 9 + .../tests/runtime/test_bootstrap.py | 39 + .../tests/runtime/test_bundle.py | 5 +- .../tests/runtime/test_context.py | 60 + .../tests/runtime/test_controller.py | 44 + .../tests/runtime/test_entrypoint.py | 2 +- .../tests/services/test_wiring.py | 48 + .../tests/state/test_retry_collection.py | 1780 +++++++++++++++++ .../tests/state/test_store.py | 16 + .../tests/state/test_submission_recovery.py | 144 ++ scripts/test_slurm_package_install.py | 8 + 50 files changed, 6431 insertions(+), 74 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/src/data_designer/slurm/state/submission_recovery.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 create mode 100644 packages/data-designer-slurm/tests/state/test_submission_recovery.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index 5f3fa3184..31011a4dd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -79,6 +79,7 @@ class _ExecutionContext: builder: DataDesignerConfigBuilder designer: DataDesigner requested_resume: ResumeMode + retry_resume: ResumeMode | None dataset_path: Path @@ -169,10 +170,16 @@ def preflight( prepared: PreparedClientEnvironment, endpoints: Mapping[str, str], plugins: tuple[ClientPluginEntryPoint, ...], + retry_resume: ResumeMode | None = None, ) -> ClientEnvironmentManifest: """Validate packages, plugins, assets, and config without generation.""" try: - context = self._build_context(plan_path, prepared=prepared, endpoints=endpoints) + context = self._build_context( + plan_path, + prepared=prepared, + endpoints=endpoints, + retry_resume=retry_resume, + ) progress = _ProgressWriter(context, prepared, self._clock) progress.required(ClientProgressPhase.VALIDATING_PLUGINS) progress.required(ClientProgressPhase.VALIDATING_CONFIG) @@ -200,12 +207,18 @@ def run( prepared: PreparedClientEnvironment, endpoints: Mapping[str, str], plugins: tuple[ClientPluginEntryPoint, ...], + retry_resume: ResumeMode | None = None, ) -> ClientResult: """Invoke the public Data Designer generation contract and persist its result.""" context: _ExecutionContext | None = None progress: _ProgressWriter | None = None try: - context = self._build_context(plan_path, prepared=prepared, endpoints=endpoints) + context = self._build_context( + plan_path, + prepared=prepared, + endpoints=endpoints, + retry_resume=retry_resume, + ) progress = _ProgressWriter(context, prepared, self._clock, revision=2) self._validate_environment_manifest(context, prepared, plugins) progress.required(ClientProgressPhase.GENERATING, completed_records=0) @@ -245,6 +258,7 @@ def _build_context( *, prepared: PreparedClientEnvironment, endpoints: Mapping[str, str], + retry_resume: ResumeMode | None = None, ) -> _ExecutionContext: try: plan = ResolvedSlurmRunPlan.model_validate_json( @@ -283,7 +297,14 @@ def _build_context( mcp_providers = self._materialize_mcp_providers(plan) managed_assets_path = self._validate_assets(plan) requested_resume = ResumeMode(plan.invocation.authored.resume) - dataset_path = self._dataset_path(plan, shard, prepared, requested_resume) + execution_resume = requested_resume if retry_resume is None else retry_resume + if ( + retry_resume is not None + and requested_resume is not ResumeMode.IF_POSSIBLE + and retry_resume is not requested_resume + ): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "retry resume mode differs from the plan") + dataset_path = self._dataset_path(plan, shard, prepared, execution_resume) designer = self._data_designer_factory( artifact_path=dataset_path.parent, model_providers=providers, @@ -292,7 +313,7 @@ def _build_context( auto_configure_logging=False, ) designer.set_run_config(RunConfig.model_validate(plan.invocation.effective_run_config)) - return _ExecutionContext(plan, shard, builder, designer, requested_resume, dataset_path) + return _ExecutionContext(plan, shard, builder, designer, requested_resume, retry_resume, dataset_path) except ClientWorkerError: raise except Exception as error: @@ -485,9 +506,10 @@ def _prepare_dataset_workspace( ensure_private_directory(context.dataset_path.parent) elif not context.dataset_path.parent.is_dir(): raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "shard dataset workspace is unavailable") - if context.requested_resume is not ResumeMode.NEVER: + execution_resume = context.requested_resume if context.retry_resume is None else context.retry_resume + if execution_resume is not ResumeMode.NEVER: ensure_private_directory(context.dataset_path) - if context.requested_resume is ResumeMode.NEVER and context.dataset_path.exists(): + if execution_resume is ResumeMode.NEVER and context.dataset_path.exists(): if not context.dataset_path.is_dir() or any(context.dataset_path.iterdir()): raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "attempt dataset workspace is not empty") @@ -503,7 +525,7 @@ def _generate_dataset( context.builder, num_records=context.shard.requested_records, dataset_name=context.dataset_path.name, - resume=context.requested_resume, + resume=context.requested_resume if context.retry_resume is None else context.retry_resume, artifact_path=context.dataset_path.parent, on_batch_complete=progress.on_batch_complete, ), @@ -538,11 +560,16 @@ def _validate_creation_result( raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer result counts are invalid") if results.early_shutdown is None or results.effective_resume_mode is None: raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer result metadata is incomplete") - if results.requested_resume_mode is not context.requested_resume: + execution_resume = context.requested_resume if context.retry_resume is None else context.retry_resume + if results.requested_resume_mode is not execution_resume: raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer resume metadata differs") dataset_path = Path(results.dataset_path) effective_resume = results.effective_resume_mode + if context.retry_resume is not None and effective_resume is not context.retry_resume: + raise ClientWorkerError( + ClientErrorCode.OUTPUT_INVALID, "Data Designer effective resume mode differs from retry intent" + ) shared_path = Path(get_container_path(context.plan, context.shard.resume_workspace.path, require_writable=True)) expected_path = shared_path if effective_resume is ResumeMode.ALWAYS else prepared.attempt_dir / "dataset" if ( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py b/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py index 5139b10c0..d34fb5dcf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py @@ -11,6 +11,7 @@ from datetime import datetime, timezone from pathlib import Path +from data_designer.config import ResumeMode from data_designer.slurm.client.environment import ( ClientEnvironmentBuilder, PreparedClientEnvironment, @@ -47,10 +48,23 @@ def main(argv: Sequence[str] | None = None) -> int: execution_module = importlib.import_module("data_designer.slurm.client.execution") ClientWorker = getattr(execution_module, "ClientWorker") worker = ClientWorker() + retry_resume = None if arguments.resume_mode is None else ResumeMode(arguments.resume_mode) if arguments.operation == "preflight": - worker.preflight(arguments.plan, prepared=prepared, endpoints=endpoints, plugins=plugins) + worker.preflight( + arguments.plan, + prepared=prepared, + endpoints=endpoints, + plugins=plugins, + retry_resume=retry_resume, + ) else: - worker.run(arguments.plan, prepared=prepared, endpoints=endpoints, plugins=plugins) + worker.run( + arguments.plan, + prepared=prepared, + endpoints=endpoints, + plugins=plugins, + retry_resume=retry_resume, + ) return 0 except ClientWorkerError as error: if prepared is not None and error.code is ClientErrorCode.PLUGIN_LOAD_FAILED: @@ -78,6 +92,7 @@ def _parse_arguments(argv: Sequence[str] | None) -> argparse.Namespace: parser.add_argument("--shard-id", required=True) parser.add_argument("--attempt-id", required=True) parser.add_argument("--attempt-dir", required=True, type=Path) + parser.add_argument("--resume-mode", choices=("never", "always")) parser.add_argument("--endpoint", action="append", default=[]) return parser.parse_args(argv) 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 81ee4d2b3..58e203105 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,23 +5,28 @@ from __future__ import annotations +import os import re import subprocess import unicodedata from collections.abc import Mapping, Sequence from dataclasses import dataclass +from datetime import datetime, timedelta from pathlib import Path from data_designer.slurm.contracts import Identifier -from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError from data_designer.slurm.launcher.models import ( SlurmAccountingEntry, SlurmJobSubmissionReceipt, + SlurmNamedJobEntry, SlurmQueueEntry, + SlurmSubmissionMatch, ) from data_designer.slurm.launcher.parsing import ( parse_accounting, parse_gpu_counts, + parse_named_jobs, parse_queue, parse_submission, ) @@ -104,12 +109,21 @@ def submit_script( raise ValueError("batch script text must be non-empty UTF-8 text without NUL") hold_arguments = ("--hold",) if hold else () export_argument, environment = _format_export_environment(export_environment) - output = self._run( - (self._executables.sbatch, "--parsable", *hold_arguments, export_argument), - input_text=script, - environment=environment, - ) - return parse_submission(output) + try: + output = self._run( + (self._executables.sbatch, "--parsable", *hold_arguments, export_argument), + input_text=script, + environment=environment, + ) + 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.""" @@ -160,6 +174,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=JobID,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))) @@ -194,13 +249,26 @@ def _run( ) else: completed = self._runner.run(command, input_text=input_text, environment=environment) - 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}") @@ -288,3 +356,26 @@ def _format_error_detail(error: BaseException) -> str: def _is_unknown_queue_job_error(error: SlurmCommandError) -> bool: message = str(error) return any(message.endswith(f": {detail}") for detail in _UNKNOWN_QUEUE_JOB_DETAILS) + + +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 new file mode 100644 index 000000000..5072d14d9 --- /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", collection_plan.submission_job_name), + ("account", resolved_plan.submission.account), + ("partition", resolved_plan.selected_profile.profile.image_build.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/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 ad71b0af5..efc16549a 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 _SYSTEM_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" @@ -25,13 +27,10 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi if plan.selected_profile.profile.gpu_request_mode == "visible": directive_text = f"{directive_text}\n#SBATCH --exclusive" attempt = f"{attempt_ordinal:04d}" - scheduler_bin_path = plan.selected_profile.profile.scheduler.bin_path - command_path = _SYSTEM_PATH if scheduler_bin_path is None else f"{scheduler_bin_path}:{_SYSTEM_PATH}" - return f"""#!/usr/bin/env bash {directive_text} set -Eeuo pipefail -export PATH={quote_shell_value(command_path)} +export PATH={quote_shell_value(_get_command_path(plan))} readonly DD_RUNTIME_ARCHIVE={quote_shell_value(plan.runtime_bundle.path)} readonly DD_RUNTIME_SHA256={quote_shell_value(plan.runtime_bundle.sha256)} @@ -65,26 +64,135 @@ 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, 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 + ) + return f"""#!/usr/bin/env bash +{directives} +set -Eeuo pipefail +export PATH={quote_shell_value(_get_command_path(plan))} + +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_RETRY_ID={quote_shell_value(retry.retry_id)} +readonly DD_RETRY_PLAN_SHA256={quote_shell_value(retry.compute_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}}" "${{DD_RETRY_ID}}" \ + "${{DD_RETRY_PLAN_SHA256}}" "${{DD_EFFECTIVE_RESUME_MODE}}" +""" + + +def _get_command_path(plan: ResolvedSlurmRunPlan) -> str: + scheduler_bin_path = plan.selected_profile.profile.scheduler.bin_path + return _SYSTEM_PATH if scheduler_bin_path is None else f"{scheduler_bin_path}:{_SYSTEM_PATH}" + + +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, *(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), + ("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)), ("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/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index 54a053231..264bfdd4a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -203,6 +203,7 @@ def _build_client_step( log_directory: Path, ) -> RuntimeStepSpec: plan = context.plan + retry_resume_mode = None if context.retry_plan is None else context.retry_plan.effective_resume_mode command = build_client_command( "preflight" if operation == "preflight" else "run", plan, @@ -210,15 +211,35 @@ def _build_client_step( context.attempt, context.attempt_directory, endpoints, + retry_resume_mode, ) if operation == "client": + endpoint_arguments = tuple( + argument + for endpoint in endpoints + for argument in ("--endpoint", f"{endpoint.model_alias}=http://{endpoint.host}:{endpoint.port}/v1") + ) + retry_binding = ( + () + if context.retry_plan is None + else ( + "--retry-id", + context.retry_plan.retry_id, + "--retry-plan-sha256", + context.retry_plan.compute_sha256(), + "--effective-resume-mode", + context.retry_plan.effective_resume_mode, + ) + ) command = ( "python3", "-m", "data_designer.slurm.runtime.entrypoint", "client", *command[4:6], - *command[10:], + *command[10:12], + *retry_binding, + *endpoint_arguments, ) secret_names = collect_secret_environment_names( (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py index 6817300b4..cbb32b181 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -14,7 +14,7 @@ from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.paths import get_container_path -from data_designer.slurm.state import SlurmStateWriter +from data_designer.slurm.state import RetryPlan, SlurmStateError, SlurmStateWriter from data_designer.slurm.state.filesystem import open_verified_directory, read_regular_text _MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 @@ -24,6 +24,10 @@ def load_allocation_context( plan_path: Path, attempt_directory: Path, environment: Mapping[str, str], + *, + retry_id: str | None = None, + retry_plan_sha256: str | None = None, + effective_resume_mode: str | None = None, ) -> tuple[AllocationContext, SlurmStateWriter]: """Load one scheduler-selected shard attempt through its container paths.""" writer = _load_state_writer(plan_path, attempt_directory) @@ -50,7 +54,40 @@ def load_allocation_context( SlurmRuntimeErrorCode.INVALID_CONTEXT, "scheduler array job does not match the persisted attempt", ) - return AllocationContext(plan, shard, attempt, host_attempt_directory), writer + retry_plan = _load_retry_plan( + writer, + retry_id=retry_id, + retry_plan_sha256=retry_plan_sha256, + effective_resume_mode=effective_resume_mode, + ) + return AllocationContext(plan, shard, attempt, host_attempt_directory, retry_plan), writer + + +def _load_retry_plan( + writer: SlurmStateWriter, + *, + retry_id: str | None, + retry_plan_sha256: str | None, + effective_resume_mode: str | None, +) -> RetryPlan | None: + arguments = (retry_id, retry_plan_sha256, effective_resume_mode) + if all(argument is None for argument in arguments): + return None + if any(argument is None for argument in arguments): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime retry binding is incomplete") + assert retry_id is not None and retry_plan_sha256 is not None and effective_resume_mode is not None + if effective_resume_mode not in {"never", "always"}: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime retry resume mode is invalid") + try: + retry_plan = writer.load_retry_plan(retry_id) + except SlurmStateError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime retry plan is unavailable") from error + if retry_plan.compute_sha256() != retry_plan_sha256 or retry_plan.effective_resume_mode != effective_resume_mode: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime retry binding differs from persisted state", + ) + return retry_plan def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWriter: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index e90cbc555..c8abeda28 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -295,6 +295,7 @@ def _prepare_runtime(self) -> _RuntimeTopology: self._context.attempt_directory, endpoints, self._environment, + retry_resume_mode=self._retry_resume_mode, ) ) self._supervisor.wait(self._supervisor.start(client_preflight)) @@ -346,17 +347,26 @@ def _run_client( self._context.attempt_directory, endpoints, self._environment, + retry_resume_mode=self._retry_resume_mode, ) ) + workspace_mode = self._retry_resume_mode + if workspace_mode is None: + workspace_mode = "never" if self._context.plan.invocation.authored.resume == "never" else "always" with self._state.acquire_dataset_workspace( self._attempt.shard_id, self._attempt.attempt_id, - self._context.plan.invocation.authored.resume, + workspace_mode, ): generation_started_at = self._now() self._supervisor.wait(self._supervisor.start(generation), required=required_processes) self._supervisor.require_running(required_processes) client_result, candidate = load_complete_client_candidate(self._context, self._attempt) + if self._retry_resume_mode is not None and client_result.effective_resume_mode != self._retry_resume_mode: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client effective resume mode differs from the persisted retry plan", + ) self._validate_client_timestamps(candidate.created_at, client_result.completed_at, generation_started_at) candidate_reference = client_result.candidate_output_manifest if candidate_reference is None: # pragma: no cover - the record contract requires this for complete results @@ -368,6 +378,11 @@ def _run_client( self._attempt = _copy_attempt(self._attempt, candidate_output=candidate_reference) return candidate_reference, client_result.completed_at + @property + def _retry_resume_mode(self) -> Literal["never", "always"] | None: + retry_plan = self._context.retry_plan + return None if retry_plan is None else retry_plan.effective_resume_mode + def _validate_client_timestamps( self, candidate_created_at: datetime, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index 66ee59f21..db0462981 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -85,10 +85,13 @@ def _parse_arguments(arguments: Sequence[str] | None) -> argparse.Namespace: def _add_context_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument("--plan", required=True, type=Path) parser.add_argument("--attempt-dir", required=True, type=Path) + parser.add_argument("--retry-id") + parser.add_argument("--retry-plan-sha256") + parser.add_argument("--effective-resume-mode", choices=("never", "always")) def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + context, writer = _load_context(arguments, environment) _validate_attempt_is_executable(context.attempt) SystemAllocationPreflight.verify_attempt_directory(arguments.attempt_dir) SystemAllocationPreflight.verify_ports(context, environment) @@ -113,7 +116,7 @@ def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + context, writer = _load_context(arguments, environment) previous = writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) timestamp = _now(context.attempt, previous) deployments = _resolve_deployments(context, environment) @@ -143,12 +146,17 @@ def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> Non def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + context, writer = _load_context(arguments, environment) generation_started_at = _now(context.attempt, _load_optional_readiness(context, writer)) + resume_mode = ( + context.plan.invocation.authored.resume + if context.retry_plan is None + else context.retry_plan.effective_resume_mode + ) with writer.acquire_dataset_workspace( context.shard.shard_id, context.attempt.attempt_id, - context.plan.invocation.authored.resume, + resume_mode, ): return_code = client_worker_main( ( @@ -161,6 +169,7 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No context.attempt.attempt_id, "--attempt-dir", arguments.attempt_dir.as_posix(), + *(() if context.retry_plan is None else ("--resume-mode", context.retry_plan.effective_resume_mode)), *(argument for endpoint in arguments.endpoint for argument in ("--endpoint", endpoint)), ) ) @@ -171,6 +180,14 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No context.attempt, attempt_directory=arguments.attempt_dir, ) + if ( + context.retry_plan is not None + and client_result.effective_resume_mode != context.retry_plan.effective_resume_mode + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client effective resume mode differs from the persisted retry plan", + ) completed_at = client_result.completed_at if candidate.created_at < generation_started_at or completed_at < generation_started_at: raise SlurmRuntimeError( @@ -186,7 +203,7 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No def _succeed(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + context, writer = _load_context(arguments, environment) stopped_at = _write_stopped_readiness(context, writer) attempt = writer.load_attempt(context.shard.shard_id, context.attempt.attempt_id) if attempt.candidate_output is None: @@ -204,7 +221,7 @@ def _succeed(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N def _fail(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + context, writer = _load_context(arguments, environment) attempt = writer.load_attempt(context.shard.shard_id, context.attempt.attempt_id) if attempt.state in {AttemptLifecycleState.SUCCEEDED, AttemptLifecycleState.FAILED}: return @@ -220,6 +237,20 @@ def _fail(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None ) +def _load_context( + arguments: argparse.Namespace, + environment: Mapping[str, str], +) -> tuple[AllocationContext, SlurmStateWriter]: + return load_allocation_context( + arguments.plan, + arguments.attempt_dir, + environment, + retry_id=arguments.retry_id, + retry_plan_sha256=arguments.retry_plan_sha256, + effective_resume_mode=arguments.effective_resume_mode, + ) + + def _begin_attempt( context: AllocationContext, writer: SlurmStateWriter, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh index dc15efd34..cbba85438 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -10,15 +10,24 @@ source "${DD_RUNTIME_DIR}/cleanup.sh" DD_RUNTIME_PREPARED=0 DD_RUNTIME_FINALIZED=0 +DD_RUNTIME_RETRY_ARGUMENTS=() dd_slurm_run_allocation() { - if [[ $# -ne 2 ]]; then - printf '%s\n' 'allocation runtime requires plan and attempt directory arguments' >&2 + if [[ $# -ne 2 && $# -ne 5 ]]; then + printf '%s\n' 'allocation runtime requires plan and attempt with optional retry binding' >&2 return 64 fi DD_PLAN_PATH=$1 DD_ATTEMPT_PATH=$2 DD_RUNTIME_MANIFEST=${DD_ATTEMPT_PATH}/runtime-manifest.json + DD_RUNTIME_RETRY_ARGUMENTS=() + if [[ $# -eq 5 ]]; then + DD_RUNTIME_RETRY_ARGUMENTS=( + --retry-id "$3" + --retry-plan-sha256 "$4" + --effective-resume-mode "$5" + ) + fi dd_read_control_plan "${DD_PLAN_PATH}" dd_read_plan_secret_names "${DD_PLAN_PATH}" dd_read_container_path "${DD_PLAN_PATH}" "${DD_PLAN_PATH}" false @@ -38,7 +47,7 @@ dd_slurm_run_allocation() { trap 'exit 130' INT TERM DD_RUNTIME_PREPARED=1 - dd_run_control_phase prepare \ + dd_run_bound_control_phase prepare \ --runtime-root "${DD_RUNTIME_DIR}" \ --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" dd_verify_runtime_manifest \ @@ -57,7 +66,7 @@ dd_slurm_run_allocation() { dd_start_endpoints dd_wait_for_role_readiness endpoint dd_require_running - dd_run_control_phase ready + dd_run_bound_control_phase ready dd_read_step_ids "${DD_RUNTIME_MANIFEST}" client ((${#DD_STEP_IDS[@]} == 1)) @@ -68,10 +77,19 @@ dd_slurm_run_allocation() { dd_require_running dd_cleanup_steps - dd_run_control_phase succeed + dd_run_bound_control_phase succeed DD_RUNTIME_FINALIZED=1 } +dd_run_bound_control_phase() { + local operation=$1 + shift + dd_run_control_phase \ + "${operation}" \ + "${DD_RUNTIME_RETRY_ARGUMENTS[@]+"${DD_RUNTIME_RETRY_ARGUMENTS[@]}"}" \ + "$@" +} + dd_verify_host_context() { local tool for tool in bash sha256sum tar jq srun scontrol getent curl; do @@ -173,7 +191,7 @@ dd_runtime_exit() { set +e dd_cleanup_steps if ((DD_RUNTIME_PREPARED == 1 && DD_RUNTIME_FINALIZED == 0)); then - dd_run_control_phase fail >/dev/null + dd_run_bound_control_phase fail >/dev/null fi dd_stop_runtime_timer exit "${status}" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py index 50d2557fe..2aa93ab16 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py @@ -16,7 +16,7 @@ from data_designer.slurm.contracts import Identifier from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode -from data_designer.slurm.state import AttemptManifest +from data_designer.slurm.state import AttemptManifest, RetryPlan _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) @@ -117,6 +117,7 @@ class AllocationContext: shard: PlannedShard attempt: AttemptManifest attempt_directory: Path + retry_plan: RetryPlan | None = None def __post_init__(self) -> None: if not isinstance(self.plan, ResolvedSlurmRunPlan): @@ -147,3 +148,25 @@ def __post_init__(self) -> None: SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation attempt does not match its planned scheduler task", ) + if self.retry_plan is not None: + selected = tuple( + retry_shard + for retry_shard in self.retry_plan.planned_shards + if retry_shard.shard_id == self.shard.shard_id + and retry_shard.attempt_id == self.attempt.attempt_id + and retry_shard.array_task_index == self.shard.array_task_index + ) + if ( + self.retry_plan.run_id != self.plan.run_id + or self.retry_plan.resolved_plan != self.attempt.resolved_plan + or ( + self.plan.invocation.authored.resume != "if_possible" + and self.retry_plan.effective_resume_mode != self.plan.invocation.authored.resume + ) + or len(selected) != 1 + or selected[0].attempt_ordinal != self.attempt.attempt_ordinal + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "allocation attempt does not match its persisted retry plan", + ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py index 8aeb96eb2..74ca803ea 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -8,7 +8,7 @@ import os from collections.abc import Mapping from pathlib import Path -from typing import Protocol +from typing import Literal, Protocol from data_designer.slurm.config.environment import ( LiteralEnvironmentBinding, @@ -55,6 +55,8 @@ def build_preflight_step( attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], source_environment: Mapping[str, str], + *, + retry_resume_mode: Literal["never", "always"] | None = None, ) -> RuntimeStep: """Build the zero-GPU client preflight step.""" ... @@ -67,6 +69,8 @@ def build_generation_step( attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], source_environment: Mapping[str, str], + *, + retry_resume_mode: Literal["never", "always"] | None = None, ) -> RuntimeStep: """Build the zero-GPU client generation step.""" ... @@ -83,6 +87,8 @@ def build_preflight_step( attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], source_environment: Mapping[str, str], + *, + retry_resume_mode: Literal["never", "always"] | None = None, ) -> RuntimeStep: """Build a deterministic client-worker preflight command.""" return self._build_step( @@ -95,6 +101,7 @@ def build_preflight_step( attempt_directory, endpoints, source_environment, + retry_resume_mode, ) def build_generation_step( @@ -105,6 +112,8 @@ def build_generation_step( attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], source_environment: Mapping[str, str], + *, + retry_resume_mode: Literal["never", "always"] | None = None, ) -> RuntimeStep: """Build a deterministic client-worker generation command.""" return self._build_step( @@ -117,6 +126,7 @@ def build_generation_step( attempt_directory, endpoints, source_environment, + retry_resume_mode, ) @staticmethod @@ -130,8 +140,17 @@ def _build_step( attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], source_environment: Mapping[str, str], + retry_resume_mode: Literal["never", "always"] | None, ) -> RuntimeStep: - command = build_client_command(operation, plan, shard, attempt, attempt_directory, endpoints) + command = build_client_command( + operation, + plan, + shard, + attempt, + attempt_directory, + endpoints, + retry_resume_mode, + ) secret_names, environment = _build_client_environment(plan, source_environment) allocation_environment = ("SLURM_JOB_GPUS",) if plan.selected_profile.profile.gpu_request_mode == "gres" else () return _build_srun_step( @@ -153,6 +172,7 @@ def build_client_command( attempt: AttemptManifest, attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], + retry_resume_mode: Literal["never", "always"] | None, ) -> tuple[str, ...]: endpoint_arguments = tuple( argument @@ -172,6 +192,7 @@ def build_client_command( attempt.attempt_id, "--attempt-dir", get_container_path(plan, attempt_directory.as_posix(), require_writable=True), + *(() if retry_resume_mode is None else ("--resume-mode", retry_resume_mode)), *endpoint_arguments, ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index df3ab4709..b13834f56 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -548,17 +548,12 @@ def _reconcile_attempts( ) if not active: return - observed_at = max(self._clock(), *(attempt.updated_at for attempt in active)) - for attempt in active: - readiness = _load_optional(lambda: writer.load_readiness(attempt.shard_id, attempt.attempt_id)) - if readiness is not None: - observed_at = max(observed_at, readiness.updated_at) try: reconciled = SlurmStateReconciler( self._profile.profile.workspace_root, writer.load_run().run_id, self._launcher, - ).refresh(observed_at=observed_at) + ).refresh(observed_at=self._clock()) except SlurmStateError as error: if isinstance(error.__cause__, SlurmLauncherError): return @@ -567,16 +562,18 @@ def _reconcile_attempts( for shard in reconciled.shards: for status in shard.attempts: if (status.attempt.shard_id, status.attempt.attempt_id) in active_identities: - self._update_reconciled_attempt(writer, status, observed_at) + self._update_reconciled_attempt(writer, status) @staticmethod def _update_reconciled_attempt( writer: SlurmStateWriter, status: AttemptStatus, - observed_at: datetime, ) -> None: attempt = status.attempt - update: dict[str, object] = {"updated_at": observed_at} + scheduler = status.scheduler + if scheduler is None: # pragma: no cover - active attempts always have scheduler evidence + raise AssertionError("active attempt has no reconciled scheduler evidence") + update: dict[str, object] = {"updated_at": scheduler.observed_at} if status.effective_state is EffectiveAttemptState.PENDING and attempt.state is AttemptLifecycleState.SUBMITTED: update["state"] = AttemptLifecycleState.PENDING elif status.effective_state is EffectiveAttemptState.RUNNING and attempt.state in { @@ -585,11 +582,10 @@ def _update_reconciled_attempt( }: update["state"] = AttemptLifecycleState.RUNNING elif status.effective_state is EffectiveAttemptState.FAILED: - scheduler_state = SchedulerState.UNKNOWN if status.scheduler is None else status.scheduler.state update.update( state=AttemptLifecycleState.FAILED, terminal_classification=_FAILURE_CLASSIFICATIONS.get( - scheduler_state, + scheduler.state, AttemptTerminalClassification.UNKNOWN, ), ) 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 c397e50f6..c68bb8738 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 Callable, Protocol +from typing import BinaryIO, Callable, Protocol import data_designer.lazy_heavy_imports as lazy from data_designer.slurm.state.filesystem import ( @@ -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.""" @@ -101,6 +109,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] + _artifacts: tuple[_ArtifactSnapshot, ...] + + class CandidateArtifactVerifier: """Verify one manifest-bounded candidate and lease its files through publication.""" @@ -127,6 +145,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.""" + 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=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.""" + actual = self.inspect(candidate) + 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.""" @@ -141,7 +210,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] @@ -165,14 +234,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 @@ -186,7 +304,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) @@ -213,6 +332,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..396ccdcf5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -0,0 +1,385 @@ +# 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, SlurmSubmissionMatch +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 +from data_designer.slurm.state.submission_recovery import ( + SUBMISSION_VISIBILITY_WINDOW, + PreparedSubmission, + resolve_prepared_submission, +) + +_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.""" + ... + + 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.""" + + 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: + current = self._reconcile_prepared_collection(current_plan, current, timestamp) + 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, + reconciliation_deadline=timestamp + SUBMISSION_VISIBILITY_WINDOW, + ) + 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.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, + 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, + reconciliation_deadline=None, + ) + 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, + 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]) + + 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..b8774661f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py @@ -0,0 +1,112 @@ +# 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_optional_utc_timestamp, + 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 + 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: + 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..7da068a42 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py @@ -0,0 +1,339 @@ +# 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" +_WORKER_LOCK = "worker.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 + + @contextmanager + def acquire_worker_lock(self, collection_id: Identifier) -> Iterator[None]: + """Serialize bulk I/O owners without blocking collection metadata refresh.""" + with self._open_collection_directory(collection_id) as descriptor: + with acquire_file_lock( + descriptor, + _WORKER_LOCK, + self.get_collection_root(collection_id) / _WORKER_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..369b4cf6d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py @@ -0,0 +1,113 @@ +# 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} + _require(len(schema_digests) == 1, "collection candidates have incompatible schemas") + _require( + all(candidate.provenance_digest == resolved_plan.compute_sha256() for candidate in candidates), + "collection candidate provenance does not match the resolved plan", + ) + + 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..8f464ee6c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py @@ -0,0 +1,292 @@ +# 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 time import sleep + +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) +_SCHEDULER_BINDING_WAIT_ATTEMPTS = 300 + + +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.""" + try: + self._wait_for_scheduler_binding() + started_at = _utc_now() if completed_at is None else completed_at + with self._collections.acquire_worker_lock(self._collection_id): + plan, status = self._begin_collection(started_at) + if status.state is CollectionState.SUCCEEDED: + return self._load_valid_result(plan, status) + resolved_plan, candidates = self._inputs.resolve(plan) + self._destinations.validate_persisted(resolved_plan, plan) + recovered = self._load_optional_result(plan) + if recovered is not None: + return self._publish_success_status(plan, recovered, recovered.completed_at) + return self._execute_collection(plan, resolved_plan, candidates, 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 _begin_collection(self, started_at: datetime) -> tuple[CollectionPlan, CollectionStatus]: + 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) + if status.state is CollectionState.SUCCEEDED: + return plan, status + return plan, self._advance_status(status, CollectionState.RUNNING, started_at) + + def _execute_collection( + self, + plan: CollectionPlan, + 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), + derive_collection_staging_directory(plan), + 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, completed_at, error) + return self._publish_success_status(plan, result, result.completed_at) + + def _recover_or_fail( + self, + plan: CollectionPlan, + completed_at: datetime | None, + error: Exception, + ) -> CollectionResult: + try: + recovered = self._load_optional_result(plan) + except (SlurmStateError, OSError, ValueError) as recovery_error: + self._publish_failed_status(self._completion_time(completed_at)) + raise recovery_error from error + if recovered is not None: + return self._publish_success_status(plan, recovered, recovered.completed_at) + self._publish_failed_status(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, + result: CollectionResult, + updated_at: datetime, + ) -> CollectionResult: + with self._collections.acquire_lock(): + previous = self._collections.read_status(self._collection_id) + bound_plan = self._load_bound_plan(previous) + if bound_plan != plan: + raise StateCorruptionError("collection plan changed during worker execution") + self._require_scheduler_job(previous) + result_reference = self._collections.get_result_reference(plan, result) + if previous.state is CollectionState.SUCCEEDED: + if previous.result != result_reference: + raise StateCorruptionError("collection status does not bind its published result") + return result + 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=max(updated_at, previous.updated_at), + state=CollectionState.SUCCEEDED, + scheduler=previous.scheduler, + scheduler_observation=previous.scheduler_observation, + result=result_reference, + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + return result + + def _publish_failed_status(self, updated_at: datetime) -> None: + with self._collections.acquire_lock(): + previous = self._collections.read_status(self._collection_id) + self._load_bound_plan(previous) + self._require_scheduler_job(previous) + self._advance_status(previous, CollectionState.FAILED, updated_at) + + def _advance_status( + self, + previous: CollectionStatus, + state: CollectionState, + updated_at: datetime, + ) -> CollectionStatus: + current = _updated_status( + previous, + revision=previous.revision + 1, + updated_at=max(updated_at, previous.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 _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") + + 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..78c8a2bc4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -0,0 +1,92 @@ +# 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 one authorized output path 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 + workspace_root = plan.selected_profile.profile.workspace_root + workspace_mount = ContainerMount(source=workspace_root, target=workspace_root, read_only=False) + authorized = { + (mount.source, mount.target): mount + for mount in (workspace_mount, *plan.container_mounts) + if not mount.read_only + } + writable = tuple( + mount + for mount in authorized.values() + if 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(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, collection_plan.host_destination) + 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/observation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py index e8da67c8d..a3102658e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py @@ -62,6 +62,7 @@ def collect( *, observed_at: datetime, previous: Mapping[SchedulerJobIdentity, SchedulerObservation | None] | None = None, + observation_floors: Mapping[SchedulerJobIdentity, datetime] | None = None, ) -> tuple[SchedulerObservation, ...]: """Return one deterministic observation for every requested identity.""" validate_utc_timestamp(observed_at) @@ -69,13 +70,14 @@ def collect( if not requested: return () prior = {} if previous is None else previous + observation_times = _resolve_observation_times(requested, observed_at, observation_floors) queue, accounting = self._query_scheduler(requested) queue_by_identity = self._index_records(queue, requested, source="active queue") accounting_by_identity = self._index_records(accounting, requested, source="accounting") return tuple( self._resolve_observation( identity, - observed_at, + observation_times[identity], queue_by_identity.get(identity), accounting_by_identity.get(identity), prior.get(identity), @@ -148,6 +150,22 @@ def _resolve_observation( return observation +def _resolve_observation_times( + requested: tuple[SchedulerJobIdentity, ...], + observed_at: datetime, + observation_floors: Mapping[SchedulerJobIdentity, datetime] | None, +) -> dict[SchedulerJobIdentity, datetime]: + floors = {} if observation_floors is None else observation_floors + if not set(floors).issubset(requested): + raise SlurmStateError("scheduler observation floors contain an unrequested identity") + resolved: dict[SchedulerJobIdentity, datetime] = {} + for identity in requested: + floor = floors.get(identity, observed_at) + validate_utc_timestamp(floor) + resolved[identity] = max(observed_at, floor) + return resolved + + def _select_observed_state( queue_state: SchedulerState | None, accounting_state: SchedulerState | None, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py index eeef347c5..4b018d726 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -53,7 +53,6 @@ class _ShardSnapshot: class _ObservationBatch: previous: dict[SchedulerIdentity, SchedulerObservation | None] current: dict[SchedulerJobIdentity, SchedulerObservation] - observed_at: datetime class SlurmStateReconciler: @@ -101,11 +100,15 @@ def refresh(self, *, observed_at: datetime | None = None) -> RunStatus: attempts_by_shard = self._reader.load_validated_attempts(run, plan, shards) previous = self._load_previous_observations(attempts_by_shard) selectors = tuple(previous.keys()) - current = self._collector.collect(selectors, observed_at=timestamp, previous=previous) + current = self._collector.collect( + selectors, + observed_at=timestamp, + previous=previous, + observation_floors=self._load_observation_floors(plan, attempts_by_shard, previous), + ) batch = _ObservationBatch( previous=previous, current={observation.scheduler: observation for observation in current}, - observed_at=timestamp, ) shard_statuses = tuple( self._refresh_shard( @@ -143,6 +146,28 @@ def _load_previous_observations( previous[attempt.scheduler] = self._reader.load_optional_scheduler_observation(attempt) return previous + def _load_observation_floors( + self, + plan: ResolvedSlurmRunPlan, + attempts_by_shard: dict[ShardId, tuple[AttemptManifest, ...]], + previous: dict[SchedulerIdentity, SchedulerObservation | None], + ) -> dict[SchedulerIdentity, datetime]: + floors: dict[SchedulerIdentity, datetime] = {} + for attempts in attempts_by_shard.values(): + for attempt in attempts: + scheduler = attempt.scheduler + if scheduler is None: + continue + floor = attempt.updated_at + prior = previous[scheduler] + if prior is not None: + floor = max(floor, prior.observed_at) + readiness = self._reader.load_optional_readiness(plan, attempt) + if readiness is not None: + floor = max(floor, readiness.updated_at) + floors[scheduler] = floor + return floors + def _refresh_shard( self, expected: _ShardSnapshot, @@ -200,7 +225,7 @@ def _build_attempt_status( attempt, readiness, scheduler, - current_time=batch.observed_at, + current_time=scheduler.observed_at, ) else: if attempt.state is not AttemptLifecycleState.CREATED: 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..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 @@ -115,6 +115,45 @@ 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) + + @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) + 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.""" @@ -131,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 new file mode 100644 index 000000000..607ab3fc9 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -0,0 +1,481 @@ +# 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, 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 +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.submission_recovery import ( + SUBMISSION_VISIBILITY_WINDOW, + PreparedSubmission, + resolve_prepared_submission, +) +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.""" + ... + + 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.""" + + 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_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: + 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_pending_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: + 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 + 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 _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, + 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, + reconciliation_deadline=timestamp + SUBMISSION_VISIBILITY_WINDOW, + ) + ) + + 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 + 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, + 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=array_job_id, + ) + validate_retry_status_transition(self._retries.read_status(plan.retry_id), submitted) + self._retries.replace_status(submitted) + return submitted + + 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..d5045c20c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py @@ -0,0 +1,76 @@ +# 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_optional_utc_timestamp, 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 + 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: + 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 + + +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/src/data_designer/slurm/state/store.py b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py index 9076e1955..6ab827ffe 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -25,12 +25,13 @@ ) from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest from data_designer.slurm.state.finalization import WinnerFinalizer -from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner +from data_designer.slurm.state.outputs import CandidateOutputManifest, RetryPlan, ShardWinner from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError from data_designer.slurm.state.reader import StateReader from data_designer.slurm.state.readiness import AttemptReadiness from data_designer.slurm.state.reconciliation import validate_readiness_transition from data_designer.slurm.state.results import AttemptResultPublisher +from data_designer.slurm.state.retry_storage import RetryStorage from data_designer.slurm.state.storage import StateStorage from data_designer.slurm.state.validation import ( StateContractError, @@ -84,6 +85,7 @@ def __init__( self._reader = StateReader(self._storage, normalized_run_id) self._results = AttemptResultPublisher(self._storage, self._reader) self._finalizer = WinnerFinalizer(self._storage, self._reader) + self._retries = RetryStorage(self._storage) self._run_id = normalized_run_id @property @@ -124,6 +126,22 @@ def load_resolved_plan(self) -> ResolvedSlurmRunPlan: """Load and digest-verify the run's immutable resolved plan.""" return self._reader.load_resolved_plan() + def load_retry_plan(self, retry_id: Identifier) -> RetryPlan: + """Load one retry plan and verify its status binds the same immutable record.""" + try: + normalized_retry_id = _IDENTIFIER_ADAPTER.validate_python(retry_id, strict=True) + plan = self._retries.read_plan(normalized_retry_id) + status = self._retries.read_status(normalized_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 + except (StateCorruptionError, StateNotFoundError): + raise + except (FileNotFoundError, ValidationError) as error: + raise StateNotFoundError(f"retry {retry_id!r} is unavailable") from error + except OSError as error: + raise StateCorruptionError(f"retry {retry_id!r} is unsafe or unreadable") from error + def load_shards(self) -> tuple[ShardManifest, ...]: """Load and validate the complete ordered shard set.""" return self._reader.load_shards() 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..99a5a2d04 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py @@ -0,0 +1,85 @@ +# 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: + 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", + "SubmissionLookup", + "resolve_prepared_submission", +] diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index c2175577b..7da326e0c 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -135,6 +135,8 @@ def test_main_preserves_equals_in_model_alias( client_worker_case.prepared.attempt_id, "--attempt-dir", client_worker_case.attempt_dir.as_posix(), + "--resume-mode", + "always", "--endpoint", f"judge=v2={endpoint}", ] @@ -142,6 +144,7 @@ def test_main_preserves_equals_in_model_alias( assert result == 0 assert worker.preflight.call_args.kwargs["endpoints"] == {"judge=v2": endpoint} + assert worker.preflight.call_args.kwargs["retry_resume"] is ResumeMode.ALWAYS @pytest.mark.parametrize( @@ -348,6 +351,77 @@ def test_if_possible_without_workspace_uses_attempt_dataset(client_worker_case: assert result.dataset_path == (client_worker_case.attempt_dir / "dataset").as_posix() +@pytest.mark.parametrize("retry_resume", [ResumeMode.NEVER, ResumeMode.ALWAYS]) +def test_retry_resume_mode_controls_if_possible_generation( + client_worker_case: ClientWorkerCase, + retry_resume: ResumeMode, +) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["resume"] = "if_possible" + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + if retry_resume is ResumeMode.ALWAYS: + resume_path = Path(plan.shards[0].resume_workspace.path) + resume_path.mkdir() + (resume_path / "partial").touch() + worker = ClientWorker( + data_designer_factory=partial(FakeDataDesigner, effective_resume=retry_resume), + ) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + retry_resume=retry_resume, + ) + + result = worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + retry_resume=retry_resume, + ) + + expected_path = ( + Path(plan.shards[0].resume_workspace.path) + if retry_resume is ResumeMode.ALWAYS + else client_worker_case.attempt_dir / "dataset" + ) + assert result.requested_resume_mode == "if_possible" + assert result.effective_resume_mode == retry_resume.value + assert result.dataset_path == expected_path.as_posix() + + +def test_retry_rejects_effective_resume_mode_drift(client_worker_case: ClientWorkerCase) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["resume"] = "if_possible" + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + resume_path = Path(plan.shards[0].resume_workspace.path) + resume_path.mkdir() + (resume_path / "partial").touch() + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + retry_resume=ResumeMode.ALWAYS, + ) + + with pytest.raises(ClientWorkerError) as error: + worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + retry_resume=ResumeMode.ALWAYS, + ) + + assert error.value.code is ClientErrorCode.OUTPUT_INVALID + + def test_if_possible_interruption_preserves_workspace_for_retry(client_worker_case: ClientWorkerCase) -> None: payload = client_worker_case.plan.model_dump(mode="json") payload["invocation"]["authored"]["resume"] = "if_possible" diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 52951158a..d17e5eb25 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -5,12 +5,13 @@ import subprocess from collections.abc import Mapping, Sequence +from datetime import datetime, timezone import pytest from slurm_test_fakes import FakeCommandResponse, FakeSlurmJob, FakeSlurmRunner from data_designer.slurm.launcher.client import SlurmCommandClient, SlurmExecutables -from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -72,6 +73,23 @@ def test_client_exports_only_explicit_environment_names_without_values_in_argv() assert "secret-value" not in " ".join(runner.command) +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") @@ -97,6 +115,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=JobID,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") @@ -364,7 +427,14 @@ def run( class _TimeoutRunner: - def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + del input_text, environment 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..a79a8187a --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -0,0 +1,205 @@ +# 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 pathlib import Path + +import pytest + +from data_designer.slurm.config import ContainerMount, injected_profile +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 +from data_designer.slurm.state.errors import StateConflictError + + +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 f"#SBATCH --job-name={collection.submission_job_name}" in script + assert "#SBATCH --partition=cpu" in script + assert "#SBATCH --partition=batch" not 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 ( + '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 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 + assert f'readonly DD_RETRY_ID="{retry.retry_id}"' in script + assert f'readonly DD_RETRY_PLAN_SHA256="{retry.compute_sha256()}"' in script + assert 'readonly DD_EFFECTIVE_RESUME_MODE="never"' 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 '"${DD_RETRY_PLAN_SHA256}" "${DD_EFFECTIVE_RESUME_MODE}"' in script + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +def test_retry_renderer_uses_profile_slurm_bin_path(multi_node_plan: ResolvedSlurmRunPlan) -> None: + scheduler = multi_node_plan.selected_profile.profile.scheduler.model_copy(update={"bin_path": "/opt/slurm/bin"}) + profile = multi_node_plan.selected_profile.profile.model_copy(update={"scheduler": scheduler}) + plan = multi_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + retry = RetryPlan( + schema_version=1, + retry_id="retry-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=( + 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(plan, retry) + + assert 'export PATH="/opt/slurm/bin:/usr/local/sbin:' in script + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +def test_destination_reauthorizes_explicit_path_through_workspace_mapping( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + workspace_root = multi_node_plan.selected_profile.profile.workspace_root + requested = (Path(workspace_root) / "explicit" / "collected").as_posix() + plan = multi_node_plan.model_copy(update={"container_mounts": ()}) + resolver = CollectionDestinationResolver() + destination = resolver.resolve(plan, requested) + 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=requested, + container_destination=requested, + num_partitions=plan.output.partitions, + ) + + assert destination.mount == ContainerMount(source=workspace_root, target=workspace_root, read_only=False) + assert resolver.validate_persisted(plan, collection) == destination + + +def test_destination_requires_one_unique_most_specific_mapping( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + requested = "/host/exports/run-001" + mounts = ( + ContainerMount(source="/host/exports", target="/container/a", read_only=False), + ContainerMount(source="/host/exports", target="/container/b", read_only=False), + ) + plan = multi_node_plan.model_copy(update={"container_mounts": mounts}) + + with pytest.raises(StateConflictError, match="ambiguous writable mount"): + CollectionDestinationResolver().resolve(plan, requested) diff --git a/packages/data-designer-slurm/tests/runtime/conftest.py b/packages/data-designer-slurm/tests/runtime/conftest.py index 475a56ff5..253f7a771 100644 --- a/packages/data-designer-slurm/tests/runtime/conftest.py +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -151,6 +151,9 @@ def verify(self, context: AllocationContext, environment: object) -> None: class FakeClientStepBuilder: + def __init__(self) -> None: + self.retry_resume_modes: list[Literal["never", "always"] | None] = [] + def build_preflight_step( self, plan: ResolvedSlurmRunPlan, @@ -159,8 +162,11 @@ def build_preflight_step( attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], source_environment: object, + *, + retry_resume_mode: Literal["never", "always"] | None = None, ) -> RuntimeStep: del plan, shard, attempt, endpoints, source_environment + self.retry_resume_modes.append(retry_resume_mode) return _step("client-preflight", RuntimeStepRole.CLIENT_PREFLIGHT, attempt_directory) def build_generation_step( @@ -171,8 +177,11 @@ def build_generation_step( attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], source_environment: object, + *, + retry_resume_mode: Literal["never", "always"] | None = None, ) -> RuntimeStep: del plan, shard, attempt, endpoints, source_environment + self.retry_resume_modes.append(retry_resume_mode) return _step("client-generation", RuntimeStepRole.CLIENT, attempt_directory) diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 51a56aa2c..10fb6c277 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -3,10 +3,13 @@ from __future__ import annotations +from dataclasses import replace + from conftest import RuntimeCase from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, build_runtime_manifest from data_designer.slurm.runtime.models import RuntimeStepRole +from data_designer.slurm.state import RetryPlan, RetryShard def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(runtime_case: RuntimeCase) -> None: @@ -41,3 +44,39 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru assert "--attempt-id" not in manifest.steps[-1].command assert "--plan" in manifest.steps[-1].command assert "--attempt-dir" in manifest.steps[-1].command + + +def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=context.plan.run_id, + created_at=runtime_case.created_at, + resolved_plan=context.attempt.resolved_plan, + planned_shards=( + RetryShard( + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + attempt_ordinal=context.attempt.attempt_ordinal, + array_task_index=context.shard.array_task_index, + ), + ), + effective_resume_mode="never", + ) + retry_context = replace(context, retry_plan=retry) + + manifest = build_runtime_manifest( + retry_context, + {"SLURM_JOB_GPUS": "0"}, + runtime_root=context.attempt_directory / "runtime", + log_directory=context.attempt_directory / "logs/execution-00000002", + ) + + preflight = manifest.steps[0].command + client = manifest.steps[-1].command + assert ("--resume-mode", "never") == preflight[preflight.index("--resume-mode") :][:2] + assert ("--retry-id", retry.retry_id) == client[client.index("--retry-id") :][:2] + assert ("--retry-plan-sha256", retry.compute_sha256()) == client[client.index("--retry-plan-sha256") :][:2] + assert ("--effective-resume-mode", "never") == client[client.index("--effective-resume-mode") :][:2] + assert "--resume-mode" not in client diff --git a/packages/data-designer-slurm/tests/runtime/test_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py index 40174aba1..5b73ad071 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bundle.py +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -47,7 +47,10 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p assert all(archive.getmember(name).uid == 0 for name in names) entrypoint = archive.extractfile("entrypoint.sh") assert entrypoint is not None - assert b"python3 -m data_designer.slurm.runtime.entrypoint" not in entrypoint.read() + entrypoint_content = entrypoint.read() + assert b"python3 -m data_designer.slurm.runtime.entrypoint" not in entrypoint_content + assert b"--retry-plan-sha256" in entrypoint_content + assert b"--effective-resume-mode" in entrypoint_content def test_runtime_bundle_recursively_collects_and_imports_nested_packages( diff --git a/packages/data-designer-slurm/tests/runtime/test_context.py b/packages/data-designer-slurm/tests/runtime/test_context.py index a0d0e7d19..ee16de9f2 100644 --- a/packages/data-designer-slurm/tests/runtime/test_context.py +++ b/packages/data-designer-slurm/tests/runtime/test_context.py @@ -7,14 +7,21 @@ from datetime import datetime, timezone from pathlib import Path from typing import cast +from unittest.mock import Mock +import pytest + +import data_designer.slurm.runtime.context as runtime_context from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.runtime.context import load_allocation_context +from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, + RetryPlan, + RetryShard, RunManifest, SchedulerIdentity, ShardManifest, @@ -22,6 +29,59 @@ ) +def test_runtime_retry_binding_requires_all_arguments() -> None: + with pytest.raises(SlurmRuntimeError, match="binding is incomplete"): + runtime_context._load_retry_plan( + Mock(), + retry_id="retry-0001", + retry_plan_sha256=None, + effective_resume_mode="never", + ) + + +@pytest.mark.parametrize( + ("digest", "resume_mode"), + [ + pytest.param("f" * 64, "never", id="digest"), + pytest.param(None, "always", id="resume-mode"), + ], +) +def test_runtime_retry_binding_rejects_tampered_arguments( + single_node_plan: ResolvedSlurmRunPlan, + digest: str | None, + resume_mode: str, +) -> None: + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=single_node_plan.run_id, + created_at=datetime(2026, 9, 2, tzinfo=timezone.utc), + resolved_plan=ArtifactReference( + path="/workspace/primary/runs/run-single/resolved-plan.json", + sha256=single_node_plan.compute_sha256(), + ), + planned_shards=( + RetryShard( + shard_id="shard-00000", + attempt_id="attempt-0002", + attempt_ordinal=2, + array_task_index=0, + ), + ), + effective_resume_mode="never", + ) + writer = Mock() + writer.load_retry_plan.return_value = retry + + with pytest.raises(SlurmRuntimeError, match="binding differs"): + runtime_context._load_retry_plan( + writer, + retry_id=retry.retry_id, + retry_plan_sha256=retry.compute_sha256() if digest is None else digest, + effective_resume_mode=resume_mode, + ) + + def test_allocation_context_reads_and_updates_state_through_remapped_workspace( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index 45ff94b19..df3439944 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -32,6 +32,8 @@ DeploymentReadiness, EndpointPublicationState, ReadinessState, + RetryPlan, + RetryShard, RunManifest, ShardManifest, SlurmStateWriter, @@ -252,6 +254,48 @@ def finalize_winner(self, *args: object, **kwargs: object) -> None: assert state.winners == [] +def test_controller_binds_retry_resume_mode_to_client_and_workspace(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=context.plan.run_id, + created_at=runtime_case.created_at, + resolved_plan=context.attempt.resolved_plan, + planned_shards=( + RetryShard( + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + attempt_ordinal=context.attempt.attempt_ordinal, + array_task_index=context.shard.array_task_index, + ), + ), + effective_resume_mode="never", + ) + retry_context = replace(context, retry_plan=retry) + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(context.attempt) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + client_steps = FakeClientStepBuilder() + controller = OneNodeAllocationController( + retry_context, + runtime_proxy_path=context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=client_steps, + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + result = controller.run() + + assert result.state is AttemptLifecycleState.SUCCEEDED + assert client_steps.retry_resume_modes == ["never", "never"] + assert state.dataset_workspace_modes == ["never"] + + def test_preflight_failure_starts_no_process_and_fails_attempt(runtime_case: RuntimeCase) -> None: clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) state = FakeStateStore(runtime_case.context.attempt) diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 86439bcc8..17c12005d 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -179,7 +179,7 @@ def _patch_runtime_context( runtime_case: RuntimeCase, state: FakeStateStore, ) -> None: - monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args: (runtime_case.context, state)) + monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args, **kwargs: (runtime_case.context, state)) monkeypatch.setattr(entrypoint, "get_container_path", lambda plan, path, **kwargs: path) monkeypatch.setenv("SLURM_JOB_GPUS", "0") diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 03dbd24df..bbe2b39d1 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -436,6 +436,54 @@ def test_status_expires_unrequeued_preemption_and_cancel_skips_terminal_job( assert launcher.cancellations == [] +def test_status_does_not_expire_requeue_window_from_another_attempt_clock( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={"array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2})} + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + current_time = [datetime(2026, 9, 8, tzinfo=timezone.utc)] + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: current_time[0], + package_version="0.9.2", + ) + result = service.execute(authored, source_root=tmp_path) + preempted = SchedulerIdentity(array_job_id=42, array_task_id=0) + running = SchedulerIdentity(array_job_id=42, array_task_id=1) + launcher.queue_entries = (SlurmQueueEntry(job_identity=running, state=SchedulerState.RUNNING),) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=preempted, + state=SchedulerState.PREEMPTED, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=0), + ), + ) + service.status(result.run_id) + writer = SlurmStateWriter(tmp_path, result.run_id) + other_attempt = writer.load_attempt("shard-00001", "attempt-0001") + writer.update_attempt(other_attempt.model_copy(update={"updated_at": current_time[0] + timedelta(minutes=10)})) + + current_time[0] += timedelta(minutes=1) + still_requeueing = service.status(result.run_id) + current_time[0] += timedelta(minutes=1) + launcher.queue_entries = ( + SlurmQueueEntry(job_identity=preempted, state=SchedulerState.PENDING), + SlurmQueueEntry(job_identity=running, state=SchedulerState.RUNNING), + ) + requeued = service.status(result.run_id) + + assert still_requeueing.shards[0].attempts[0].attempt.state is AttemptLifecycleState.PENDING + assert requeued.shards[0].attempts[0].attempt.state is AttemptLifecycleState.PENDING + + def test_auto_gpu_resolution_rejects_mixed_node_shapes( tmp_path: Path, profile_catalog: SlurmProfileCatalog, 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..441744846 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -0,0 +1,1780 @@ +# 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 threading import Event +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 +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 +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, + StateContractError, + 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_inputs import CollectionInputResolver +from data_designer.slurm.state.collection_storage import CollectionStorage +from data_designer.slurm.state.collection_validation import validate_collection_inputs +from data_designer.slurm.state.collection_worker import SlurmCollectionWorker +from data_designer.slurm.state.reader import StateReader +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) + retry_plan = case.writer.load_retry_plan("retry-0001") + assert retry_plan.effective_resume_mode == "never" + assert retry_plan.planned_shards[0].attempt_id == attempts[0].attempt_id + 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_waits_for_scheduler_visibility_without_a_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 + 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, + 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_stable_plan_provenance_across_distinct_shards( + 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)) + state = StateStorage(case.workspace, case.plan.run_id) + collection_plan = CollectionStorage(state).read_plan(submitted.collection_id) + _, candidates = CollectionInputResolver(state, StateReader(state, case.plan.run_id)).resolve(collection_plan) + + assert len({candidate.shard_id for candidate in candidates}) == len(candidates) + assert len({candidate.files[0].sha256 for candidate in candidates}) == len(candidates) + assert {candidate.provenance_digest for candidate in candidates} == {case.plan.compute_sha256()} + drifted = tuple(candidate.model_copy(update={"provenance_digest": "f" * 64}) for candidate in candidates) + with pytest.raises(StateContractError, match="does not match the resolved plan"): + validate_collection_inputs(case.plan, collection_plan, drifted) + + +def test_collection_refresh_observes_running_during_bulk_merge( + 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) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + runner.set_job_state(5101, queue_state="RUNNING", accounting_state=None) + merge_started = Event() + release_merge = Event() + original_merge = collection_merge.CollectionMerger.merge + + def pause_merge(*args: object, **kwargs: object) -> object: + merge_started.set() + assert release_merge.wait(timeout=3) + return original_merge(*args, **kwargs) + + monkeypatch.setattr(collection_merge.CollectionMerger, "merge", pause_merge) + worker = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + worker_result = executor.submit(worker.run, completed_at=case.created_at + timedelta(minutes=11)) + assert merge_started.wait(timeout=3) + try: + refresh_result = executor.submit( + coordinator.refresh, + observed_at=case.created_at + timedelta(minutes=12), + ).result(timeout=2) + finally: + release_merge.set() + result = worker_result.result(timeout=5) + + assert refresh_result.state is CollectionState.RUNNING + assert result.actual_records == case.plan.invocation.authored.num_records + assert ( + CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id).state + is CollectionState.SUCCEEDED + ) + + +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_waits_for_scheduler_visibility_without_a_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) + runner = FakeSlurmRunner() + scheduler = SlurmCommandClient(runner) + 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 + 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, + 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/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 9a59c86cd..a7865afc4 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -1633,6 +1633,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())]) 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..5e29db0d0 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_submission_recovery.py @@ -0,0 +1,144 @@ +# 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, 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), + ) + + +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 + ) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 17e535c3b..c2c696633 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