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..5f968fe73 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -0,0 +1,280 @@ +# 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 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.""" + + +class PlanStateValidator: + """Validate records against one resolved plan with reusable derived state. + + This is an in-process validation service, not a persisted contract record. + """ + + 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(), + ) + 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, + 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.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, + "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 _require(condition: bool, message: str) -> None: + if not condition: + raise IntegrationContractError(message) + + +__all__ = [ + "IntegrationContractError", + "PlanStateValidator", +] 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..17e99ea2c --- /dev/null +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -0,0 +1,708 @@ +# 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 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, ResumeWorkspace +from data_designer.slurm.integration import IntegrationContractError, PlanStateValidator +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, + CandidateOutcome, + 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=timezone.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_and_extra_state(records: IntegrationRecords) -> None: + with pytest.raises(IntegrationContractError, match="exactly the run shard count"): + records.validator.validate_plan_shards(records.run, ()) + with pytest.raises(IntegrationContractError, match="exactly the run shard count"): + records.validator.validate_plan_shards(records.run, records.shards + records.shards) + + +@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): + 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"): + validator.validate_plan_shards(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"): + 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"): + 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"): + validator.validate_initial_readiness(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): + 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"): + records.validator.validate_initial_readiness( + records.attempt, + records.readiness.model_copy(update={"revision": 2}), + ) + with pytest.raises(IntegrationContractError, match="must be pending"): + records.validator.validate_initial_readiness( + 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"), + ("unsubmitted_attempt_with_scheduler", "submitted"), + ], +) +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}) + elif mutation == "unsubmitted_attempt": + attempt = attempt.model_copy( + update={ + "state": AttemptLifecycleState.CREATED, + "scheduler": None, + "terminal_classification": None, + "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): + 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"): + 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"): + 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"): + records.validator.validate_planned_attempt(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"): + 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"): + 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"): + 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: + 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"), + ("wrong_terminal_classification", "successfully classified"), + ("ineligible_candidate", "complete candidate outputs"), + ("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 == "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": + 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): + records.validator.validate_finalization_chain( + records.plan.shards[0], + attempt, + client_result, + candidate, + winner, + ) + + +@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): + records.validator.validate_finalization_chain( + 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=mutation): + records.validator.validate_planned_attempt(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} + ) + validator = PlanStateValidator(plan) + + assert ( + validator.validate_finalization_chain( + 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 + ), + ) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index bf2cffd28..f6ba85e19 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