From b2f85c85c3983e578ad2bdeeb91dcd951580b59d Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 18 Aug 2026 12:02:44 -0600 Subject: [PATCH 1/4] feat: add Slurm state record contracts Define strict, versioned execution, readiness, output, reconciliation, and scheduler records. Enforce cross-record identity, digest, shard-set, and terminal-state invariants with deterministic golden fixtures and built-wheel coverage. Closes #865 --- packages/data-designer-slurm/pyproject.toml | 5 +- .../src/data_designer/slurm/state/__init__.py | 96 +++++ .../src/data_designer/slurm/state/base.py | 127 +++++++ .../data_designer/slurm/state/execution.py | 125 +++++++ .../src/data_designer/slurm/state/outputs.py | 137 +++++++ .../data_designer/slurm/state/readiness.py | 142 ++++++++ .../slurm/state/reconciliation.py | 158 +++++++++ .../data_designer/slurm/state/scheduler.py | 63 ++++ .../data_designer/slurm/state/validation.py | 158 +++++++++ .../tests/state/golden/accounting_lag.json | 10 + .../tests/state/golden/candidate_output.json | 23 ++ .../tests/state/golden/collection_plan.json | 23 ++ .../tests/state/golden/failed_attempt.json | 20 ++ .../state/golden/multi_node_readiness.json | 39 ++ .../tests/state/golden/run_manifest.json | 14 + .../tests/state/golden/shard_manifest.json | 16 + .../tests/state/golden/shard_winner.json | 12 + .../state/golden/single_node_readiness.json | 25 ++ .../tests/state/golden/stale_readiness.json | 25 ++ .../state/golden/successful_attempt.json | 23 ++ .../tests/state/test_golden_records.py | 48 +++ .../tests/state/test_records.py | 214 +++++++++++ .../tests/state/test_validation.py | 334 ++++++++++++++++++ scripts/test_slurm_package_install.py | 4 + uv.lock | 6 +- 25 files changed, 1845 insertions(+), 2 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/base.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/execution.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/validation.py create mode 100644 packages/data-designer-slurm/tests/state/golden/accounting_lag.json create mode 100644 packages/data-designer-slurm/tests/state/golden/candidate_output.json create mode 100644 packages/data-designer-slurm/tests/state/golden/collection_plan.json create mode 100644 packages/data-designer-slurm/tests/state/golden/failed_attempt.json create mode 100644 packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/run_manifest.json create mode 100644 packages/data-designer-slurm/tests/state/golden/shard_manifest.json create mode 100644 packages/data-designer-slurm/tests/state/golden/shard_winner.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/stale_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/successful_attempt.json create mode 100644 packages/data-designer-slurm/tests/state/test_golden_records.py create mode 100644 packages/data-designer-slurm/tests/state/test_records.py create mode 100644 packages/data-designer-slurm/tests/state/test_validation.py 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/state/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py new file mode 100644 index 000000000..d9607afb0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -0,0 +1,96 @@ +# 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.state.base import ( + ArtifactReference, + Identifier, + SchedulerIdentity, + Sha256Digest, + StateRecord, + StateValue, +) +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RecordRange, + 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_collection_plan, + validate_shard_manifest, + validate_shard_set, + validate_shard_winner, +) + +__all__ = [ + "ArtifactReference", + "AttemptLifecycleState", + "AttemptManifest", + "AttemptReadiness", + "AttemptTerminalClassification", + "CandidateOutcome", + "CandidateOutputFile", + "CandidateOutputManifest", + "CollectionPlan", + "CollectionShard", + "DeploymentReadiness", + "EffectiveAttemptState", + "EndpointPublicationState", + "Identifier", + "ProbeEvidence", + "ProbeOutcome", + "ReadinessState", + "ReasonCode", + "RecordRange", + "RunManifest", + "SchedulerIdentity", + "SchedulerObservation", + "SchedulerState", + "Sha256Digest", + "ShardManifest", + "ShardWinner", + "StateContractError", + "StateRecord", + "StateValue", + "reconcile_attempt_observation", + "validate_attempt_manifest", + "validate_collection_plan", + "validate_readiness_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..8523c8883 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import posixpath +from datetime import datetime, timedelta +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, NonNegativeInt, PositiveInt, StringConstraints, field_validator + +Identifier = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", + ), +] +Sha256Digest = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + + +class StateValue(BaseModel): + """Base for strict, immutable values nested within state records.""" + + model_config = ConfigDict( + extra="forbid", + frozen=True, + protected_namespaces=(), + strict=True, + validate_default=True, + ) + + +class StateRecord(StateValue): + """Base for immutable, strictly versioned Slurm state records.""" + + schema_version: Literal[1] = 1 + + def serialize_canonical_json(self) -> bytes: + """Serialize the record to stable bytes suitable for hashing.""" + return json.dumps( + self.model_dump(mode="json"), + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + def serialize_json(self) -> str: + """Serialize the record to deterministic, human-readable JSON.""" + return ( + json.dumps( + self.model_dump(mode="json"), + allow_nan=False, + ensure_ascii=False, + indent=2, + sort_keys=True, + ) + + "\n" + ) + + def compute_sha256(self) -> Sha256Digest: + """Compute the digest of the canonical JSON representation.""" + return hashlib.sha256(self.serialize_canonical_json()).hexdigest() + + +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) + + +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 == "/": + 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(StateValue): + """Immutable reference to an on-disk artifact and its content digest.""" + + path: str + sha256: Sha256Digest + + _path_is_safe = field_validator("path")(validate_absolute_path) + + +class SchedulerIdentity(StateValue): + """Slurm array job and task identity assigned to one attempt.""" + + array_job_id: PositiveInt + array_task_id: NonNegativeInt 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..05a9b656d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/execution.py @@ -0,0 +1,125 @@ +# 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.state.base import ( + ArtifactReference, + Identifier, + SchedulerIdentity, + StateRecord, + StateValue, + validate_utc_timestamp, +) + + +class RecordRange(StateValue): + """Half-open global record range assigned to a 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 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: Identifier + shard_index: NonNegativeInt + record_range: RecordRange + input_partition: ArtifactReference | None = None + resume_workspace_id: Identifier + 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: Identifier + attempt_id: Identifier + 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..b1ebc3449 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -0,0 +1,137 @@ +# 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.state.base import ( + ArtifactReference, + Identifier, + Sha256Digest, + StateRecord, + StateValue, + validate_absolute_path, + validate_relative_path, + 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: Identifier + attempt_id: Identifier + attempt_ordinal: PositiveInt + created_at: datetime + dataset_path: str + requested_records: PositiveInt + actual_records: NonNegativeInt + require_exact_record_count: bool + 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 policy permits publishing this candidate as the shard winner.""" + return self.actual_records > 0 and ( + not self.require_exact_record_count or self.actual_records == self.requested_records + ) + + @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: Identifier + attempt_id: Identifier + 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: Identifier + 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..36182c2e2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -0,0 +1,142 @@ +# 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.state.base import Identifier, 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_name: Identifier + model_alias: Identifier + 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: Identifier + attempt_id: Identifier + 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_names = [deployment.deployment_name for deployment in self.deployments] + model_aliases = [deployment.model_alias for deployment in self.deployments] + if len(deployment_names) != len(set(deployment_names)): + raise ValueError("deployment names must be unique") + if len(model_aliases) != len(set(model_aliases)): + raise ValueError("model aliases must be unique") + + 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..356c82c40 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -0,0 +1,158 @@ +# 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_name == new_deployment.deployment_name, + "readiness deployment order or name 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_name!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_name!r} endpoint publication 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(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") + + 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.PENDING: + return EffectiveAttemptState.PENDING + if scheduler.state is SchedulerState.RUNNING: + return EffectiveAttemptState.RUNNING + + if readiness.state is ReadinessState.FAILED: + return EffectiveAttemptState.FAILED + if readiness.state is ReadinessState.PENDING: + return EffectiveAttemptState.PENDING + if readiness.state in {ReadinessState.STARTING, ReadinessState.READY}: + return EffectiveAttemptState.RUNNING + return EffectiveAttemptState.UNKNOWN + + +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..50c8ffee8 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -0,0 +1,158 @@ +# 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, +) + + +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 the exact, ordered set of shards belonging to a run.""" + _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) + shard_indices = tuple(shard.shard_index for shard in shards) + _require(len(set(shard_ids)) == len(shards), "shard IDs must be unique") + _require( + shard_indices == tuple(range(run.shard_count)), + "shards must be ordered by a complete zero-based shard index", + ) + 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_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..5cf9e1c1d --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/candidate_output.json @@ -0,0 +1,23 @@ +{ + "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-0000/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, + "require_exact_record_count": true, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-0000" +} 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..4dcf8b6ab --- /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-0000", + "winner_manifest": { + "path": "/workspace/runs/run-0001/shards/shard-0000/winner.json", + "sha256": "089139c096a31edcd65aa06a1bc79734860bef3fdb5253dda972b1ca805f79b5" + } + } + ], + "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..500b5fbe7 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/failed_attempt.json @@ -0,0 +1,20 @@ +{ + "attempt_id": "attempt-failed", + "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-0000", + "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..34d245d8a --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json @@ -0,0 +1,39 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_name": "generator-model", + "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_name": "judge-model", + "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-0000", + "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..805623e19 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json @@ -0,0 +1,16 @@ +{ + "created_at": "2026-08-18T12:00:01Z", + "input_partition": { + "path": "/workspace/runs/run-0001/shards/shard-0000/input-partition.json", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "record_range": { + "end_index_exclusive": 100, + "start_index": 0 + }, + "resume_workspace_id": "resume-shard-0000", + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-0000", + "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..5dda1da4a --- /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-0000/attempts/attempt-0001/candidate-output.json", + "sha256": "f1ddf26dbe8af3c6a8f486646ea9c8a42102709b15c10fe43adc1d69f4b952af" + }, + "published_at": "2026-08-18T12:05:02Z", + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-0000" +} 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..b5626e499 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_name": "primary-model", + "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-0000", + "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..78cfffb80 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/stale_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-failed", + "deployments": [ + { + "deployment_name": "primary-model", + "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-0000", + "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..9068082d1 --- /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-0000/attempts/attempt-0001/candidate-output.json", + "sha256": "f1ddf26dbe8af3c6a8f486646ea9c8a42102709b15c10fe43adc1d69f4b952af" + }, + "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-0000", + "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..11dce6ba6 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_golden_records.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +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_json()) == record + assert len(record.compute_sha256()) == 64 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..a84047d8b --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_records.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 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( + 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_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_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_candidate_output_counts_and_policy_are_explicit() -> None: + payload = _golden_payload("candidate_output.json") + payload["actual_records"] = 99 + payload["outcome"] = "partial" + payload["files"][0]["record_count"] = 99 + + exact_candidate = CandidateOutputManifest.model_validate_json(json.dumps(payload)) + assert not exact_candidate.winner_eligible + + payload["require_exact_record_count"] = False + partial_candidate = CandidateOutputManifest.model_validate_json(json.dumps(payload)) + assert 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..573092020 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -0,0 +1,334 @@ +# 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, + reconcile_attempt_observation, + validate_attempt_manifest, + validate_collection_plan, + validate_readiness_transition, + validate_shard_manifest, + validate_shard_set, + validate_shard_winner, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" +RecordT = TypeVar("RecordT", bound=StateRecord) + + +def _load(model: type[RecordT], filename: str) -> RecordT: + return model.model_validate_json((GOLDEN_DIRECTORY / filename).read_text()) + + +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, shard.model_copy(update={"run_id": "another-run"})) + + wrong_reference = ArtifactReference( + path=winner.candidate_manifest.path, + sha256="0" * 64, + ) + wrong_winner = winner.model_copy(update={"candidate_manifest": wrong_reference}) + with pytest.raises(StateContractError, match="candidate digest"): + validate_shard_winner(run, shard, attempt, candidate, wrong_winner) + + wrong_count = candidate.model_copy(update={"requested_records": 101, "require_exact_record_count": False}) + with pytest.raises(StateContractError, match="record count"): + validate_shard_winner(run, shard, attempt, wrong_count, winner) + + +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 = attempt.model_copy( + update={ + "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)) + with pytest.raises(StateContractError, match="winner policy"): + validate_shard_winner(run, shard, attempt, 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 = current.deployments[0].model_copy( + update={ + "state": ReadinessState.STARTING, + "ready_backends": 0, + "endpoint_publication": EndpointPublicationState.PENDING, + } + ) + previous = current.model_copy( + update={ + "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, current.model_copy(update={"revision": 2})) + + multi = _load(AttemptReadiness, "multi_node_readiness.json") + reordered = multi.model_copy( + update={ + "revision": 5, + "deployments": tuple(reversed(multi.deployments)), + } + ) + with pytest.raises(StateContractError, match="order or name"): + 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 = ready.deployments[0].model_copy( + update={ + "state": ReadinessState.STARTING, + "ready_backends": 0, + "endpoint_publication": EndpointPublicationState.PENDING, + } + ) + backward = ready.model_copy( + update={ + "revision": 4, + "state": ReadinessState.STARTING, + "deployments": (starting_deployment,), + } + ) + with pytest.raises(StateContractError, match="cannot move"): + validate_readiness_transition(ready, backward) + + changed_count = ready.model_copy( + update={ + "revision": 4, + "deployments": (ready.deployments[0].model_copy(update={"expected_backends": 2}),), + } + ) + with pytest.raises(StateContractError, match="backend count"): + validate_readiness_transition(ready, changed_count) + + published_starting = ready.model_copy( + update={ + "state": ReadinessState.STARTING, + "deployments": (ready.deployments[0].model_copy(update={"state": ReadinessState.STARTING}),), + } + ) + unpublished = published_starting.model_copy( + update={ + "revision": published_starting.revision + 1, + "deployments": ( + published_starting.deployments[0].model_copy( + update={"endpoint_publication": EndpointPublicationState.PENDING} + ), + ), + } + ) + with pytest.raises(StateContractError, match="endpoint publication"): + validate_readiness_transition(published_starting, unpublished) + + +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 = plan.planned_shards[0].model_copy(update={"winner_manifest": wrong_reference}) + wrong_plan = plan.model_copy(update={"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 = _load(ShardWinner, "shard_winner.json").model_copy(update={"shard_id": "invented-shard"}) + winner_reference = ArtifactReference( + path=plan.planned_shards[0].winner_manifest.path, + sha256=winner.compute_sha256(), + ) + planned_shard = plan.planned_shards[0].model_copy( + update={"shard_id": "invented-shard", "winner_manifest": winner_reference} + ) + altered_plan = plan.model_copy(update={"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( + 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( + 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 + ) + + +def test_accounting_lag_is_nonterminal_until_its_deadline() -> None: + attempt = _load(AttemptManifest, "successful_attempt.json").model_copy( + update={ + "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_readiness_never_declares_success() -> None: + attempt = _load(AttemptManifest, "successful_attempt.json").model_copy( + update={ + "state": AttemptLifecycleState.RUNNING, + "terminal_classification": None, + "candidate_output": None, + } + ) + readiness = _load(AttemptReadiness, "single_node_readiness.json") + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 2, tzinfo=timezone.utc), + state=SchedulerState.UNKNOWN, + ) + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 2, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.RUNNING + ) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index ff6fd6b6a..f3df40237 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -125,6 +125,8 @@ 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.state import RunManifest +assert RunManifest.__name__ == "RunManifest" """ run([str(python), "-c", statement], cwd=cwd) @@ -170,8 +172,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" From 3dd8b36dcd81cef835bc13248adbcba11fab2bf3 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 18 Aug 2026 14:58:43 -0600 Subject: [PATCH 2/4] fix: harden Slurm state contracts Reject partial winners and conflicting shard ownership. Require explicit versions, enforce evidence chronology and readiness precedence, and cover every state-contract path with validated fixtures. Refs #865 --- .../src/data_designer/slurm/state/base.py | 2 +- .../src/data_designer/slurm/state/outputs.py | 7 +- .../data_designer/slurm/state/readiness.py | 5 + .../slurm/state/reconciliation.py | 19 +- .../data_designer/slurm/state/validation.py | 7 + .../tests/state/golden/candidate_output.json | 1 - .../tests/state/golden/collection_plan.json | 2 +- .../tests/state/golden/shard_winner.json | 2 +- .../state/golden/successful_attempt.json | 2 +- .../tests/state/test_model_edges.py | 212 ++++++++++ .../tests/state/test_records.py | 17 +- .../tests/state/test_validation.py | 386 ++++++++++++++---- 12 files changed, 562 insertions(+), 100 deletions(-) create mode 100644 packages/data-designer-slurm/tests/state/test_model_edges.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py index 8523c8883..442dd6f3c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -37,7 +37,7 @@ class StateValue(BaseModel): class StateRecord(StateValue): """Base for immutable, strictly versioned Slurm state records.""" - schema_version: Literal[1] = 1 + schema_version: Literal[1] def serialize_canonical_json(self) -> bytes: """Serialize the record to stable bytes suitable for hashing.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py index b1ebc3449..d1959ab7b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -49,7 +49,6 @@ class CandidateOutputManifest(StateRecord): dataset_path: str requested_records: PositiveInt actual_records: NonNegativeInt - require_exact_record_count: bool outcome: CandidateOutcome files: tuple[CandidateOutputFile, ...] dataset_schema_digest: Sha256Digest @@ -60,10 +59,8 @@ class CandidateOutputManifest(StateRecord): @property def winner_eligible(self) -> bool: - """Whether policy permits publishing this candidate as the shard winner.""" - return self.actual_records > 0 and ( - not self.require_exact_record_count or self.actual_records == self.requested_records - ) + """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: 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 index 36182c2e2..19b544890 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -121,6 +121,11 @@ def validate_deployments(self) -> AttemptReadiness: raise ValueError("deployment names 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: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index 356c82c40..db2c3a896 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -98,6 +98,17 @@ def validate_readiness_transition( new_deployment.endpoint_publication in _ALLOWED_ENDPOINT_TRANSITIONS[old_deployment.endpoint_publication], f"deployment {old_deployment.deployment_name!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_name!r} probe evidence cannot be removed" + ) + _require( + new_probe.observed_at >= old_probe.observed_at, + f"deployment {old_deployment.deployment_name!r} probe observation cannot move backward", + ) return current @@ -111,11 +122,15 @@ def reconcile_attempt_observation( """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 @@ -132,13 +147,13 @@ def reconcile_attempt_observation( if current_time <= deadline: return EffectiveAttemptState.ACCOUNTING_LAG 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 - if readiness.state is ReadinessState.FAILED: - return EffectiveAttemptState.FAILED if readiness.state is ReadinessState.PENDING: return EffectiveAttemptState.PENDING if readiness.state in {ReadinessState.STARTING, ReadinessState.READY}: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py index 50c8ffee8..08f66e33b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -39,12 +39,19 @@ def validate_shard_set( validate_shard_manifest(run, shard) shard_ids = tuple(shard.shard_id for shard in shards) + resume_workspace_ids = tuple(shard.resume_workspace_id 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_ids)) == len(shards), "resume workspace IDs 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 diff --git a/packages/data-designer-slurm/tests/state/golden/candidate_output.json b/packages/data-designer-slurm/tests/state/golden/candidate_output.json index 5cf9e1c1d..17366b3d1 100644 --- a/packages/data-designer-slurm/tests/state/golden/candidate_output.json +++ b/packages/data-designer-slurm/tests/state/golden/candidate_output.json @@ -16,7 +16,6 @@ "outcome": "complete", "provenance_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", "requested_records": 100, - "require_exact_record_count": true, "run_id": "run-0001", "schema_version": 1, "shard_id": "shard-0000" diff --git a/packages/data-designer-slurm/tests/state/golden/collection_plan.json b/packages/data-designer-slurm/tests/state/golden/collection_plan.json index 4dcf8b6ab..a9d5f36bf 100644 --- a/packages/data-designer-slurm/tests/state/golden/collection_plan.json +++ b/packages/data-designer-slurm/tests/state/golden/collection_plan.json @@ -10,7 +10,7 @@ "shard_id": "shard-0000", "winner_manifest": { "path": "/workspace/runs/run-0001/shards/shard-0000/winner.json", - "sha256": "089139c096a31edcd65aa06a1bc79734860bef3fdb5253dda972b1ca805f79b5" + "sha256": "40ea4f22e70929b7311dfd81a3d67bd83277a035a5d4390872a5c3bd1ef29f74" } } ], diff --git a/packages/data-designer-slurm/tests/state/golden/shard_winner.json b/packages/data-designer-slurm/tests/state/golden/shard_winner.json index 5dda1da4a..255665bd6 100644 --- a/packages/data-designer-slurm/tests/state/golden/shard_winner.json +++ b/packages/data-designer-slurm/tests/state/golden/shard_winner.json @@ -3,7 +3,7 @@ "attempt_ordinal": 1, "candidate_manifest": { "path": "/workspace/runs/run-0001/shards/shard-0000/attempts/attempt-0001/candidate-output.json", - "sha256": "f1ddf26dbe8af3c6a8f486646ea9c8a42102709b15c10fe43adc1d69f4b952af" + "sha256": "01eb18e264e6d2c39b1e5807e5a9cdab78ca30554cc94c1e4ac316feed02f9b6" }, "published_at": "2026-08-18T12:05:02Z", "run_id": "run-0001", diff --git a/packages/data-designer-slurm/tests/state/golden/successful_attempt.json b/packages/data-designer-slurm/tests/state/golden/successful_attempt.json index 9068082d1..af652f7b4 100644 --- a/packages/data-designer-slurm/tests/state/golden/successful_attempt.json +++ b/packages/data-designer-slurm/tests/state/golden/successful_attempt.json @@ -3,7 +3,7 @@ "attempt_ordinal": 1, "candidate_output": { "path": "/workspace/runs/run-0001/shards/shard-0000/attempts/attempt-0001/candidate-output.json", - "sha256": "f1ddf26dbe8af3c6a8f486646ea9c8a42102709b15c10fe43adc1d69f4b952af" + "sha256": "01eb18e264e6d2c39b1e5807e5a9cdab78ca30554cc94c1e4ac316feed02f9b6" }, "created_at": "2026-08-18T12:00:02Z", "resolved_plan": { 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..010565852 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_model_edges.py @@ -0,0 +1,212 @@ +# 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", + ( + "/", + "/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-0001/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-0001" + 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_names() -> None: + payload = _golden_payload("multi_node_readiness.json") + payload["deployments"][1]["deployment_name"] = payload["deployments"][0]["deployment_name"] + + with pytest.raises(ValidationError, match="deployment names 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 index a84047d8b..c08f3af48 100644 --- a/packages/data-designer-slurm/tests/state/test_records.py +++ b/packages/data-designer-slurm/tests/state/test_records.py @@ -29,6 +29,7 @@ def _golden_payload(filename: str) -> dict[str, Any]: 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), @@ -67,6 +68,14 @@ def test_records_reject_unknown_fields() -> None: 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) @@ -187,18 +196,14 @@ def test_pending_and_stopped_deployments_reject_ready_backends( AttemptReadiness.model_validate_json(json.dumps(payload)) -def test_candidate_output_counts_and_policy_are_explicit() -> None: +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 - exact_candidate = CandidateOutputManifest.model_validate_json(json.dumps(payload)) - assert not exact_candidate.winner_eligible - - payload["require_exact_record_count"] = False partial_candidate = CandidateOutputManifest.model_validate_json(json.dumps(payload)) - assert partial_candidate.winner_eligible + assert not partial_candidate.winner_eligible def test_candidate_output_rejects_count_and_file_mismatches() -> None: diff --git a/packages/data-designer-slurm/tests/state/test_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py index 573092020..8a21f08eb 100644 --- a/packages/data-designer-slurm/tests/state/test_validation.py +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -28,6 +28,7 @@ ShardWinner, StateContractError, StateRecord, + StateValue, reconcile_attempt_observation, validate_attempt_manifest, validate_collection_plan, @@ -39,12 +40,29 @@ 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") @@ -65,19 +83,54 @@ def test_cross_record_identity_and_digest_mismatches_fail() -> None: winner = _load(ShardWinner, "shard_winner.json") with pytest.raises(StateContractError, match="run_id"): - validate_shard_manifest(run, shard.model_copy(update={"run_id": "another-run"})) + validate_shard_manifest(run, _validated_copy(shard, run_id="another-run")) wrong_reference = ArtifactReference( path=winner.candidate_manifest.path, sha256="0" * 64, ) - wrong_winner = winner.model_copy(update={"candidate_manifest": wrong_reference}) + 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_count = candidate.model_copy(update={"requested_records": 101, "require_exact_record_count": False}) + 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, shard, attempt, wrong_count, winner) + 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-0001", + shard_index=1, + record_range=second_range, + resume_workspace_id="resume-shard-0001", + ) + + 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_id=first.resume_workspace_id) + with pytest.raises(StateContractError, match="workspace IDs must be unique"): + validate_shard_set(run, (first, shared_workspace)) def test_failed_partial_and_existing_winners_are_rejected() -> None: @@ -87,21 +140,31 @@ def test_failed_partial_and_existing_winners_are_rejected() -> None: candidate = _load(CandidateOutputManifest, "candidate_output.json") winner = _load(ShardWinner, "shard_winner.json") - failed_attempt = attempt.model_copy( - update={ - "state": AttemptLifecycleState.FAILED, - "terminal_classification": AttemptTerminalClassification.NODE_FAILED, - } + 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.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, attempt, partial, winner) + validate_shard_winner(run, shard, partial_attempt, partial, partial_winner) with pytest.raises(StateContractError, match="immutable"): validate_shard_winner( @@ -116,32 +179,30 @@ def test_failed_partial_and_existing_winners_are_rejected() -> None: def test_readiness_revisions_are_monotonic_and_preserve_authored_order() -> None: current = _load(AttemptReadiness, "single_node_readiness.json") - starting_deployment = current.deployments[0].model_copy( - update={ - "state": ReadinessState.STARTING, - "ready_backends": 0, - "endpoint_publication": EndpointPublicationState.PENDING, - } + starting_deployment = _validated_copy( + current.deployments[0], + state=ReadinessState.STARTING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + last_probe=None, ) - previous = current.model_copy( - update={ - "revision": 2, - "updated_at": datetime(2026, 8, 18, 12, 1, tzinfo=timezone.utc), - "state": ReadinessState.STARTING, - "deployments": (starting_deployment,), - } + 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, current.model_copy(update={"revision": 2})) + validate_readiness_transition(previous, _validated_copy(current, revision=2)) multi = _load(AttemptReadiness, "multi_node_readiness.json") - reordered = multi.model_copy( - update={ - "revision": 5, - "deployments": tuple(reversed(multi.deployments)), - } + reordered = _validated_copy( + multi, + revision=5, + deployments=tuple(reversed(multi.deployments)), ) with pytest.raises(StateContractError, match="order or name"): validate_readiness_transition(multi, reordered) @@ -149,50 +210,78 @@ def test_readiness_revisions_are_monotonic_and_preserve_authored_order() -> None def test_readiness_cannot_move_backward_or_change_backend_count() -> None: ready = _load(AttemptReadiness, "single_node_readiness.json") - starting_deployment = ready.deployments[0].model_copy( - update={ - "state": ReadinessState.STARTING, - "ready_backends": 0, - "endpoint_publication": EndpointPublicationState.PENDING, - } + starting_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.STARTING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, ) - backward = ready.model_copy( - update={ - "revision": 4, - "state": ReadinessState.STARTING, - "deployments": (starting_deployment,), - } + 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_count = ready.model_copy( - update={ - "revision": 4, - "deployments": (ready.deployments[0].model_copy(update={"expected_backends": 2}),), - } + 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) - published_starting = ready.model_copy( - update={ - "state": ReadinessState.STARTING, - "deployments": (ready.deployments[0].model_copy(update={"state": ReadinessState.STARTING}),), - } + stopped_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.STOPPED, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, ) - unpublished = published_starting.model_copy( - update={ - "revision": published_starting.revision + 1, - "deployments": ( - published_starting.deployments[0].model_copy( - update={"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(published_starting, unpublished) + 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: @@ -210,8 +299,8 @@ def test_collection_requires_exact_winner_set_and_digests() -> None: path=plan.planned_shards[0].winner_manifest.path, sha256="0" * 64, ) - wrong_planned_shard = plan.planned_shards[0].model_copy(update={"winner_manifest": wrong_reference}) - wrong_plan = plan.model_copy(update={"planned_shards": (wrong_planned_shard,)}) + 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,)) @@ -220,15 +309,17 @@ def test_collection_rejects_an_invented_shard_even_when_its_digest_matches() -> run = _load(RunManifest, "run_manifest.json") shard = _load(ShardManifest, "shard_manifest.json") plan = _load(CollectionPlan, "collection_plan.json") - winner = _load(ShardWinner, "shard_winner.json").model_copy(update={"shard_id": "invented-shard"}) + winner = _validated_copy(_load(ShardWinner, "shard_winner.json"), shard_id="invented-shard") winner_reference = ArtifactReference( path=plan.planned_shards[0].winner_manifest.path, sha256=winner.compute_sha256(), ) - planned_shard = plan.planned_shards[0].model_copy( - update={"shard_id": "invented-shard", "winner_manifest": winner_reference} + planned_shard = _validated_copy( + plan.planned_shards[0], + shard_id="invented-shard", + winner_manifest=winner_reference, ) - altered_plan = plan.model_copy(update={"planned_shards": (planned_shard,)}) + altered_plan = _validated_copy(plan, planned_shards=(planned_shard,)) with pytest.raises(StateContractError, match="planned shards"): validate_collection_plan(run, altered_plan, (shard,), (winner,)) @@ -239,6 +330,7 @@ def test_terminal_attempt_evidence_overrides_stale_readiness() -> None: 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, @@ -260,6 +352,7 @@ def test_terminal_scheduler_failure_overrides_successful_attempt() -> 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, 6, tzinfo=timezone.utc), state=SchedulerState.NODE_FAILED, @@ -276,13 +369,142 @@ def test_terminal_scheduler_failure_overrides_successful_attempt() -> None: ) +@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.PENDING), + (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 = _load(AttemptManifest, "successful_attempt.json").model_copy( - update={ - "state": AttemptLifecycleState.RUNNING, - "terminal_classification": None, - "candidate_output": 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") @@ -308,18 +530,18 @@ def test_accounting_lag_is_nonterminal_until_its_deadline() -> None: def test_readiness_never_declares_success() -> None: - attempt = _load(AttemptManifest, "successful_attempt.json").model_copy( - update={ - "state": AttemptLifecycleState.RUNNING, - "terminal_classification": None, - "candidate_output": 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, 2, tzinfo=timezone.utc), + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), state=SchedulerState.UNKNOWN, ) @@ -328,7 +550,7 @@ def test_readiness_never_declares_success() -> None: attempt, readiness, scheduler, - current_time=datetime(2026, 8, 18, 12, 2, tzinfo=timezone.utc), + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), ) is EffectiveAttemptState.RUNNING ) From 7a4b1e0438fba8ed6e615ca8d486204eb6c96d1b Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 19 Aug 2026 09:31:19 -0600 Subject: [PATCH 3/4] fix(slurm): harden state contract invariants Hash persisted record bytes, align state identities with the resolved plan vocabulary, and preserve public model aliases. Add monotonic attempt and scheduler transition checks, including fixed accounting-lag deadlines and UNKNOWN precedence. Refs #865 --- .../src/data_designer/slurm/state/__init__.py | 6 + .../src/data_designer/slurm/state/base.py | 8 +- .../data_designer/slurm/state/readiness.py | 10 +- .../slurm/state/reconciliation.py | 21 ++-- .../data_designer/slurm/state/validation.py | 104 ++++++++++++++++- .../tests/state/golden/candidate_output.json | 4 +- .../tests/state/golden/collection_plan.json | 6 +- .../tests/state/golden/failed_attempt.json | 4 +- .../state/golden/multi_node_readiness.json | 6 +- .../tests/state/golden/shard_manifest.json | 6 +- .../tests/state/golden/shard_winner.json | 6 +- .../state/golden/single_node_readiness.json | 4 +- .../tests/state/golden/stale_readiness.json | 6 +- .../state/golden/successful_attempt.json | 6 +- .../tests/state/test_golden_records.py | 17 ++- .../tests/state/test_model_edges.py | 12 +- .../tests/state/test_records.py | 10 ++ .../tests/state/test_validation.py | 110 ++++++++++++++++-- 18 files changed, 287 insertions(+), 59 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py index d9607afb0..9c3baa3c9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -50,7 +50,10 @@ 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, @@ -88,8 +91,11 @@ "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 index 442dd6f3c..6c5cebe23 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -40,7 +40,7 @@ class StateRecord(StateValue): schema_version: Literal[1] def serialize_canonical_json(self) -> bytes: - """Serialize the record to stable bytes suitable for hashing.""" + """Serialize the record to a stable compact JSON representation.""" return json.dumps( self.model_dump(mode="json"), allow_nan=False, @@ -63,8 +63,8 @@ def serialize_json(self) -> str: ) def compute_sha256(self) -> Sha256Digest: - """Compute the digest of the canonical JSON representation.""" - return hashlib.sha256(self.serialize_canonical_json()).hexdigest() + """Compute the digest of the exact bytes written by ``serialize_json``.""" + return hashlib.sha256(self.serialize_json().encode("utf-8")).hexdigest() def validate_utc_timestamp(value: datetime) -> datetime: @@ -87,6 +87,8 @@ 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): 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 index 19b544890..d6cdb6c87 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -61,8 +61,8 @@ def validate_redacted_message(cls, value: str) -> str: class DeploymentReadiness(StateValue): """Readiness state for one authored-order model deployment.""" - deployment_name: Identifier - model_alias: Identifier + deployment_id: Identifier + model_alias: str state: ReadinessState expected_backends: PositiveInt ready_backends: NonNegativeInt @@ -115,10 +115,10 @@ class AttemptReadiness(StateRecord): @model_validator(mode="after") def validate_deployments(self) -> AttemptReadiness: - deployment_names = [deployment.deployment_name for deployment in self.deployments] + deployment_ids = [deployment.deployment_id for deployment in self.deployments] model_aliases = [deployment.model_alias for deployment in self.deployments] - if len(deployment_names) != len(set(deployment_names)): - raise ValueError("deployment names must be unique") + 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( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index db2c3a896..aacfacc65 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -76,8 +76,8 @@ def validate_readiness_transition( for old_deployment, new_deployment in zip(previous.deployments, current.deployments, strict=True): _require( - old_deployment.deployment_name == new_deployment.deployment_name, - "readiness deployment order or name cannot change", + old_deployment.deployment_id == new_deployment.deployment_id, + "readiness deployment order or ID cannot change", ) _require( old_deployment.model_alias == new_deployment.model_alias, @@ -90,24 +90,24 @@ def validate_readiness_transition( _require( new_deployment.state in _ALLOWED_READINESS_TRANSITIONS[old_deployment.state], ( - f"deployment {old_deployment.deployment_name!r} cannot move from " + 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_name!r} endpoint publication cannot move backward", + 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_name!r} probe evidence cannot be removed" + 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_name!r} probe observation cannot move backward", + f"deployment {old_deployment.deployment_id!r} probe observation cannot move backward", ) return current @@ -147,18 +147,15 @@ def reconcile_attempt_observation( 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 - - if readiness.state is ReadinessState.PENDING: - return EffectiveAttemptState.PENDING - if readiness.state in {ReadinessState.STARTING, ReadinessState.READY}: - return EffectiveAttemptState.RUNNING - return EffectiveAttemptState.UNKNOWN + raise AssertionError(f"unhandled scheduler state: {scheduler.state}") # pragma: no cover def _require(condition: bool, message: str) -> None: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py index 08f66e33b..999c32308 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -15,6 +15,20 @@ 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): @@ -33,7 +47,7 @@ def validate_shard_set( run: RunManifest, shards: tuple[ShardManifest, ...], ) -> tuple[ShardManifest, ...]: - """Validate the exact, ordered set of shards belonging to a run.""" + """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) @@ -72,6 +86,94 @@ def validate_attempt_manifest( 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, diff --git a/packages/data-designer-slurm/tests/state/golden/candidate_output.json b/packages/data-designer-slurm/tests/state/golden/candidate_output.json index 17366b3d1..81e9a63f4 100644 --- a/packages/data-designer-slurm/tests/state/golden/candidate_output.json +++ b/packages/data-designer-slurm/tests/state/golden/candidate_output.json @@ -3,7 +3,7 @@ "attempt_id": "attempt-0001", "attempt_ordinal": 1, "created_at": "2026-08-18T12:05:00Z", - "dataset_path": "/workspace/runs/run-0001/shards/shard-0000/attempts/attempt-0001/dataset", + "dataset_path": "/workspace/runs/run-0001/shards/shard-00000/attempts/attempt-0001/dataset", "dataset_schema_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", "files": [ { @@ -18,5 +18,5 @@ "requested_records": 100, "run_id": "run-0001", "schema_version": 1, - "shard_id": "shard-0000" + "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 index a9d5f36bf..1d528f0d1 100644 --- a/packages/data-designer-slurm/tests/state/golden/collection_plan.json +++ b/packages/data-designer-slurm/tests/state/golden/collection_plan.json @@ -7,10 +7,10 @@ "overwrite": false, "planned_shards": [ { - "shard_id": "shard-0000", + "shard_id": "shard-00000", "winner_manifest": { - "path": "/workspace/runs/run-0001/shards/shard-0000/winner.json", - "sha256": "40ea4f22e70929b7311dfd81a3d67bd83277a035a5d4390872a5c3bd1ef29f74" + "path": "/workspace/runs/run-0001/shards/shard-00000/winner.json", + "sha256": "e6efb5fd97b97c36e7dd1c4e5ab61dc954b611fd68d5e85956e357b36d2a8930" } } ], diff --git a/packages/data-designer-slurm/tests/state/golden/failed_attempt.json b/packages/data-designer-slurm/tests/state/golden/failed_attempt.json index 500b5fbe7..244b7977d 100644 --- a/packages/data-designer-slurm/tests/state/golden/failed_attempt.json +++ b/packages/data-designer-slurm/tests/state/golden/failed_attempt.json @@ -1,5 +1,5 @@ { - "attempt_id": "attempt-failed", + "attempt_id": "attempt-0002", "attempt_ordinal": 2, "candidate_output": null, "created_at": "2026-08-18T13:00:00Z", @@ -13,7 +13,7 @@ "array_task_id": 0 }, "schema_version": 1, - "shard_id": "shard-0000", + "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 index 34d245d8a..444351540 100644 --- a/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json +++ b/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json @@ -2,7 +2,7 @@ "attempt_id": "attempt-0001", "deployments": [ { - "deployment_name": "generator-model", + "deployment_id": "deployment-00000", "endpoint_publication": "pending", "expected_backends": 2, "last_probe": { @@ -16,7 +16,7 @@ "state": "starting" }, { - "deployment_name": "judge-model", + "deployment_id": "deployment-00001", "endpoint_publication": "published", "expected_backends": 2, "last_probe": { @@ -33,7 +33,7 @@ "revision": 4, "run_id": "run-multi", "schema_version": 1, - "shard_id": "shard-0000", + "shard_id": "shard-00000", "state": "starting", "updated_at": "2026-08-18T14:02:00Z" } diff --git a/packages/data-designer-slurm/tests/state/golden/shard_manifest.json b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json index 805623e19..2e4eadccc 100644 --- a/packages/data-designer-slurm/tests/state/golden/shard_manifest.json +++ b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json @@ -1,16 +1,16 @@ { "created_at": "2026-08-18T12:00:01Z", "input_partition": { - "path": "/workspace/runs/run-0001/shards/shard-0000/input-partition.json", + "path": "/workspace/runs/run-0001/shards/shard-00000/input-partition.json", "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" }, "record_range": { "end_index_exclusive": 100, "start_index": 0 }, - "resume_workspace_id": "resume-shard-0000", + "resume_workspace_id": "resume-shard-00000", "run_id": "run-0001", "schema_version": 1, - "shard_id": "shard-0000", + "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 index 255665bd6..20a903c35 100644 --- a/packages/data-designer-slurm/tests/state/golden/shard_winner.json +++ b/packages/data-designer-slurm/tests/state/golden/shard_winner.json @@ -2,11 +2,11 @@ "attempt_id": "attempt-0001", "attempt_ordinal": 1, "candidate_manifest": { - "path": "/workspace/runs/run-0001/shards/shard-0000/attempts/attempt-0001/candidate-output.json", - "sha256": "01eb18e264e6d2c39b1e5807e5a9cdab78ca30554cc94c1e4ac316feed02f9b6" + "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-0000" + "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 index b5626e499..9337193d5 100644 --- a/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json +++ b/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json @@ -2,7 +2,7 @@ "attempt_id": "attempt-0001", "deployments": [ { - "deployment_name": "primary-model", + "deployment_id": "deployment-00000", "endpoint_publication": "published", "expected_backends": 1, "last_probe": { @@ -19,7 +19,7 @@ "revision": 3, "run_id": "run-0001", "schema_version": 1, - "shard_id": "shard-0000", + "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 index 78cfffb80..511337eb8 100644 --- a/packages/data-designer-slurm/tests/state/golden/stale_readiness.json +++ b/packages/data-designer-slurm/tests/state/golden/stale_readiness.json @@ -1,8 +1,8 @@ { - "attempt_id": "attempt-failed", + "attempt_id": "attempt-0002", "deployments": [ { - "deployment_name": "primary-model", + "deployment_id": "deployment-00000", "endpoint_publication": "published", "expected_backends": 1, "last_probe": { @@ -19,7 +19,7 @@ "revision": 2, "run_id": "run-failed", "schema_version": 1, - "shard_id": "shard-0000", + "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 index af652f7b4..03d28f620 100644 --- a/packages/data-designer-slurm/tests/state/golden/successful_attempt.json +++ b/packages/data-designer-slurm/tests/state/golden/successful_attempt.json @@ -2,8 +2,8 @@ "attempt_id": "attempt-0001", "attempt_ordinal": 1, "candidate_output": { - "path": "/workspace/runs/run-0001/shards/shard-0000/attempts/attempt-0001/candidate-output.json", - "sha256": "01eb18e264e6d2c39b1e5807e5a9cdab78ca30554cc94c1e4ac316feed02f9b6" + "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": { @@ -16,7 +16,7 @@ "array_task_id": 0 }, "schema_version": 1, - "shard_id": "shard-0000", + "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 index 11dce6ba6..fa31859db 100644 --- a/packages/data-designer-slurm/tests/state/test_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_golden_records.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib from pathlib import Path import pytest @@ -44,5 +45,19 @@ def test_golden_record_round_trip_is_deterministic(filename: str, model: type[St 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 len(record.compute_sha256()) == 64 + 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 index 010565852..0babd1f76 100644 --- a/packages/data-designer-slurm/tests/state/test_model_edges.py +++ b/packages/data-designer-slurm/tests/state/test_model_edges.py @@ -32,6 +32,8 @@ def _golden_payload(filename: str) -> dict[str, Any]: "path", ( "/", + "//", + "//host/path", "/workspace/control\ncharacter", "/workspace//unnormalized", ), @@ -122,14 +124,14 @@ def test_candidate_output_rejects_remaining_invalid_inventory_forms() -> None: 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-0001/winner.json" + 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-0001" + 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)) @@ -174,11 +176,11 @@ def test_readiness_rejects_remaining_contradictory_states( AttemptReadiness.model_validate_json(json.dumps(payload)) -def test_readiness_rejects_duplicate_deployment_names() -> None: +def test_readiness_rejects_duplicate_deployment_ids() -> None: payload = _golden_payload("multi_node_readiness.json") - payload["deployments"][1]["deployment_name"] = payload["deployments"][0]["deployment_name"] + payload["deployments"][1]["deployment_id"] = payload["deployments"][0]["deployment_id"] - with pytest.raises(ValidationError, match="deployment names must be unique"): + with pytest.raises(ValidationError, match="deployment IDs must be unique"): AttemptReadiness.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 index c08f3af48..2eac315bf 100644 --- a/packages/data-designer-slurm/tests/state/test_records.py +++ b/packages/data-designer-slurm/tests/state/test_records.py @@ -142,6 +142,16 @@ def test_readiness_rejects_duplicate_deployment_identity() -> None: 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" diff --git a/packages/data-designer-slurm/tests/state/test_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py index 8a21f08eb..8094a36d1 100644 --- a/packages/data-designer-slurm/tests/state/test_validation.py +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -31,8 +31,11 @@ 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, @@ -112,10 +115,10 @@ def test_shard_set_rejects_overlapping_ranges_and_shared_resume_workspace() -> N ) second = _validated_copy( first, - shard_id="shard-0001", + shard_id="shard-00001", shard_index=1, record_range=second_range, - resume_workspace_id="resume-shard-0001", + resume_workspace_id="resume-shard-00001", ) assert validate_shard_set(run, (first, second)) == (first, second) @@ -133,6 +136,70 @@ def test_shard_set_rejects_overlapping_ranges_and_shared_resume_workspace() -> N 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_id="resume-shard-00001", + ) + 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") @@ -204,7 +271,7 @@ def test_readiness_revisions_are_monotonic_and_preserve_authored_order() -> None revision=5, deployments=tuple(reversed(multi.deployments)), ) - with pytest.raises(StateContractError, match="order or name"): + with pytest.raises(StateContractError, match="order or ID"): validate_readiness_transition(multi, reordered) @@ -477,7 +544,7 @@ def test_reconciliation_covers_nonterminal_and_fallback_states() -> None: (attempt, ready, SchedulerState.PENDING, EffectiveAttemptState.PENDING), (attempt, ready, SchedulerState.RUNNING, EffectiveAttemptState.RUNNING), (attempt, ready, SchedulerState.COMPLETED, EffectiveAttemptState.FAILED), - (attempt, pending, SchedulerState.UNKNOWN, EffectiveAttemptState.PENDING), + (attempt, pending, SchedulerState.UNKNOWN, EffectiveAttemptState.UNKNOWN), (attempt, stopped, SchedulerState.UNKNOWN, EffectiveAttemptState.UNKNOWN), (successful_attempt, ready, SchedulerState.UNKNOWN, EffectiveAttemptState.SUCCEEDED), ) @@ -529,7 +596,34 @@ def test_accounting_lag_is_nonterminal_until_its_deadline() -> None: ) -def test_readiness_never_declares_success() -> None: +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, @@ -541,7 +635,7 @@ def test_readiness_never_declares_success() -> None: scheduler = SchedulerObservation( schema_version=1, scheduler=attempt.scheduler, - observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + observed_at=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), state=SchedulerState.UNKNOWN, ) @@ -550,7 +644,7 @@ def test_readiness_never_declares_success() -> None: attempt, readiness, scheduler, - current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + current_time=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), ) - is EffectiveAttemptState.RUNNING + is EffectiveAttemptState.UNKNOWN ) From 2ecc3f7a88e402827cdf5bb25e98087b9bf36f36 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 19 Aug 2026 09:58:24 -0600 Subject: [PATCH 4/4] refactor(slurm): share contract primitives Make planning and runtime state consume one public immutable model family. Align shard, attempt, range, artifact, and resume workspace shapes so the configuration-plan branch can rebase without preserving duplicate Pydantic types.\n\nRefs #865 --- .../src/data_designer/slurm/contracts.py | 180 ++++++++++++++++++ .../src/data_designer/slurm/state/__init__.py | 20 +- .../src/data_designer/slurm/state/base.py | 123 +++--------- .../data_designer/slurm/state/execution.py | 34 ++-- .../src/data_designer/slurm/state/outputs.py | 20 +- .../data_designer/slurm/state/readiness.py | 9 +- .../data_designer/slurm/state/validation.py | 4 +- .../tests/state/golden/shard_manifest.json | 4 +- .../tests/state/test_validation.py | 12 +- .../tests/test_contracts.py | 65 +++++++ scripts/test_slurm_package_install.py | 6 + 11 files changed, 333 insertions(+), 144 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/contracts.py create mode 100644 packages/data-designer-slurm/tests/test_contracts.py 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 index 9c3baa3c9..2d7eb81d8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -5,11 +5,20 @@ from __future__ import annotations -from data_designer.slurm.state.base import ( +from data_designer.slurm.contracts import ( ArtifactReference, + AttemptId, + ContractRecord, + ContractValue, Identifier, - SchedulerIdentity, + ModelAlias, + RecordRange, + ResumeWorkspace, Sha256Digest, + ShardId, +) +from data_designer.slurm.state.base import ( + SchedulerIdentity, StateRecord, StateValue, ) @@ -17,7 +26,6 @@ AttemptLifecycleState, AttemptManifest, AttemptTerminalClassification, - RecordRange, RunManifest, ShardManifest, ) @@ -63,6 +71,7 @@ "ArtifactReference", "AttemptLifecycleState", "AttemptManifest", + "AttemptId", "AttemptReadiness", "AttemptTerminalClassification", "CandidateOutcome", @@ -70,21 +79,26 @@ "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", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py index 6c5cebe23..e9825bd95 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -3,68 +3,22 @@ from __future__ import annotations -import hashlib -import json -import posixpath from datetime import datetime, timedelta -from typing import Annotated, Literal -from pydantic import BaseModel, ConfigDict, NonNegativeInt, PositiveInt, StringConstraints, field_validator +from pydantic import NonNegativeInt, PositiveInt -Identifier = Annotated[ - str, - StringConstraints( - min_length=1, - max_length=128, - pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", - ), -] -Sha256Digest = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] - - -class StateValue(BaseModel): - """Base for strict, immutable values nested within state records.""" - - model_config = ConfigDict( - extra="forbid", - frozen=True, - protected_namespaces=(), - strict=True, - validate_default=True, - ) - - -class StateRecord(StateValue): - """Base for immutable, strictly versioned Slurm state records.""" +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractRecord, + ContractValue, + Identifier, + Sha256Digest, + validate_absolute_path, + validate_relative_path, +) - schema_version: Literal[1] - - def serialize_canonical_json(self) -> bytes: - """Serialize the record to a stable compact JSON representation.""" - return json.dumps( - self.model_dump(mode="json"), - allow_nan=False, - ensure_ascii=False, - separators=(",", ":"), - sort_keys=True, - ).encode("utf-8") - - def serialize_json(self) -> str: - """Serialize the record to deterministic, human-readable JSON.""" - return ( - json.dumps( - self.model_dump(mode="json"), - allow_nan=False, - ensure_ascii=False, - indent=2, - sort_keys=True, - ) - + "\n" - ) - - 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() +StateValue = ContractValue +StateRecord = ContractRecord def validate_utc_timestamp(value: datetime) -> datetime: @@ -83,47 +37,22 @@ def validate_optional_utc_timestamp(value: datetime | None) -> datetime | None: return validate_utc_timestamp(value) -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(StateValue): - """Immutable reference to an on-disk artifact and its content digest.""" - - path: str - sha256: Sha256Digest - - _path_is_safe = field_validator("path")(validate_absolute_path) - - 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 index 05a9b656d..927401001 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/execution.py @@ -8,33 +8,21 @@ from pydantic import NonNegativeInt, PositiveInt, field_validator, model_validator -from data_designer.slurm.state.base import ( +from data_designer.slurm.contracts import ( ArtifactReference, + AttemptId, Identifier, + RecordRange, + ResumeWorkspace, + ShardId, +) +from data_designer.slurm.state.base import ( SchedulerIdentity, StateRecord, - StateValue, validate_utc_timestamp, ) -class RecordRange(StateValue): - """Half-open global record range assigned to a 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 RunManifest(StateRecord): """Identity and immutable authored/resolved inputs for a Slurm run.""" @@ -51,11 +39,11 @@ class ShardManifest(StateRecord): """Stable shard identity and planner-owned input partition reference.""" run_id: Identifier - shard_id: Identifier + shard_id: ShardId shard_index: NonNegativeInt record_range: RecordRange input_partition: ArtifactReference | None = None - resume_workspace_id: Identifier + resume_workspace: ResumeWorkspace created_at: datetime _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) @@ -86,8 +74,8 @@ class AttemptManifest(StateRecord): """Attempt identity, lifecycle, scheduler identity, and output reference.""" run_id: Identifier - shard_id: Identifier - attempt_id: Identifier + shard_id: ShardId + attempt_id: AttemptId attempt_ordinal: PositiveInt resolved_plan: ArtifactReference state: AttemptLifecycleState diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py index d1959ab7b..756e00b44 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -9,14 +9,18 @@ from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator -from data_designer.slurm.state.base import ( +from data_designer.slurm.contracts import ( ArtifactReference, + AttemptId, Identifier, Sha256Digest, - StateRecord, - StateValue, + ShardId, validate_absolute_path, validate_relative_path, +) +from data_designer.slurm.state.base import ( + StateRecord, + StateValue, validate_utc_timestamp, ) @@ -42,8 +46,8 @@ class CandidateOutputManifest(StateRecord): """Attempt-local output that may become the immutable shard winner.""" run_id: Identifier - shard_id: Identifier - attempt_id: Identifier + shard_id: ShardId + attempt_id: AttemptId attempt_ordinal: PositiveInt created_at: datetime dataset_path: str @@ -91,8 +95,8 @@ class ShardWinner(StateRecord): """Immutable pointer selecting exactly one candidate for a shard.""" run_id: Identifier - shard_id: Identifier - attempt_id: Identifier + shard_id: ShardId + attempt_id: AttemptId attempt_ordinal: PositiveInt candidate_manifest: ArtifactReference published_at: datetime @@ -103,7 +107,7 @@ class ShardWinner(StateRecord): class CollectionShard(StateValue): """Winner manifest selected for one shard in a collection plan.""" - shard_id: Identifier + shard_id: ShardId winner_manifest: ArtifactReference 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 index d6cdb6c87..a7f538165 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -9,7 +9,8 @@ from pydantic import Field, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator -from data_designer.slurm.state.base import Identifier, StateRecord, StateValue, validate_utc_timestamp +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, @@ -62,7 +63,7 @@ class DeploymentReadiness(StateValue): """Readiness state for one authored-order model deployment.""" deployment_id: Identifier - model_alias: str + model_alias: ModelAlias state: ReadinessState expected_backends: PositiveInt ready_backends: NonNegativeInt @@ -104,8 +105,8 @@ class AttemptReadiness(StateRecord): """Revisioned readiness snapshot for all deployments in one attempt.""" run_id: Identifier - shard_id: Identifier - attempt_id: Identifier + shard_id: ShardId + attempt_id: AttemptId revision: PositiveInt updated_at: datetime state: ReadinessState diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py index 999c32308..30fdbc066 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -53,10 +53,10 @@ def validate_shard_set( validate_shard_manifest(run, shard) shard_ids = tuple(shard.shard_id for shard in shards) - resume_workspace_ids = tuple(shard.resume_workspace_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_ids)) == len(shards), "resume workspace 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", diff --git a/packages/data-designer-slurm/tests/state/golden/shard_manifest.json b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json index 2e4eadccc..5e885b143 100644 --- a/packages/data-designer-slurm/tests/state/golden/shard_manifest.json +++ b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json @@ -8,7 +8,9 @@ "end_index_exclusive": 100, "start_index": 0 }, - "resume_workspace_id": "resume-shard-00000", + "resume_workspace": { + "path": "/workspace/runs/run-0001/shards/shard-00000/dataset" + }, "run_id": "run-0001", "schema_version": 1, "shard_id": "shard-00000", diff --git a/packages/data-designer-slurm/tests/state/test_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py index 8094a36d1..037785899 100644 --- a/packages/data-designer-slurm/tests/state/test_validation.py +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -118,7 +118,7 @@ def test_shard_set_rejects_overlapping_ranges_and_shared_resume_workspace() -> N shard_id="shard-00001", shard_index=1, record_range=second_range, - resume_workspace_id="resume-shard-00001", + resume_workspace={"path": "/workspace/runs/run-0001/shards/shard-00001/dataset"}, ) assert validate_shard_set(run, (first, second)) == (first, second) @@ -131,8 +131,8 @@ def test_shard_set_rejects_overlapping_ranges_and_shared_resume_workspace() -> N with pytest.raises(StateContractError, match="must not overlap"): validate_shard_set(run, (first, overlapping)) - shared_workspace = _validated_copy(second, resume_workspace_id=first.resume_workspace_id) - with pytest.raises(StateContractError, match="workspace IDs must be unique"): + 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)) @@ -144,7 +144,7 @@ def test_attempt_set_rejects_scheduler_identity_collisions() -> None: shard_id="shard-00001", shard_index=1, record_range=_validated_copy(first_shard.record_range, start_index=100, end_index_exclusive=200), - resume_workspace_id="resume-shard-00001", + resume_workspace={"path": "/workspace/runs/run-0001/shards/shard-00001/dataset"}, ) first_attempt = _validated_copy( _load(AttemptManifest, "successful_attempt.json"), @@ -376,14 +376,14 @@ def test_collection_rejects_an_invented_shard_even_when_its_digest_matches() -> 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="invented-shard") + 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="invented-shard", + shard_id="shard-99999", winner_manifest=winner_reference, ) altered_plan = _validated_copy(plan, planned_shards=(planned_shard,)) 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 f3df40237..d244200d8 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -125,8 +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)