From e85b0eacd5c63df3f0d58386a0134a64ed8b5a31 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 24 Aug 2026 13:31:17 -0600 Subject: [PATCH 1/5] feat: integrate Slurm plan and state Validate plan-aware shards, readiness, attempts, and finalization through a reusable immutable context. Add deterministic goldens and negative coverage for mismatch and failure paths. Closes #880 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/contracts.py | 3 + .../src/data_designer/slurm/integration.py | 328 +++++++++++++ .../golden/finalization_chain.json | 128 +++++ .../test_integration_validation.py | 452 ++++++++++++++++++ 4 files changed, 911 insertions(+) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/integration.py create mode 100644 packages/data-designer-slurm/tests/integration/golden/finalization_chain.json create mode 100644 packages/data-designer-slurm/tests/integration/test_integration_validation.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 index f07063417..fc47926b1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -199,6 +199,7 @@ def validate_relative_path(value: str) -> str: def validate_local_config_path(value: str) -> str: + """Validate a local JSON or YAML configuration path.""" validate_plain_text(value, field_name="path") if "://" in value: raise ValueError("builder and config sources must be local paths") @@ -211,6 +212,7 @@ def validate_local_config_path(value: str) -> str: def validate_plain_text(value: str, *, field_name: str) -> str: + """Reject empty text and control characters at persisted boundaries.""" if not value: raise ValueError(f"{field_name} must not be empty") if any(ord(character) < 32 or ord(character) == 127 for character in value): @@ -219,6 +221,7 @@ def validate_plain_text(value: str, *, field_name: str) -> str: def validate_url(value: str, *, field_name: str) -> str: + """Validate an HTTP(S) URL with a valid host and port.""" validate_plain_text(value, field_name=field_name) try: parsed = urlsplit(value) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/integration.py b/packages/data-designer-slurm/src/data_designer/slurm/integration.py new file mode 100644 index 000000000..33e151841 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure validation across Slurm execution-plan and runtime-state contracts.""" + +from __future__ import annotations + +import posixpath +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType + +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.contracts import ArtifactReference, ShardId +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutputManifest, + ReadinessState, + RunManifest, + ShardManifest, + ShardWinner, + StateContractError, + validate_shard_set, +) + + +class IntegrationContractError(ValueError): + """Raised when reviewed plan and state records disagree.""" + + +@dataclass(frozen=True, slots=True) +class PlanStateValidator: + """Validate records against one resolved plan with reusable derived state.""" + + plan: ResolvedSlurmRunPlan + _plan_reference: ArtifactReference = field(init=False, repr=False) + _shards_by_id: Mapping[ShardId, PlannedShard] = field(init=False, repr=False) + + def __post_init__(self) -> None: + run_root = posixpath.dirname(self.plan.authored_config.path) + object.__setattr__( + self, + "_plan_reference", + ArtifactReference( + path=posixpath.join(run_root, "resolved-plan.json"), + sha256=self.plan.compute_sha256(), + ), + ) + object.__setattr__( + self, + "_shards_by_id", + MappingProxyType({planned_shard.shard_id: planned_shard for planned_shard in self.plan.shards}), + ) + + def validate_plan_shards( + self, + run: RunManifest, + shards: tuple[ShardManifest, ...], + ) -> tuple[ShardManifest, ...]: + """Validate the complete ordered state shard set against planned intent.""" + _require(run.run_id == self.plan.run_id, "run manifest identity does not match the resolved plan") + _require( + run.authored_config == self.plan.authored_config, + "run authored config does not match the resolved plan", + ) + self._validate_plan_reference(run.resolved_plan) + _require(run.shard_count == len(self.plan.shards), "run shard count does not match the resolved plan") + try: + validate_shard_set(run, shards) + except StateContractError as exc: + raise IntegrationContractError(str(exc)) from exc + + _require(len(shards) == len(self.plan.shards), "state shards must exactly match the planned shard count") + for planned, persisted in zip(self.plan.shards, shards, strict=True): + _require(persisted.shard_id == planned.shard_id, "state shard identity does not match planned order") + _require(persisted.shard_index == planned.shard_index, "state shard index does not match planned order") + _require(persisted.record_range == planned.record_range, "state shard record range does not match the plan") + _require( + persisted.input_partition == planned.input_partition, + "state shard input partition does not match the plan", + ) + _require( + persisted.resume_workspace == planned.resume_workspace, + "state shard resume workspace does not match the plan", + ) + return shards + + def validate_initial_readiness( + self, + attempt: AttemptManifest, + readiness: AttemptReadiness, + ) -> AttemptReadiness: + """Anchor the first readiness snapshot to a fully validated planned attempt.""" + planned_shard = self._get_planned_shard(attempt.shard_id) + self.validate_planned_attempt(planned_shard, attempt) + _require(readiness.run_id == attempt.run_id, "readiness run_id does not match the attempt") + _require(readiness.shard_id == attempt.shard_id, "readiness shard_id does not match the attempt") + _require(readiness.attempt_id == attempt.attempt_id, "readiness attempt_id does not match the attempt") + _require(readiness.revision == 1, "initial readiness must use revision 1") + _require(readiness.state is ReadinessState.PENDING, "initial readiness must be pending") + _require(readiness.updated_at >= attempt.created_at, "initial readiness cannot precede attempt creation") + + expected = tuple( + ( + deployment.deployment_id, + deployment.authored.model_alias, + deployment.topology.replica_count, + ) + for deployment in self.plan.deployments + ) + actual = tuple( + (deployment.deployment_id, deployment.model_alias, deployment.expected_backends) + for deployment in readiness.deployments + ) + _require(actual == expected, "initial readiness deployments do not match the resolved plan") + return readiness + + def validate_planned_attempt( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + ) -> AttemptManifest: + """Validate attempt identity and scheduler task ownership against one shard.""" + _require(attempt.run_id == self.plan.run_id, "attempt run_id does not match the resolved plan") + self._validate_plan_reference(attempt.resolved_plan) + canonical_shard = self._get_planned_shard(attempt.shard_id) + _require(canonical_shard == planned_shard, "planned shard is not the canonical shard for the attempt") + expected_attempt_id = f"attempt-{attempt.attempt_ordinal:04d}" + _require(attempt.attempt_id == expected_attempt_id, "attempt ID does not match its ordinal") + _require(attempt.scheduler is not None, "planned attempts require scheduler array-task identity") + _require( + attempt.scheduler.array_task_id == planned_shard.array_task_index, + "attempt scheduler array task does not match the planned shard", + ) + return attempt + + def validate_finalization_chain( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + winner: ShardWinner, + ) -> ShardWinner: + """Validate a complete semantic result through immutable winner publication.""" + self.validate_planned_attempt(planned_shard, attempt) + _require(attempt.state is AttemptLifecycleState.SUCCEEDED, "only successful attempts may be finalized") + _require( + attempt.terminal_classification is AttemptTerminalClassification.SUCCEEDED, + "only successfully classified attempts may be finalized", + ) + _require(client_result.outcome is ClientOutcome.COMPLETE, "only complete client results may be finalized") + _require(candidate.winner_eligible, "only complete candidate outputs may be finalized") + + for record_name, run_id in ( + ("client result", client_result.run_id), + ("candidate", candidate.run_id), + ("winner", winner.run_id), + ): + _require(run_id == self.plan.run_id, f"{record_name} run_id does not match the resolved plan") + for record_name, shard_id in ( + ("client result", client_result.shard_id), + ("candidate", candidate.shard_id), + ("winner", winner.shard_id), + ): + _require(shard_id == planned_shard.shard_id, f"{record_name} shard_id does not match the planned shard") + for record_name, attempt_id in ( + ("client result", client_result.attempt_id), + ("candidate", candidate.attempt_id), + ("winner", winner.attempt_id), + ): + _require(attempt_id == attempt.attempt_id, f"{record_name} attempt_id does not match the attempt") + + _require( + candidate.attempt_ordinal == attempt.attempt_ordinal == winner.attempt_ordinal, + "candidate and winner attempt ordinals must match the attempt", + ) + _require( + client_result.requested_records == candidate.requested_records == planned_shard.requested_records, + "client and candidate requested records must match the planned shard", + ) + _require( + client_result.actual_records == candidate.actual_records == planned_shard.requested_records, + "client and candidate actual records must complete the planned shard", + ) + _require( + client_result.requested_resume_mode == self.plan.invocation.authored.resume, + "client requested resume mode does not match the resolved plan", + ) + + expected_dataset_path = self._get_expected_dataset_path(planned_shard, attempt, client_result) + _require( + client_result.dataset_path == expected_dataset_path, "client dataset path does not match planned intent" + ) + _require( + candidate.dataset_path == expected_dataset_path, "candidate dataset path does not match planned intent" + ) + + expected_manifest_path = posixpath.join( + posixpath.dirname(planned_shard.resume_workspace.path), + "attempts", + attempt.attempt_id, + "output-manifest.json", + ) + candidate_reference = client_result.candidate_output_manifest + _require(candidate_reference is not None, "complete client result has no candidate manifest reference") + _require( + candidate_reference.path == expected_manifest_path, + "candidate manifest path does not match planned intent", + ) + _require( + candidate_reference.sha256 == candidate.compute_sha256(), + "client candidate digest does not match the manifest", + ) + _require( + attempt.candidate_output == candidate_reference, + "attempt candidate reference does not match client result", + ) + _require( + winner.candidate_manifest == candidate_reference, + "winner candidate reference does not match client result", + ) + + _require(candidate.created_at >= attempt.created_at, "candidate creation cannot precede attempt creation") + _require( + client_result.completed_at >= candidate.created_at, + "client completion cannot precede candidate creation", + ) + _require( + attempt.updated_at >= client_result.completed_at, "attempt completion cannot precede client completion" + ) + _require(winner.published_at >= attempt.updated_at, "winner publication cannot precede attempt completion") + return winner + + def _get_planned_shard(self, shard_id: ShardId) -> PlannedShard: + planned_shard = self._shards_by_id.get(shard_id) + if planned_shard is None: + raise IntegrationContractError("attempt shard_id does not identify a planned shard") + return planned_shard + + def _validate_plan_reference(self, reference: ArtifactReference) -> None: + _require( + reference.path == self._plan_reference.path, + "resolved plan reference path does not match planned intent", + ) + _require( + reference.sha256 == self._plan_reference.sha256, + "resolved plan reference digest does not match plan bytes", + ) + + @staticmethod + def _get_expected_dataset_path( + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + ) -> str: + if client_result.effective_resume_mode == "always": + return planned_shard.resume_workspace.path + return posixpath.join( + posixpath.dirname(planned_shard.resume_workspace.path), + "attempts", + attempt.attempt_id, + "dataset", + ) + + +def validate_plan_shards( + plan: ResolvedSlurmRunPlan, + run: RunManifest, + shards: tuple[ShardManifest, ...], +) -> tuple[ShardManifest, ...]: + """Validate state shards for a one-off plan; reuse ``PlanStateValidator`` for batches.""" + return PlanStateValidator(plan).validate_plan_shards(run, shards) + + +def validate_initial_readiness( + plan: ResolvedSlurmRunPlan, + attempt: AttemptManifest, + readiness: AttemptReadiness, +) -> AttemptReadiness: + """Validate initial readiness for a one-off plan; reuse ``PlanStateValidator`` for batches.""" + return PlanStateValidator(plan).validate_initial_readiness(attempt, readiness) + + +def validate_planned_attempt( + plan: ResolvedSlurmRunPlan, + planned_shard: PlannedShard, + attempt: AttemptManifest, +) -> AttemptManifest: + """Validate one planned attempt; reuse ``PlanStateValidator`` for batches.""" + return PlanStateValidator(plan).validate_planned_attempt(planned_shard, attempt) + + +def validate_finalization_chain( + plan: ResolvedSlurmRunPlan, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + winner: ShardWinner, +) -> ShardWinner: + """Validate one finalization chain; reuse ``PlanStateValidator`` for batches.""" + return PlanStateValidator(plan).validate_finalization_chain( + planned_shard, + attempt, + client_result, + candidate, + winner, + ) + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise IntegrationContractError(message) + + +__all__ = [ + "IntegrationContractError", + "PlanStateValidator", + "validate_finalization_chain", + "validate_initial_readiness", + "validate_plan_shards", + "validate_planned_attempt", +] diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json new file mode 100644 index 000000000..517a45ac5 --- /dev/null +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -0,0 +1,128 @@ +{ + "attempt": { + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "candidate_output": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + }, + "created_at": "2026-08-19T12:00:02Z", + "resolved_plan": { + "path": "/workspace/primary/runs/run-single/resolved-plan.json", + "sha256": "d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" + }, + "run_id": "run-single", + "scheduler": { + "array_job_id": 4101, + "array_task_id": 0 + }, + "schema_version": 1, + "shard_id": "shard-00000", + "state": "succeeded", + "terminal_classification": "succeeded", + "updated_at": "2026-08-19T12:05:02Z" + }, + "candidate": { + "actual_records": 8, + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "created_at": "2026-08-19T12:05:00Z", + "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", + "dataset_schema_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + "files": [ + { + "byte_size": 4096, + "record_count": 8, + "relative_path": "part-00000.parquet", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + } + ], + "outcome": "complete", + "provenance_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "requested_records": 8, + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000" + }, + "client_result": { + "actual_records": 8, + "attempt_id": "attempt-0001", + "candidate_output_manifest": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + }, + "completed_at": "2026-08-19T12:05:01Z", + "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", + "early_shutdown": false, + "effective_resume_mode": "never", + "outcome": "complete", + "requested_records": 8, + "requested_resume_mode": "never", + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000" + }, + "readiness": { + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 1, + "last_probe": null, + "model_alias": "generator", + "ready_backends": 0, + "state": "pending" + } + ], + "revision": 1, + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "pending", + "updated_at": "2026-08-19T12:00:03Z" + }, + "run": { + "authored_config": { + "path": "/workspace/primary/runs/run-single/authored-config.json", + "sha256": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" + }, + "created_at": "2026-08-19T12:00:00Z", + "resolved_plan": { + "path": "/workspace/primary/runs/run-single/resolved-plan.json", + "sha256": "d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" + }, + "run_id": "run-single", + "schema_version": 1, + "shard_count": 1 + }, + "shards": [ + { + "created_at": "2026-08-19T12:00:01Z", + "input_partition": null, + "record_range": { + "end_index_exclusive": 8, + "start_index": 0 + }, + "resume_workspace": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/dataset" + }, + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000", + "shard_index": 0 + } + ], + "winner": { + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "candidate_manifest": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + }, + "published_at": "2026-08-19T12:05:03Z", + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000" + } +} diff --git a/packages/data-designer-slurm/tests/integration/test_integration_validation.py b/packages/data-designer-slurm/tests/integration/test_integration_validation.py new file mode 100644 index 000000000..a440ec96f --- /dev/null +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -0,0 +1,452 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import posixpath +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import TypeVar + +import pytest + +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.contracts import ArtifactReference, ContractValue, RecordRange +from data_designer.slurm.integration import ( + IntegrationContractError, + PlanStateValidator, + validate_finalization_chain, + validate_initial_readiness, + validate_plan_shards, + validate_planned_attempt, +) +from data_designer.slurm.planning import ( + ArtifactReference as PlanningArtifactReference, +) +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + ArtifactReference as StateArtifactReference, +) +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutputManifest, + DeploymentReadiness, + EndpointPublicationState, + ReadinessState, + RunManifest, + SchedulerIdentity, + ShardManifest, + ShardWinner, +) + +TEST_ROOT = Path(__file__).parents[1] +CONTRACT_GOLDEN_DIR = TEST_ROOT / "contracts" / "golden" +INTEGRATION_GOLDEN_DIR = Path(__file__).parent / "golden" +CREATED_AT = datetime(2026, 8, 19, 12, 0, tzinfo=UTC) +_RecordT = TypeVar("_RecordT", bound=ContractValue) + + +@dataclass(frozen=True) +class IntegrationRecords: + plan: ResolvedSlurmRunPlan + validator: PlanStateValidator + run: RunManifest + shards: tuple[ShardManifest, ...] + attempt: AttemptManifest + readiness: AttemptReadiness + client_result: ClientResult + candidate: CandidateOutputManifest + winner: ShardWinner + + +@pytest.fixture +def records() -> IntegrationRecords: + plan = _load_plan("single_node_plan.json") + payload = json.loads((INTEGRATION_GOLDEN_DIR / "finalization_chain.json").read_text()) + return IntegrationRecords( + plan=plan, + validator=PlanStateValidator(plan), + run=_load_record(RunManifest, payload["run"]), + shards=tuple(_load_record(ShardManifest, shard) for shard in payload["shards"]), + attempt=_load_record(AttemptManifest, payload["attempt"]), + readiness=_load_record(AttemptReadiness, payload["readiness"]), + client_result=_load_record(ClientResult, payload["client_result"]), + candidate=_load_record(CandidateOutputManifest, payload["candidate"]), + winner=_load_record(ShardWinner, payload["winner"]), + ) + + +def test_shared_contract_types_retain_exact_identity() -> None: + assert PlanningArtifactReference is ArtifactReference + assert StateArtifactReference is ArtifactReference + + +def test_golden_records_validate_every_plan_state_join(records: IntegrationRecords) -> None: + planned_shard = records.plan.shards[0] + + assert records.validator.validate_plan_shards(records.run, records.shards) is records.shards + assert records.validator.validate_initial_readiness(records.attempt, records.readiness) is records.readiness + assert records.validator.validate_planned_attempt(planned_shard, records.attempt) is records.attempt + assert ( + records.validator.validate_finalization_chain( + planned_shard, + records.attempt, + records.client_result, + records.candidate, + records.winner, + ) + is records.winner + ) + + +def test_plan_shards_reject_missing_extra_and_mismatched_state(records: IntegrationRecords) -> None: + with pytest.raises(IntegrationContractError, match="exactly the run shard count"): + validate_plan_shards(records.plan, records.run, ()) + with pytest.raises(IntegrationContractError, match="exactly the run shard count"): + validate_plan_shards(records.plan, records.run, records.shards + records.shards) + + mismatched = records.shards[0].model_copy( + update={"record_range": RecordRange(start_index=1, end_index_exclusive=8)} + ) + with pytest.raises(IntegrationContractError, match="record range"): + validate_plan_shards(records.plan, records.run, (mismatched,)) + + +def test_plan_shards_reject_reordered_state() -> None: + plan = _load_plan("multi_node_plan.json") + run, shards = _state_shards_for_plan(plan) + + with pytest.raises(IntegrationContractError, match="ordered"): + validate_plan_shards(plan, run, tuple(reversed(shards))) + + +def test_initial_readiness_rejects_plan_order_alias_and_backend_count(records: IntegrationRecords) -> None: + deployment = records.readiness.deployments[0] + wrong_alias = deployment.model_copy(update={"model_alias": "other"}) + readiness = records.readiness.model_copy(update={"deployments": (wrong_alias,)}) + with pytest.raises(IntegrationContractError, match="deployments"): + validate_initial_readiness(records.plan, records.attempt, readiness) + + wrong_count = deployment.model_copy(update={"expected_backends": 2}) + readiness = records.readiness.model_copy(update={"deployments": (wrong_count,)}) + with pytest.raises(IntegrationContractError, match="deployments"): + validate_initial_readiness(records.plan, records.attempt, readiness) + + multi_plan = _load_plan("multi_node_plan.json") + attempt = _attempt_for_plan(multi_plan) + readiness = _pending_readiness(multi_plan, attempt) + reordered = readiness.model_copy(update={"deployments": tuple(reversed(readiness.deployments))}) + with pytest.raises(IntegrationContractError, match="deployments"): + validate_initial_readiness(multi_plan, attempt, reordered) + + +def test_initial_readiness_requires_first_pending_revision(records: IntegrationRecords) -> None: + with pytest.raises(IntegrationContractError, match="revision 1"): + validate_initial_readiness( + records.plan, + records.attempt, + records.readiness.model_copy(update={"revision": 2}), + ) + with pytest.raises(IntegrationContractError, match="must be pending"): + validate_initial_readiness( + records.plan, + records.attempt, + records.readiness.model_copy(update={"state": ReadinessState.READY}), + ) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("unplanned_shard", "planned shard"), + ("wrong_array_task", "array task"), + ("nondeterministic_attempt_id", "ordinal"), + ("unsubmitted_attempt", "scheduler"), + ], +) +def test_initial_readiness_rejects_invalid_planned_attempt( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + attempt = records.attempt + readiness = records.readiness + if mutation == "unplanned_shard": + attempt = attempt.model_copy(update={"shard_id": "shard-99999"}) + readiness = readiness.model_copy(update={"shard_id": attempt.shard_id}) + elif mutation == "wrong_array_task": + attempt = attempt.model_copy(update={"scheduler": SchedulerIdentity(array_job_id=4101, array_task_id=1)}) + elif mutation == "nondeterministic_attempt_id": + attempt = attempt.model_copy(update={"attempt_id": "attempt-0002"}) + readiness = readiness.model_copy(update={"attempt_id": attempt.attempt_id}) + else: + attempt = attempt.model_copy( + update={ + "state": AttemptLifecycleState.CREATED, + "scheduler": None, + "terminal_classification": None, + "candidate_output": None, + } + ) + + with pytest.raises(IntegrationContractError, match=message): + validate_initial_readiness(records.plan, attempt, readiness) + + +def test_planned_attempt_rejects_task_and_deterministic_identity_mismatch(records: IntegrationRecords) -> None: + planned_shard = records.plan.shards[0] + wrong_task = records.attempt.model_copy(update={"scheduler": SchedulerIdentity(array_job_id=4101, array_task_id=1)}) + with pytest.raises(IntegrationContractError, match="array task"): + validate_planned_attempt(records.plan, planned_shard, wrong_task) + + wrong_id = records.attempt.model_copy(update={"attempt_id": "attempt-0002"}) + with pytest.raises(IntegrationContractError, match="ordinal"): + validate_planned_attempt(records.plan, planned_shard, wrong_id) + + no_scheduler = records.attempt.model_copy( + update={ + "state": AttemptLifecycleState.CREATED, + "scheduler": None, + "terminal_classification": None, + "candidate_output": None, + } + ) + with pytest.raises(IntegrationContractError, match="scheduler"): + validate_planned_attempt(records.plan, planned_shard, no_scheduler) + + +def test_plan_state_validator_reuses_one_context_for_an_attempt_batch() -> None: + plan = _load_plan("multi_node_plan.json") + validator = PlanStateValidator(plan) + attempts = tuple(_attempt_for_planned_shard(plan, planned_shard) for planned_shard in plan.shards) + + assert ( + tuple( + validator.validate_planned_attempt(planned_shard, attempt) + for planned_shard, attempt in zip(plan.shards, attempts, strict=True) + ) + == attempts + ) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("partial_result", "complete client results"), + ("failed_result", "complete client results"), + ("failed_attempt", "successful attempts"), + ("stale_winner", "winner attempt_id"), + ("digest_mismatch", "digest"), + ("dataset_path_mismatch", "dataset path"), + ("client_before_candidate", "client completion"), + ("winner_before_attempt", "winner publication"), + ], +) +def test_finalization_rejects_invalid_chains( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + attempt = records.attempt + client_result = records.client_result + candidate = records.candidate + winner = records.winner + if mutation == "partial_result": + client_result = client_result.model_copy( + update={"outcome": ClientOutcome.PARTIAL, "actual_records": 4, "early_shutdown": True} + ) + elif mutation == "failed_result": + client_result = ClientResult.model_validate( + client_result.model_dump(mode="python") + | { + "actual_records": None, + "outcome": ClientOutcome.FAILED, + "dataset_path": None, + "early_shutdown": None, + "effective_resume_mode": None, + "candidate_output_manifest": None, + "error_code": "generation_failed", + "redacted_message": "generation failed", + } + ) + elif mutation == "failed_attempt": + attempt = AttemptManifest.model_validate( + attempt.model_dump(mode="python") + | { + "state": AttemptLifecycleState.FAILED, + "terminal_classification": AttemptTerminalClassification.FAILED, + "candidate_output": None, + } + ) + elif mutation == "stale_winner": + winner = winner.model_copy(update={"attempt_id": "attempt-0002"}) + elif mutation == "digest_mismatch": + candidate = candidate.model_copy(update={"provenance_digest": "c" * 64}) + elif mutation == "dataset_path_mismatch": + candidate = candidate.model_copy(update={"dataset_path": "/workspace/other/dataset"}) + elif mutation == "client_before_candidate": + client_result = client_result.model_copy(update={"completed_at": candidate.created_at - timedelta(seconds=1)}) + else: + winner = winner.model_copy(update={"published_at": records.attempt.updated_at - timedelta(seconds=1)}) + + with pytest.raises(IntegrationContractError, match=message): + validate_finalization_chain( + records.plan, + records.plan.shards[0], + attempt, + client_result, + candidate, + winner, + ) + + +def test_finalization_rejects_plan_reference_drift(records: IntegrationRecords) -> None: + drifted = records.attempt.model_copy( + update={ + "resolved_plan": ArtifactReference( + path=records.attempt.resolved_plan.path, + sha256="a" * 64, + ) + } + ) + + with pytest.raises(IntegrationContractError, match="digest"): + validate_planned_attempt(records.plan, records.plan.shards[0], drifted) + + +def test_finalization_accepts_effective_resume_always(records: IntegrationRecords) -> None: + plan_payload = records.plan.model_dump(mode="python") + plan_payload["invocation"]["authored"]["resume"] = "always" + plan = ResolvedSlurmRunPlan.model_validate(plan_payload) + planned_shard = plan.shards[0] + candidate = CandidateOutputManifest.model_validate( + records.candidate.model_dump(mode="python") | {"dataset_path": planned_shard.resume_workspace.path} + ) + original_candidate_reference = records.client_result.candidate_output_manifest + assert original_candidate_reference is not None + candidate_reference = ArtifactReference( + path=original_candidate_reference.path, + sha256=candidate.compute_sha256(), + ) + attempt = AttemptManifest.model_validate( + records.attempt.model_dump(mode="python") + | {"candidate_output": candidate_reference, "resolved_plan": _plan_reference(plan)} + ) + client_result = ClientResult.model_validate( + records.client_result.model_dump(mode="python") + | { + "requested_resume_mode": "always", + "effective_resume_mode": "always", + "dataset_path": planned_shard.resume_workspace.path, + "candidate_output_manifest": candidate_reference, + } + ) + winner = ShardWinner.model_validate( + records.winner.model_dump(mode="python") | {"candidate_manifest": candidate_reference} + ) + + assert ( + validate_finalization_chain( + plan, + planned_shard, + attempt, + client_result, + candidate, + winner, + ) + is winner + ) + + +def test_integration_golden_contains_no_environment_specific_values() -> None: + payload = (INTEGRATION_GOLDEN_DIR / "finalization_chain.json").read_text().casefold() + + for forbidden in ('"token":', '"password":', '"account":', '"partition":', "/users/", "/home/"): + assert forbidden not in payload + + +def _load_plan(name: str) -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIR / name).read_text()) + + +def _load_record(record_type: type[_RecordT], payload: object) -> _RecordT: + return record_type.model_validate_json(json.dumps(payload)) + + +def _plan_reference(plan: ResolvedSlurmRunPlan) -> ArtifactReference: + return ArtifactReference( + path=posixpath.join(posixpath.dirname(plan.authored_config.path), "resolved-plan.json"), + sha256=plan.compute_sha256(), + ) + + +def _state_shards_for_plan(plan: ResolvedSlurmRunPlan) -> tuple[RunManifest, tuple[ShardManifest, ...]]: + run = RunManifest( + schema_version=1, + run_id=plan.run_id, + created_at=CREATED_AT, + authored_config=plan.authored_config, + resolved_plan=_plan_reference(plan), + shard_count=len(plan.shards), + ) + shards = tuple( + ShardManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=planned.shard_id, + shard_index=planned.shard_index, + record_range=planned.record_range, + input_partition=planned.input_partition, + resume_workspace=planned.resume_workspace, + created_at=CREATED_AT + timedelta(seconds=index + 1), + ) + for index, planned in enumerate(plan.shards) + ) + return run, shards + + +def _attempt_for_plan(plan: ResolvedSlurmRunPlan) -> AttemptManifest: + return _attempt_for_planned_shard(plan, plan.shards[0]) + + +def _attempt_for_planned_shard(plan: ResolvedSlurmRunPlan, shard: PlannedShard) -> AttemptManifest: + return AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=_plan_reference(plan), + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=shard.array_task_index), + created_at=CREATED_AT, + updated_at=CREATED_AT, + ) + + +def _pending_readiness(plan: ResolvedSlurmRunPlan, attempt: AttemptManifest) -> AttemptReadiness: + return AttemptReadiness( + schema_version=1, + run_id=attempt.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + revision=1, + updated_at=attempt.created_at, + state=ReadinessState.PENDING, + deployments=tuple( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.PENDING, + expected_backends=deployment.topology.replica_count, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + for deployment in plan.deployments + ), + ) From 993058669844e2144acec1aaf423047b62aa9375 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 24 Aug 2026 13:55:52 -0600 Subject: [PATCH 2/5] fix: reject unsubmitted planned attempts Require scheduler-bound plan validation to start only after an attempt leaves the created state. Cover created attempts that already carry scheduler metadata and verify the integration validator imports from an isolated Slurm wheel. Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/integration.py | 9 ++++++++- .../test_integration_validation.py | 20 +++++++++++++++++-- scripts/test_slurm_package_install.py | 2 ++ 3 files changed, 28 insertions(+), 3 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/integration.py b/packages/data-designer-slurm/src/data_designer/slurm/integration.py index 33e151841..44414553e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/integration.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -34,7 +34,10 @@ class IntegrationContractError(ValueError): @dataclass(frozen=True, slots=True) class PlanStateValidator: - """Validate records against one resolved plan with reusable derived state.""" + """Validate records against one resolved plan with reusable derived state. + + This is an in-process validation service, not a persisted contract record. + """ plan: ResolvedSlurmRunPlan _plan_reference: ArtifactReference = field(init=False, repr=False) @@ -131,6 +134,10 @@ def validate_planned_attempt( _require(canonical_shard == planned_shard, "planned shard is not the canonical shard for the attempt") expected_attempt_id = f"attempt-{attempt.attempt_ordinal:04d}" _require(attempt.attempt_id == expected_attempt_id, "attempt ID does not match its ordinal") + _require( + attempt.state is not AttemptLifecycleState.CREATED, + "planned attempts must be submitted before scheduler task validation", + ) _require(attempt.scheduler is not None, "planned attempts require scheduler array-task identity") _require( attempt.scheduler.array_task_id == planned_shard.array_task_index, diff --git a/packages/data-designer-slurm/tests/integration/test_integration_validation.py b/packages/data-designer-slurm/tests/integration/test_integration_validation.py index a440ec96f..3003c91ba 100644 --- a/packages/data-designer-slurm/tests/integration/test_integration_validation.py +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -167,6 +167,7 @@ def test_initial_readiness_requires_first_pending_revision(records: IntegrationR ("wrong_array_task", "array task"), ("nondeterministic_attempt_id", "ordinal"), ("unsubmitted_attempt", "scheduler"), + ("unsubmitted_attempt_with_scheduler", "submitted"), ], ) def test_initial_readiness_rejects_invalid_planned_attempt( @@ -184,7 +185,7 @@ def test_initial_readiness_rejects_invalid_planned_attempt( elif mutation == "nondeterministic_attempt_id": attempt = attempt.model_copy(update={"attempt_id": "attempt-0002"}) readiness = readiness.model_copy(update={"attempt_id": attempt.attempt_id}) - else: + elif mutation == "unsubmitted_attempt": attempt = attempt.model_copy( update={ "state": AttemptLifecycleState.CREATED, @@ -193,12 +194,21 @@ def test_initial_readiness_rejects_invalid_planned_attempt( "candidate_output": None, } ) + else: + attempt = AttemptManifest.model_validate( + attempt.model_dump(mode="python") + | { + "state": AttemptLifecycleState.CREATED, + "terminal_classification": None, + "candidate_output": None, + } + ) with pytest.raises(IntegrationContractError, match=message): validate_initial_readiness(records.plan, attempt, readiness) -def test_planned_attempt_rejects_task_and_deterministic_identity_mismatch(records: IntegrationRecords) -> None: +def test_planned_attempt_rejects_task_identity_and_unsubmitted_state(records: IntegrationRecords) -> None: planned_shard = records.plan.shards[0] wrong_task = records.attempt.model_copy(update={"scheduler": SchedulerIdentity(array_job_id=4101, array_task_id=1)}) with pytest.raises(IntegrationContractError, match="array task"): @@ -219,6 +229,12 @@ def test_planned_attempt_rejects_task_and_deterministic_identity_mismatch(record with pytest.raises(IntegrationContractError, match="scheduler"): validate_planned_attempt(records.plan, planned_shard, no_scheduler) + created_with_scheduler = AttemptManifest.model_validate( + no_scheduler.model_dump(mode="python") | {"scheduler": records.attempt.scheduler} + ) + with pytest.raises(IntegrationContractError, match="submitted"): + validate_planned_attempt(records.plan, planned_shard, created_with_scheduler) + def test_plan_state_validator_reuses_one_context_for_an_attempt_batch() -> None: plan = _load_plan("multi_node_plan.json") diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index bf2cffd28..ef24f7ced 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -128,6 +128,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.contracts import ArtifactReference as ContractArtifactReference from data_designer.slurm.contracts import RecordRange as ContractRecordRange from data_designer.slurm.contracts import ResumeWorkspace as ContractResumeWorkspace +from data_designer.slurm.integration import PlanStateValidator from data_designer.slurm.planning import ArtifactReference as PlanningArtifactReference from data_designer.slurm.planning import RecordRange as PlanningRecordRange from data_designer.slurm.planning import ResumeWorkspace as PlanningResumeWorkspace @@ -136,6 +137,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.state import ResumeWorkspace as StateResumeWorkspace from data_designer.slurm.state import RunManifest assert RunManifest.__name__ == "RunManifest" +assert PlanStateValidator.__name__ == "PlanStateValidator" assert PlanningArtifactReference is ContractArtifactReference assert PlanningRecordRange is ContractRecordRange assert PlanningResumeWorkspace is ContractResumeWorkspace From b4ab6cc7c9b5a2996288925cfdf30fa68fb0f63c Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 24 Aug 2026 15:06:24 -0600 Subject: [PATCH 3/5] test: harden Slurm integration join coverage Signed-off-by: Nabin Mulepati --- .../test_integration_validation.py | 288 ++++++++++++++++-- scripts/test_slurm_package_install.py | 1 - 2 files changed, 268 insertions(+), 21 deletions(-) diff --git a/packages/data-designer-slurm/tests/integration/test_integration_validation.py b/packages/data-designer-slurm/tests/integration/test_integration_validation.py index 3003c91ba..b6b25981d 100644 --- a/packages/data-designer-slurm/tests/integration/test_integration_validation.py +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -6,14 +6,14 @@ import json import posixpath from dataclasses import dataclass -from datetime import UTC, datetime, timedelta +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import TypeVar import pytest from data_designer.slurm.client import ClientOutcome, ClientResult -from data_designer.slurm.contracts import ArtifactReference, ContractValue, RecordRange +from data_designer.slurm.contracts import ArtifactReference, ContractValue, RecordRange, ResumeWorkspace from data_designer.slurm.integration import ( IntegrationContractError, PlanStateValidator, @@ -34,6 +34,7 @@ AttemptManifest, AttemptReadiness, AttemptTerminalClassification, + CandidateOutcome, CandidateOutputManifest, DeploymentReadiness, EndpointPublicationState, @@ -47,7 +48,7 @@ TEST_ROOT = Path(__file__).parents[1] CONTRACT_GOLDEN_DIR = TEST_ROOT / "contracts" / "golden" INTEGRATION_GOLDEN_DIR = Path(__file__).parent / "golden" -CREATED_AT = datetime(2026, 8, 19, 12, 0, tzinfo=UTC) +CREATED_AT = datetime(2026, 8, 19, 12, 0, tzinfo=timezone.utc) _RecordT = TypeVar("_RecordT", bound=ContractValue) @@ -104,17 +105,91 @@ def test_golden_records_validate_every_plan_state_join(records: IntegrationRecor ) -def test_plan_shards_reject_missing_extra_and_mismatched_state(records: IntegrationRecords) -> None: +def test_plan_shards_reject_missing_and_extra_state(records: IntegrationRecords) -> None: with pytest.raises(IntegrationContractError, match="exactly the run shard count"): validate_plan_shards(records.plan, records.run, ()) with pytest.raises(IntegrationContractError, match="exactly the run shard count"): validate_plan_shards(records.plan, records.run, records.shards + records.shards) - mismatched = records.shards[0].model_copy( - update={"record_range": RecordRange(start_index=1, end_index_exclusive=8)} - ) - with pytest.raises(IntegrationContractError, match="record range"): - validate_plan_shards(records.plan, records.run, (mismatched,)) + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("run_id", "run manifest identity"), + ("authored_config", "authored config"), + ("plan_path", "reference path"), + ("plan_digest", "reference digest"), + ("shard_count", "run shard count"), + ("shard_id", "shard identity"), + ("shard_index", "outside the run shard count"), + ("record_range", "record range"), + ("input_partition", "input partition"), + ("resume_workspace", "resume workspace"), + ("created_at", "cannot precede"), + ], +) +def test_plan_shards_reject_one_field_drift( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + run = records.run + shard = records.shards[0] + if mutation == "run_id": + run = run.model_copy(update={"run_id": "run-other"}) + elif mutation == "authored_config": + run = run.model_copy( + update={ + "authored_config": ArtifactReference( + path=run.authored_config.path, + sha256="a" * 64, + ) + } + ) + elif mutation == "plan_path": + run = run.model_copy( + update={ + "resolved_plan": ArtifactReference( + path=posixpath.join(posixpath.dirname(run.resolved_plan.path), "other-plan.json"), + sha256=run.resolved_plan.sha256, + ) + } + ) + elif mutation == "plan_digest": + run = run.model_copy( + update={ + "resolved_plan": ArtifactReference( + path=run.resolved_plan.path, + sha256="a" * 64, + ) + } + ) + elif mutation == "shard_count": + run = run.model_copy(update={"shard_count": 2}) + elif mutation == "shard_id": + shard = shard.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "shard_index": + shard = shard.model_copy(update={"shard_index": 1}) + elif mutation == "record_range": + shard = shard.model_copy(update={"record_range": RecordRange(start_index=1, end_index_exclusive=8)}) + elif mutation == "input_partition": + shard = shard.model_copy( + update={ + "input_partition": ArtifactReference( + path=posixpath.join(posixpath.dirname(shard.resume_workspace.path), "input-partition.json"), + sha256="a" * 64, + ) + } + ) + elif mutation == "resume_workspace": + shard = shard.model_copy( + update={"resume_workspace": ResumeWorkspace(path=f"{shard.resume_workspace.path}-other")} + ) + else: + shard = shard.model_copy(update={"created_at": run.created_at - timedelta(seconds=1)}) + + with pytest.raises(IntegrationContractError, match=message): + validate_plan_shards(records.plan, run, (shard,)) def test_plan_shards_reject_reordered_state() -> None: @@ -145,6 +220,38 @@ def test_initial_readiness_rejects_plan_order_alias_and_backend_count(records: I validate_initial_readiness(multi_plan, attempt, reordered) +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("run_id", "readiness run_id"), + ("shard_id", "readiness shard_id"), + ("attempt_id", "readiness attempt_id"), + ("updated_at", "cannot precede"), + ("deployment_id", "deployments"), + ], +) +def test_initial_readiness_rejects_one_field_identity_and_time_drift( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + readiness = records.readiness + if mutation == "run_id": + readiness = readiness.model_copy(update={"run_id": "run-other"}) + elif mutation == "shard_id": + readiness = readiness.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "attempt_id": + readiness = readiness.model_copy(update={"attempt_id": "attempt-0002"}) + elif mutation == "updated_at": + readiness = readiness.model_copy(update={"updated_at": records.attempt.created_at - timedelta(seconds=1)}) + else: + deployment = readiness.deployments[0].model_copy(update={"deployment_id": "deployment-99999"}) + readiness = readiness.model_copy(update={"deployments": (deployment,)}) + + with pytest.raises(IntegrationContractError, match=message): + validate_initial_readiness(records.plan, records.attempt, readiness) + + def test_initial_readiness_requires_first_pending_revision(records: IntegrationRecords) -> None: with pytest.raises(IntegrationContractError, match="revision 1"): validate_initial_readiness( @@ -208,8 +315,12 @@ def test_initial_readiness_rejects_invalid_planned_attempt( validate_initial_readiness(records.plan, attempt, readiness) -def test_planned_attempt_rejects_task_identity_and_unsubmitted_state(records: IntegrationRecords) -> None: +def test_planned_attempt_rejects_identity_task_and_unsubmitted_state(records: IntegrationRecords) -> None: planned_shard = records.plan.shards[0] + wrong_run = records.attempt.model_copy(update={"run_id": "run-other"}) + with pytest.raises(IntegrationContractError, match="attempt run_id"): + validate_planned_attempt(records.plan, planned_shard, wrong_run) + wrong_task = records.attempt.model_copy(update={"scheduler": SchedulerIdentity(array_job_id=4101, array_task_id=1)}) with pytest.raises(IntegrationContractError, match="array task"): validate_planned_attempt(records.plan, planned_shard, wrong_task) @@ -235,6 +346,10 @@ def test_planned_attempt_rejects_task_identity_and_unsubmitted_state(records: In with pytest.raises(IntegrationContractError, match="submitted"): validate_planned_attempt(records.plan, planned_shard, created_with_scheduler) + multi_plan = _load_plan("multi_node_plan.json") + with pytest.raises(IntegrationContractError, match="canonical shard"): + validate_planned_attempt(multi_plan, multi_plan.shards[1], _attempt_for_plan(multi_plan)) + def test_plan_state_validator_reuses_one_context_for_an_attempt_batch() -> None: plan = _load_plan("multi_node_plan.json") @@ -256,6 +371,8 @@ def test_plan_state_validator_reuses_one_context_for_an_attempt_batch() -> None: ("partial_result", "complete client results"), ("failed_result", "complete client results"), ("failed_attempt", "successful attempts"), + ("wrong_terminal_classification", "successfully classified"), + ("ineligible_candidate", "complete candidate outputs"), ("stale_winner", "winner attempt_id"), ("digest_mismatch", "digest"), ("dataset_path_mismatch", "dataset path"), @@ -299,6 +416,10 @@ def test_finalization_rejects_invalid_chains( "candidate_output": None, } ) + elif mutation == "wrong_terminal_classification": + attempt = attempt.model_copy(update={"terminal_classification": AttemptTerminalClassification.FAILED}) + elif mutation == "ineligible_candidate": + candidate = candidate.model_copy(update={"outcome": CandidateOutcome.PARTIAL}) elif mutation == "stale_winner": winner = winner.model_copy(update={"attempt_id": "attempt-0002"}) elif mutation == "digest_mismatch": @@ -321,17 +442,144 @@ def test_finalization_rejects_invalid_chains( ) -def test_finalization_rejects_plan_reference_drift(records: IntegrationRecords) -> None: - drifted = records.attempt.model_copy( - update={ - "resolved_plan": ArtifactReference( - path=records.attempt.resolved_plan.path, - sha256="a" * 64, - ) - } - ) +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("client_run_id", "client result run_id"), + ("candidate_run_id", "candidate run_id"), + ("winner_run_id", "winner run_id"), + ("client_shard_id", "client result shard_id"), + ("candidate_shard_id", "candidate shard_id"), + ("winner_shard_id", "winner shard_id"), + ("client_attempt_id", "client result attempt_id"), + ("candidate_attempt_id", "candidate attempt_id"), + ("candidate_ordinal", "attempt ordinals"), + ("winner_ordinal", "attempt ordinals"), + ("client_requested_records", "requested records"), + ("candidate_requested_records", "requested records"), + ("client_actual_records", "actual records"), + ("candidate_actual_records", "actual records"), + ("requested_resume_mode", "resume mode"), + ("effective_resume_mode", "client dataset path"), + ("client_dataset_path", "client dataset path"), + ("missing_candidate_reference", "no candidate manifest"), + ("candidate_reference_path", "candidate manifest path"), + ("attempt_candidate_reference", "attempt candidate reference"), + ("winner_candidate_reference", "winner candidate reference"), + ("candidate_before_attempt", "candidate creation"), + ("attempt_before_client", "attempt completion"), + ], +) +def test_finalization_rejects_one_field_join_drift( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + attempt = records.attempt + client_result = records.client_result + candidate = records.candidate + winner = records.winner + candidate_reference = client_result.candidate_output_manifest + assert candidate_reference is not None + if mutation == "client_run_id": + client_result = client_result.model_copy(update={"run_id": "run-other"}) + elif mutation == "candidate_run_id": + candidate = candidate.model_copy(update={"run_id": "run-other"}) + elif mutation == "winner_run_id": + winner = winner.model_copy(update={"run_id": "run-other"}) + elif mutation == "client_shard_id": + client_result = client_result.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "candidate_shard_id": + candidate = candidate.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "winner_shard_id": + winner = winner.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "client_attempt_id": + client_result = client_result.model_copy(update={"attempt_id": "attempt-0002"}) + elif mutation == "candidate_attempt_id": + candidate = candidate.model_copy(update={"attempt_id": "attempt-0002"}) + elif mutation == "candidate_ordinal": + candidate = candidate.model_copy(update={"attempt_ordinal": 2}) + elif mutation == "winner_ordinal": + winner = winner.model_copy(update={"attempt_ordinal": 2}) + elif mutation == "client_requested_records": + client_result = client_result.model_copy(update={"requested_records": 7}) + elif mutation == "candidate_requested_records": + candidate = candidate.model_copy(update={"requested_records": 7}) + elif mutation == "client_actual_records": + client_result = client_result.model_copy(update={"actual_records": 7}) + elif mutation == "candidate_actual_records": + candidate = candidate.model_copy(update={"actual_records": 7}) + elif mutation == "requested_resume_mode": + client_result = client_result.model_copy(update={"requested_resume_mode": "if_possible"}) + elif mutation == "effective_resume_mode": + client_result = client_result.model_copy(update={"effective_resume_mode": "always"}) + elif mutation == "client_dataset_path": + client_result = client_result.model_copy(update={"dataset_path": "/workspace/other/dataset"}) + elif mutation == "missing_candidate_reference": + client_result = client_result.model_copy(update={"candidate_output_manifest": None}) + elif mutation == "candidate_reference_path": + client_result = client_result.model_copy( + update={ + "candidate_output_manifest": ArtifactReference( + path=posixpath.join(posixpath.dirname(candidate_reference.path), "other-manifest.json"), + sha256=candidate_reference.sha256, + ) + } + ) + elif mutation == "attempt_candidate_reference": + attempt = attempt.model_copy( + update={ + "candidate_output": ArtifactReference( + path=candidate_reference.path, + sha256="a" * 64, + ) + } + ) + elif mutation == "winner_candidate_reference": + winner = winner.model_copy( + update={ + "candidate_manifest": ArtifactReference( + path=candidate_reference.path, + sha256="a" * 64, + ) + } + ) + elif mutation == "candidate_before_attempt": + candidate = candidate.model_copy(update={"created_at": attempt.created_at - timedelta(seconds=1)}) + candidate_reference = ArtifactReference( + path=candidate_reference.path, + sha256=candidate.compute_sha256(), + ) + client_result = client_result.model_copy(update={"candidate_output_manifest": candidate_reference}) + attempt = attempt.model_copy(update={"candidate_output": candidate_reference}) + winner = winner.model_copy(update={"candidate_manifest": candidate_reference}) + else: + attempt = attempt.model_copy(update={"updated_at": client_result.completed_at - timedelta(seconds=1)}) + + with pytest.raises(IntegrationContractError, match=message): + validate_finalization_chain( + records.plan, + records.plan.shards[0], + attempt, + client_result, + candidate, + winner, + ) + + +@pytest.mark.parametrize("mutation", ["path", "digest"]) +def test_planned_attempt_rejects_plan_reference_drift(records: IntegrationRecords, mutation: str) -> None: + reference = records.attempt.resolved_plan + if mutation == "path": + reference = ArtifactReference( + path=posixpath.join(posixpath.dirname(reference.path), "other-plan.json"), + sha256=reference.sha256, + ) + else: + reference = ArtifactReference(path=reference.path, sha256="a" * 64) + drifted = records.attempt.model_copy(update={"resolved_plan": reference}) - with pytest.raises(IntegrationContractError, match="digest"): + with pytest.raises(IntegrationContractError, match=mutation): validate_planned_attempt(records.plan, records.plan.shards[0], drifted) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index ef24f7ced..f6ba85e19 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -137,7 +137,6 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.state import ResumeWorkspace as StateResumeWorkspace from data_designer.slurm.state import RunManifest assert RunManifest.__name__ == "RunManifest" -assert PlanStateValidator.__name__ == "PlanStateValidator" assert PlanningArtifactReference is ContractArtifactReference assert PlanningRecordRange is ContractRecordRange assert PlanningResumeWorkspace is ContractResumeWorkspace From 3a91aa9a87caf0457cfe4085fa858b70bba93c95 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 24 Aug 2026 15:20:26 -0600 Subject: [PATCH 4/5] refactor: simplify plan state validator Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/integration.py | 32 ++++++++----------- 1 file changed, 13 insertions(+), 19 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/integration.py b/packages/data-designer-slurm/src/data_designer/slurm/integration.py index 44414553e..58f33c3bf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/integration.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -7,7 +7,6 @@ import posixpath from collections.abc import Mapping -from dataclasses import dataclass, field from types import MappingProxyType from data_designer.slurm.client import ClientOutcome, ClientResult @@ -32,33 +31,28 @@ class IntegrationContractError(ValueError): """Raised when reviewed plan and state records disagree.""" -@dataclass(frozen=True, slots=True) class PlanStateValidator: """Validate records against one resolved plan with reusable derived state. This is an in-process validation service, not a persisted contract record. """ - plan: ResolvedSlurmRunPlan - _plan_reference: ArtifactReference = field(init=False, repr=False) - _shards_by_id: Mapping[ShardId, PlannedShard] = field(init=False, repr=False) - - def __post_init__(self) -> None: - run_root = posixpath.dirname(self.plan.authored_config.path) - object.__setattr__( - self, - "_plan_reference", - ArtifactReference( - path=posixpath.join(run_root, "resolved-plan.json"), - sha256=self.plan.compute_sha256(), - ), + def __init__(self, plan: ResolvedSlurmRunPlan) -> None: + self._plan: ResolvedSlurmRunPlan = plan + run_root = posixpath.dirname(plan.authored_config.path) + self._plan_reference: ArtifactReference = ArtifactReference( + path=posixpath.join(run_root, "resolved-plan.json"), + sha256=plan.compute_sha256(), ) - object.__setattr__( - self, - "_shards_by_id", - MappingProxyType({planned_shard.shard_id: planned_shard for planned_shard in self.plan.shards}), + self._shards_by_id: Mapping[ShardId, PlannedShard] = MappingProxyType( + {planned_shard.shard_id: planned_shard for planned_shard in plan.shards} ) + @property + def plan(self) -> ResolvedSlurmRunPlan: + """The resolved plan used to build this validation context.""" + return self._plan + def validate_plan_shards( self, run: RunManifest, From 4f0fea1d2198d0088790fa21408585d3e065d21f Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 24 Aug 2026 15:27:58 -0600 Subject: [PATCH 5/5] refactor: remove validator wrapper functions Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/integration.py | 49 --------------- .../test_integration_validation.py | 60 ++++++++----------- 2 files changed, 26 insertions(+), 83 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/integration.py b/packages/data-designer-slurm/src/data_designer/slurm/integration.py index 58f33c3bf..5f968fe73 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/integration.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -269,51 +269,6 @@ def _get_expected_dataset_path( ) -def validate_plan_shards( - plan: ResolvedSlurmRunPlan, - run: RunManifest, - shards: tuple[ShardManifest, ...], -) -> tuple[ShardManifest, ...]: - """Validate state shards for a one-off plan; reuse ``PlanStateValidator`` for batches.""" - return PlanStateValidator(plan).validate_plan_shards(run, shards) - - -def validate_initial_readiness( - plan: ResolvedSlurmRunPlan, - attempt: AttemptManifest, - readiness: AttemptReadiness, -) -> AttemptReadiness: - """Validate initial readiness for a one-off plan; reuse ``PlanStateValidator`` for batches.""" - return PlanStateValidator(plan).validate_initial_readiness(attempt, readiness) - - -def validate_planned_attempt( - plan: ResolvedSlurmRunPlan, - planned_shard: PlannedShard, - attempt: AttemptManifest, -) -> AttemptManifest: - """Validate one planned attempt; reuse ``PlanStateValidator`` for batches.""" - return PlanStateValidator(plan).validate_planned_attempt(planned_shard, attempt) - - -def validate_finalization_chain( - plan: ResolvedSlurmRunPlan, - planned_shard: PlannedShard, - attempt: AttemptManifest, - client_result: ClientResult, - candidate: CandidateOutputManifest, - winner: ShardWinner, -) -> ShardWinner: - """Validate one finalization chain; reuse ``PlanStateValidator`` for batches.""" - return PlanStateValidator(plan).validate_finalization_chain( - planned_shard, - attempt, - client_result, - candidate, - winner, - ) - - def _require(condition: bool, message: str) -> None: if not condition: raise IntegrationContractError(message) @@ -322,8 +277,4 @@ def _require(condition: bool, message: str) -> None: __all__ = [ "IntegrationContractError", "PlanStateValidator", - "validate_finalization_chain", - "validate_initial_readiness", - "validate_plan_shards", - "validate_planned_attempt", ] diff --git a/packages/data-designer-slurm/tests/integration/test_integration_validation.py b/packages/data-designer-slurm/tests/integration/test_integration_validation.py index b6b25981d..17e99ea2c 100644 --- a/packages/data-designer-slurm/tests/integration/test_integration_validation.py +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -14,14 +14,7 @@ from data_designer.slurm.client import ClientOutcome, ClientResult from data_designer.slurm.contracts import ArtifactReference, ContractValue, RecordRange, ResumeWorkspace -from data_designer.slurm.integration import ( - IntegrationContractError, - PlanStateValidator, - validate_finalization_chain, - validate_initial_readiness, - validate_plan_shards, - validate_planned_attempt, -) +from data_designer.slurm.integration import IntegrationContractError, PlanStateValidator from data_designer.slurm.planning import ( ArtifactReference as PlanningArtifactReference, ) @@ -107,9 +100,9 @@ def test_golden_records_validate_every_plan_state_join(records: IntegrationRecor def test_plan_shards_reject_missing_and_extra_state(records: IntegrationRecords) -> None: with pytest.raises(IntegrationContractError, match="exactly the run shard count"): - validate_plan_shards(records.plan, records.run, ()) + records.validator.validate_plan_shards(records.run, ()) with pytest.raises(IntegrationContractError, match="exactly the run shard count"): - validate_plan_shards(records.plan, records.run, records.shards + records.shards) + records.validator.validate_plan_shards(records.run, records.shards + records.shards) @pytest.mark.parametrize( @@ -189,15 +182,16 @@ def test_plan_shards_reject_one_field_drift( shard = shard.model_copy(update={"created_at": run.created_at - timedelta(seconds=1)}) with pytest.raises(IntegrationContractError, match=message): - validate_plan_shards(records.plan, run, (shard,)) + records.validator.validate_plan_shards(run, (shard,)) def test_plan_shards_reject_reordered_state() -> None: plan = _load_plan("multi_node_plan.json") + validator = PlanStateValidator(plan) run, shards = _state_shards_for_plan(plan) with pytest.raises(IntegrationContractError, match="ordered"): - validate_plan_shards(plan, run, tuple(reversed(shards))) + validator.validate_plan_shards(run, tuple(reversed(shards))) def test_initial_readiness_rejects_plan_order_alias_and_backend_count(records: IntegrationRecords) -> None: @@ -205,19 +199,20 @@ def test_initial_readiness_rejects_plan_order_alias_and_backend_count(records: I wrong_alias = deployment.model_copy(update={"model_alias": "other"}) readiness = records.readiness.model_copy(update={"deployments": (wrong_alias,)}) with pytest.raises(IntegrationContractError, match="deployments"): - validate_initial_readiness(records.plan, records.attempt, readiness) + records.validator.validate_initial_readiness(records.attempt, readiness) wrong_count = deployment.model_copy(update={"expected_backends": 2}) readiness = records.readiness.model_copy(update={"deployments": (wrong_count,)}) with pytest.raises(IntegrationContractError, match="deployments"): - validate_initial_readiness(records.plan, records.attempt, readiness) + records.validator.validate_initial_readiness(records.attempt, readiness) multi_plan = _load_plan("multi_node_plan.json") + validator = PlanStateValidator(multi_plan) attempt = _attempt_for_plan(multi_plan) readiness = _pending_readiness(multi_plan, attempt) reordered = readiness.model_copy(update={"deployments": tuple(reversed(readiness.deployments))}) with pytest.raises(IntegrationContractError, match="deployments"): - validate_initial_readiness(multi_plan, attempt, reordered) + validator.validate_initial_readiness(attempt, reordered) @pytest.mark.parametrize( @@ -249,19 +244,17 @@ def test_initial_readiness_rejects_one_field_identity_and_time_drift( readiness = readiness.model_copy(update={"deployments": (deployment,)}) with pytest.raises(IntegrationContractError, match=message): - validate_initial_readiness(records.plan, records.attempt, readiness) + records.validator.validate_initial_readiness(records.attempt, readiness) def test_initial_readiness_requires_first_pending_revision(records: IntegrationRecords) -> None: with pytest.raises(IntegrationContractError, match="revision 1"): - validate_initial_readiness( - records.plan, + records.validator.validate_initial_readiness( records.attempt, records.readiness.model_copy(update={"revision": 2}), ) with pytest.raises(IntegrationContractError, match="must be pending"): - validate_initial_readiness( - records.plan, + records.validator.validate_initial_readiness( records.attempt, records.readiness.model_copy(update={"state": ReadinessState.READY}), ) @@ -312,22 +305,22 @@ def test_initial_readiness_rejects_invalid_planned_attempt( ) with pytest.raises(IntegrationContractError, match=message): - validate_initial_readiness(records.plan, attempt, readiness) + records.validator.validate_initial_readiness(attempt, readiness) def test_planned_attempt_rejects_identity_task_and_unsubmitted_state(records: IntegrationRecords) -> None: planned_shard = records.plan.shards[0] wrong_run = records.attempt.model_copy(update={"run_id": "run-other"}) with pytest.raises(IntegrationContractError, match="attempt run_id"): - validate_planned_attempt(records.plan, planned_shard, wrong_run) + records.validator.validate_planned_attempt(planned_shard, wrong_run) wrong_task = records.attempt.model_copy(update={"scheduler": SchedulerIdentity(array_job_id=4101, array_task_id=1)}) with pytest.raises(IntegrationContractError, match="array task"): - validate_planned_attempt(records.plan, planned_shard, wrong_task) + records.validator.validate_planned_attempt(planned_shard, wrong_task) wrong_id = records.attempt.model_copy(update={"attempt_id": "attempt-0002"}) with pytest.raises(IntegrationContractError, match="ordinal"): - validate_planned_attempt(records.plan, planned_shard, wrong_id) + records.validator.validate_planned_attempt(planned_shard, wrong_id) no_scheduler = records.attempt.model_copy( update={ @@ -338,17 +331,18 @@ def test_planned_attempt_rejects_identity_task_and_unsubmitted_state(records: In } ) with pytest.raises(IntegrationContractError, match="scheduler"): - validate_planned_attempt(records.plan, planned_shard, no_scheduler) + records.validator.validate_planned_attempt(planned_shard, no_scheduler) created_with_scheduler = AttemptManifest.model_validate( no_scheduler.model_dump(mode="python") | {"scheduler": records.attempt.scheduler} ) with pytest.raises(IntegrationContractError, match="submitted"): - validate_planned_attempt(records.plan, planned_shard, created_with_scheduler) + records.validator.validate_planned_attempt(planned_shard, created_with_scheduler) multi_plan = _load_plan("multi_node_plan.json") + validator = PlanStateValidator(multi_plan) with pytest.raises(IntegrationContractError, match="canonical shard"): - validate_planned_attempt(multi_plan, multi_plan.shards[1], _attempt_for_plan(multi_plan)) + validator.validate_planned_attempt(multi_plan.shards[1], _attempt_for_plan(multi_plan)) def test_plan_state_validator_reuses_one_context_for_an_attempt_batch() -> None: @@ -432,8 +426,7 @@ def test_finalization_rejects_invalid_chains( winner = winner.model_copy(update={"published_at": records.attempt.updated_at - timedelta(seconds=1)}) with pytest.raises(IntegrationContractError, match=message): - validate_finalization_chain( - records.plan, + records.validator.validate_finalization_chain( records.plan.shards[0], attempt, client_result, @@ -557,8 +550,7 @@ def test_finalization_rejects_one_field_join_drift( attempt = attempt.model_copy(update={"updated_at": client_result.completed_at - timedelta(seconds=1)}) with pytest.raises(IntegrationContractError, match=message): - validate_finalization_chain( - records.plan, + records.validator.validate_finalization_chain( records.plan.shards[0], attempt, client_result, @@ -580,7 +572,7 @@ def test_planned_attempt_rejects_plan_reference_drift(records: IntegrationRecord drifted = records.attempt.model_copy(update={"resolved_plan": reference}) with pytest.raises(IntegrationContractError, match=mutation): - validate_planned_attempt(records.plan, records.plan.shards[0], drifted) + records.validator.validate_planned_attempt(records.plan.shards[0], drifted) def test_finalization_accepts_effective_resume_always(records: IntegrationRecords) -> None: @@ -613,10 +605,10 @@ def test_finalization_accepts_effective_resume_always(records: IntegrationRecord winner = ShardWinner.model_validate( records.winner.model_dump(mode="python") | {"candidate_manifest": candidate_reference} ) + validator = PlanStateValidator(plan) assert ( - validate_finalization_chain( - plan, + validator.validate_finalization_chain( planned_shard, attempt, client_result,