diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 70866a129..0fd4ec5ff 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -35,7 +35,10 @@ style = "pep440" bump = true [tool.hatch.metadata.hooks.uv-dynamic-versioning] -dependencies = ["data-designer=={{ version }}"] +dependencies = [ + "data-designer=={{ version }}", + "pydantic>=2.9.2,<3", +] [tool.hatch.build.targets.wheel] packages = ["src/data_designer"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py new file mode 100644 index 000000000..087b2a678 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared immutable contracts for Slurm planning and runtime state.""" + +from __future__ import annotations + +import hashlib +import json +import posixpath +from typing import Annotated, Literal + +from pydantic import ( + BaseModel, + ConfigDict, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) + +Identifier = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", + ), +] +ModelAlias = str +ShardId = Annotated[str, StringConstraints(pattern=r"^shard-[0-9]{5,}$")] +AttemptId = Annotated[str, StringConstraints(pattern=r"^attempt-[0-9]{4,}$")] +SchemaVersion = Literal[1] +Sha256Digest = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + + +class ContractValue(BaseModel): + """Base for strict immutable values shared across Slurm boundaries.""" + + model_config = ConfigDict( + extra="forbid", + frozen=True, + allow_inf_nan=False, + protected_namespaces=(), + strict=True, + validate_default=True, + ) + + +class ContractRecord(ContractValue): + """Base for immutable, explicitly versioned Slurm records.""" + + schema_version: SchemaVersion + + def serialize_canonical_json(self) -> bytes: + """Serialize the record to stable compact UTF-8 bytes.""" + return canonical_json(self.model_dump(mode="json")) + + def serialize_json(self) -> str: + """Serialize the record to deterministic persisted text.""" + return pretty_json(self.model_dump(mode="json")) + + def compute_sha256(self) -> Sha256Digest: + """Compute the digest of the exact bytes written by ``serialize_json``.""" + return hashlib.sha256(self.serialize_json().encode("utf-8")).hexdigest() + + +def canonical_json(value: object) -> bytes: + """Serialize a JSON-compatible value to stable UTF-8 bytes.""" + return json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + +def pretty_json(value: object) -> str: + """Serialize a JSON-compatible value to deterministic persisted text.""" + return ( + json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + indent=2, + sort_keys=True, + ) + + "\n" + ) + + +def compute_sha256(value: object) -> Sha256Digest: + """Compute the canonical JSON digest of a JSON-compatible value.""" + return hashlib.sha256(canonical_json(value)).hexdigest() + + +def validate_absolute_path(value: str) -> str: + """Validate a normalized, absolute POSIX path below the filesystem root.""" + if not value.startswith("/"): + raise ValueError("path must be absolute") + if value.startswith("//"): + raise ValueError("path must have exactly one leading slash") + if value == "/": + raise ValueError("path must not be the filesystem root") + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("path must not contain control characters") + if ".." in value.split("/"): + raise ValueError("path must not contain parent-directory components") + if posixpath.normpath(value) != value: + raise ValueError("path must be normalized") + return value + + +def validate_relative_path(value: str) -> str: + """Validate a normalized relative POSIX path without parent traversal.""" + if not value or value.startswith("/"): + raise ValueError("path must be a non-empty relative path") + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("path must not contain control characters") + if ".." in value.split("/"): + raise ValueError("path must not contain parent-directory components") + if posixpath.normpath(value) != value or value == ".": + raise ValueError("path must be normalized") + return value + + +class ArtifactReference(ContractValue): + """Immutable reference to persisted file bytes and their digest.""" + + path: str + sha256: Sha256Digest + + _path_is_absolute = field_validator("path")(validate_absolute_path) + + +class RecordRange(ContractValue): + """Half-open global record range assigned to one shard.""" + + start_index: NonNegativeInt + end_index_exclusive: PositiveInt + + @property + def record_count(self) -> int: + return self.end_index_exclusive - self.start_index + + @model_validator(mode="after") + def validate_bounds(self) -> RecordRange: + if self.end_index_exclusive <= self.start_index: + raise ValueError("end_index_exclusive must be greater than start_index") + return self + + +class ResumeWorkspace(ContractValue): + """Canonical shard-owned dataset workspace.""" + + path: str + + _path_is_absolute = field_validator("path")(validate_absolute_path) + + +__all__ = [ + "ArtifactReference", + "AttemptId", + "ContractRecord", + "ContractValue", + "Identifier", + "ModelAlias", + "RecordRange", + "ResumeWorkspace", + "SchemaVersion", + "Sha256Digest", + "ShardId", + "canonical_json", + "compute_sha256", + "pretty_json", + "validate_absolute_path", + "validate_relative_path", +] 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 new file mode 100644 index 000000000..2d7eb81d8 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public state contracts for the optional Slurm execution package.""" + +from __future__ import annotations + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + ContractRecord, + ContractValue, + Identifier, + ModelAlias, + RecordRange, + ResumeWorkspace, + Sha256Digest, + ShardId, +) +from data_designer.slurm.state.base import ( + SchedulerIdentity, + StateRecord, + StateValue, +) +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + ShardManifest, +) +from data_designer.slurm.state.outputs import ( + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + CollectionPlan, + CollectionShard, + ShardWinner, +) +from data_designer.slurm.state.readiness import ( + AttemptReadiness, + DeploymentReadiness, + EndpointPublicationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, + ReasonCode, +) +from data_designer.slurm.state.reconciliation import ( + reconcile_attempt_observation, + validate_readiness_transition, +) +from data_designer.slurm.state.scheduler import ( + EffectiveAttemptState, + SchedulerObservation, + SchedulerState, +) +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_manifest, + validate_attempt_set, + validate_attempt_transition, + validate_collection_plan, + validate_scheduler_observation_transition, + validate_shard_manifest, + validate_shard_set, + validate_shard_winner, +) + +__all__ = [ + "ArtifactReference", + "AttemptLifecycleState", + "AttemptManifest", + "AttemptId", + "AttemptReadiness", + "AttemptTerminalClassification", + "CandidateOutcome", + "CandidateOutputFile", + "CandidateOutputManifest", + "CollectionPlan", + "CollectionShard", + "ContractRecord", + "ContractValue", + "DeploymentReadiness", + "EffectiveAttemptState", + "EndpointPublicationState", + "Identifier", + "ModelAlias", + "ProbeEvidence", + "ProbeOutcome", + "ReadinessState", + "ReasonCode", + "RecordRange", + "RunManifest", + "ResumeWorkspace", + "SchedulerIdentity", + "SchedulerObservation", + "SchedulerState", + "Sha256Digest", + "ShardManifest", + "ShardId", + "ShardWinner", + "StateContractError", + "StateRecord", + "StateValue", + "reconcile_attempt_observation", + "validate_attempt_manifest", + "validate_attempt_set", + "validate_attempt_transition", + "validate_collection_plan", + "validate_readiness_transition", + "validate_scheduler_observation_transition", + "validate_shard_manifest", + "validate_shard_set", + "validate_shard_winner", +] 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 new file mode 100644 index 000000000..e9825bd95 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timedelta + +from pydantic import NonNegativeInt, PositiveInt + +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractRecord, + ContractValue, + Identifier, + Sha256Digest, + validate_absolute_path, + validate_relative_path, +) + +StateValue = ContractValue +StateRecord = ContractRecord + + +def validate_utc_timestamp(value: datetime) -> datetime: + """Validate that a timestamp is timezone-aware and expressed in UTC.""" + if value.tzinfo is None or value.utcoffset() is None: + raise ValueError("timestamp must include timezone information") + if value.utcoffset() != timedelta(0): + raise ValueError("timestamp must be in UTC") + return value + + +def validate_optional_utc_timestamp(value: datetime | None) -> datetime | None: + """Validate an optional timestamp when it is present.""" + if value is None: + return None + return validate_utc_timestamp(value) + + +class SchedulerIdentity(StateValue): + """Slurm array job and task identity assigned to one attempt.""" + + array_job_id: PositiveInt + array_task_id: NonNegativeInt + + +__all__ = [ + "ArtifactReference", + "Identifier", + "SchedulerIdentity", + "Sha256Digest", + "StateRecord", + "StateValue", + "validate_absolute_path", + "validate_optional_utc_timestamp", + "validate_relative_path", + "validate_utc_timestamp", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/state/execution.py new file mode 100644 index 000000000..927401001 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/execution.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from pydantic import NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + Identifier, + RecordRange, + ResumeWorkspace, + ShardId, +) +from data_designer.slurm.state.base import ( + SchedulerIdentity, + StateRecord, + validate_utc_timestamp, +) + + +class RunManifest(StateRecord): + """Identity and immutable authored/resolved inputs for a Slurm run.""" + + run_id: Identifier + created_at: datetime + authored_config: ArtifactReference + resolved_plan: ArtifactReference + shard_count: PositiveInt + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + + +class ShardManifest(StateRecord): + """Stable shard identity and planner-owned input partition reference.""" + + run_id: Identifier + shard_id: ShardId + shard_index: NonNegativeInt + record_range: RecordRange + input_partition: ArtifactReference | None = None + resume_workspace: ResumeWorkspace + created_at: datetime + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + + +class AttemptLifecycleState(str, Enum): + CREATED = "created" + SUBMITTED = "submitted" + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class AttemptTerminalClassification(str, Enum): + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + TIMED_OUT = "timed_out" + NODE_FAILED = "node_failed" + PREEMPTED = "preempted" + REQUEUED = "requeued" + OUT_OF_MEMORY = "out_of_memory" + UNKNOWN = "unknown" + + +class AttemptManifest(StateRecord): + """Attempt identity, lifecycle, scheduler identity, and output reference.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + attempt_ordinal: PositiveInt + resolved_plan: ArtifactReference + state: AttemptLifecycleState + terminal_classification: AttemptTerminalClassification | None = None + scheduler: SchedulerIdentity | None = None + candidate_output: ArtifactReference | None = None + created_at: datetime + updated_at: datetime + + _timestamps_are_utc = field_validator("created_at", "updated_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_lifecycle(self) -> AttemptManifest: + if self.updated_at < self.created_at: + raise ValueError("updated_at must not precede created_at") + + if self.state is not AttemptLifecycleState.CREATED and self.scheduler is None: + raise ValueError("submitted and later attempts require scheduler identity") + + terminal = self.state in { + AttemptLifecycleState.SUCCEEDED, + AttemptLifecycleState.FAILED, + } + if terminal != (self.terminal_classification is not None): + raise ValueError("terminal classification must be present exactly for terminal attempts") + + if self.state is AttemptLifecycleState.SUCCEEDED: + if self.terminal_classification is not AttemptTerminalClassification.SUCCEEDED: + raise ValueError("successful attempts require a succeeded terminal classification") + if self.candidate_output is None: + raise ValueError("successful attempts require a candidate output reference") + elif self.terminal_classification is AttemptTerminalClassification.SUCCEEDED: + raise ValueError("failed attempts cannot have a succeeded terminal classification") + + return self 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 new file mode 100644 index 000000000..756e00b44 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime +from enum import Enum +from typing import Literal + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + Identifier, + Sha256Digest, + ShardId, + validate_absolute_path, + validate_relative_path, +) +from data_designer.slurm.state.base import ( + StateRecord, + StateValue, + validate_utc_timestamp, +) + + +class CandidateOutcome(str, Enum): + COMPLETE = "complete" + PARTIAL = "partial" + EMPTY = "empty" + + +class CandidateOutputFile(StateValue): + """One immutable file contained in a candidate output.""" + + relative_path: str + sha256: Sha256Digest + byte_size: NonNegativeInt + record_count: NonNegativeInt + + _relative_path_is_safe = field_validator("relative_path")(validate_relative_path) + + +class CandidateOutputManifest(StateRecord): + """Attempt-local output that may become the immutable shard winner.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + attempt_ordinal: PositiveInt + created_at: datetime + dataset_path: str + requested_records: PositiveInt + actual_records: NonNegativeInt + outcome: CandidateOutcome + files: tuple[CandidateOutputFile, ...] + dataset_schema_digest: Sha256Digest + provenance_digest: Sha256Digest + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + _dataset_path_is_safe = field_validator("dataset_path")(validate_absolute_path) + + @property + def winner_eligible(self) -> bool: + """Whether this complete candidate may be published as the shard winner.""" + return self.outcome is CandidateOutcome.COMPLETE + + @model_validator(mode="after") + def validate_output(self) -> CandidateOutputManifest: + if self.actual_records > self.requested_records: + raise ValueError("actual_records must not exceed requested_records") + + expected_outcome = ( + CandidateOutcome.EMPTY + if self.actual_records == 0 + else CandidateOutcome.COMPLETE + if self.actual_records == self.requested_records + else CandidateOutcome.PARTIAL + ) + if self.outcome is not expected_outcome: + raise ValueError(f"outcome must be {expected_outcome.value!r} for this record count") + + relative_paths = [output_file.relative_path for output_file in self.files] + if len(relative_paths) != len(set(relative_paths)): + raise ValueError("candidate output file paths must be unique") + if self.actual_records > 0 and not self.files: + raise ValueError("non-empty candidate outputs require at least one file") + if sum(output_file.record_count for output_file in self.files) != self.actual_records: + raise ValueError("candidate output file record counts must equal actual_records") + return self + + +class ShardWinner(StateRecord): + """Immutable pointer selecting exactly one candidate for a shard.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + attempt_ordinal: PositiveInt + candidate_manifest: ArtifactReference + published_at: datetime + + _published_at_is_utc = field_validator("published_at")(validate_utc_timestamp) + + +class CollectionShard(StateValue): + """Winner manifest selected for one shard in a collection plan.""" + + shard_id: ShardId + winner_manifest: ArtifactReference + + +class CollectionPlan(StateRecord): + """Immutable inputs and destinations for deterministic collection.""" + + collection_id: Identifier + run_id: Identifier + created_at: datetime + resolved_plan: ArtifactReference + planned_shards: tuple[CollectionShard, ...] = Field(min_length=1) + host_destination: str + container_destination: str + num_partitions: PositiveInt + overwrite: Literal[False] = False + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + _destinations_are_safe = field_validator("host_destination", "container_destination")(validate_absolute_path) + + @model_validator(mode="after") + def validate_shards(self) -> CollectionPlan: + shard_ids = [shard.shard_id for shard in self.planned_shards] + winner_paths = [shard.winner_manifest.path for shard in self.planned_shards] + if len(shard_ids) != len(set(shard_ids)): + raise ValueError("collection shard IDs must be unique") + if len(winner_paths) != len(set(winner_paths)): + raise ValueError("collection winner manifest paths must be unique") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py new file mode 100644 index 000000000..a7f538165 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +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 AttemptId, Identifier, ModelAlias, ShardId +from data_designer.slurm.state.base import StateRecord, StateValue, validate_utc_timestamp + +ReasonCode = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=64, + pattern=r"^[a-z][a-z0-9_]*$", + ), +] + + +class ReadinessState(str, Enum): + PENDING = "pending" + STARTING = "starting" + READY = "ready" + FAILED = "failed" + STOPPED = "stopped" + + +class EndpointPublicationState(str, Enum): + PENDING = "pending" + PUBLISHED = "published" + FAILED = "failed" + + +class ProbeOutcome(str, Enum): + SUCCESS = "success" + FAILURE = "failure" + + +class ProbeEvidence(StateValue): + """Bounded, redacted evidence from one readiness probe.""" + + observed_at: datetime + outcome: ProbeOutcome + reason_code: ReasonCode + redacted_message: Annotated[str, StringConstraints(max_length=512)] + + _observed_at_is_utc = field_validator("observed_at")(validate_utc_timestamp) + + @field_validator("redacted_message") + @classmethod + def validate_redacted_message(cls, value: str) -> str: + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("redacted_message must not contain control characters") + return value + + +class DeploymentReadiness(StateValue): + """Readiness state for one authored-order model deployment.""" + + deployment_id: Identifier + model_alias: ModelAlias + state: ReadinessState + expected_backends: PositiveInt + ready_backends: NonNegativeInt + endpoint_publication: EndpointPublicationState + last_probe: ProbeEvidence | None = None + + @model_validator(mode="after") + def validate_counts_and_state(self) -> DeploymentReadiness: + if self.ready_backends > self.expected_backends: + raise ValueError("ready_backends must not exceed expected_backends") + if self.endpoint_publication is EndpointPublicationState.FAILED and self.state not in { + ReadinessState.FAILED, + ReadinessState.STOPPED, + }: + raise ValueError("failed endpoint publication requires a failed or stopped deployment") + + if self.state is ReadinessState.PENDING: + if self.ready_backends != 0: + raise ValueError("pending deployments cannot have ready backends") + if self.endpoint_publication is not EndpointPublicationState.PENDING: + raise ValueError("pending deployments require pending endpoint publication") + elif self.state is ReadinessState.STARTING: + if ( + self.ready_backends == self.expected_backends + and self.endpoint_publication is EndpointPublicationState.PUBLISHED + ): + raise ValueError("fully ready published deployments must use the ready state") + elif self.state is ReadinessState.READY: + if self.ready_backends != self.expected_backends: + raise ValueError("ready deployments require every expected backend") + if self.endpoint_publication is not EndpointPublicationState.PUBLISHED: + raise ValueError("ready deployments require a published endpoint") + elif self.state is ReadinessState.STOPPED and self.ready_backends != 0: + raise ValueError("stopped deployments cannot have ready backends") + return self + + +class AttemptReadiness(StateRecord): + """Revisioned readiness snapshot for all deployments in one attempt.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + revision: PositiveInt + updated_at: datetime + state: ReadinessState + deployments: tuple[DeploymentReadiness, ...] = Field(min_length=1) + + _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_deployments(self) -> AttemptReadiness: + deployment_ids = [deployment.deployment_id for deployment in self.deployments] + model_aliases = [deployment.model_alias for deployment in self.deployments] + if len(deployment_ids) != len(set(deployment_ids)): + raise ValueError("deployment IDs must be unique") + if len(model_aliases) != len(set(model_aliases)): + raise ValueError("model aliases must be unique") + if any( + deployment.last_probe is not None and deployment.last_probe.observed_at > self.updated_at + for deployment in self.deployments + ): + raise ValueError("probe observations must not be later than the readiness snapshot") + + deployment_states = tuple(deployment.state for deployment in self.deployments) + if self.state is ReadinessState.PENDING: + if any(state is not ReadinessState.PENDING for state in deployment_states): + raise ValueError("pending attempts require every deployment to be pending") + elif self.state is ReadinessState.STARTING: + if any(state in {ReadinessState.FAILED, ReadinessState.STOPPED} for state in deployment_states): + raise ValueError("starting attempts cannot contain failed or stopped deployments") + if all(state is ReadinessState.READY for state in deployment_states): + raise ValueError("attempts with every deployment ready must use the ready state") + elif self.state is ReadinessState.READY: + if any(state is not ReadinessState.READY for state in deployment_states): + raise ValueError("ready attempts require every deployment to be ready") + elif self.state is ReadinessState.FAILED: + if not any(state is ReadinessState.FAILED for state in deployment_states): + raise ValueError("failed attempts require at least one failed deployment") + elif any(state is not ReadinessState.STOPPED for state in deployment_states): + raise ValueError("stopped attempts require every deployment to be stopped") + return 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 new file mode 100644 index 000000000..aacfacc65 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timedelta + +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest +from data_designer.slurm.state.readiness import ( + AttemptReadiness, + EndpointPublicationState, + ReadinessState, +) +from data_designer.slurm.state.scheduler import ( + EffectiveAttemptState, + SchedulerObservation, + SchedulerState, +) +from data_designer.slurm.state.validation import StateContractError + +_ALLOWED_READINESS_TRANSITIONS: dict[ReadinessState, frozenset[ReadinessState]] = { + ReadinessState.PENDING: frozenset( + {ReadinessState.PENDING, ReadinessState.STARTING, ReadinessState.FAILED, ReadinessState.STOPPED} + ), + ReadinessState.STARTING: frozenset( + {ReadinessState.STARTING, ReadinessState.READY, ReadinessState.FAILED, ReadinessState.STOPPED} + ), + ReadinessState.READY: frozenset({ReadinessState.READY, ReadinessState.FAILED, ReadinessState.STOPPED}), + ReadinessState.FAILED: frozenset({ReadinessState.FAILED, ReadinessState.STOPPED}), + ReadinessState.STOPPED: frozenset({ReadinessState.STOPPED}), +} + +_ALLOWED_ENDPOINT_TRANSITIONS: dict[EndpointPublicationState, frozenset[EndpointPublicationState]] = { + EndpointPublicationState.PENDING: frozenset( + { + EndpointPublicationState.PENDING, + EndpointPublicationState.PUBLISHED, + EndpointPublicationState.FAILED, + } + ), + EndpointPublicationState.PUBLISHED: frozenset({EndpointPublicationState.PUBLISHED}), + 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, + current: AttemptReadiness, +) -> AttemptReadiness: + """Validate monotonic revision, identity, order, and readiness transitions.""" + _require(previous.run_id == current.run_id, "readiness run_id cannot change") + _require(previous.shard_id == current.shard_id, "readiness shard_id cannot change") + _require(previous.attempt_id == current.attempt_id, "readiness attempt_id cannot change") + _require(current.revision == previous.revision + 1, "readiness revision must increase by exactly one") + _require(current.updated_at >= previous.updated_at, "readiness updated_at cannot move backward") + _require( + current.state in _ALLOWED_READINESS_TRANSITIONS[previous.state], + f"attempt readiness cannot move from {previous.state.value} to {current.state.value}", + ) + _require( + len(previous.deployments) == len(current.deployments), + "readiness deployment count cannot change", + ) + + for old_deployment, new_deployment in zip(previous.deployments, current.deployments, strict=True): + _require( + old_deployment.deployment_id == new_deployment.deployment_id, + "readiness deployment order or ID cannot change", + ) + _require( + old_deployment.model_alias == new_deployment.model_alias, + "readiness deployment model alias cannot change", + ) + _require( + old_deployment.expected_backends == new_deployment.expected_backends, + "readiness expected backend count cannot change", + ) + _require( + new_deployment.state in _ALLOWED_READINESS_TRANSITIONS[old_deployment.state], + ( + f"deployment {old_deployment.deployment_id!r} cannot move from " + f"{old_deployment.state.value} to {new_deployment.state.value}" + ), + ) + _require( + new_deployment.endpoint_publication in _ALLOWED_ENDPOINT_TRANSITIONS[old_deployment.endpoint_publication], + f"deployment {old_deployment.deployment_id!r} endpoint publication cannot move backward", + ) + old_probe = old_deployment.last_probe + new_probe = new_deployment.last_probe + if old_probe is not None: + if new_probe is None: + raise StateContractError( + f"deployment {old_deployment.deployment_id!r} probe evidence cannot be removed" + ) + _require( + new_probe.observed_at >= old_probe.observed_at, + f"deployment {old_deployment.deployment_id!r} probe observation cannot move backward", + ) + return current + + +def reconcile_attempt_observation( + attempt: AttemptManifest, + readiness: AttemptReadiness, + scheduler: SchedulerObservation, + *, + current_time: datetime, +) -> EffectiveAttemptState: + """Apply scheduler terminal precedence without treating readiness as success.""" + _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") + _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 scheduler.state in _SCHEDULER_FAILURE_STATES: + return EffectiveAttemptState.FAILED + if attempt.state is AttemptLifecycleState.FAILED: + return EffectiveAttemptState.FAILED + if attempt.state is AttemptLifecycleState.SUCCEEDED: + return EffectiveAttemptState.SUCCEEDED + + if scheduler.state is SchedulerState.COMPLETED: + return EffectiveAttemptState.FAILED + if scheduler.state is SchedulerState.ACCOUNTING_LAG: + deadline = scheduler.reconciliation_deadline + _require(deadline is not None, "accounting lag has no reconciliation deadline") + if current_time <= deadline: + return EffectiveAttemptState.ACCOUNTING_LAG + return EffectiveAttemptState.UNKNOWN + if scheduler.state is SchedulerState.UNKNOWN: + return EffectiveAttemptState.UNKNOWN + if readiness.state is ReadinessState.FAILED: + return EffectiveAttemptState.FAILED + if scheduler.state is SchedulerState.PENDING: + return EffectiveAttemptState.PENDING + if scheduler.state is SchedulerState.RUNNING: + return EffectiveAttemptState.RUNNING + raise AssertionError(f"unhandled scheduler state: {scheduler.state}") # pragma: no cover + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise StateContractError(message) + + +def _require_utc(value: datetime, field_name: str) -> None: + _require( + value.tzinfo is not None and value.utcoffset() == timedelta(0), + f"{field_name} must be in UTC", + ) 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 new file mode 100644 index 000000000..3245c75fa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from pydantic import field_validator, model_validator + +from data_designer.slurm.state.base import ( + SchedulerIdentity, + StateRecord, + validate_optional_utc_timestamp, + validate_utc_timestamp, +) + + +class SchedulerState(str, Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + TIMED_OUT = "timed_out" + NODE_FAILED = "node_failed" + PREEMPTED = "preempted" + REQUEUED = "requeued" + OUT_OF_MEMORY = "out_of_memory" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" + + +class SchedulerObservation(StateRecord): + """Normalized scheduler observation used for deterministic reconciliation.""" + + scheduler: SchedulerIdentity + observed_at: datetime + state: SchedulerState + reconciliation_deadline: datetime | None = None + + _observed_at_is_utc = field_validator("observed_at")(validate_utc_timestamp) + _reconciliation_deadline_is_utc = field_validator("reconciliation_deadline")(validate_optional_utc_timestamp) + + @model_validator(mode="after") + 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") + return self + + +class EffectiveAttemptState(str, Enum): + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" 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 new file mode 100644 index 000000000..30fdbc066 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + ShardManifest, +) +from data_designer.slurm.state.outputs import ( + CandidateOutputManifest, + CollectionPlan, + ShardWinner, +) +from data_designer.slurm.state.scheduler import SchedulerObservation, SchedulerState + +_ATTEMPT_STATE_ORDER = { + AttemptLifecycleState.CREATED: 0, + AttemptLifecycleState.SUBMITTED: 1, + AttemptLifecycleState.PENDING: 2, + AttemptLifecycleState.RUNNING: 3, +} +_TERMINAL_ATTEMPT_STATES = frozenset( + { + AttemptLifecycleState.SUCCEEDED, + AttemptLifecycleState.FAILED, + } +) + + +class StateContractError(ValueError): + """Raised when related state records violate a cross-record contract.""" + + +def validate_shard_manifest(run: RunManifest, shard: ShardManifest) -> ShardManifest: + """Validate a shard against the immutable run identity and bounds.""" + _require(shard.run_id == run.run_id, "shard run_id does not match run manifest") + _require(shard.shard_index < run.shard_count, "shard_index is outside the run shard count") + _require(shard.created_at >= run.created_at, "shard creation cannot precede run creation") + return shard + + +def validate_shard_set( + run: RunManifest, + shards: tuple[ShardManifest, ...], +) -> tuple[ShardManifest, ...]: + """Validate ordered run state shards without duplicating resolved-plan coverage.""" + _require(len(shards) == run.shard_count, "shard set must include exactly the run shard count") + for shard in shards: + validate_shard_manifest(run, shard) + + shard_ids = tuple(shard.shard_id for shard in shards) + resume_workspace_paths = tuple(shard.resume_workspace.path for shard in shards) + shard_indices = tuple(shard.shard_index for shard in shards) + _require(len(set(shard_ids)) == len(shards), "shard IDs must be unique") + _require(len(set(resume_workspace_paths)) == len(shards), "resume workspace paths must be unique") + _require( + shard_indices == tuple(range(run.shard_count)), + "shards must be ordered by a complete zero-based shard index", + ) + for previous, current in zip(shards, shards[1:]): + _require( + current.record_range.start_index >= previous.record_range.end_index_exclusive, + "shard record ranges must not overlap", + ) + return shards + + +def validate_attempt_manifest( + run: RunManifest, + shard: ShardManifest, + attempt: AttemptManifest, +) -> AttemptManifest: + """Validate an attempt against its run, shard, and resolved plan.""" + validate_shard_manifest(run, shard) + _require(attempt.run_id == run.run_id, "attempt run_id does not match run manifest") + _require(attempt.shard_id == shard.shard_id, "attempt shard_id does not match shard manifest") + _require( + attempt.resolved_plan == run.resolved_plan, + "attempt resolved plan does not match run manifest", + ) + _require(attempt.created_at >= shard.created_at, "attempt creation cannot precede shard creation") + return attempt + + +def validate_attempt_set( + run: RunManifest, + shards: tuple[ShardManifest, ...], + attempts: tuple[AttemptManifest, ...], +) -> tuple[AttemptManifest, ...]: + """Validate attempt identities and scheduler ownership across a run.""" + validate_shard_set(run, shards) + shard_by_id = {shard.shard_id: shard for shard in shards} + + for attempt in attempts: + shard = shard_by_id.get(attempt.shard_id) + _require(shard is not None, f"attempt references unknown shard {attempt.shard_id!r}") + validate_attempt_manifest(run, shard, attempt) + + attempt_ids = tuple((attempt.shard_id, attempt.attempt_id) for attempt in attempts) + shard_ordinals = tuple((attempt.shard_id, attempt.attempt_ordinal) for attempt in attempts) + scheduler_identities = tuple(attempt.scheduler for attempt in attempts if attempt.scheduler is not None) + _require(len(set(attempt_ids)) == len(attempts), "attempt IDs must be unique within each shard") + _require( + len(set(shard_ordinals)) == len(attempts), + "attempt ordinals must be unique within each shard", + ) + _require( + len(set(scheduler_identities)) == len(scheduler_identities), + "scheduler identities must be unique across attempts", + ) + return attempts + + +def validate_attempt_transition( + previous: AttemptManifest, + current: AttemptManifest, +) -> AttemptManifest: + """Validate immutable identity and monotonic lifecycle updates for an attempt.""" + for field_name in ( + "run_id", + "shard_id", + "attempt_id", + "attempt_ordinal", + "resolved_plan", + "created_at", + ): + _require( + getattr(previous, field_name) == getattr(current, field_name), + f"attempt {field_name} cannot change", + ) + _require(current.updated_at >= previous.updated_at, "attempt updated_at cannot move backward") + + if previous.state in _TERMINAL_ATTEMPT_STATES: + _require(current == previous, "terminal attempt manifests are immutable") + return current + + if current.state not in _TERMINAL_ATTEMPT_STATES: + _require( + _ATTEMPT_STATE_ORDER[current.state] >= _ATTEMPT_STATE_ORDER[previous.state], + f"attempt state cannot move from {previous.state.value} to {current.state.value}", + ) + if previous.scheduler is not None: + _require(current.scheduler == previous.scheduler, "attempt scheduler identity cannot change") + if previous.candidate_output is not None: + _require(current.candidate_output == previous.candidate_output, "attempt candidate output cannot change") + return current + + +def validate_scheduler_observation_transition( + previous: SchedulerObservation, + current: SchedulerObservation, +) -> SchedulerObservation: + """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 previous.state is SchedulerState.ACCOUNTING_LAG: + deadline = previous.reconciliation_deadline + _require(deadline is not None, "accounting lag has no reconciliation deadline") + if current.state is SchedulerState.ACCOUNTING_LAG: + _require( + current.reconciliation_deadline == deadline, + "accounting-lag reconciliation deadline cannot change", + ) + elif current.state is SchedulerState.UNKNOWN: + _require( + current.observed_at > deadline, + "accounting lag cannot become unknown before its reconciliation deadline expires", + ) + return current + + +def validate_shard_winner( + run: RunManifest, + shard: ShardManifest, + attempt: AttemptManifest, + candidate: CandidateOutputManifest, + winner: ShardWinner, + *, + existing_winner: ShardWinner | None = None, +) -> ShardWinner: + """Validate first-writer winner publication without performing persistence.""" + validate_attempt_manifest(run, shard, attempt) + _require(existing_winner is None, "a shard winner is immutable once published") + + for record_name, record_run_id in ( + ("candidate", candidate.run_id), + ("winner", winner.run_id), + ): + _require(record_run_id == run.run_id, f"{record_name} run_id does not match run manifest") + for record_name, record_shard_id in ( + ("candidate", candidate.shard_id), + ("winner", winner.shard_id), + ): + _require(record_shard_id == shard.shard_id, f"{record_name} shard_id does not match shard manifest") + + _require(candidate.attempt_id == attempt.attempt_id, "candidate attempt_id does not match attempt") + _require(winner.attempt_id == attempt.attempt_id, "winner attempt_id does not match attempt") + _require( + candidate.attempt_ordinal == attempt.attempt_ordinal == winner.attempt_ordinal, + "candidate and winner attempt ordinals must match the attempt", + ) + _require(attempt.state is AttemptLifecycleState.SUCCEEDED, "only successful attempts may win") + _require( + attempt.terminal_classification is AttemptTerminalClassification.SUCCEEDED, + "only successfully classified attempts may win", + ) + _require(candidate.winner_eligible, "candidate output does not satisfy winner policy") + _require( + candidate.requested_records == shard.record_range.record_count, + "candidate requested record count does not match shard record range", + ) + _require(candidate.created_at >= attempt.created_at, "candidate creation cannot precede attempt creation") + _require(attempt.updated_at >= candidate.created_at, "attempt completion cannot precede candidate creation") + _require(winner.published_at >= attempt.updated_at, "winner publication cannot precede attempt completion") + _require( + winner.candidate_manifest.sha256 == candidate.compute_sha256(), + "winner candidate digest does not match candidate manifest", + ) + _require( + attempt.candidate_output == winner.candidate_manifest, + "attempt candidate output reference does not match winner", + ) + return winner + + +def validate_collection_plan( + run: RunManifest, + plan: CollectionPlan, + shards: tuple[ShardManifest, ...], + winners: tuple[ShardWinner, ...], +) -> CollectionPlan: + """Validate collection against the exact run shard set and winner digests.""" + validate_shard_set(run, shards) + _require(plan.run_id == run.run_id, "collection run_id does not match run manifest") + _require(plan.resolved_plan == run.resolved_plan, "collection resolved plan does not match run manifest") + _require(plan.created_at >= run.created_at, "collection creation cannot precede run creation") + + expected_shard_ids = tuple(shard.shard_id for shard in shards) + planned_shard_ids = tuple(shard.shard_id for shard in plan.planned_shards) + _require( + planned_shard_ids == expected_shard_ids, + "collection planned shards must exactly match the ordered run shard set", + ) + _require(len(winners) == len(shards), "winner set must include exactly the run shard count") + + winner_by_shard = {winner.shard_id: winner for winner in winners} + _require(len(winner_by_shard) == len(winners), "winner shard IDs must be unique") + _require(set(winner_by_shard) == set(expected_shard_ids), "winner shard IDs must exactly match the run shard set") + for planned_shard in plan.planned_shards: + winner = winner_by_shard[planned_shard.shard_id] + _require(winner.run_id == run.run_id, "winner run_id does not match run manifest") + _require(plan.created_at >= winner.published_at, "collection creation cannot precede winner publication") + _require( + planned_shard.winner_manifest.sha256 == winner.compute_sha256(), + f"winner digest mismatch for shard {planned_shard.shard_id!r}", + ) + return plan + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise StateContractError(message) diff --git a/packages/data-designer-slurm/tests/state/golden/accounting_lag.json b/packages/data-designer-slurm/tests/state/golden/accounting_lag.json new file mode 100644 index 000000000..9c97822f4 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/accounting_lag.json @@ -0,0 +1,10 @@ +{ + "observed_at": "2026-08-18T12:05:02Z", + "reconciliation_deadline": "2026-08-18T12:10:02Z", + "scheduler": { + "array_job_id": 4101, + "array_task_id": 0 + }, + "schema_version": 1, + "state": "accounting_lag" +} diff --git a/packages/data-designer-slurm/tests/state/golden/candidate_output.json b/packages/data-designer-slurm/tests/state/golden/candidate_output.json new file mode 100644 index 000000000..81e9a63f4 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/candidate_output.json @@ -0,0 +1,22 @@ +{ + "actual_records": 100, + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "created_at": "2026-08-18T12:05:00Z", + "dataset_path": "/workspace/runs/run-0001/shards/shard-00000/attempts/attempt-0001/dataset", + "dataset_schema_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + "files": [ + { + "byte_size": 4096, + "record_count": 100, + "relative_path": "part-00000.parquet", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + } + ], + "outcome": "complete", + "provenance_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "requested_records": 100, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000" +} diff --git a/packages/data-designer-slurm/tests/state/golden/collection_plan.json b/packages/data-designer-slurm/tests/state/golden/collection_plan.json new file mode 100644 index 000000000..1d528f0d1 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/collection_plan.json @@ -0,0 +1,23 @@ +{ + "collection_id": "collection-0001", + "container_destination": "/mnt/data/run-0001.parquet", + "created_at": "2026-08-18T12:06:00Z", + "host_destination": "/workspace/runs/run-0001/collection/run-0001.parquet", + "num_partitions": 1, + "overwrite": false, + "planned_shards": [ + { + "shard_id": "shard-00000", + "winner_manifest": { + "path": "/workspace/runs/run-0001/shards/shard-00000/winner.json", + "sha256": "e6efb5fd97b97c36e7dd1c4e5ab61dc954b611fd68d5e85956e357b36d2a8930" + } + } + ], + "resolved_plan": { + "path": "/workspace/runs/run-0001/resolved-plan.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "run_id": "run-0001", + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/state/golden/failed_attempt.json b/packages/data-designer-slurm/tests/state/golden/failed_attempt.json new file mode 100644 index 000000000..244b7977d --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/failed_attempt.json @@ -0,0 +1,20 @@ +{ + "attempt_id": "attempt-0002", + "attempt_ordinal": 2, + "candidate_output": null, + "created_at": "2026-08-18T13:00:00Z", + "resolved_plan": { + "path": "/workspace/runs/run-failed/resolved-plan.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "run_id": "run-failed", + "scheduler": { + "array_job_id": 4201, + "array_task_id": 0 + }, + "schema_version": 1, + "shard_id": "shard-00000", + "state": "failed", + "terminal_classification": "node_failed", + "updated_at": "2026-08-18T13:05:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json b/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json new file mode 100644 index 000000000..444351540 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json @@ -0,0 +1,39 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 2, + "last_probe": { + "observed_at": "2026-08-18T14:02:00Z", + "outcome": "failure", + "reason_code": "backend_starting", + "redacted_message": "One of two backends is ready" + }, + "model_alias": "generator", + "ready_backends": 1, + "state": "starting" + }, + { + "deployment_id": "deployment-00001", + "endpoint_publication": "published", + "expected_backends": 2, + "last_probe": { + "observed_at": "2026-08-18T14:02:00Z", + "outcome": "success", + "reason_code": "backends_ready", + "redacted_message": "All backends are ready" + }, + "model_alias": "judge", + "ready_backends": 2, + "state": "ready" + } + ], + "revision": 4, + "run_id": "run-multi", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "starting", + "updated_at": "2026-08-18T14:02:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/run_manifest.json b/packages/data-designer-slurm/tests/state/golden/run_manifest.json new file mode 100644 index 000000000..38825f78a --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/run_manifest.json @@ -0,0 +1,14 @@ +{ + "authored_config": { + "path": "/workspace/runs/run-0001/authored-config.json", + "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + }, + "created_at": "2026-08-18T12:00:00Z", + "resolved_plan": { + "path": "/workspace/runs/run-0001/resolved-plan.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "run_id": "run-0001", + "schema_version": 1, + "shard_count": 1 +} diff --git a/packages/data-designer-slurm/tests/state/golden/shard_manifest.json b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json new file mode 100644 index 000000000..5e885b143 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json @@ -0,0 +1,18 @@ +{ + "created_at": "2026-08-18T12:00:01Z", + "input_partition": { + "path": "/workspace/runs/run-0001/shards/shard-00000/input-partition.json", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "record_range": { + "end_index_exclusive": 100, + "start_index": 0 + }, + "resume_workspace": { + "path": "/workspace/runs/run-0001/shards/shard-00000/dataset" + }, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "shard_index": 0 +} diff --git a/packages/data-designer-slurm/tests/state/golden/shard_winner.json b/packages/data-designer-slurm/tests/state/golden/shard_winner.json new file mode 100644 index 000000000..20a903c35 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/shard_winner.json @@ -0,0 +1,12 @@ +{ + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "candidate_manifest": { + "path": "/workspace/runs/run-0001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "c639099f85b3cdd2759d3b74c792f06dc4a61a8550622b50ba2a3512c10fdb68" + }, + "published_at": "2026-08-18T12:05:02Z", + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json new file mode 100644 index 000000000..9337193d5 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:01:30Z", + "outcome": "success", + "reason_code": "backend_ready", + "redacted_message": "Backend is ready" + }, + "model_alias": "primary", + "ready_backends": 1, + "state": "ready" + } + ], + "revision": 3, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "ready", + "updated_at": "2026-08-18T12:01:30Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/stale_readiness.json b/packages/data-designer-slurm/tests/state/golden/stale_readiness.json new file mode 100644 index 000000000..511337eb8 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/stale_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0002", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T13:04:00Z", + "outcome": "success", + "reason_code": "backend_ready", + "redacted_message": "Backend was ready before the terminal failure" + }, + "model_alias": "primary", + "ready_backends": 1, + "state": "ready" + } + ], + "revision": 2, + "run_id": "run-failed", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "ready", + "updated_at": "2026-08-18T13:04:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/successful_attempt.json b/packages/data-designer-slurm/tests/state/golden/successful_attempt.json new file mode 100644 index 000000000..03d28f620 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/successful_attempt.json @@ -0,0 +1,23 @@ +{ + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "candidate_output": { + "path": "/workspace/runs/run-0001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "c639099f85b3cdd2759d3b74c792f06dc4a61a8550622b50ba2a3512c10fdb68" + }, + "created_at": "2026-08-18T12:00:02Z", + "resolved_plan": { + "path": "/workspace/runs/run-0001/resolved-plan.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "run_id": "run-0001", + "scheduler": { + "array_job_id": 4101, + "array_task_id": 0 + }, + "schema_version": 1, + "shard_id": "shard-00000", + "state": "succeeded", + "terminal_classification": "succeeded", + "updated_at": "2026-08-18T12:05:01Z" +} diff --git a/packages/data-designer-slurm/tests/state/test_golden_records.py b/packages/data-designer-slurm/tests/state/test_golden_records.py new file mode 100644 index 000000000..fa31859db --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_golden_records.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest + +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerObservation, + ShardManifest, + ShardWinner, + StateRecord, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" +GOLDEN_MODELS: tuple[tuple[str, type[StateRecord]], ...] = ( + ("run_manifest.json", RunManifest), + ("shard_manifest.json", ShardManifest), + ("successful_attempt.json", AttemptManifest), + ("single_node_readiness.json", AttemptReadiness), + ("multi_node_readiness.json", AttemptReadiness), + ("failed_attempt.json", AttemptManifest), + ("stale_readiness.json", AttemptReadiness), + ("accounting_lag.json", SchedulerObservation), + ("candidate_output.json", CandidateOutputManifest), + ("shard_winner.json", ShardWinner), + ("collection_plan.json", CollectionPlan), +) + + +@pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) +def test_golden_record_round_trip_is_deterministic(filename: str, model: type[StateRecord]) -> None: + serialized = (GOLDEN_DIRECTORY / filename).read_text() + + record = model.model_validate_json(serialized) + direct_record = model(**record.model_dump(mode="python")) + + assert direct_record == record + assert record.serialize_json() == serialized + assert model.model_validate_json(record.serialize_canonical_json()) == record + assert model.model_validate_json(record.serialize_json()) == record + assert record.compute_sha256() == hashlib.sha256(serialized.encode("utf-8")).hexdigest() + + +def test_golden_artifact_references_hash_exact_persisted_bytes() -> None: + candidate_bytes = (GOLDEN_DIRECTORY / "candidate_output.json").read_bytes() + winner_bytes = (GOLDEN_DIRECTORY / "shard_winner.json").read_bytes() + attempt = AttemptManifest.model_validate_json((GOLDEN_DIRECTORY / "successful_attempt.json").read_text()) + winner = ShardWinner.model_validate_json(winner_bytes) + collection = CollectionPlan.model_validate_json((GOLDEN_DIRECTORY / "collection_plan.json").read_text()) + + assert attempt.candidate_output is not None + assert attempt.candidate_output.sha256 == hashlib.sha256(candidate_bytes).hexdigest() + assert winner.candidate_manifest.sha256 == hashlib.sha256(candidate_bytes).hexdigest() + assert collection.planned_shards[0].winner_manifest.sha256 == hashlib.sha256(winner_bytes).hexdigest() diff --git a/packages/data-designer-slurm/tests/state/test_model_edges.py b/packages/data-designer-slurm/tests/state/test_model_edges.py new file mode 100644 index 000000000..0babd1f76 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_model_edges.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.state import ( + ArtifactReference, + AttemptManifest, + AttemptReadiness, + CandidateOutputFile, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerObservation, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" + + +def _golden_payload(filename: str) -> dict[str, Any]: + return json.loads((GOLDEN_DIRECTORY / filename).read_text()) + + +@pytest.mark.parametrize( + "path", + ( + "/", + "//", + "//host/path", + "/workspace/control\ncharacter", + "/workspace//unnormalized", + ), +) +def test_absolute_artifact_paths_reject_unsafe_forms(path: str) -> None: + with pytest.raises(ValidationError): + ArtifactReference(path=path, sha256="a" * 64) + + +@pytest.mark.parametrize( + "path", + ( + "", + "/absolute", + "control\ncharacter", + ".", + "directory//unnormalized", + ), +) +def test_relative_output_paths_reject_unsafe_forms(path: str) -> None: + with pytest.raises(ValidationError): + CandidateOutputFile( + relative_path=path, + sha256="a" * 64, + byte_size=0, + record_count=0, + ) + + +@pytest.mark.parametrize( + "created_at", + ( + "2026-08-18T12:00:00", + "2026-08-18T13:00:00+01:00", + ), +) +def test_run_manifest_requires_utc_timestamp(created_at: str) -> None: + payload = _golden_payload("run_manifest.json") + payload["created_at"] = created_at + + with pytest.raises(ValidationError): + RunManifest.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("mutation", "message"), + ( + ({"updated_at": "2026-08-18T12:00:01Z"}, "must not precede"), + ({"terminal_classification": "failed"}, "succeeded terminal classification"), + ({"state": "failed", "terminal_classification": "succeeded"}, "failed attempts"), + ), +) +def test_attempt_manifest_rejects_invalid_terminal_details( + mutation: dict[str, object], + message: str, +) -> None: + payload = _golden_payload("successful_attempt.json") + payload.update(mutation) + + with pytest.raises(ValidationError, match=message): + AttemptManifest.model_validate_json(json.dumps(payload)) + + +def test_candidate_output_rejects_remaining_invalid_inventory_forms() -> None: + payload = _golden_payload("candidate_output.json") + payload["actual_records"] = 101 + with pytest.raises(ValidationError, match="must not exceed"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + payload["outcome"] = "partial" + with pytest.raises(ValidationError, match="outcome must be"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + duplicate_file = dict(payload["files"][0]) + duplicate_file["record_count"] = 0 + payload["files"].append(duplicate_file) + with pytest.raises(ValidationError, match="paths must be unique"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + payload["files"] = [] + with pytest.raises(ValidationError, match="require at least one file"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + +def test_collection_plan_rejects_duplicate_shards_and_winner_paths() -> None: + payload = _golden_payload("collection_plan.json") + duplicate_shard = json.loads(json.dumps(payload["planned_shards"][0])) + duplicate_shard["winner_manifest"]["path"] = "/workspace/runs/run-0001/shards/shard-00001/winner.json" + payload["planned_shards"].append(duplicate_shard) + with pytest.raises(ValidationError, match="shard IDs must be unique"): + CollectionPlan.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("collection_plan.json") + duplicate_winner = json.loads(json.dumps(payload["planned_shards"][0])) + duplicate_winner["shard_id"] = "shard-00001" + payload["planned_shards"].append(duplicate_winner) + with pytest.raises(ValidationError, match="winner manifest paths must be unique"): + CollectionPlan.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("attempt_state", "deployment_updates", "message"), + ( + ("ready", {"ready_backends": 2}, "must not exceed"), + ("ready", {"endpoint_publication": "failed"}, "failed endpoint publication"), + ( + "pending", + {"state": "pending", "ready_backends": 0, "endpoint_publication": "published"}, + "pending endpoint publication", + ), + ("starting", {"state": "starting"}, "must use the ready state"), + ("ready", {"endpoint_publication": "pending"}, "published endpoint"), + ( + "starting", + {"state": "stopped", "ready_backends": 0}, + "cannot contain failed or stopped", + ), + ("starting", {}, "every deployment ready"), + ( + "ready", + {"state": "starting", "ready_backends": 0, "endpoint_publication": "pending"}, + "every deployment to be ready", + ), + ("failed", {}, "at least one failed deployment"), + ), +) +def test_readiness_rejects_remaining_contradictory_states( + attempt_state: str, + deployment_updates: dict[str, object], + message: str, +) -> None: + payload = _golden_payload("single_node_readiness.json") + payload["state"] = attempt_state + payload["deployments"][0].update(deployment_updates) + + with pytest.raises(ValidationError, match=message): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_readiness_rejects_duplicate_deployment_ids() -> None: + payload = _golden_payload("multi_node_readiness.json") + payload["deployments"][1]["deployment_id"] = payload["deployments"][0]["deployment_id"] + + with pytest.raises(ValidationError, match="deployment IDs must be unique"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_readiness_rejects_probe_observation_after_snapshot() -> None: + payload = _golden_payload("single_node_readiness.json") + payload["deployments"][0]["last_probe"]["observed_at"] = "2026-08-18T12:03:00Z" + + with pytest.raises(ValidationError, match="later than the readiness snapshot"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_scheduler_observation_requires_consistent_deadline() -> None: + payload = _golden_payload("accounting_lag.json") + payload["reconciliation_deadline"] = None + with pytest.raises(ValidationError, match="requires a reconciliation deadline"): + SchedulerObservation.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("accounting_lag.json") + payload["reconciliation_deadline"] = "2026-08-18T12:05:01Z" + with pytest.raises(ValidationError, match="must not precede"): + SchedulerObservation.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("accounting_lag.json") + payload["state"] = "running" + with pytest.raises(ValidationError, match="only accounting lag"): + SchedulerObservation.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("accounting_lag.json") + payload["reconciliation_deadline"] = "2026-08-18T13:10:02+01:00" + with pytest.raises(ValidationError, match="must be in UTC"): + SchedulerObservation.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/state/test_records.py b/packages/data-designer-slurm/tests/state/test_records.py new file mode 100644 index 000000000..2eac315bf --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_records.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.state import ( + ArtifactReference, + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + RecordRange, + RunManifest, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" + + +def _golden_payload(filename: str) -> dict[str, Any]: + return json.loads((GOLDEN_DIRECTORY / filename).read_text()) + + +def test_direct_construction_and_json_loading_produce_identical_record() -> None: + record = RunManifest( + schema_version=1, + run_id="run-direct", + created_at=datetime(2026, 8, 18, 12, tzinfo=timezone.utc), + authored_config=ArtifactReference(path="/workspace/run/authored.json", sha256="a" * 64), + resolved_plan=ArtifactReference(path="/workspace/run/plan.json", sha256="b" * 64), + shard_count=2, + ) + + assert RunManifest.model_validate_json(record.serialize_json()) == record + assert record.serialize_json() == RunManifest.model_validate_json(record.serialize_json()).serialize_json() + + +@pytest.mark.parametrize( + ("field_name", "invalid_value"), + ( + ("schema_version", 2), + ("run_id", "contains/slash"), + ("run_id", " leading-space"), + ), +) +def test_run_manifest_rejects_invalid_version_and_identity( + field_name: str, + invalid_value: object, +) -> None: + payload = _golden_payload("run_manifest.json") + payload[field_name] = invalid_value + + with pytest.raises(ValidationError): + RunManifest.model_validate_json(json.dumps(payload)) + + +def test_records_reject_unknown_fields() -> None: + payload = _golden_payload("run_manifest.json") + payload["unreviewed"] = True + + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + RunManifest.model_validate_json(json.dumps(payload)) + + +def test_records_require_explicit_schema_version() -> None: + payload = _golden_payload("run_manifest.json") + payload.pop("schema_version") + + with pytest.raises(ValidationError, match="Field required"): + RunManifest.model_validate_json(json.dumps(payload)) + + +def test_artifact_reference_rejects_invalid_digest_and_path() -> None: + with pytest.raises(ValidationError): + ArtifactReference(path="/workspace/plan.json", sha256="A" * 64) + with pytest.raises(ValidationError, match="parent-directory"): + ArtifactReference(path="/workspace/../private/plan.json", sha256="a" * 64) + with pytest.raises(ValidationError, match="absolute"): + ArtifactReference(path="relative/plan.json", sha256="a" * 64) + + +def test_records_are_frozen() -> None: + record = RunManifest.model_validate_json((GOLDEN_DIRECTORY / "run_manifest.json").read_text()) + + with pytest.raises(ValidationError, match="Instance is frozen"): + record.run_id = "another-run" + + +def test_record_range_is_half_open_and_non_empty() -> None: + record_range = RecordRange(start_index=10, end_index_exclusive=15) + + assert record_range.record_count == 5 + with pytest.raises(ValidationError, match="greater than start_index"): + RecordRange(start_index=10, end_index_exclusive=10) + + +@pytest.mark.parametrize( + ("filename", "mutation"), + ( + ( + "successful_attempt.json", + {"candidate_output": None}, + ), + ( + "successful_attempt.json", + {"state": "running", "terminal_classification": "succeeded"}, + ), + ( + "successful_attempt.json", + {"state": "running", "terminal_classification": None, "scheduler": None}, + ), + ), +) +def test_attempt_lifecycle_invariants(filename: str, mutation: dict[str, object]) -> None: + payload = _golden_payload(filename) + payload.update(mutation) + + with pytest.raises(ValidationError): + AttemptManifest.model_validate_json(json.dumps(payload)) + + +def test_ready_deployment_requires_all_backends_and_published_endpoint() -> None: + payload = _golden_payload("single_node_readiness.json") + deployment = payload["deployments"][0] + deployment["ready_backends"] = 0 + + with pytest.raises(ValidationError, match="every expected backend"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_readiness_rejects_duplicate_deployment_identity() -> None: + payload = _golden_payload("multi_node_readiness.json") + payload["deployments"][1]["model_alias"] = payload["deployments"][0]["model_alias"] + + with pytest.raises(ValidationError, match="model aliases must be unique"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_readiness_preserves_the_public_model_alias_contract() -> None: + payload = _golden_payload("single_node_readiness.json") + model_alias = "model/模型 alias " + "x" * 140 + payload["deployments"][0]["model_alias"] = model_alias + + readiness = AttemptReadiness.model_validate_json(json.dumps(payload)) + + assert readiness.deployments[0].model_alias == model_alias + + +def test_probe_evidence_is_bounded_and_single_line() -> None: + payload = _golden_payload("single_node_readiness.json") + payload["deployments"][0]["last_probe"]["redacted_message"] = "line one\nline two" + + with pytest.raises(ValidationError, match="control characters"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + payload["deployments"][0]["last_probe"]["redacted_message"] = "x" * 513 + with pytest.raises(ValidationError): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("attempt_state", "deployment_state"), + ( + ("stopped", "ready"), + ("pending", "ready"), + ("starting", "stopped"), + ), +) +def test_attempt_readiness_rejects_contradictory_deployment_states( + attempt_state: str, + deployment_state: str, +) -> None: + payload = _golden_payload("single_node_readiness.json") + payload["state"] = attempt_state + payload["deployments"][0]["state"] = deployment_state + + with pytest.raises(ValidationError): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("state", "endpoint_publication"), + (("pending", "pending"), ("stopped", "published")), +) +def test_pending_and_stopped_deployments_reject_ready_backends( + state: str, + endpoint_publication: str, +) -> None: + payload = _golden_payload("single_node_readiness.json") + payload["state"] = state + payload["deployments"][0].update( + { + "state": state, + "endpoint_publication": endpoint_publication, + "ready_backends": 1, + } + ) + + with pytest.raises(ValidationError, match="cannot have ready backends"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_partial_candidate_is_never_winner_eligible() -> None: + payload = _golden_payload("candidate_output.json") + payload["actual_records"] = 99 + payload["outcome"] = "partial" + payload["files"][0]["record_count"] = 99 + + partial_candidate = CandidateOutputManifest.model_validate_json(json.dumps(payload)) + assert not partial_candidate.winner_eligible + + +def test_candidate_output_rejects_count_and_file_mismatches() -> None: + payload = _golden_payload("candidate_output.json") + payload["files"][0]["record_count"] = 99 + + with pytest.raises(ValidationError, match="record counts"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + payload["files"][0]["relative_path"] = "../escaped.parquet" + with pytest.raises(ValidationError, match="parent-directory"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/state/test_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py new file mode 100644 index 000000000..037785899 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -0,0 +1,650 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import TypeVar + +import pytest + +from data_designer.slurm.state import ( + ArtifactReference, + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutputManifest, + CollectionPlan, + EffectiveAttemptState, + EndpointPublicationState, + ReadinessState, + RunManifest, + SchedulerObservation, + SchedulerState, + ShardManifest, + ShardWinner, + StateContractError, + StateRecord, + StateValue, + reconcile_attempt_observation, + validate_attempt_manifest, + validate_attempt_set, + validate_attempt_transition, + validate_collection_plan, + validate_readiness_transition, + validate_scheduler_observation_transition, + validate_shard_manifest, + validate_shard_set, + validate_shard_winner, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" +RecordT = TypeVar("RecordT", bound=StateRecord) +ValueT = TypeVar("ValueT", bound=StateValue) + + +def _load(model: type[RecordT], filename: str) -> RecordT: + return model.model_validate_json((GOLDEN_DIRECTORY / filename).read_text()) + + +def _json_value(value: object) -> object: + if isinstance(value, datetime): + return value.isoformat() + if isinstance(value, StateValue): + return value.model_dump(mode="json") + if isinstance(value, tuple): + return [_json_value(item) for item in value] + return value + + +def _validated_copy(record: ValueT, **updates: object) -> ValueT: + payload = record.model_dump(mode="json") + payload.update({key: _json_value(value) for key, value in updates.items()}) + return type(record).model_validate_json(json.dumps(payload)) + + +def test_run_shard_attempt_and_winner_contracts_match() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + attempt = _load(AttemptManifest, "successful_attempt.json") + candidate = _load(CandidateOutputManifest, "candidate_output.json") + winner = _load(ShardWinner, "shard_winner.json") + + assert validate_shard_manifest(run, shard) is shard + assert validate_attempt_manifest(run, shard, attempt) is attempt + assert validate_shard_winner(run, shard, attempt, candidate, winner) is winner + + +def test_cross_record_identity_and_digest_mismatches_fail() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + attempt = _load(AttemptManifest, "successful_attempt.json") + candidate = _load(CandidateOutputManifest, "candidate_output.json") + winner = _load(ShardWinner, "shard_winner.json") + + with pytest.raises(StateContractError, match="run_id"): + validate_shard_manifest(run, _validated_copy(shard, run_id="another-run")) + + wrong_reference = ArtifactReference( + path=winner.candidate_manifest.path, + sha256="0" * 64, + ) + wrong_winner = _validated_copy(winner, candidate_manifest=wrong_reference) + with pytest.raises(StateContractError, match="candidate digest"): + validate_shard_winner(run, shard, attempt, candidate, wrong_winner) + + wrong_range = _validated_copy( + shard.record_range, + end_index_exclusive=shard.record_range.end_index_exclusive + 1, + ) + wrong_shard = _validated_copy(shard, record_range=wrong_range) + with pytest.raises(StateContractError, match="record count"): + validate_shard_winner(run, wrong_shard, attempt, candidate, winner) + + +def test_shard_set_rejects_overlapping_ranges_and_shared_resume_workspace() -> None: + run = _validated_copy(_load(RunManifest, "run_manifest.json"), shard_count=2) + first = _load(ShardManifest, "shard_manifest.json") + second_range = _validated_copy( + first.record_range, + start_index=first.record_range.end_index_exclusive, + end_index_exclusive=first.record_range.end_index_exclusive + first.record_range.record_count, + ) + second = _validated_copy( + first, + shard_id="shard-00001", + shard_index=1, + record_range=second_range, + resume_workspace={"path": "/workspace/runs/run-0001/shards/shard-00001/dataset"}, + ) + + assert validate_shard_set(run, (first, second)) == (first, second) + + overlapping_range = _validated_copy( + second.record_range, + start_index=first.record_range.end_index_exclusive - 1, + ) + overlapping = _validated_copy(second, record_range=overlapping_range) + with pytest.raises(StateContractError, match="must not overlap"): + validate_shard_set(run, (first, overlapping)) + + shared_workspace = _validated_copy(second, resume_workspace=first.resume_workspace) + with pytest.raises(StateContractError, match="workspace paths must be unique"): + validate_shard_set(run, (first, shared_workspace)) + + +def test_attempt_set_rejects_scheduler_identity_collisions() -> None: + run = _validated_copy(_load(RunManifest, "run_manifest.json"), shard_count=2) + first_shard = _load(ShardManifest, "shard_manifest.json") + second_shard = _validated_copy( + first_shard, + shard_id="shard-00001", + shard_index=1, + record_range=_validated_copy(first_shard.record_range, start_index=100, end_index_exclusive=200), + resume_workspace={"path": "/workspace/runs/run-0001/shards/shard-00001/dataset"}, + ) + first_attempt = _validated_copy( + _load(AttemptManifest, "successful_attempt.json"), + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + second_attempt = _validated_copy( + first_attempt, + shard_id=second_shard.shard_id, + ) + assert first_attempt.scheduler is not None + + with pytest.raises(StateContractError, match="scheduler identities must be unique"): + validate_attempt_set(run, (first_shard, second_shard), (first_attempt, second_attempt)) + + unique_scheduler = _validated_copy(first_attempt.scheduler, array_task_id=1) + unique_second_attempt = _validated_copy(second_attempt, scheduler=unique_scheduler) + assert validate_attempt_set( + run, + (first_shard, second_shard), + (first_attempt, unique_second_attempt), + ) == (first_attempt, unique_second_attempt) + + +def test_attempt_transitions_preserve_identity_and_terminal_state() -> None: + succeeded = _load(AttemptManifest, "successful_attempt.json") + running = _validated_copy( + succeeded, + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + updated_at=datetime(2026, 8, 18, 12, 4, tzinfo=timezone.utc), + ) + + assert validate_attempt_transition(running, succeeded) is succeeded + assert validate_attempt_transition(succeeded, succeeded) is succeeded + with pytest.raises(StateContractError, match="state cannot move"): + validate_attempt_transition(running, _validated_copy(running, state=AttemptLifecycleState.PENDING)) + + changed_scheduler = _validated_copy(running.scheduler, array_job_id=9999) + with pytest.raises(StateContractError, match="scheduler identity cannot change"): + validate_attempt_transition(running, _validated_copy(running, scheduler=changed_scheduler)) + + candidate_running = _validated_copy(running, candidate_output=succeeded.candidate_output) + with pytest.raises(StateContractError, match="candidate output cannot change"): + validate_attempt_transition(candidate_running, _validated_copy(candidate_running, candidate_output=None)) + + with pytest.raises(StateContractError, match="terminal attempt manifests are immutable"): + validate_attempt_transition( + succeeded, + _validated_copy(succeeded, updated_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc)), + ) + + +def test_failed_partial_and_existing_winners_are_rejected() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + attempt = _load(AttemptManifest, "successful_attempt.json") + candidate = _load(CandidateOutputManifest, "candidate_output.json") + winner = _load(ShardWinner, "shard_winner.json") + + failed_attempt = _validated_copy( + attempt, + state=AttemptLifecycleState.FAILED, + terminal_classification=AttemptTerminalClassification.NODE_FAILED, + ) + with pytest.raises(StateContractError, match="successful attempts"): + validate_shard_winner(run, shard, failed_attempt, candidate, winner) + + partial_payload = json.loads((GOLDEN_DIRECTORY / "candidate_output.json").read_text()) + partial_payload.update( + { + "actual_records": 99, + "outcome": "partial", + } + ) + partial_payload["files"][0]["record_count"] = 99 + partial = CandidateOutputManifest.model_validate_json(json.dumps(partial_payload)) + partial_reference = ArtifactReference( + path=winner.candidate_manifest.path, + sha256=partial.compute_sha256(), + ) + partial_attempt = _validated_copy(attempt, candidate_output=partial_reference) + partial_winner = _validated_copy(winner, candidate_manifest=partial_reference) + with pytest.raises(StateContractError, match="winner policy"): + validate_shard_winner(run, shard, partial_attempt, partial, partial_winner) + + with pytest.raises(StateContractError, match="immutable"): + validate_shard_winner( + run, + shard, + attempt, + candidate, + winner, + existing_winner=winner, + ) + + +def test_readiness_revisions_are_monotonic_and_preserve_authored_order() -> None: + current = _load(AttemptReadiness, "single_node_readiness.json") + starting_deployment = _validated_copy( + current.deployments[0], + state=ReadinessState.STARTING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + last_probe=None, + ) + previous = _validated_copy( + current, + revision=2, + updated_at=datetime(2026, 8, 18, 12, 1, tzinfo=timezone.utc), + state=ReadinessState.STARTING, + deployments=(starting_deployment,), + ) + + assert validate_readiness_transition(previous, current) is current + with pytest.raises(StateContractError, match="revision"): + validate_readiness_transition(previous, _validated_copy(current, revision=2)) + + multi = _load(AttemptReadiness, "multi_node_readiness.json") + reordered = _validated_copy( + multi, + revision=5, + deployments=tuple(reversed(multi.deployments)), + ) + with pytest.raises(StateContractError, match="order or ID"): + validate_readiness_transition(multi, reordered) + + +def test_readiness_cannot_move_backward_or_change_backend_count() -> None: + ready = _load(AttemptReadiness, "single_node_readiness.json") + starting_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.STARTING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + backward = _validated_copy( + ready, + revision=4, + state=ReadinessState.STARTING, + deployments=(starting_deployment,), + ) + with pytest.raises(StateContractError, match="cannot move"): + validate_readiness_transition(ready, backward) + + changed_deployment = _validated_copy( + ready.deployments[0], + expected_backends=2, + ready_backends=2, + ) + changed_count = _validated_copy( + ready, + revision=4, + deployments=(changed_deployment,), + ) + with pytest.raises(StateContractError, match="backend count"): + validate_readiness_transition(ready, changed_count) + + stopped_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.STOPPED, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + unpublished = _validated_copy( + ready, + revision=ready.revision + 1, + state=ReadinessState.STOPPED, + deployments=(stopped_deployment,), + ) + with pytest.raises(StateContractError, match="endpoint publication"): + validate_readiness_transition(ready, unpublished) + + +def test_readiness_probe_evidence_cannot_regress_or_disappear() -> None: + previous = _load(AttemptReadiness, "single_node_readiness.json") + previous_probe = previous.deployments[0].last_probe + assert previous_probe is not None + + regressed_probe = _validated_copy( + previous_probe, + observed_at=datetime(2026, 8, 18, 12, 1, tzinfo=timezone.utc), + ) + regressed_deployment = _validated_copy(previous.deployments[0], last_probe=regressed_probe) + regressed = _validated_copy( + previous, + revision=previous.revision + 1, + updated_at=datetime(2026, 8, 18, 12, 3, tzinfo=timezone.utc), + deployments=(regressed_deployment,), + ) + with pytest.raises(StateContractError, match="probe observation cannot move backward"): + validate_readiness_transition(previous, regressed) + + removed_deployment = _validated_copy(previous.deployments[0], last_probe=None) + removed = _validated_copy( + previous, + revision=previous.revision + 1, + updated_at=datetime(2026, 8, 18, 12, 3, tzinfo=timezone.utc), + deployments=(removed_deployment,), + ) + with pytest.raises(StateContractError, match="probe evidence cannot be removed"): + validate_readiness_transition(previous, removed) + + +def test_collection_requires_exact_winner_set_and_digests() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + plan = _load(CollectionPlan, "collection_plan.json") + winner = _load(ShardWinner, "shard_winner.json") + + assert validate_shard_set(run, (shard,)) == (shard,) + assert validate_collection_plan(run, plan, (shard,), (winner,)) is plan + with pytest.raises(StateContractError, match="winner set"): + validate_collection_plan(run, plan, (shard,), ()) + + wrong_reference = ArtifactReference( + path=plan.planned_shards[0].winner_manifest.path, + sha256="0" * 64, + ) + wrong_planned_shard = _validated_copy(plan.planned_shards[0], winner_manifest=wrong_reference) + wrong_plan = _validated_copy(plan, planned_shards=(wrong_planned_shard,)) + with pytest.raises(StateContractError, match="digest mismatch"): + validate_collection_plan(run, wrong_plan, (shard,), (winner,)) + + +def test_collection_rejects_an_invented_shard_even_when_its_digest_matches() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + plan = _load(CollectionPlan, "collection_plan.json") + winner = _validated_copy(_load(ShardWinner, "shard_winner.json"), shard_id="shard-99999") + winner_reference = ArtifactReference( + path=plan.planned_shards[0].winner_manifest.path, + sha256=winner.compute_sha256(), + ) + planned_shard = _validated_copy( + plan.planned_shards[0], + shard_id="shard-99999", + winner_manifest=winner_reference, + ) + altered_plan = _validated_copy(plan, planned_shards=(planned_shard,)) + + with pytest.raises(StateContractError, match="planned shards"): + validate_collection_plan(run, altered_plan, (shard,), (winner,)) + + +def test_terminal_attempt_evidence_overrides_stale_readiness() -> None: + attempt = _load(AttemptManifest, "failed_attempt.json") + readiness = _load(AttemptReadiness, "stale_readiness.json") + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 13, 5, tzinfo=timezone.utc), + state=SchedulerState.RUNNING, + ) + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 13, 5, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.FAILED + ) + + +def test_terminal_scheduler_failure_overrides_successful_attempt() -> None: + attempt = _load(AttemptManifest, "successful_attempt.json") + readiness = _load(AttemptReadiness, "single_node_readiness.json") + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + state=SchedulerState.NODE_FAILED, + ) + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.FAILED + ) + + +@pytest.mark.parametrize("scheduler_state", (SchedulerState.PENDING, SchedulerState.RUNNING)) +def test_failed_readiness_overrides_nonterminal_scheduler_state(scheduler_state: SchedulerState) -> None: + attempt = _validated_copy( + _load(AttemptManifest, "successful_attempt.json"), + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + readiness = _load(AttemptReadiness, "single_node_readiness.json") + failed_deployment = _validated_copy( + readiness.deployments[0], + state=ReadinessState.FAILED, + ready_backends=0, + endpoint_publication=EndpointPublicationState.FAILED, + ) + failed_readiness = _validated_copy( + readiness, + state=ReadinessState.FAILED, + deployments=(failed_deployment,), + ) + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + state=scheduler_state, + ) + + assert ( + reconcile_attempt_observation( + attempt, + failed_readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.FAILED + ) + + +def test_reconciliation_rejects_stale_or_future_evidence() -> None: + attempt = _load(AttemptManifest, "successful_attempt.json") + readiness = _load(AttemptReadiness, "single_node_readiness.json") + assert attempt.scheduler is not None + + stale_scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 0, 1, tzinfo=timezone.utc), + state=SchedulerState.NODE_FAILED, + ) + with pytest.raises(StateContractError, match="cannot precede attempt creation"): + reconcile_attempt_observation( + attempt, + readiness, + stale_scheduler, + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + + current_scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + state=SchedulerState.UNKNOWN, + ) + future_readiness = _validated_copy( + readiness, + updated_at=datetime(2026, 8, 18, 12, 7, tzinfo=timezone.utc), + ) + with pytest.raises(StateContractError, match="cannot precede readiness update"): + reconcile_attempt_observation( + attempt, + future_readiness, + current_scheduler, + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + + +def test_reconciliation_covers_nonterminal_and_fallback_states() -> None: + successful_attempt = _load(AttemptManifest, "successful_attempt.json") + attempt = _validated_copy( + successful_attempt, + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + ready = _load(AttemptReadiness, "single_node_readiness.json") + pending_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.PENDING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + last_probe=None, + ) + pending = _validated_copy(ready, state=ReadinessState.PENDING, deployments=(pending_deployment,)) + stopped_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.STOPPED, + ready_backends=0, + ) + stopped = _validated_copy(ready, state=ReadinessState.STOPPED, deployments=(stopped_deployment,)) + assert attempt.scheduler is not None + observed_at = datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc) + current_time = observed_at + + cases = ( + (attempt, ready, SchedulerState.PENDING, EffectiveAttemptState.PENDING), + (attempt, ready, SchedulerState.RUNNING, EffectiveAttemptState.RUNNING), + (attempt, ready, SchedulerState.COMPLETED, EffectiveAttemptState.FAILED), + (attempt, pending, SchedulerState.UNKNOWN, EffectiveAttemptState.UNKNOWN), + (attempt, stopped, SchedulerState.UNKNOWN, EffectiveAttemptState.UNKNOWN), + (successful_attempt, ready, SchedulerState.UNKNOWN, EffectiveAttemptState.SUCCEEDED), + ) + for case_attempt, case_readiness, scheduler_state, expected in cases: + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=observed_at, + state=scheduler_state, + ) + assert ( + reconcile_attempt_observation( + case_attempt, + case_readiness, + scheduler, + current_time=current_time, + ) + is expected + ) + + +def test_accounting_lag_is_nonterminal_until_its_deadline() -> None: + attempt = _validated_copy( + _load(AttemptManifest, "successful_attempt.json"), + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + readiness = _load(AttemptReadiness, "single_node_readiness.json") + scheduler = _load(SchedulerObservation, "accounting_lag.json") + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 7, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.ACCOUNTING_LAG + ) + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.UNKNOWN + ) + + +def test_accounting_lag_deadline_is_fixed_across_observations() -> None: + previous = _load(SchedulerObservation, "accounting_lag.json") + continued = _validated_copy( + previous, + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + refreshed = _validated_copy( + continued, + reconciliation_deadline=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), + ) + expired = SchedulerObservation( + schema_version=1, + scheduler=previous.scheduler, + observed_at=datetime(2026, 8, 18, 12, 10, 3, tzinfo=timezone.utc), + state=SchedulerState.UNKNOWN, + ) + + assert validate_scheduler_observation_transition(previous, continued) is continued + assert validate_scheduler_observation_transition(previous, expired) is expired + with pytest.raises(StateContractError, match="deadline cannot change"): + validate_scheduler_observation_transition(previous, refreshed) + + early_unknown = _validated_copy(expired, observed_at=datetime(2026, 8, 18, 12, 9, tzinfo=timezone.utc)) + with pytest.raises(StateContractError, match="before its reconciliation deadline expires"): + validate_scheduler_observation_transition(previous, early_unknown) + + +def test_unknown_scheduler_state_overrides_stale_ready_readiness() -> None: + attempt = _validated_copy( + _load(AttemptManifest, "successful_attempt.json"), + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + readiness = _load(AttemptReadiness, "single_node_readiness.json") + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), + state=SchedulerState.UNKNOWN, + ) + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.UNKNOWN + ) diff --git a/packages/data-designer-slurm/tests/test_contracts.py b/packages/data-designer-slurm/tests/test_contracts.py new file mode 100644 index 000000000..3319d16e2 --- /dev/null +++ b/packages/data-designer-slurm/tests/test_contracts.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.contracts import ( + ArtifactReference as ContractArtifactReference, +) +from data_designer.slurm.contracts import ( + ContractRecord, + ContractValue, + ResumeWorkspace, + canonical_json, + compute_sha256, + pretty_json, +) +from data_designer.slurm.contracts import ( + RecordRange as ContractRecordRange, +) +from data_designer.slurm.state import ( + ArtifactReference as StateArtifactReference, +) +from data_designer.slurm.state import ( + ContractRecord as StateContractRecord, +) +from data_designer.slurm.state import ( + ContractValue as StateContractValue, +) +from data_designer.slurm.state import ( + RecordRange as StateRecordRange, +) +from data_designer.slurm.state import ( + StateRecord, + StateValue, +) + + +def test_state_exports_exact_shared_contract_types() -> None: + assert StateArtifactReference is ContractArtifactReference + assert StateRecordRange is ContractRecordRange + assert StateContractValue is ContractValue + assert StateContractRecord is ContractRecord + assert StateValue is ContractValue + assert StateRecord is ContractRecord + + +def test_shared_json_helpers_are_deterministic() -> None: + value = {"unicode": "模型", "number": 1} + + assert canonical_json(value) == b'{"number":1,"unicode":"\xe6\xa8\xa1\xe5\x9e\x8b"}' + assert pretty_json(value) == '{\n "number": 1,\n "unicode": "模型"\n}\n' + assert compute_sha256(value) == hashlib.sha256(canonical_json(value)).hexdigest() + + +def test_resume_workspace_requires_a_safe_absolute_path() -> None: + workspace = ResumeWorkspace(path="/workspace/runs/run-0001/shards/shard-00000/dataset") + + assert workspace.path.endswith("/dataset") + with pytest.raises(ValidationError, match="exactly one leading slash"): + ResumeWorkspace(path="//workspace/dataset") diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index ff6fd6b6a..d244200d8 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -125,6 +125,14 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non assert slurm_help_result.exit_code == 0, (slurm_help_result.output, repr(slurm_help_result.exception)) assert "data_designer.slurm.cli" in sys.modules assert version("data-designer-slurm") == {version!r} +from data_designer.slurm.contracts import ArtifactReference as ContractArtifactReference +from data_designer.slurm.contracts import RecordRange as ContractRecordRange +from data_designer.slurm.state import ArtifactReference as StateArtifactReference +from data_designer.slurm.state import RecordRange as StateRecordRange +from data_designer.slurm.state import RunManifest +assert RunManifest.__name__ == "RunManifest" +assert StateArtifactReference is ContractArtifactReference +assert StateRecordRange is ContractRecordRange """ run([str(python), "-c", statement], cwd=cwd) @@ -170,8 +178,10 @@ def main() -> None: base_leaf_requirement = requirement(base_metadata, "data-designer-slurm") leaf_base_requirement = requirement(leaf_metadata, "data-designer") + leaf_pydantic_requirement = requirement(leaf_metadata, "pydantic") assert str(base_leaf_requirement.specifier) == f"=={version}" assert str(leaf_base_requirement.specifier) == f"=={version}" + assert leaf_pydantic_requirement.specifier == Requirement("pydantic>=2.9.2,<3").specifier assert base_leaf_requirement.marker is not None assert base_leaf_requirement.marker.evaluate({"extra": "slurm"}) assert not base_leaf_requirement.marker.evaluate({"extra": ""}) diff --git a/uv.lock b/uv.lock index 28264c968..6e0db450d 100644 --- a/uv.lock +++ b/uv.lock @@ -973,10 +973,14 @@ name = "data-designer-slurm" source = { editable = "packages/data-designer-slurm" } dependencies = [ { name = "data-designer" }, + { name = "pydantic" }, ] [package.metadata] -requires-dist = [{ name = "data-designer", editable = "packages/data-designer" }] +requires-dist = [ + { name = "data-designer", editable = "packages/data-designer" }, + { name = "pydantic", specifier = ">=2.9.2,<3" }, +] [[package]] name = "data-designer-workspace"