From 506de4896653c3da090453ec3163a0613dc0625e Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 2 Sep 2026 16:26:12 -0600 Subject: [PATCH] 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 | 31 +- .../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, 1403 insertions(+), 57 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 47d061ad1..f21e7c282 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 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}$") _MAX_SLURM_INTEGER = (1 << 32) - 1 @@ -86,7 +83,7 @@ def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: ) 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) @@ -107,7 +104,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) @@ -130,7 +127,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))) @@ -162,13 +159,13 @@ def _run(self, command: Sequence[str], *, input_text: str | None = None) -> str: return stdout -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: @@ -184,16 +181,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 57f0fe4a5..39a1062ca 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 0d00e7b3b..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, @@ -145,7 +134,7 @@ def validate_readiness_transition( def reconcile_attempt_observation( attempt: AttemptManifest, - readiness: AttemptReadiness, + readiness: AttemptReadiness | None, scheduler: SchedulerObservation, *, current_time: datetime, @@ -154,16 +143,18 @@ def reconcile_attempt_observation( _require_utc(current_time, "current_time") _require(current_time >= scheduler.observed_at, "current_time cannot precede scheduler observation") _require(current_time >= attempt.updated_at, "current_time cannot precede attempt update") - _require(current_time >= readiness.updated_at, "current_time cannot precede readiness update") - _require(readiness.run_id == attempt.run_id, "readiness run_id does not match attempt") - _require(readiness.shard_id == attempt.shard_id, "readiness shard_id does not match attempt") - _require(readiness.attempt_id == attempt.attempt_id, "readiness attempt_id does not match attempt") + if readiness is not None: + _require(current_time >= readiness.updated_at, "current_time cannot precede readiness update") + _require(readiness.run_id == attempt.run_id, "readiness run_id does not match attempt") + _require(readiness.shard_id == attempt.shard_id, "readiness shard_id does not match attempt") + _require(readiness.attempt_id == attempt.attempt_id, "readiness attempt_id does not match attempt") _require(attempt.scheduler is not None, "attempt has no scheduler identity") _require(scheduler.scheduler == attempt.scheduler, "scheduler identity does not match attempt") _require(scheduler.observed_at >= attempt.created_at, "scheduler observation cannot precede attempt creation") - _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") + 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 @@ -180,7 +171,7 @@ def reconcile_attempt_observation( return EffectiveAttemptState.UNKNOWN if scheduler.state is SchedulerState.UNKNOWN: return EffectiveAttemptState.UNKNOWN - if readiness.state is ReadinessState.FAILED: + if readiness is not None and readiness.state is ReadinessState.FAILED: return EffectiveAttemptState.FAILED if scheduler.state is SchedulerState.PENDING: return EffectiveAttemptState.PENDING 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 618c156ba..fa8dac9e7 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" @@ -82,6 +84,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 @@ -270,6 +275,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 f80247355..8375421d3 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