feat: define Slurm runtime and state record contracts

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"]

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",
+]

packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Public state contracts for the optional Slurm execution package."""
+
+from __future__ import annotations
+
+from data_designer.slurm.contracts import (
+    ArtifactReference,
+    AttemptId,
+    ContractRecord,
+    ContractValue,
+    Identifier,
+    ModelAlias,
+    RecordRange,
+    ResumeWorkspace,
+    Sha256Digest,
+    ShardId,
+)
+from data_designer.slurm.state.base import (
+    SchedulerIdentity,
+    StateRecord,
+    StateValue,
+)
+from data_designer.slurm.state.execution import (
+    AttemptLifecycleState,
+    AttemptManifest,
+    AttemptTerminalClassification,
+    RunManifest,
+    ShardManifest,
+)
+from data_designer.slurm.state.outputs import (
+    CandidateOutcome,
+    CandidateOutputFile,
+    CandidateOutputManifest,
+    CollectionPlan,
+    CollectionShard,
+    ShardWinner,
+)
+from data_designer.slurm.state.readiness import (
+    AttemptReadiness,
+    DeploymentReadiness,
+    EndpointPublicationState,
+    ProbeEvidence,
+    ProbeOutcome,
+    ReadinessState,
+    ReasonCode,
+)
+from data_designer.slurm.state.reconciliation import (
+    reconcile_attempt_observation,
+    validate_readiness_transition,
+)
+from data_designer.slurm.state.scheduler import (
+    EffectiveAttemptState,
+    SchedulerObservation,
+    SchedulerState,
+)
+from data_designer.slurm.state.validation import (
+    StateContractError,
+    validate_attempt_manifest,
+    validate_attempt_set,
+    validate_attempt_transition,
+    validate_collection_plan,
+    validate_scheduler_observation_transition,
+    validate_shard_manifest,
+    validate_shard_set,
+    validate_shard_winner,
+)
+
+__all__ = [
+    "ArtifactReference",
+    "AttemptLifecycleState",
+    "AttemptManifest",
+    "AttemptId",
+    "AttemptReadiness",
+    "AttemptTerminalClassification",
+    "CandidateOutcome",
+    "CandidateOutputFile",
+    "CandidateOutputManifest",
+    "CollectionPlan",
+    "CollectionShard",
+    "ContractRecord",
+    "ContractValue",
+    "DeploymentReadiness",
+    "EffectiveAttemptState",
+    "EndpointPublicationState",
+    "Identifier",
+    "ModelAlias",
+    "ProbeEvidence",
+    "ProbeOutcome",
+    "ReadinessState",
+    "ReasonCode",
+    "RecordRange",
+    "RunManifest",
+    "ResumeWorkspace",
+    "SchedulerIdentity",
+    "SchedulerObservation",
+    "SchedulerState",
+    "Sha256Digest",
+    "ShardManifest",
+    "ShardId",
+    "ShardWinner",
+    "StateContractError",
+    "StateRecord",
+    "StateValue",
+    "reconcile_attempt_observation",
+    "validate_attempt_manifest",
+    "validate_attempt_set",
+    "validate_attempt_transition",
+    "validate_collection_plan",
+    "validate_readiness_transition",
+    "validate_scheduler_observation_transition",
+    "validate_shard_manifest",
+    "validate_shard_set",
+    "validate_shard_winner",
+]

packages/data-designer-slurm/src/data_designer/slurm/state/base.py

@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+
+from pydantic import NonNegativeInt, PositiveInt
+
+from data_designer.slurm.contracts import (
+    ArtifactReference,
+    ContractRecord,
+    ContractValue,
+    Identifier,
+    Sha256Digest,
+    validate_absolute_path,
+    validate_relative_path,
+)
+
+StateValue = ContractValue
+StateRecord = ContractRecord
+
+
+def validate_utc_timestamp(value: datetime) -> datetime:
+    """Validate that a timestamp is timezone-aware and expressed in UTC."""
+    if value.tzinfo is None or value.utcoffset() is None:
+        raise ValueError("timestamp must include timezone information")
+    if value.utcoffset() != timedelta(0):
+        raise ValueError("timestamp must be in UTC")
+    return value
+
+
+def validate_optional_utc_timestamp(value: datetime | None) -> datetime | None:
+    """Validate an optional timestamp when it is present."""
+    if value is None:
+        return None
+    return validate_utc_timestamp(value)
+
+
+class SchedulerIdentity(StateValue):
+    """Slurm array job and task identity assigned to one attempt."""
+
+    array_job_id: PositiveInt
+    array_task_id: NonNegativeInt
+
+
+__all__ = [
+    "ArtifactReference",
+    "Identifier",
+    "SchedulerIdentity",
+    "Sha256Digest",
+    "StateRecord",
+    "StateValue",
+    "validate_absolute_path",
+    "validate_optional_utc_timestamp",
+    "validate_relative_path",
+    "validate_utc_timestamp",
+]