diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 0fd4ec5ff..340c8f289 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -37,6 +37,7 @@ bump = true [tool.hatch.metadata.hooks.uv-dynamic-versioning] dependencies = [ "data-designer=={{ version }}", + "packaging>=25,<27", "pydantic>=2.9.2,<3", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py new file mode 100644 index 000000000..40acb8f5b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Immutable benchmark records for Data Designer Slurm.""" + +from __future__ import annotations + +from data_designer.slurm.benchmark.records import ( + BenchmarkCaseResult, + BenchmarkChildRun, + BenchmarkManifest, + BenchmarkOutcome, + BenchmarkRecommendation, + BenchmarkRecommendationKind, + BenchmarkReport, +) + +__all__ = [ + "BenchmarkCaseResult", + "BenchmarkChildRun", + "BenchmarkManifest", + "BenchmarkOutcome", + "BenchmarkRecommendation", + "BenchmarkRecommendationKind", + "BenchmarkReport", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py new file mode 100644 index 000000000..11c797527 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timedelta +from enum import Enum +from typing import Annotated + +from pydantic import ( + Field, + NonNegativeFloat, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) + +from data_designer.slurm.contracts import ArtifactReference, ContractRecord, ContractValue, Identifier + + +class BenchmarkChildRun(ContractValue): + case_id: Identifier + child_run_id: Identifier + child_authored_config: ArtifactReference + + @model_validator(mode="after") + def validate_authored_config(self) -> BenchmarkChildRun: + expected_suffix = f"/runs/{self.child_run_id}/authored-config.json" + if not self.child_authored_config.path.endswith(expected_suffix): + raise ValueError("child authored config path must match the child run identity") + return self + + +class BenchmarkManifest(ContractRecord): + """Stable mapping from benchmark cases to ordinary child runs.""" + + benchmark_id: Identifier + benchmark_config: ArtifactReference + children: tuple[BenchmarkChildRun, ...] = Field(min_length=1) + + @model_validator(mode="after") + def validate_children(self) -> BenchmarkManifest: + case_ids = tuple(child.case_id for child in self.children) + run_ids = tuple(child.child_run_id for child in self.children) + if len(case_ids) != len(set(case_ids)): + raise ValueError("benchmark case IDs must be unique") + if len(run_ids) != len(set(run_ids)): + raise ValueError("benchmark child run IDs must be unique") + return self + + +class BenchmarkOutcome(str, Enum): + PENDING = "pending" + ACCOUNTING_LAG = "accounting_lag" + SUCCEEDED = "succeeded" + FAILED = "failed" + INCOMPLETE = "incomplete" + + +class BenchmarkCaseResult(ContractValue): + case_id: Identifier + child_run_id: Identifier + outcome: BenchmarkOutcome + topology_digest: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + requested_records: PositiveInt + actual_records: NonNegativeInt | None = None + boot_seconds: NonNegativeFloat | None = None + generation_seconds: NonNegativeFloat | None = None + wall_seconds: NonNegativeFloat | None = None + rows_per_second: NonNegativeFloat | None = None + request_count: NonNegativeInt | None = None + token_count: NonNegativeInt | None = None + gpus_per_job: PositiveInt + nodes_per_job: PositiveInt + gpu_hours_per_job: NonNegativeFloat | None = None + total_gpu_hours: NonNegativeFloat | None = None + target_jobs: PositiveInt | None = None + feasible: bool | None = None + + @model_validator(mode="after") + def validate_metrics(self) -> BenchmarkCaseResult: + if self.actual_records is not None and self.actual_records > self.requested_records: + raise ValueError("benchmark actual_records must not exceed requested_records") + required = ( + self.actual_records, + self.boot_seconds, + self.generation_seconds, + self.wall_seconds, + self.rows_per_second, + self.gpu_hours_per_job, + self.total_gpu_hours, + self.target_jobs, + self.feasible, + ) + if self.outcome is BenchmarkOutcome.SUCCEEDED and any(value is None for value in required): + raise ValueError("successful benchmark cases require complete timing and feasibility metrics") + if self.outcome is BenchmarkOutcome.SUCCEEDED: + if self.actual_records != self.requested_records: + raise ValueError("successful benchmark cases require the requested record count") + if self.generation_seconds == 0 or self.wall_seconds == 0 or self.rows_per_second == 0: + raise ValueError("successful benchmark generation, wall time, and throughput must be positive") + return self + + +class BenchmarkRecommendationKind(str, Enum): + PARETO = "pareto" + MINIMUM_JOBS = "minimum_jobs" + MINIMUM_GPU_HOURS = "minimum_gpu_hours" + + +class BenchmarkRecommendation(ContractValue): + kind: BenchmarkRecommendationKind + case_id: Identifier + + +class BenchmarkReport(ContractRecord): + """Atomic point-in-time benchmark analysis output.""" + + benchmark_id: Identifier + analysis_id: Identifier + benchmark_manifest: ArtifactReference + created_at: datetime + cases: tuple[BenchmarkCaseResult, ...] = Field(min_length=1) + recommendations: tuple[BenchmarkRecommendation, ...] = () + + @field_validator("created_at") + @classmethod + def validate_created_at(cls, value: datetime) -> datetime: + if value.tzinfo is None or value.utcoffset() != timedelta(0): + raise ValueError("created_at must be timezone-aware UTC") + return value + + @model_validator(mode="after") + def validate_report(self) -> BenchmarkReport: + case_ids = tuple(case.case_id for case in self.cases) + child_run_ids = tuple(case.child_run_id for case in self.cases) + if len(case_ids) != len(set(case_ids)): + raise ValueError("benchmark report case IDs must be unique") + if len(child_run_ids) != len(set(child_run_ids)): + raise ValueError("benchmark report child run IDs must be unique") + unknown = {recommendation.case_id for recommendation in self.recommendations}.difference(case_ids) + if unknown: + raise ValueError(f"recommendations reference unknown cases: {', '.join(sorted(unknown))}") + recommendable = { + case.case_id for case in self.cases if case.outcome is BenchmarkOutcome.SUCCEEDED and case.feasible is True + } + identities: set[tuple[BenchmarkRecommendationKind, str]] = set() + singleton_kinds: set[BenchmarkRecommendationKind] = set() + for recommendation in self.recommendations: + if recommendation.case_id not in recommendable: + raise ValueError("benchmark recommendations must reference successful feasible cases") + identity = (recommendation.kind, recommendation.case_id) + if identity in identities: + raise ValueError("benchmark recommendations must be unique") + identities.add(identity) + if recommendation.kind is not BenchmarkRecommendationKind.PARETO: + if recommendation.kind in singleton_kinds: + raise ValueError("minimum benchmark recommendation kinds must be unique") + singleton_kinds.add(recommendation.kind) + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py new file mode 100644 index 000000000..96fe5effa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Semantic client records shared with Slurm state consumers.""" + +from __future__ import annotations + +from data_designer.slurm.client.records import ClientOutcome, ClientResult + +__all__ = ["ClientOutcome", "ClientResult"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/records.py b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py new file mode 100644 index 000000000..a76290cfc --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timedelta +from enum import Enum +from typing import Annotated, Literal + +from pydantic import NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + ContractRecord, + Identifier, + ShardId, + validate_absolute_path, +) + + +class ClientOutcome(str, Enum): + COMPLETE = "complete" + PARTIAL = "partial" + FAILED = "failed" + + +class ClientResult(ContractRecord): + """Semantic Data Designer outcome independent of engine-internal result types.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + completed_at: datetime + requested_records: PositiveInt + actual_records: NonNegativeInt | None + outcome: ClientOutcome + dataset_path: str | None = None + early_shutdown: bool | None = None + requested_resume_mode: Literal["never", "always", "if_possible"] + effective_resume_mode: Literal["never", "always"] | None = None + candidate_output_manifest: ArtifactReference | None = None + error_code: Identifier | None = None + redacted_message: Annotated[str, StringConstraints(max_length=512)] | None = None + + @field_validator("completed_at") + @classmethod + def validate_completed_at(cls, value: datetime) -> datetime: + if value.tzinfo is None or value.utcoffset() != timedelta(0): + raise ValueError("completed_at must be timezone-aware UTC") + return value + + @field_validator("dataset_path") + @classmethod + def validate_dataset_path(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + @field_validator("redacted_message") + @classmethod + def validate_message(cls, value: str | None) -> str | None: + if value is not None and any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("redacted_message must not contain control characters") + return value + + @model_validator(mode="after") + def validate_outcome(self) -> ClientResult: + if self.actual_records is not None and self.actual_records > self.requested_records: + raise ValueError("actual_records must not exceed requested_records") + if self.requested_resume_mode != "if_possible" and self.effective_resume_mode not in { + None, + self.requested_resume_mode, + }: + raise ValueError("effective resume mode must match a fixed requested mode") + if self.outcome is not ClientOutcome.FAILED: + if self.early_shutdown is None or self.effective_resume_mode is None: + raise ValueError("non-failed client results require resume and early-shutdown facts") + if self.outcome is ClientOutcome.COMPLETE: + if self.actual_records != self.requested_records: + raise ValueError("complete client results require the requested record count") + if self.early_shutdown: + raise ValueError("complete client results cannot report early shutdown") + self._require_success_artifacts() + elif self.outcome is ClientOutcome.PARTIAL: + if self.actual_records is None or self.actual_records >= self.requested_records: + raise ValueError("partial client results require fewer than the requested record count") + self._require_success_artifacts() + else: + if self.candidate_output_manifest is not None: + raise ValueError("failed client results cannot reference a candidate output manifest") + if self.error_code is None: + raise ValueError("failed client results require error_code") + return self + + def _require_success_artifacts(self) -> None: + if self.dataset_path is None or self.candidate_output_manifest is None: + raise ValueError("successful client results require dataset and candidate manifest paths") + if self.error_code is not None or self.redacted_message is not None: + raise ValueError("successful client results cannot contain failure details") + shard_root = f"/runs/{self.run_id}/shards/{self.shard_id}" + if self.effective_resume_mode == "never": + expected_dataset = f"{shard_root}/attempts/{self.attempt_id}/dataset" + else: + expected_dataset = f"{shard_root}/dataset" + if not self.dataset_path.endswith(expected_dataset): + raise ValueError("dataset path must match the run, shard, attempt, and resume policy") + expected_manifest = f"{shard_root}/attempts/{self.attempt_id}/output-manifest.json" + if not self.candidate_output_manifest.path.endswith(expected_manifest): + raise ValueError("candidate output reference must match the run, shard, and attempt") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py new file mode 100644 index 000000000..5cda08684 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public authored configuration contracts for Data Designer Slurm.""" + +from __future__ import annotations + +from data_designer.slurm.config.benchmark import ( + AdaptiveRecordPolicy, + BenchmarkAnalysisTargets, + BenchmarkBaseRun, + BenchmarkDeploymentCase, + BenchmarkDeploymentOverride, + DataDesignerSlurmBenchmarkConfig, + FixedRecordPolicy, +) +from data_designer.slurm.config.images import ( + ClientImageInspection, + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, + ServingImageInspection, +) +from data_designer.slurm.config.profiles import ( + ContainerMount, + GpuRequestMode, + ImageBuildProfile, + ProfileSelectionSource, + SchedulerProfile, + SelectedSlurmProfile, + SlurmProfile, + SlurmProfileCatalog, + injected_profile, + select_profile, + validate_selected_profile, +) +from data_designer.slurm.config.run import ( + ArrayTasksConfig, + BuilderInput, + ClientConfig, + ClientDependencies, + DataDesignerSlurmConfig, + DeploymentResources, + DeploymentTopology, + InputBindings, + InvocationConfig, + InvocationDiagnostics, + LiteralEnvironmentBinding, + LocalStdioMCPProviderConfig, + OutputConfig, + QueueBackpressureConfig, + RemoteMCPProviderConfig, + SecretRef, + ServerDeploymentConfig, + SubmissionConfig, + VllmServerConfig, +) + +__all__ = [ + "AdaptiveRecordPolicy", + "ArrayTasksConfig", + "BenchmarkAnalysisTargets", + "BenchmarkBaseRun", + "BenchmarkDeploymentCase", + "BenchmarkDeploymentOverride", + "BuilderInput", + "ClientConfig", + "ClientDependencies", + "ClientImageInspection", + "ContainerMount", + "DataDesignerSlurmBenchmarkConfig", + "DataDesignerSlurmConfig", + "DeploymentResources", + "DeploymentTopology", + "FixedRecordPolicy", + "GpuRequestMode", + "ImageBuildProfile", + "ImageBuildRequest", + "ImageInspectionRecord", + "ImageKind", + "ImageRef", + "InputBindings", + "InstalledDistribution", + "InvocationConfig", + "InvocationDiagnostics", + "LiteralEnvironmentBinding", + "LocalStdioMCPProviderConfig", + "OutputConfig", + "ProfileSelectionSource", + "QueueBackpressureConfig", + "RemoteMCPProviderConfig", + "SchedulerProfile", + "SecretRef", + "SelectedSlurmProfile", + "ServerDeploymentConfig", + "ServingImageInspection", + "SlurmProfile", + "SlurmProfileCatalog", + "SubmissionConfig", + "VllmServerConfig", + "injected_profile", + "select_profile", + "validate_selected_profile", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py new file mode 100644 index 000000000..e7c465c2e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Annotated, Literal + +from pydantic import Field, PositiveFloat, PositiveInt, field_validator, model_validator + +from data_designer.slurm.config.run import DataDesignerSlurmConfig +from data_designer.slurm.contracts import ( + AuthoredConfig, + Duration, + Identifier, + ModelAlias, + SchemaVersion, + validate_local_config_path, +) + + +class BenchmarkBaseRun(AuthoredConfig): + source: str | None = None + inline: DataDesignerSlurmConfig | None = None + + @model_validator(mode="before") + @classmethod + def normalize_source(cls, value: object) -> object: + if isinstance(value, str): + return {"source": value} + return value + + @field_validator("source") + @classmethod + def validate_source(cls, value: str | None) -> str | None: + return None if value is None else validate_local_config_path(value) + + @model_validator(mode="after") + def validate_base_run(self) -> BenchmarkBaseRun: + if (self.source is None) == (self.inline is None): + raise ValueError("base_run requires exactly one of source or inline") + return self + + +class BenchmarkDeploymentOverride(AuthoredConfig): + nodes: PositiveInt + nodes_per_replica: PositiveInt + + @model_validator(mode="after") + def validate_topology(self) -> BenchmarkDeploymentOverride: + if self.nodes % self.nodes_per_replica: + raise ValueError("nodes_per_replica must divide benchmark deployment nodes") + return self + + +class BenchmarkDeploymentCase(AuthoredConfig): + name: Identifier + deployments: dict[ModelAlias, BenchmarkDeploymentOverride] = Field(min_length=1) + + +class FixedRecordPolicy(AuthoredConfig): + type: Literal["fixed"] + records: PositiveInt + + +class AdaptiveRecordPolicy(AuthoredConfig): + type: Literal["adaptive"] + base_records: PositiveInt + max_records: PositiveInt + records_per_concurrency: PositiveFloat + + @model_validator(mode="after") + def validate_bounds(self) -> AdaptiveRecordPolicy: + if self.max_records < self.base_records: + raise ValueError("adaptive max_records must not be less than base_records") + return self + + +BenchmarkRecordPolicy = Annotated[FixedRecordPolicy | AdaptiveRecordPolicy, Field(discriminator="type")] + + +class BenchmarkAnalysisTargets(AuthoredConfig): + target_total_records: PositiveInt + target_runtime: Duration + + +class DataDesignerSlurmBenchmarkConfig(AuthoredConfig): + """Authored benchmark intent expanded into ordinary Slurm run configs.""" + + schema_version: SchemaVersion + name: Identifier + base_run: BenchmarkBaseRun + model_aliases: Literal["all"] | list[ModelAlias] + concurrency_values: list[PositiveInt] = Field(min_length=1) + deployment_cases: list[BenchmarkDeploymentCase] = Field(min_length=1) + record_policy: BenchmarkRecordPolicy + analysis: BenchmarkAnalysisTargets + + @model_validator(mode="after") + def validate_benchmark(self) -> DataDesignerSlurmBenchmarkConfig: + if isinstance(self.model_aliases, list): + if not self.model_aliases: + raise ValueError("model_aliases must not be empty") + if len(self.model_aliases) != len(set(self.model_aliases)): + raise ValueError("benchmark model aliases must be unique") + if len(self.concurrency_values) != len(set(self.concurrency_values)): + raise ValueError("benchmark concurrency values must be unique") + case_names = [case.name for case in self.deployment_cases] + if len(case_names) != len(set(case_names)): + raise ValueError("benchmark deployment case names must be unique") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/images.py b/packages/data-designer-slurm/src/data_designer/slurm/config/images.py new file mode 100644 index 000000000..0b4e983de --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/images.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import re +from enum import Enum +from typing import Annotated, Literal + +from pydantic import Field, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ( + AuthoredConfig, + ContractRecord, + ContractValue, + Identifier, + Sha256Digest, + validate_absolute_path, + validate_plain_text, +) + +DistributionName = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[a-z0-9]+(?:-[a-z0-9]+)*$", + ), +] + + +class ImageKind(str, Enum): + CLIENT = "client" + SERVING = "serving" + + +class ImageRef(AuthoredConfig): + """Authored reference to one registered image alias or SQSH path.""" + + name: Identifier | None = None + path: str | None = None + + @field_validator("path") + @classmethod + def validate_path(cls, value: str | None) -> str | None: + if value is None: + return None + validate_absolute_path(value) + if not value.endswith(".sqsh"): + raise ValueError("image path must end in .sqsh") + return value + + @model_validator(mode="after") + def validate_reference(self) -> ImageRef: + if (self.name is None) == (self.path is None): + raise ValueError("image reference requires exactly one of name or path") + return self + + +class ImageBuildRequest(AuthoredConfig): + """Typed input for one image import or existing-SQSH registration.""" + + name: Identifier + kind: Literal["client", "serving"] + source: str + + @field_validator("source") + @classmethod + def validate_source(cls, value: str) -> str: + validate_plain_text(value, field_name="source") + if value.endswith(".sqsh"): + return validate_absolute_path(value) + if not re.fullmatch(r"[^\s]+@sha256:[0-9a-f]{64}", value): + raise ValueError("OCI image source must be digest-qualified") + return value + + +class InstalledDistribution(ContractValue): + name: DistributionName + version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + + +class ClientImageInspection(ContractValue): + kind: Literal[ImageKind.CLIENT] + python_implementation: Identifier + python_version: Annotated[str, StringConstraints(pattern=r"^[0-9]+\.[0-9]+\.[0-9]+$")] + python_abi: Identifier + distributions: tuple[InstalledDistribution, ...] + installer_path: str + installer_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + + _installer_path_is_absolute = field_validator("installer_path")(validate_absolute_path) + + @model_validator(mode="after") + def validate_distributions(self) -> ClientImageInspection: + names = tuple(distribution.name for distribution in self.distributions) + if len(names) != len(set(names)): + raise ValueError("installed distribution names must be unique") + return self + + +class ServingImageInspection(ContractValue): + kind: Literal[ImageKind.SERVING] + server_type: Literal["vllm"] + runtime_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + executable_path: str + + _executable_path_is_absolute = field_validator("executable_path")(validate_absolute_path) + + +ImageInspection = Annotated[ClientImageInspection | ServingImageInspection, Field(discriminator="kind")] + + +class ImageInspectionRecord(ContractRecord): + """Digest-bound factual inspection output produced inside an SQSH.""" + + inspector_version: Identifier + sqsh_sha256: Sha256Digest + inspection: ImageInspection diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py new file mode 100644 index 000000000..87412d84d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from enum import Enum +from fnmatch import fnmatchcase +from typing import Annotated, Literal + +from pydantic import Field, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ( + AuthoredConfig, + ContractRecord, + Identifier, + SchemaVersion, + Sha256Digest, + compute_sha256, + validate_absolute_path, + validate_plain_text, +) + + +class GpuRequestMode(str, Enum): + GRES = "gres" + VISIBLE = "visible" + + +class SchedulerProfile(AuthoredConfig): + account: Identifier | None = None + partition: Identifier | None = None + mem_per_gpu: Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:K|M|G|T)$")] | None = None + + +class ImageBuildProfile(AuthoredConfig): + partition: Identifier + + +class ContainerMount(AuthoredConfig): + source: str + target: str + read_only: bool = False + + _paths_are_absolute = field_validator("source", "target")(validate_absolute_path) + + +class SlurmProfile(AuthoredConfig): + """Strict facts for one Slurm cluster.""" + + schema_version: SchemaVersion + host_patterns: list[str] = Field(default_factory=list) + scheduler: SchedulerProfile = Field(default_factory=SchedulerProfile) + gpus_per_node: PositiveInt | Literal["auto"] + workspace_root: str + image_build: ImageBuildProfile + gpu_request_mode: Literal["gres", "visible"] = "gres" + container_mounts: list[ContainerMount] = Field(default_factory=list) + + _workspace_root_is_absolute = field_validator("workspace_root")(validate_absolute_path) + + @field_validator("host_patterns") + @classmethod + def validate_host_patterns(cls, values: list[str]) -> list[str]: + normalized: set[str] = set() + for value in values: + validate_plain_text(value, field_name="host pattern") + _validate_hostname_glob(value) + pattern = value.casefold() + if pattern in normalized: + raise ValueError(f"duplicate hostname glob: {value!r}") + normalized.add(pattern) + return values + + @model_validator(mode="after") + def validate_mounts(self) -> SlurmProfile: + targets = [mount.target for mount in self.container_mounts] + if len(targets) != len(set(targets)): + raise ValueError("container mount targets must be unique") + return self + + +class SlurmProfileCatalog(AuthoredConfig): + """Versioned catalog of independently complete cluster profiles.""" + + schema_version: SchemaVersion + default_cluster: Identifier + clusters: dict[Identifier, SlurmProfile] = Field(min_length=1) + + @model_validator(mode="after") + def validate_catalog(self) -> SlurmProfileCatalog: + if self.default_cluster not in self.clusters: + raise ValueError("default_cluster must name a configured cluster") + + patterns: dict[str, str] = {} + for cluster_name, profile in self.clusters.items(): + for pattern in profile.host_patterns: + normalized = pattern.casefold() + if normalized in patterns: + raise ValueError( + f"hostname glob {pattern!r} is duplicated by clusters " + f"{patterns[normalized]!r} and {cluster_name!r}" + ) + patterns[normalized] = cluster_name + return self + + +class ProfileSelectionSource(str, Enum): + EXPLICIT = "explicit" + HOSTNAME = "hostname" + DEFAULT = "default" + INJECTED = "injected" + + +class SelectedSlurmProfile(ContractRecord): + """Selected profile and provenance persisted in a resolved run plan.""" + + cluster_name: Identifier | None = None + selection_source: ProfileSelectionSource + matched_pattern: str | None = None + catalog_path: str | None = None + catalog_sha256: Sha256Digest | None = None + profile_sha256: Sha256Digest + profile: SlurmProfile + + @field_validator("catalog_path") + @classmethod + def validate_catalog_path(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + @model_validator(mode="after") + def validate_selection(self) -> SelectedSlurmProfile: + if self.profile_sha256 != _profile_digest(self.profile): + raise ValueError("profile_sha256 does not match the selected profile") + + catalog_fields = (self.cluster_name, self.catalog_sha256) + if self.selection_source is ProfileSelectionSource.INJECTED: + if any(value is not None for value in (*catalog_fields, self.catalog_path, self.matched_pattern)): + raise ValueError("injected profiles must not contain catalog selection fields") + else: + if any(value is None for value in catalog_fields): + raise ValueError("catalog selections require cluster_name and catalog_sha256") + if self.selection_source is ProfileSelectionSource.HOSTNAME: + if self.matched_pattern is None: + raise ValueError("hostname selection requires matched_pattern") + elif self.matched_pattern is not None: + raise ValueError("only hostname selection may contain matched_pattern") + return self + + +def select_profile( + catalog: SlurmProfileCatalog, + *, + cluster: str | None = None, + hostnames: tuple[str, ...] = (), + catalog_path: str | None = None, +) -> SelectedSlurmProfile: + """Select a profile with explicit, hostname, then default precedence.""" + catalog_sha256 = compute_sha256(catalog.model_dump(mode="json")) + if cluster is not None: + if cluster not in catalog.clusters: + raise ValueError(f"unknown cluster {cluster!r}") + return _selection( + catalog, + cluster, + ProfileSelectionSource.EXPLICIT, + catalog_sha256, + catalog_path=catalog_path, + ) + + normalized_hosts = {hostname.casefold() for hostname in hostnames if hostname} + matches: dict[str, list[str]] = {} + for cluster_name, profile in catalog.clusters.items(): + matching_patterns = sorted( + pattern + for pattern in profile.host_patterns + if any(fnmatchcase(hostname, pattern.casefold()) for hostname in normalized_hosts) + ) + if matching_patterns: + matches[cluster_name] = matching_patterns + + if len(matches) > 1: + raise ValueError(f"hostname matches multiple clusters: {', '.join(sorted(matches))}") + if matches: + selected_name = next(iter(matches)) + return _selection( + catalog, + selected_name, + ProfileSelectionSource.HOSTNAME, + catalog_sha256, + catalog_path=catalog_path, + matched_pattern=matches[selected_name][0], + ) + return _selection( + catalog, + catalog.default_cluster, + ProfileSelectionSource.DEFAULT, + catalog_sha256, + catalog_path=catalog_path, + ) + + +def injected_profile(profile: SlurmProfile) -> SelectedSlurmProfile: + """Record a directly injected effective profile.""" + return SelectedSlurmProfile( + schema_version=1, + selection_source=ProfileSelectionSource.INJECTED, + profile_sha256=_profile_digest(profile), + profile=profile, + ) + + +def validate_selected_profile( + catalog: SlurmProfileCatalog, + selected: SelectedSlurmProfile, +) -> SelectedSlurmProfile: + """Validate a persisted catalog selection against its source catalog.""" + if selected.selection_source is ProfileSelectionSource.INJECTED: + raise ValueError("injected profile selection has no source catalog") + if selected.catalog_sha256 != compute_sha256(catalog.model_dump(mode="json")): + raise ValueError("selected profile catalog digest does not match the catalog") + if selected.cluster_name not in catalog.clusters: + raise ValueError("selected cluster is absent from the catalog") + if selected.profile != catalog.clusters[selected.cluster_name]: + raise ValueError("selected profile does not match its catalog entry") + if selected.selection_source is ProfileSelectionSource.DEFAULT and selected.cluster_name != catalog.default_cluster: + raise ValueError("default profile selection does not match the catalog default") + if ( + selected.selection_source is ProfileSelectionSource.HOSTNAME + and selected.matched_pattern not in selected.profile.host_patterns + ): + raise ValueError("hostname selection pattern is absent from the selected profile") + return selected + + +def _selection( + catalog: SlurmProfileCatalog, + cluster_name: str, + source: ProfileSelectionSource, + catalog_sha256: Sha256Digest, + *, + catalog_path: str | None, + matched_pattern: str | None = None, +) -> SelectedSlurmProfile: + profile = catalog.clusters[cluster_name] + return SelectedSlurmProfile( + schema_version=1, + cluster_name=cluster_name, + selection_source=source, + matched_pattern=matched_pattern, + catalog_path=catalog_path, + catalog_sha256=catalog_sha256, + profile_sha256=_profile_digest(profile), + profile=profile, + ) + + +def _profile_digest(profile: SlurmProfile) -> Sha256Digest: + return compute_sha256(profile.model_dump(mode="json")) + + +def _validate_hostname_glob(value: str) -> None: + if "/" in value or any(character.isspace() for character in value): + raise ValueError(f"invalid hostname glob: {value!r}") + open_class: int | None = None + for index, character in enumerate(value): + if character == "[": + if open_class is not None: + raise ValueError(f"invalid hostname glob: {value!r}") + open_class = index + elif character == "]": + if open_class is None or index == open_class + 1: + raise ValueError(f"invalid hostname glob: {value!r}") + open_class = None + if open_class is not None: + raise ValueError(f"invalid hostname glob: {value!r}") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py new file mode 100644 index 000000000..c6f7d69d4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import posixpath +import re +from collections.abc import Mapping +from typing import Annotated, Literal +from urllib.parse import urlsplit + +from packaging.requirements import InvalidRequirement, Requirement +from packaging.utils import canonicalize_name +from pydantic import ( + Field, + JsonValue, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) + +from data_designer.config import RunConfig +from data_designer.slurm.config.images import ImageRef +from data_designer.slurm.contracts import ( + AuthoredConfig, + Duration, + EnvironmentName, + Identifier, + ModelAlias, + SchemaVersion, + validate_absolute_path, + validate_local_config_path, + validate_plain_text, + validate_url, +) + +_OWNED_VLLM_FLAGS = { + "--api-key", + "--distributed-executor-backend", + "--distributed-init-address", + "--enable-expert-parallel", + "--headless", + "--host", + "--middleware", + "--model", + "--pipeline-parallel-size", + "--port", + "--served-model-name", + "--tensor-parallel-size", +} +_DURATION_FACTORS = {"s": 1, "m": 60, "h": 3600, "d": 86400} +_NON_SECRET_PAYLOAD_KEYS = frozenset({"idempotency_key", "partition_key", "primary_key", "sort_key"}) +_SECRET_NAME_PARTS = frozenset({"credential", "credentials", "password", "secret", "token"}) + + +def _duration_seconds(value: Duration) -> int: + return int(value[:-1]) * _DURATION_FACTORS[value[-1]] + + +def _secret_name_segments(value: str) -> list[str]: + snake_case = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", value) + normalized = re.sub(r"[^a-z0-9]+", "_", snake_case.casefold()).strip("_") + return normalized.split("_") + + +def _is_secret_name(value: str) -> bool: + segments = _secret_name_segments(value) + return bool( + _SECRET_NAME_PARTS.intersection(segments) + or {"access", "key"}.issubset(segments) + or segments[-1] in {"auth", "key"} + ) + + +def _is_secret_payload_name(value: str) -> bool: + segments = _secret_name_segments(value) + normalized = "_".join(segments) + return normalized not in _NON_SECRET_PAYLOAD_KEYS and _is_secret_name(value) + + +def _option_flag(value: str) -> str: + return re.split(r"[=\s]", value.lstrip(), maxsplit=1)[0] + + +def _contains_secret_key(value: object) -> bool: + if isinstance(value, Mapping): + return any( + (_is_secret_payload_name(str(key)) and item is not None) or _contains_secret_key(item) + for key, item in value.items() + ) + if isinstance(value, list | tuple): + return any(_contains_secret_key(item) for item in value) + return False + + +def _validate_environment_bindings( + values: dict[EnvironmentName, EnvironmentBinding], +) -> dict[EnvironmentName, EnvironmentBinding]: + literal_secrets = [ + name + for name, binding in values.items() + if _is_secret_name(name) and isinstance(binding, LiteralEnvironmentBinding) + ] + if literal_secrets: + raise ValueError("secret-shaped environment names require external secret references") + return values + + +class LiteralEnvironmentBinding(AuthoredConfig): + type: Literal["literal"] + value: Annotated[str, StringConstraints(max_length=4096)] + + @field_validator("value") + @classmethod + def validate_value(cls, value: str) -> str: + return validate_plain_text(value, field_name="environment value") + + +class SecretRef(AuthoredConfig): + type: Literal["secret"] + environment: EnvironmentName + + +EnvironmentBinding = Annotated[ + LiteralEnvironmentBinding | SecretRef, + Field(discriminator="type"), +] + + +class BuilderInput(AuthoredConfig): + source: str | None = None + inline: dict[str, JsonValue] | None = None + + @field_validator("source") + @classmethod + def validate_source(cls, value: str | None) -> str | None: + return None if value is None else validate_local_config_path(value) + + @model_validator(mode="after") + def validate_input(self) -> BuilderInput: + if (self.source is None) == (self.inline is None): + raise ValueError("builder requires exactly one of source or inline") + if self.inline is not None: + if not self.inline: + raise ValueError("inline builder input must not be empty") + retired = {"dependencies", "sandbox_config", "server_configs"}.intersection(self.inline) + if retired: + raise ValueError(f"builder input contains retired Big Iron fields: {', '.join(sorted(retired))}") + if "data_designer" in self.inline: + unknown = set(self.inline).difference({"data_designer", "library_version"}) + library_version = self.inline.get("library_version") + valid = not unknown and isinstance(self.inline["data_designer"], dict) + valid = valid and (library_version is None or isinstance(library_version, str)) + else: + valid = isinstance(self.inline.get("columns"), list) + if not valid: + raise ValueError("inline builder input must be one complete serialized Data Designer config") + if _contains_secret_key(self.inline): + raise ValueError("inline builder input must not contain secret values") + return self + + +class InputBindings(AuthoredConfig): + seed_path: str | None = None + managed_assets_path: str | None = None + + @field_validator("seed_path", "managed_assets_path") + @classmethod + def validate_paths(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + +class RemoteMCPProviderConfig(AuthoredConfig): + provider_type: Literal["sse", "streamable_http"] + name: Identifier + endpoint: str + api_key: SecretRef | None = None + + @field_validator("endpoint") + @classmethod + def validate_endpoint(cls, value: str) -> str: + validate_url(value, field_name="MCP endpoint") + parsed = urlsplit(value) + if parsed.username is not None or parsed.password is not None or parsed.query or parsed.fragment: + raise ValueError("MCP endpoint must not embed credentials, query parameters, or fragments") + return value + + +class LocalStdioMCPProviderConfig(AuthoredConfig): + provider_type: Literal["stdio"] + name: Identifier + command: str + args: list[str] = Field(default_factory=list) + environment: dict[EnvironmentName, EnvironmentBinding] = Field(default_factory=dict) + + @field_validator("command") + @classmethod + def validate_command(cls, value: str) -> str: + validate_plain_text(value, field_name="MCP command") + if any(character.isspace() for character in value): + raise ValueError("MCP command must be one executable token") + return value + + @field_validator("args") + @classmethod + def validate_args(cls, values: list[str]) -> list[str]: + for value in values: + validate_plain_text(value, field_name="MCP argument") + option = _option_flag(value).lstrip("-") + if _is_secret_name(option): + raise ValueError("secret-shaped MCP arguments must use an environment secret reference") + return values + + _environment_uses_secret_references = field_validator("environment")(_validate_environment_bindings) + + +MCPProviderConfig = Annotated[ + RemoteMCPProviderConfig | LocalStdioMCPProviderConfig, + Field(discriminator="provider_type"), +] + + +class InvocationDiagnostics(AuthoredConfig): + log_requests: bool = False + + +class InvocationConfig(AuthoredConfig): + num_records: PositiveInt + dataset_name: Identifier + resume: Literal["never", "always", "if_possible"] = "never" + run_config: dict[str, JsonValue] = Field(default_factory=dict) + input_bindings: InputBindings = Field(default_factory=InputBindings) + mcp_providers: list[MCPProviderConfig] = Field(default_factory=list) + model_concurrency: dict[ModelAlias, PositiveInt] = Field(default_factory=dict) + diagnostics: InvocationDiagnostics = Field(default_factory=InvocationDiagnostics) + + @field_validator("run_config") + @classmethod + def validate_run_config_keys(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: + unknown = set(value).difference(RunConfig.model_fields) + if unknown: + raise ValueError(f"unknown Data Designer RunConfig fields: {', '.join(sorted(unknown))}") + RunConfig.model_validate(value) + return value + + @model_validator(mode="after") + def validate_mcp_providers(self) -> InvocationConfig: + names = [provider.name for provider in self.mcp_providers] + if len(names) != len(set(names)): + raise ValueError("MCP provider names must be unique") + return self + + +class ClientDependencies(AuthoredConfig): + requirements: list[str] | None = Field(default_factory=list) + lock_file: str | None = None + index_credentials: dict[str, SecretRef] = Field(default_factory=dict) + + @field_validator("requirements") + @classmethod + def validate_requirements(cls, values: list[str] | None) -> list[str] | None: + if values is None: + return None + names: list[str] = [] + for value in values: + validate_plain_text(value, field_name="dependency requirement") + if value != value.strip() or value.startswith(("-e ", "/", "./", "../")) or "git+" in value: + raise ValueError(f"dependency requirement must identify a package or immutable wheel: {value!r}") + try: + requirement = Requirement(value) + except InvalidRequirement as error: + raise ValueError(f"invalid dependency requirement: {value!r}") from error + if requirement.marker is not None: + raise ValueError("dependency requirement environment markers are not supported") + if requirement.url is not None: + validate_url(requirement.url, field_name="direct dependency URL") + parsed = urlsplit(requirement.url) + valid_wheel = ( + parsed.scheme == "https" + and parsed.hostname is not None + and parsed.username is None + and parsed.password is None + and not parsed.query + and parsed.path.endswith(".whl") + and re.fullmatch(r"sha256=[0-9a-f]{64}", parsed.fragment) is not None + ) + if not valid_wheel: + raise ValueError("direct dependency URLs must be HTTPS wheels with a SHA-256 fragment") + names.append(canonicalize_name(requirement.name)) + if len(names) != len(set(names)): + raise ValueError("dependency requirements must have unique normalized names") + return values + + @field_validator("lock_file") + @classmethod + def validate_lock_file(cls, value: str | None) -> str | None: + if value is None: + return None + validate_plain_text(value, field_name="dependency lock path") + if "://" in value or ".." in value.split("/"): + raise ValueError("dependency lock must be a normalized local path") + normalized = posixpath.normpath(value) + if not normalized.endswith(".json"): + raise ValueError("dependency lock path must end in .json") + return normalized + + @model_validator(mode="after") + def validate_source(self) -> ClientDependencies: + if self.lock_file is None and self.requirements is None: + raise ValueError("client dependencies require requirements or lock_file") + if self.lock_file is not None and self.requirements is not None: + raise ValueError("client dependencies cannot contain both requirements and lock_file") + return self + + +class ClientConfig(AuthoredConfig): + cpus: PositiveInt = 32 + image: ImageRef + dependencies: ClientDependencies = Field(default_factory=ClientDependencies) + + +class QueueBackpressureConfig(AuthoredConfig): + max_waiting_requests: NonNegativeInt = 128 + retry_after_seconds: NonNegativeInt | None = 1 + + +class VllmServerConfig(AuthoredConfig): + type: Literal["vllm"] + image: ImageRef + startup_timeout: Duration = "15m" + distributed_init_timeout: Duration = "10m" + readiness_path: str = "/health" + enable_expert_parallel: bool = False + queue_backpressure: QueueBackpressureConfig = Field(default_factory=QueueBackpressureConfig) + extra_args: list[str] = Field(default_factory=list) + environment: dict[EnvironmentName, EnvironmentBinding] = Field(default_factory=dict) + + @field_validator("readiness_path") + @classmethod + def validate_readiness_path(cls, value: str) -> str: + validate_plain_text(value, field_name="readiness path") + if not value.startswith("/") or "?" in value or "#" in value: + raise ValueError("readiness_path must be an absolute URL path without query or fragment") + return value + + @field_validator("extra_args") + @classmethod + def validate_extra_args(cls, values: list[str]) -> list[str]: + for value in values: + validate_plain_text(value, field_name="vLLM argument") + flag = _option_flag(value) + if flag in _OWNED_VLLM_FLAGS: + raise ValueError(f"vLLM argument {flag!r} is owned by the compiler or runtime") + if _is_secret_name(flag.lstrip("-")): + raise ValueError("secret-shaped vLLM arguments must use an environment secret reference") + return values + + _environment_uses_secret_references = field_validator("environment")(_validate_environment_bindings) + + @model_validator(mode="after") + def validate_timeouts(self) -> VllmServerConfig: + if _duration_seconds(self.distributed_init_timeout) > _duration_seconds(self.startup_timeout): + raise ValueError("distributed_init_timeout must not exceed startup_timeout") + return self + + +class DeploymentResources(AuthoredConfig): + nodes: PositiveInt = 1 + + +class DeploymentTopology(AuthoredConfig): + tensor_parallel: PositiveInt = 1 + nodes_per_replica: PositiveInt = 1 + + +class ServerDeploymentConfig(AuthoredConfig): + model_alias: ModelAlias + served_model_name: str | None = None + model: str + server: VllmServerConfig + resources: DeploymentResources = Field(default_factory=DeploymentResources) + topology: DeploymentTopology = Field(default_factory=DeploymentTopology) + + @field_validator("model") + @classmethod + def validate_model(cls, value: str) -> str: + validate_plain_text(value, field_name="model") + if value.startswith("/"): + return validate_absolute_path(value) + if any(character.isspace() for character in value): + raise ValueError("Hugging Face model identifiers must not contain whitespace") + return value + + @field_validator("served_model_name") + @classmethod + def validate_served_model_name(cls, value: str | None) -> str | None: + return None if value is None else validate_plain_text(value, field_name="served model name") + + @model_validator(mode="after") + def validate_topology(self) -> ServerDeploymentConfig: + if self.resources.nodes % self.topology.nodes_per_replica: + raise ValueError("nodes_per_replica must divide deployment nodes") + if self.server.enable_expert_parallel and self.topology.nodes_per_replica > 1: + raise ValueError("multi-node expert parallel is not supported in v1") + return self + + +class ArrayTasksConfig(AuthoredConfig): + count: PositiveInt = 1 + max_concurrent: PositiveInt = 1 + + @model_validator(mode="after") + def validate_concurrency(self) -> ArrayTasksConfig: + if self.max_concurrent > self.count: + raise ValueError("array task concurrency must not exceed task count") + return self + + +class SubmissionConfig(AuthoredConfig): + account: Identifier | None = None + partition: Identifier | None = None + job_name: Identifier = "data-designer" + time_limit: Annotated[str, StringConstraints(pattern=r"^(?:[0-9]+-)?[0-9]{2}:[0-9]{2}:[0-9]{2}$")] = "03:55:00" + comment: Annotated[str, StringConstraints(max_length=256)] | None = None + + @field_validator("time_limit") + @classmethod + def validate_time_limit(cls, value: str) -> str: + clock = value.rsplit("-", maxsplit=1)[-1] + _, minutes, seconds = (int(part) for part in clock.split(":")) + if minutes >= 60 or seconds >= 60: + raise ValueError("time_limit minutes and seconds must be below 60") + return value + + @field_validator("comment") + @classmethod + def validate_comment(cls, value: str | None) -> str | None: + return None if value is None else validate_plain_text(value, field_name="submission comment") + + +class OutputConfig(AuthoredConfig): + root: str | None = None + format: Literal["parquet", "jsonl", "csv"] = "parquet" + partitions: PositiveInt = 1 + require_exact_record_count: bool = False + + @field_validator("root") + @classmethod + def validate_root(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + +class DataDesignerSlurmConfig(AuthoredConfig): + """Complete portable intent for one Data Designer Slurm run.""" + + schema_version: SchemaVersion + name: Identifier + builder: BuilderInput + invocation: InvocationConfig + client: ClientConfig + deployments: list[ServerDeploymentConfig] = Field(min_length=1) + array_tasks: ArrayTasksConfig = Field(default_factory=ArrayTasksConfig) + submission: SubmissionConfig = Field(default_factory=SubmissionConfig) + output: OutputConfig = Field(default_factory=OutputConfig) + + @model_validator(mode="after") + def validate_run(self) -> DataDesignerSlurmConfig: + if self.array_tasks.count > self.invocation.num_records: + raise ValueError("array task count must not exceed requested records") + aliases = [deployment.model_alias for deployment in self.deployments] + if len(aliases) != len(set(aliases)): + raise ValueError("deployment model aliases must be unique") + unknown_concurrency = set(self.invocation.model_concurrency).difference(aliases) + if unknown_concurrency: + raise ValueError( + f"model concurrency references undeclared aliases: {', '.join(sorted(unknown_concurrency))}" + ) + return self 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 087b2a678..f07063417 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -8,7 +8,9 @@ import hashlib import json import posixpath -from typing import Annotated, Literal +from collections.abc import Mapping +from typing import Annotated, Literal, TypeVar +from urllib.parse import urlsplit from pydantic import ( BaseModel, @@ -32,7 +34,60 @@ ShardId = Annotated[str, StringConstraints(pattern=r"^shard-[0-9]{5,}$")] AttemptId = Annotated[str, StringConstraints(pattern=r"^attempt-[0-9]{4,}$")] SchemaVersion = Literal[1] +EnvironmentName = Annotated[str, StringConstraints(pattern=r"^[A-Za-z_][A-Za-z0-9_]*$")] Sha256Digest = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] +Duration = Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:s|m|h|d)$")] + +_Key = TypeVar("_Key") +_Value = TypeVar("_Value") + + +class _FrozenList(list[_Value]): + """List that retains JSON compatibility without exposing mutation.""" + + def _immutable(self, *args: object, **kwargs: object) -> None: + del args, kwargs + raise TypeError("frozen list cannot be modified") + + __delitem__ = _immutable + __iadd__ = _immutable + __imul__ = _immutable + __setitem__ = _immutable + append = _immutable + clear = _immutable + extend = _immutable + insert = _immutable + pop = _immutable + remove = _immutable + reverse = _immutable + sort = _immutable + + +class _FrozenDict(dict[_Key, _Value]): + """Dictionary that retains JSON compatibility without exposing mutation.""" + + def _immutable(self, *args: object, **kwargs: object) -> None: + del args, kwargs + raise TypeError("frozen dictionary cannot be modified") + + __delitem__ = _immutable + __ior__ = _immutable + __setitem__ = _immutable + clear = _immutable + pop = _immutable + popitem = _immutable + setdefault = _immutable + update = _immutable + + +def _freeze_collections(value: object) -> object: + if isinstance(value, Mapping): + return _FrozenDict({key: _freeze_collections(item) for key, item in value.items()}) + if isinstance(value, list): + return _FrozenList(_freeze_collections(item) for item in value) + if isinstance(value, tuple): + return tuple(_freeze_collections(item) for item in value) + return value class ContractValue(BaseModel): @@ -47,6 +102,24 @@ class ContractValue(BaseModel): validate_default=True, ) + @field_validator("*", mode="after") + @classmethod + def freeze_collections(cls, value: object) -> object: + return _freeze_collections(value) + + +class AuthoredConfig(ContractValue): + """Base for strict authored configuration values.""" + + def serialize_canonical_json(self) -> bytes: + return canonical_json(self.model_dump(mode="json")) + + def serialize_json(self) -> str: + return pretty_json(self.model_dump(mode="json")) + + def compute_sha256(self) -> Sha256Digest: + return hashlib.sha256(self.serialize_json().encode("utf-8")).hexdigest() + class ContractRecord(ContractValue): """Base for immutable, explicitly versioned Slurm records.""" @@ -104,8 +177,7 @@ def validate_absolute_path(value: str) -> str: raise ValueError("path must have exactly one leading slash") if value == "/": raise ValueError("path must not be the filesystem root") - if any(ord(character) < 32 or ord(character) == 127 for character in value): - raise ValueError("path must not contain control characters") + validate_plain_text(value, field_name="path") if ".." in value.split("/"): raise ValueError("path must not contain parent-directory components") if posixpath.normpath(value) != value: @@ -126,6 +198,42 @@ def validate_relative_path(value: str) -> str: return value +def validate_local_config_path(value: str) -> str: + validate_plain_text(value, field_name="path") + if "://" in value: + raise ValueError("builder and config sources must be local paths") + if ".." in value.split("/"): + raise ValueError("path must not contain parent-directory components") + normalized = posixpath.normpath(value) + if posixpath.splitext(normalized)[1] not in {".json", ".yaml", ".yml"}: + raise ValueError("config path must end in .json, .yaml, or .yml") + return normalized + + +def validate_plain_text(value: str, *, field_name: str) -> str: + 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): + raise ValueError(f"{field_name} must not contain control characters") + return value + + +def validate_url(value: str, *, field_name: str) -> str: + validate_plain_text(value, field_name=field_name) + try: + parsed = urlsplit(value) + parsed.port + except ValueError as error: + raise ValueError(f"{field_name} must be an HTTP(S) URL with a valid host and port") from error + if ( + parsed.scheme not in {"http", "https"} + or parsed.hostname is None + or any(character.isspace() for character in value) + ): + raise ValueError(f"{field_name} must be an HTTP(S) URL") + return value + + class ArtifactReference(ContractValue): """Immutable reference to persisted file bytes and their digest.""" @@ -163,8 +271,11 @@ class ResumeWorkspace(ContractValue): __all__ = [ "ArtifactReference", "AttemptId", + "AuthoredConfig", "ContractRecord", "ContractValue", + "Duration", + "EnvironmentName", "Identifier", "ModelAlias", "RecordRange", @@ -176,5 +287,8 @@ class ResumeWorkspace(ContractValue): "compute_sha256", "pretty_json", "validate_absolute_path", + "validate_local_config_path", + "validate_plain_text", "validate_relative_path", + "validate_url", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py new file mode 100644 index 000000000..160b7ef77 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolved execution-plan contracts for Data Designer Slurm.""" + +from __future__ import annotations + +from data_designer.slurm.planning.models import ( + ArtifactReference, + LockedPackage, + PlannedShard, + PortClaim, + RecordRange, + ResolvedBuilderInput, + ResolvedClient, + ResolvedDependencyLock, + ResolvedDeployment, + ResolvedImage, + ResolvedInvocation, + ResolvedOutput, + ResolvedSlurmRunPlan, + ResolvedSubmission, + ResolvedTopology, + ResumeWorkspace, +) +from data_designer.slurm.planning.validation import PlanContractError, validate_resolved_plan + +__all__ = [ + "ArtifactReference", + "LockedPackage", + "PlanContractError", + "PlannedShard", + "PortClaim", + "RecordRange", + "ResolvedBuilderInput", + "ResolvedClient", + "ResolvedDependencyLock", + "ResolvedDeployment", + "ResolvedImage", + "ResolvedInvocation", + "ResolvedOutput", + "ResolvedSlurmRunPlan", + "ResolvedSubmission", + "ResolvedTopology", + "ResumeWorkspace", + "validate_resolved_plan", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py new file mode 100644 index 000000000..d1eff7ab6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -0,0 +1,522 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import posixpath +from typing import Annotated, Literal +from urllib.parse import urlsplit + +from packaging.requirements import Requirement +from packaging.utils import canonicalize_name +from packaging.version import InvalidVersion +from pydantic import Field, JsonValue, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.config import RunConfig +from data_designer.slurm.config.images import ( + DistributionName, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, +) +from data_designer.slurm.config.profiles import ContainerMount, SelectedSlurmProfile +from data_designer.slurm.config.run import ( + ArrayTasksConfig, + ClientConfig, + ClientDependencies, + InvocationConfig, + ServerDeploymentConfig, + SubmissionConfig, +) +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractRecord, + ContractValue, + Identifier, + ModelAlias, + RecordRange, + ResumeWorkspace, + Sha256Digest, + ShardId, + compute_sha256, + validate_absolute_path, + validate_local_config_path, + validate_plain_text, +) + + +class ResolvedImage(ContractValue): + """Immutable SQSH path and the digest-bound inspection that approved it.""" + + authored_ref: ImageRef + path: str + sha256: Sha256Digest + inspection: ImageInspectionRecord + + @field_validator("path") + @classmethod + def validate_path(cls, value: str) -> str: + validate_absolute_path(value) + if not value.endswith(".sqsh"): + raise ValueError("resolved image path must end in .sqsh") + return value + + @model_validator(mode="after") + def validate_image(self) -> ResolvedImage: + if self.inspection.sqsh_sha256 != self.sha256: + raise ValueError("image inspection digest does not match the resolved SQSH") + if self.authored_ref.path is not None and self.authored_ref.path != self.path: + raise ValueError("resolved image path does not match the authored path") + return self + + @property + def kind(self) -> ImageKind: + return self.inspection.inspection.kind + + +class LockedPackage(ContractValue): + name: DistributionName + version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + artifact: ArtifactReference + + @model_validator(mode="after") + def validate_artifact(self) -> LockedPackage: + if not self.artifact.path.endswith(".whl"): + raise ValueError("locked overlay artifacts must be wheels") + return self + + +class ResolvedDependencyLock(ContractRecord): + """Immutable client dependency resolution against one fixed image inventory.""" + + resolver_version: Identifier + python_abi: Identifier + client_image_sha256: Sha256Digest + authored_requirements: tuple[str, ...] + authored_source: str | None = None + source: ArtifactReference | None = None + image_distributions: tuple[InstalledDistribution, ...] + overlay_packages: tuple[LockedPackage, ...] + + @field_validator("authored_source") + @classmethod + def validate_authored_source(cls, value: str | None) -> str | None: + if value is None: + return None + normalized = validate_local_config_path(value) + if not normalized.endswith(".json"): + raise ValueError("dependency lock source must end in .json") + return normalized + + @model_validator(mode="after") + def validate_packages(self) -> ResolvedDependencyLock: + if (self.authored_source is None) != (self.source is None): + raise ValueError("dependency lock authored and resolved sources must be provided together") + ClientDependencies(requirements=list(self.authored_requirements)) + image_names = tuple(distribution.name for distribution in self.image_distributions) + overlay_names = tuple(package.name for package in self.overlay_packages) + if image_names != tuple(sorted(image_names)) or overlay_names != tuple(sorted(overlay_names)): + raise ValueError("dependency lock distributions must be sorted by normalized name") + if len(image_names) != len(set(image_names)) or len(overlay_names) != len(set(overlay_names)): + raise ValueError("dependency lock distribution names must be unique") + overlap = set(image_names).intersection(overlay_names) + if overlap: + raise ValueError(f"overlay packages overlap image-owned distributions: {', '.join(sorted(overlap))}") + image_packages = {distribution.name: distribution for distribution in self.image_distributions} + overlay_packages = {package.name: package for package in self.overlay_packages} + for value in self.authored_requirements: + requirement = Requirement(value) + name = canonicalize_name(requirement.name) + if requirement.url is not None: + package = overlay_packages.get(name) + digest = urlsplit(requirement.url).fragment.removeprefix("sha256=") + if package is None or package.artifact.sha256 != digest: + raise ValueError(f"direct requirement {name!r} must match one locked overlay artifact") + continue + package = overlay_packages.get(name) or image_packages.get(name) + if package is None: + raise ValueError(f"authored requirement {name!r} is missing from the dependency lock") + try: + satisfied = requirement.specifier.contains(package.version, prereleases=True) + except InvalidVersion as error: + raise ValueError(f"locked package {name!r} has an invalid version") from error + if not satisfied: + raise ValueError(f"locked package {name!r} does not satisfy its authored requirement") + return self + + +class ResolvedBuilderInput(ContractValue): + authored_source: str | None = None + source: ArtifactReference | None = None + inline: dict[str, JsonValue] | None = None + content_sha256: Sha256Digest + model_aliases: tuple[ModelAlias, ...] + referenced_model_aliases: tuple[ModelAlias, ...] = () + + @model_validator(mode="after") + def validate_input(self) -> ResolvedBuilderInput: + if (self.source is None) == (self.inline is None): + raise ValueError("resolved builder requires exactly one of source or inline") + if self.source is None: + if self.authored_source is not None: + raise ValueError("inline builder input cannot contain authored_source") + expected_digest = compute_sha256(self.inline) + model_aliases, referenced_aliases = _extract_builder_aliases(self.inline) + if self.model_aliases != model_aliases: + raise ValueError("resolved model aliases do not match the inline builder") + if self.referenced_model_aliases != referenced_aliases: + raise ValueError("resolved referenced aliases do not match the inline builder") + else: + if self.authored_source is None: + raise ValueError("resolved builder source requires authored_source") + expected_digest = self.source.sha256 + if len(self.model_aliases) != len(set(self.model_aliases)): + raise ValueError("resolved builder model aliases must be unique") + if len(self.referenced_model_aliases) != len(set(self.referenced_model_aliases)): + raise ValueError("resolved builder referenced aliases must be unique") + if self.content_sha256 != expected_digest: + raise ValueError("builder content digest does not match the resolved input") + return self + + +class ResolvedInvocation(ContractValue): + authored: InvocationConfig + effective_run_config: dict[str, JsonValue] + + @field_validator("effective_run_config") + @classmethod + def validate_run_config_is_materialized(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: + materialized = RunConfig.model_validate(value).model_dump(mode="json") + if value != materialized: + raise ValueError("effective_run_config must contain the fully materialized Data Designer RunConfig") + return value + + +class PortClaim(ContractValue): + name: Identifier + role: Literal["http", "rendezvous", "logical_endpoint"] + node_index: NonNegativeInt + port: Annotated[int, Field(ge=1024, le=65535)] + + +class ResolvedTopology(ContractValue): + tensor_parallel: PositiveInt + nodes_per_replica: PositiveInt + pipeline_parallel: PositiveInt + node_group_count: PositiveInt + replicas_per_node_group: PositiveInt + replica_count: PositiveInt + gpus_per_replica: PositiveInt + + +class ResolvedDeployment(ContractValue): + deployment_id: Identifier + authored: ServerDeploymentConfig + served_model_name: str + image: ResolvedImage + node_indices: tuple[NonNegativeInt, ...] = Field(min_length=1) + gpus_per_node: PositiveInt + topology: ResolvedTopology + ports: tuple[PortClaim, ...] = () + + @model_validator(mode="after") + def validate_deployment(self) -> ResolvedDeployment: + validate_plain_text(self.served_model_name, field_name="served model name") + if self.served_model_name != (self.authored.served_model_name or self.authored.model): + raise ValueError("resolved served model name does not match the authored deployment") + if self.image.kind is not ImageKind.SERVING: + raise ValueError("server deployments require serving images") + if self.image.authored_ref != self.authored.server.image: + raise ValueError("resolved serving image does not match the authored image reference") + if len(self.node_indices) != self.authored.resources.nodes: + raise ValueError("deployment placement must contain exactly the requested node count") + if self.node_indices != tuple(sorted(set(self.node_indices))): + raise ValueError("deployment node indices must be sorted and unique") + if self.gpus_per_node % self.authored.topology.tensor_parallel: + raise ValueError("tensor_parallel must divide resolved GPUs per node") + expected = ResolvedTopology( + tensor_parallel=self.authored.topology.tensor_parallel, + nodes_per_replica=self.authored.topology.nodes_per_replica, + pipeline_parallel=self.authored.topology.nodes_per_replica, + node_group_count=self.authored.resources.nodes // self.authored.topology.nodes_per_replica, + replicas_per_node_group=self.gpus_per_node // self.authored.topology.tensor_parallel, + replica_count=(self.authored.resources.nodes // self.authored.topology.nodes_per_replica) + * (self.gpus_per_node // self.authored.topology.tensor_parallel), + gpus_per_replica=self.authored.topology.tensor_parallel * self.authored.topology.nodes_per_replica, + ) + if self.topology != expected: + raise ValueError("resolved topology does not match deployment resources") + if any(port.node_index not in self.node_indices for port in self.ports): + raise ValueError("deployment port claims must use deployment nodes") + names = tuple(port.name for port in self.ports) + if len(names) != len(set(names)): + raise ValueError("deployment port claim names must be unique") + if any(not name.startswith(f"{self.deployment_id}-") for name in names): + raise ValueError("deployment port claim names must use the deployment ID") + if any(port.role == "logical_endpoint" for port in self.ports): + raise ValueError("logical endpoint ports belong to the resolved client") + + group_heads = self.node_indices[:: self.topology.nodes_per_replica] + expected_http_nodes = tuple(head for head in group_heads for _ in range(self.topology.replicas_per_node_group)) + http_ports = tuple(port for port in self.ports if port.role == "http") + http_nodes = tuple(port.node_index for port in http_ports) + if http_nodes != expected_http_nodes: + raise ValueError("deployment requires one ordered HTTP port claim per replica lane") + expected_http_names = tuple(f"{self.deployment_id}-http-{index:05d}" for index in range(len(http_ports))) + if tuple(port.name for port in http_ports) != expected_http_names: + raise ValueError("deployment HTTP port names must match their ordered replica lane") + + expected_rendezvous_nodes = ( + tuple(head for head in group_heads for _ in range(self.topology.replicas_per_node_group)) + if self.topology.nodes_per_replica > 1 + else () + ) + rendezvous_ports = tuple(port for port in self.ports if port.role == "rendezvous") + rendezvous_nodes = tuple(port.node_index for port in rendezvous_ports) + if rendezvous_nodes != expected_rendezvous_nodes: + raise ValueError("deployment requires one ordered rendezvous port claim per multi-node replica lane") + expected_rendezvous_names = tuple( + f"{self.deployment_id}-rendezvous-{index:05d}" for index in range(len(rendezvous_ports)) + ) + if tuple(port.name for port in rendezvous_ports) != expected_rendezvous_names: + raise ValueError("deployment rendezvous port names must match their ordered replica lane") + return self + + +class ResolvedClient(ContractValue): + authored: ClientConfig + image: ResolvedImage + dependency_lock: ArtifactReference + host_node_index: NonNegativeInt + gpu_count: Literal[0] + ports: tuple[PortClaim, ...] = () + + @model_validator(mode="after") + def validate_client(self) -> ResolvedClient: + if self.image.kind is not ImageKind.CLIENT: + raise ValueError("Data Designer client requires a client image") + if self.image.authored_ref != self.authored.image: + raise ValueError("resolved client image does not match the authored image reference") + if any(port.role != "logical_endpoint" for port in self.ports): + raise ValueError("resolved client ports must be logical endpoints") + if any(port.node_index != self.host_node_index for port in self.ports): + raise ValueError("logical endpoint ports must use the client host") + names = tuple(port.name for port in self.ports) + if len(names) != len(set(names)): + raise ValueError("logical endpoint port claim names must be unique") + return self + + +class PlannedShard(ContractValue): + shard_id: ShardId + shard_index: NonNegativeInt + array_task_index: NonNegativeInt + record_range: RecordRange + input_partition: ArtifactReference | None = None + resume_workspace: ResumeWorkspace + + @property + def requested_records(self) -> int: + return self.record_range.record_count + + +class ResolvedSubmission(ContractValue): + account: Identifier | None = None + partition: Identifier | None = None + job_name: Identifier + time_limit: str + comment: str | None = None + + @model_validator(mode="after") + def validate_submission(self) -> ResolvedSubmission: + SubmissionConfig.model_validate(self.model_dump(mode="python")) + return self + + +class ResolvedOutput(ContractValue): + root: str + format: Literal["parquet", "jsonl", "csv"] + partitions: PositiveInt + require_exact_record_count: bool + + _root_is_absolute = field_validator("root")(validate_absolute_path) + + +class ResolvedSlurmRunPlan(ContractRecord): + """Immutable allocation input consumed without ambient configuration.""" + + run_id: Identifier + package_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + authored_config: ArtifactReference + selected_profile: SelectedSlurmProfile + resolved_gpus_per_node: PositiveInt + builder: ResolvedBuilderInput + invocation: ResolvedInvocation + client: ResolvedClient + deployments: tuple[ResolvedDeployment, ...] = Field(min_length=1) + array_tasks: ArrayTasksConfig + shards: tuple[PlannedShard, ...] = Field(min_length=1) + submission: ResolvedSubmission + output: ResolvedOutput + container_mounts: tuple[ContainerMount, ...] = () + runtime_bundle: ArtifactReference + + @model_validator(mode="after") + def validate_plan(self) -> ResolvedSlurmRunPlan: + profile = self.selected_profile.profile + if profile.gpus_per_node != "auto" and profile.gpus_per_node != self.resolved_gpus_per_node: + raise ValueError("resolved GPU count does not match the selected profile") + if any(deployment.gpus_per_node != self.resolved_gpus_per_node for deployment in self.deployments): + raise ValueError("every deployment must use the resolved profile GPU count") + if tuple(profile.container_mounts) != self.container_mounts: + raise ValueError("plan mount mappings must match the selected profile") + + deployment_ids = tuple(deployment.deployment_id for deployment in self.deployments) + aliases = tuple(deployment.authored.model_alias for deployment in self.deployments) + expected_deployment_ids = tuple(f"deployment-{index:05d}" for index in range(len(self.deployments))) + if deployment_ids != expected_deployment_ids: + raise ValueError("resolved deployment IDs must use complete ordered zero-based identities") + if len(aliases) != len(set(aliases)): + raise ValueError("resolved deployment aliases must be unique") + if not set(aliases).issubset(self.builder.model_aliases): + raise ValueError("each deployment alias must match a resolved Data Designer model alias") + if not set(self.builder.referenced_model_aliases).issubset(aliases): + raise ValueError("each referenced Data Designer model alias requires a deployment") + + node_indices = tuple(index for deployment in self.deployments for index in deployment.node_indices) + if node_indices != tuple(range(len(node_indices))): + raise ValueError("deployment nodes must be disjoint and contiguous in authored order") + if self.client.host_node_index != self.deployments[0].node_indices[0]: + raise ValueError("client must be colocated on the first node of the first deployment") + if "non_inference_max_parallel_workers" not in self.invocation.authored.run_config: + workers = self.invocation.effective_run_config["non_inference_max_parallel_workers"] + if workers != RunConfig().non_inference_max_parallel_workers: + raise ValueError("default non-inference worker count must match the Data Designer RunConfig default") + + expected_logical_names = tuple( + f"{deployment.deployment_id}-logical-endpoint" for deployment in self.deployments + ) + logical_names = tuple(port.name for port in self.client.ports) + if logical_names != expected_logical_names: + raise ValueError("client requires one ordered logical endpoint port per deployment") + + ports = self.client.ports + tuple(port for deployment in self.deployments for port in deployment.ports) + port_keys = tuple((port.node_index, port.port) for port in ports) + if len(port_keys) != len(set(port_keys)): + raise ValueError("plan port claims must be unique per node") + port_names = tuple(port.name for port in ports) + if len(port_names) != len(set(port_names)): + raise ValueError("plan port claim names must be unique") + otel_port = self.invocation.effective_run_config.get("otel_metrics_port") + if type(otel_port) is int and any( + port.node_index == self.client.host_node_index and port.port == otel_port for port in ports + ): + raise ValueError("client OTEL metrics port collides with a plan port claim") + + run_root = posixpath.join(profile.workspace_root, "runs", self.run_id) + if self.authored_config.path != posixpath.join(run_root, "authored-config.json"): + raise ValueError("authored config reference must use the plan run root") + if self.client.dependency_lock.path != posixpath.join(run_root, "dependency-lock.json"): + raise ValueError("dependency lock reference must use the plan run root") + self._validate_shards(run_root) + if not _is_below(self.output.root, profile.workspace_root): + raise ValueError("resolved output root must be below the selected workspace_root") + shards_root = posixpath.join(run_root, "shards") + if ( + self.output.root == shards_root + or _is_below(self.output.root, shards_root) + or _is_below(shards_root, self.output.root) + ): + raise ValueError("resolved output root must not overlap the run shard workspace") + return self + + def _validate_shards(self, run_root: str) -> None: + if len(self.shards) != self.array_tasks.count: + raise ValueError("plan must contain exactly one shard per array task") + requested_records = self.invocation.authored.num_records + floor_count = requested_records // self.array_tasks.count + expected_start = 0 + shard_ids: list[ShardId] = [] + workspace_paths: list[str] = [] + partition_paths: list[str] = [] + requires_partition = self.invocation.authored.input_bindings.seed_path is not None + for index, shard in enumerate(self.shards): + if shard.shard_index != index or shard.array_task_index != index: + raise ValueError("shards must use complete ordered zero-based identities") + if shard.shard_id != f"shard-{index:05d}": + raise ValueError("shard IDs must match their zero-based shard index") + if shard.record_range.start_index != expected_start: + raise ValueError("shard record ranges must be contiguous") + expected_count = ( + requested_records - floor_count * (self.array_tasks.count - 1) + if index == self.array_tasks.count - 1 + else floor_count + ) + if shard.requested_records != expected_count: + raise ValueError("shards must use deterministic floor/remainder record counts") + expected_workspace = posixpath.join(run_root, "shards", shard.shard_id, "dataset") + if shard.resume_workspace.path != expected_workspace: + raise ValueError("shard resume workspace must match the run and shard identity") + if (shard.input_partition is not None) != requires_partition: + raise ValueError("shard input partition presence must match the authored seed input") + if shard.input_partition is not None: + expected_partition = posixpath.join(run_root, "shards", shard.shard_id, "input-partition.json") + if shard.input_partition.path != expected_partition: + raise ValueError("shard input partition must match the run and shard identity") + partition_paths.append(shard.input_partition.path) + shard_ids.append(shard.shard_id) + workspace_paths.append(shard.resume_workspace.path) + expected_start = shard.record_range.end_index_exclusive + if expected_start != requested_records: + raise ValueError("shard record ranges must cover the requested records") + if len(shard_ids) != len(set(shard_ids)): + raise ValueError("shard IDs must be unique") + if len(workspace_paths) != len(set(workspace_paths)): + raise ValueError("shard resume workspaces must be unique") + if len(partition_paths) != len(set(partition_paths)): + raise ValueError("shard input partitions must be unique") + + +def _is_below(path: str, root: str) -> bool: + return path != root and posixpath.commonpath((path, root)) == root + + +def _extract_builder_aliases(builder: dict[str, JsonValue]) -> tuple[tuple[ModelAlias, ...], tuple[ModelAlias, ...]]: + data_designer = builder.get("data_designer", builder) + if not isinstance(data_designer, dict): + raise ValueError("builder data_designer value must be an object") + model_configs = data_designer.get("model_configs") or [] + if not isinstance(model_configs, list): + raise ValueError("builder model_configs must be a list") + + model_aliases: list[ModelAlias] = [] + for model_config in model_configs: + if not isinstance(model_config, dict) or not isinstance(model_config.get("alias"), str): + raise ValueError("each builder model config must contain a string alias") + model_aliases.append(model_config["alias"]) + + referenced_aliases: list[ModelAlias] = [] + + def collect(value: JsonValue, *, key: str | None = None) -> None: + if key == "model_configs": + return + if key == "model_alias" or (key is not None and key.endswith("_model_alias")): + if not isinstance(value, str): + raise ValueError(f"builder {key} must be a string") + referenced_aliases.append(value) + return + if key == "model_aliases": + if not isinstance(value, list) or any(not isinstance(alias, str) for alias in value): + raise ValueError("builder model_aliases must be a list of strings") + referenced_aliases.extend(value) + return + if isinstance(value, dict): + for child_key, child in value.items(): + collect(child, key=child_key) + elif isinstance(value, list): + for child in value: + collect(child) + + collect(data_designer) + return tuple(model_aliases), tuple(dict.fromkeys(referenced_aliases)) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py new file mode 100644 index 000000000..c47ebc512 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pydantic import JsonValue + +from data_designer.config import RunConfig +from data_designer.slurm.config.images import ClientImageInspection +from data_designer.slurm.config.run import DataDesignerSlurmConfig +from data_designer.slurm.planning.models import ( + ResolvedDependencyLock, + ResolvedSlurmRunPlan, + _extract_builder_aliases, +) + + +class PlanContractError(ValueError): + """Raised when a resolved plan does not match its authored inputs.""" + + +def validate_resolved_plan( + authored: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + plan: ResolvedSlurmRunPlan, + *, + builder_payload: dict[str, JsonValue] | None = None, +) -> ResolvedSlurmRunPlan: + """Validate cross-record identities and digests for one resolved plan.""" + _require( + plan.authored_config.sha256 == authored.compute_sha256(), + "authored config digest does not match the resolved plan", + ) + _require(plan.invocation.authored == authored.invocation, "resolved invocation does not match authored input") + explicit_run_config = RunConfig.model_validate(authored.invocation.run_config).model_dump( + mode="json", + exclude_unset=True, + ) + _require_json_subset(plan.invocation.effective_run_config, explicit_run_config, path="run_config") + _require(plan.client.authored == authored.client, "resolved client does not match authored input") + _require( + tuple(deployment.authored for deployment in plan.deployments) == tuple(authored.deployments), + "resolved deployments do not match authored order and values", + ) + _require(plan.array_tasks == authored.array_tasks, "resolved array task policy does not match authored input") + + if authored.builder.inline is not None: + _require( + plan.builder.inline == authored.builder.inline, "resolved inline builder does not match authored input" + ) + else: + _require( + plan.builder.authored_source == authored.builder.source, + "resolved builder source does not match authored input", + ) + if builder_payload is None: + raise PlanContractError("sourced builder validation requires its resolved payload") + model_aliases, referenced_aliases = _extract_builder_aliases(builder_payload) + _require(plan.builder.model_aliases == model_aliases, "resolved model aliases do not match builder source") + _require( + plan.builder.referenced_model_aliases == referenced_aliases, + "resolved referenced aliases do not match builder source", + ) + + expected_account = authored.submission.account or plan.selected_profile.profile.scheduler.account + expected_partition = authored.submission.partition or plan.selected_profile.profile.scheduler.partition + _require(plan.submission.account == expected_account, "resolved account does not match authored/profile input") + _require( + plan.submission.partition == expected_partition, "resolved partition does not match authored/profile input" + ) + _require( + plan.submission.job_name == authored.submission.job_name, "resolved job name does not match authored input" + ) + _require( + plan.submission.time_limit == authored.submission.time_limit, + "resolved time limit does not match authored input", + ) + _require(plan.submission.comment == authored.submission.comment, "resolved comment does not match authored input") + + _require(plan.output.format == authored.output.format, "resolved output format does not match authored input") + _require( + plan.output.partitions == authored.output.partitions, "resolved output partitions do not match authored input" + ) + _require( + plan.output.require_exact_record_count == authored.output.require_exact_record_count, + "resolved exact-record policy does not match authored input", + ) + if authored.output.root is not None: + _require(plan.output.root == authored.output.root, "resolved output root does not match authored input") + + _require( + plan.client.dependency_lock.sha256 == dependency_lock.compute_sha256(), + "dependency lock digest does not match the resolved plan", + ) + _require( + dependency_lock.client_image_sha256 == plan.client.image.sha256, + "dependency lock client image digest does not match the resolved client image", + ) + inspection = plan.client.image.inspection.inspection + _require(isinstance(inspection, ClientImageInspection), "resolved client image lacks client inspection facts") + _require( + dependency_lock.python_abi == inspection.python_abi, "dependency lock Python ABI does not match client image" + ) + _require( + dependency_lock.image_distributions == inspection.distributions, + "dependency lock image inventory does not match client image inspection", + ) + authored_requirements = authored.client.dependencies.requirements + if authored_requirements is not None: + _require( + dependency_lock.authored_source is None and dependency_lock.source is None, + "dependency lock source is present for authored requirements", + ) + _require( + dependency_lock.authored_requirements == tuple(authored_requirements), + "dependency lock requirements do not match authored requirements", + ) + else: + _require( + dependency_lock.authored_source == authored.client.dependencies.lock_file + and dependency_lock.source is not None, + "dependency lock source does not match the authored lock file", + ) + return plan + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise PlanContractError(message) + + +def _require_json_subset(actual: JsonValue, expected: JsonValue, *, path: str) -> None: + if isinstance(expected, dict): + _require(isinstance(actual, dict), f"explicit {path} value does not match the resolved plan") + for key, value in expected.items(): + _require(key in actual, f"explicit {path}.{key} value is missing from the resolved plan") + _require_json_subset(actual[key], value, path=f"{path}.{key}") + return + _require(actual == expected, f"explicit {path} value does not match the resolved plan") diff --git a/packages/data-designer-slurm/tests/contracts/conftest.py b/packages/data-designer-slurm/tests/contracts/conftest.py new file mode 100644 index 000000000..44a74d99c --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/conftest.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfileCatalog +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan + +GOLDEN_DIR = Path(__file__).parent / "golden" + + +@pytest.fixture +def authored_run() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json((GOLDEN_DIR / "authored_run.json").read_text()) + + +@pytest.fixture +def authored_run_single() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json((GOLDEN_DIR / "authored_run_single.json").read_text()) + + +@pytest.fixture +def profile_catalog() -> SlurmProfileCatalog: + return SlurmProfileCatalog.model_validate_json((GOLDEN_DIR / "profile_catalog.json").read_text()) + + +@pytest.fixture +def dependency_lock() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json((GOLDEN_DIR / "dependency_lock.json").read_text()) + + +@pytest.fixture +def dependency_lock_single() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json((GOLDEN_DIR / "dependency_lock_single.json").read_text()) + + +@pytest.fixture +def single_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((GOLDEN_DIR / "single_node_plan.json").read_text()) + + +@pytest.fixture +def multi_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((GOLDEN_DIR / "multi_node_plan.json").read_text()) diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json new file mode 100644 index 000000000..28f65680a --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -0,0 +1,148 @@ +{ + "array_tasks": { + "count": 2, + "max_concurrent": 2 + }, + "builder": { + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + }, + { + "alias": "judge", + "model": "/models/judge", + "provider": "openai" + } + ] + } + }, + "source": null + }, + "client": { + "cpus": 32, + "dependencies": { + "index_credentials": { + "private-index": { + "environment": "PACKAGE_INDEX_TOKEN", + "type": "secret" + } + }, + "lock_file": null, + "requirements": [ + "data-designer-speech==0.2.0" + ] + }, + "image": { + "name": "dd-client-0.9", + "path": null + } + }, + "deployments": [ + { + "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 2 + }, + "served_model_name": null, + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": { + "HF_TOKEN": { + "environment": "HF_TOKEN", + "type": "secret" + } + }, + "extra_args": [ + "--max-model-len", + "32768" + ], + "image": { + "name": "vllm-0-21", + "path": null + }, + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 2, + "tensor_parallel": 8 + } + }, + { + "model": "/models/judge", + "model_alias": "judge", + "resources": { + "nodes": 1 + }, + "served_model_name": "judge-api", + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "image": { + "name": null, + "path": "/images/vllm-0-22.sqsh" + }, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 1 + } + } + ], + "invocation": { + "dataset_name": "training-data", + "diagnostics": { + "log_requests": false + }, + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 64, + "judge": 32 + }, + "num_records": 100, + "resume": "if_possible", + "run_config": { + "buffer_size": 8192 + } + }, + "name": "two-model-generation", + "output": { + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false, + "root": null + }, + "schema_version": 1, + "submission": { + "account": null, + "comment": null, + "job_name": "dd-two-model", + "partition": null, + "time_limit": "03:55:00" + } +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json new file mode 100644 index 000000000..c6a46897a --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -0,0 +1,96 @@ +{ + "array_tasks": { + "count": 1, + "max_concurrent": 1 + }, + "builder": { + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + } + ] + } + }, + "source": null + }, + "client": { + "cpus": 32, + "dependencies": { + "index_credentials": {}, + "lock_file": null, + "requirements": [] + }, + "image": { + "name": "dd-client-0.9", + "path": null + } + }, + "deployments": [ + { + "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 1 + }, + "served_model_name": null, + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "image": { + "name": "vllm-0-21", + "path": null + }, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 8 + } + } + ], + "invocation": { + "dataset_name": "single-node", + "diagnostics": { + "log_requests": false + }, + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 8 + }, + "num_records": 8, + "resume": "never", + "run_config": {} + }, + "name": "single-node-generation", + "output": { + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false, + "root": null + }, + "schema_version": 1, + "submission": { + "account": null, + "comment": null, + "job_name": "data-designer", + "partition": null, + "time_limit": "03:55:00" + } +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json b/packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json new file mode 100644 index 000000000..dce19282d --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json @@ -0,0 +1,46 @@ +{ + "analysis": { + "target_runtime": "4h", + "target_total_records": 1000000 + }, + "base_run": { + "inline": null, + "source": "run.yaml" + }, + "concurrency_values": [ + 32, + 64, + 128 + ], + "deployment_cases": [ + { + "deployments": { + "generator": { + "nodes": 2, + "nodes_per_replica": 1 + } + }, + "name": "two-independent-replicas" + }, + { + "deployments": { + "generator": { + "nodes": 2, + "nodes_per_replica": 2 + } + }, + "name": "one-two-node-replica" + } + ], + "model_aliases": [ + "generator" + ], + "name": "generator-scaling", + "record_policy": { + "base_records": 1000, + "max_records": 5000, + "records_per_concurrency": 1.0, + "type": "adaptive" + }, + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json b/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json new file mode 100644 index 000000000..d72acb06f --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json @@ -0,0 +1,26 @@ +{ + "benchmark_config": { + "path": "/workspace/primary/benchmarks/benchmark-001/config.json", + "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + }, + "benchmark_id": "benchmark-001", + "children": [ + { + "case_id": "two-independent-replicas-c32", + "child_authored_config": { + "path": "/workspace/primary/runs/run-benchmark-001-c32/authored-config.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "child_run_id": "run-benchmark-001-c32" + }, + { + "case_id": "one-two-node-replica-c32", + "child_authored_config": { + "path": "/workspace/primary/runs/run-benchmark-002-c32/authored-config.json", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "child_run_id": "run-benchmark-002-c32" + } + ], + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json b/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json new file mode 100644 index 000000000..99fbf17ac --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json @@ -0,0 +1,58 @@ +{ + "analysis_id": "analysis-001", + "benchmark_id": "benchmark-001", + "benchmark_manifest": { + "path": "/workspace/primary/benchmarks/benchmark-001/benchmark.json", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "cases": [ + { + "actual_records": 1000, + "boot_seconds": 60.0, + "case_id": "two-independent-replicas-c32", + "child_run_id": "run-benchmark-001-c32", + "feasible": true, + "generation_seconds": 120.0, + "gpu_hours_per_job": 0.8, + "gpus_per_job": 16, + "nodes_per_job": 2, + "outcome": "succeeded", + "request_count": 1000, + "requested_records": 1000, + "rows_per_second": 8.333333333333334, + "target_jobs": 1000, + "token_count": 200000, + "topology_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + "total_gpu_hours": 800.0, + "wall_seconds": 180.0 + }, + { + "actual_records": null, + "boot_seconds": null, + "case_id": "one-two-node-replica-c32", + "child_run_id": "run-benchmark-002-c32", + "feasible": null, + "generation_seconds": null, + "gpu_hours_per_job": null, + "gpus_per_job": 16, + "nodes_per_job": 2, + "outcome": "pending", + "request_count": null, + "requested_records": 1000, + "rows_per_second": null, + "target_jobs": null, + "token_count": null, + "topology_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "total_gpu_hours": null, + "wall_seconds": null + } + ], + "created_at": "2026-08-19T12:00:00Z", + "recommendations": [ + { + "case_id": "two-independent-replicas-c32", + "kind": "minimum_gpu_hours" + } + ], + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json new file mode 100644 index 000000000..6466182d7 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json @@ -0,0 +1,23 @@ +{ + "inspection": { + "distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "installer_path": "/usr/bin/pip", + "installer_version": "26.1", + "kind": "client", + "python_abi": "cp312", + "python_implementation": "cpython", + "python_version": "3.12.12" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/client_result.json b/packages/data-designer-slurm/tests/contracts/golden/client_result.json new file mode 100644 index 000000000..f3b5edac6 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/client_result.json @@ -0,0 +1,20 @@ +{ + "actual_records": 50, + "attempt_id": "attempt-0001", + "candidate_output_manifest": { + "path": "/workspace/primary/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + }, + "completed_at": "2026-08-19T12:00:00Z", + "dataset_path": "/workspace/primary/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset", + "early_shutdown": false, + "effective_resume_mode": "never", + "error_code": null, + "outcome": "complete", + "redacted_message": null, + "requested_records": 50, + "requested_resume_mode": "if_possible", + "run_id": "run-001", + "schema_version": 1, + "shard_id": "shard-00000" +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json new file mode 100644 index 000000000..87190b04c --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -0,0 +1,31 @@ +{ + "authored_requirements": [ + "data-designer-speech==0.2.0" + ], + "authored_source": null, + "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "image_distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "overlay_packages": [ + { + "artifact": { + "path": "/workspace/primary/runs/run-001/dependencies/data_designer_speech-0.2.0.whl", + "sha256": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff" + }, + "name": "data-designer-speech", + "version": "0.2.0" + } + ], + "python_abi": "cp312", + "resolver_version": "resolver-1", + "schema_version": 1, + "source": null +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json new file mode 100644 index 000000000..b86f29d6c --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json @@ -0,0 +1,16 @@ +{ + "authored_requirements": [], + "authored_source": null, + "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "image_distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + } + ], + "overlay_packages": [], + "python_abi": "cp312", + "resolver_version": "resolver-1", + "schema_version": 1, + "source": null +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json new file mode 100644 index 000000000..0f2fd69b1 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -0,0 +1,445 @@ +{ + "array_tasks": { + "count": 2, + "max_concurrent": 2 + }, + "authored_config": { + "path": "/workspace/primary/runs/run-001/authored-config.json", + "sha256": "9bfdbcb1ddd374a5d43eea44890d13c08798ea9e18c30ff5d5258fa0d27f539d" + }, + "builder": { + "authored_source": null, + "content_sha256": "f37227ca7c67e203abe881c0228b4308a8e741364296d293159a1201949732f2", + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + }, + { + "alias": "judge", + "model": "/models/judge", + "provider": "openai" + } + ] + } + }, + "model_aliases": [ + "generator", + "judge" + ], + "referenced_model_aliases": [], + "source": null + }, + "client": { + "authored": { + "cpus": 32, + "dependencies": { + "index_credentials": { + "private-index": { + "environment": "PACKAGE_INDEX_TOKEN", + "type": "secret" + } + }, + "lock_file": null, + "requirements": [ + "data-designer-speech==0.2.0" + ] + }, + "image": { + "name": "dd-client-0.9", + "path": null + } + }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-001/dependency-lock.json", + "sha256": "a86032b310aa6bdb95fc7f35ef606fec2e56b977ba60e36b8b9293341ac43e00" + }, + "gpu_count": 0, + "host_node_index": 0, + "image": { + "authored_ref": { + "name": "dd-client-0.9", + "path": null + }, + "inspection": { + "inspection": { + "distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "installer_path": "/usr/bin/pip", + "installer_version": "26.1", + "kind": "client", + "python_abi": "cp312", + "python_implementation": "cpython", + "python_version": "3.12.12" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "path": "/images/dd-client-0.9.sqsh", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "ports": [ + { + "name": "deployment-00000-logical-endpoint", + "node_index": 0, + "port": 17000, + "role": "logical_endpoint" + }, + { + "name": "deployment-00001-logical-endpoint", + "node_index": 0, + "port": 17001, + "role": "logical_endpoint" + } + ] + }, + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "deployments": [ + { + "authored": { + "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 2 + }, + "served_model_name": null, + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": { + "HF_TOKEN": { + "environment": "HF_TOKEN", + "type": "secret" + } + }, + "extra_args": [ + "--max-model-len", + "32768" + ], + "image": { + "name": "vllm-0-21", + "path": null + }, + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 2, + "tensor_parallel": 8 + } + }, + "deployment_id": "deployment-00000", + "gpus_per_node": 8, + "image": { + "authored_ref": { + "name": "vllm-0-21", + "path": null + }, + "inspection": { + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.21.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "node_indices": [ + 0, + 1 + ], + "ports": [ + { + "name": "deployment-00000-http-00000", + "node_index": 0, + "port": 18000, + "role": "http" + }, + { + "name": "deployment-00000-rendezvous-00000", + "node_index": 0, + "port": 19000, + "role": "rendezvous" + } + ], + "served_model_name": "example/generator", + "topology": { + "gpus_per_replica": 16, + "node_group_count": 1, + "nodes_per_replica": 2, + "pipeline_parallel": 2, + "replica_count": 1, + "replicas_per_node_group": 1, + "tensor_parallel": 8 + } + }, + { + "authored": { + "model": "/models/judge", + "model_alias": "judge", + "resources": { + "nodes": 1 + }, + "served_model_name": "judge-api", + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "image": { + "name": null, + "path": "/images/vllm-0-22.sqsh" + }, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 1 + } + }, + "deployment_id": "deployment-00001", + "gpus_per_node": 8, + "image": { + "authored_ref": { + "name": null, + "path": "/images/vllm-0-22.sqsh" + }, + "inspection": { + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.22.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "9999999999999999999999999999999999999999999999999999999999999999" + }, + "path": "/images/vllm-0-22.sqsh", + "sha256": "9999999999999999999999999999999999999999999999999999999999999999" + }, + "node_indices": [ + 2 + ], + "ports": [ + { + "name": "deployment-00001-http-00000", + "node_index": 2, + "port": 18000, + "role": "http" + }, + { + "name": "deployment-00001-http-00001", + "node_index": 2, + "port": 18001, + "role": "http" + }, + { + "name": "deployment-00001-http-00002", + "node_index": 2, + "port": 18002, + "role": "http" + }, + { + "name": "deployment-00001-http-00003", + "node_index": 2, + "port": 18003, + "role": "http" + }, + { + "name": "deployment-00001-http-00004", + "node_index": 2, + "port": 18004, + "role": "http" + }, + { + "name": "deployment-00001-http-00005", + "node_index": 2, + "port": 18005, + "role": "http" + }, + { + "name": "deployment-00001-http-00006", + "node_index": 2, + "port": 18006, + "role": "http" + }, + { + "name": "deployment-00001-http-00007", + "node_index": 2, + "port": 18007, + "role": "http" + } + ], + "served_model_name": "judge-api", + "topology": { + "gpus_per_replica": 1, + "node_group_count": 1, + "nodes_per_replica": 1, + "pipeline_parallel": 1, + "replica_count": 8, + "replicas_per_node_group": 8, + "tensor_parallel": 1 + } + } + ], + "invocation": { + "authored": { + "dataset_name": "training-data", + "diagnostics": { + "log_requests": false + }, + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 64, + "judge": 32 + }, + "num_records": 100, + "resume": "if_possible", + "run_config": { + "buffer_size": 8192 + } + }, + "effective_run_config": { + "async_trace": false, + "buffer_size": 8192, + "disable_early_shutdown": true, + "display_tui": false, + "jinja_rendering_engine": "secure", + "max_concurrent_row_groups": 3, + "max_conversation_correction_steps": 0, + "max_conversation_restarts": 0, + "max_in_flight_tasks": 1024, + "non_inference_max_parallel_workers": 4, + "otel_metrics_port": null, + "preserve_dropped_columns": true, + "progress_interval": 5.0, + "request_admission": null, + "shutdown_error_rate": 1.0, + "shutdown_error_window": 10, + "write_scheduler_events": false + } + }, + "output": { + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false, + "root": "/workspace/primary/runs/run-001/output" + }, + "package_version": "0.9.2", + "resolved_gpus_per_node": 8, + "run_id": "run-001", + "runtime_bundle": { + "path": "/workspace/primary/runtime/runtime.tar.gz", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + }, + "schema_version": 1, + "selected_profile": { + "catalog_path": "/workspace/profile.json", + "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "cluster_name": "primary", + "matched_pattern": null, + "profile": { + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "gpu_request_mode": "gres", + "gpus_per_node": 8, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "image_build": { + "partition": "cpu" + }, + "scheduler": { + "account": "research", + "mem_per_gpu": null, + "partition": "batch" + }, + "schema_version": 1, + "workspace_root": "/workspace/primary" + }, + "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "schema_version": 1, + "selection_source": "explicit" + }, + "shards": [ + { + "array_task_index": 0, + "input_partition": null, + "record_range": { + "end_index_exclusive": 50, + "start_index": 0 + }, + "resume_workspace": { + "path": "/workspace/primary/runs/run-001/shards/shard-00000/dataset" + }, + "shard_id": "shard-00000", + "shard_index": 0 + }, + { + "array_task_index": 1, + "input_partition": null, + "record_range": { + "end_index_exclusive": 100, + "start_index": 50 + }, + "resume_workspace": { + "path": "/workspace/primary/runs/run-001/shards/shard-00001/dataset" + }, + "shard_id": "shard-00001", + "shard_index": 1 + } + ], + "submission": { + "account": "research", + "comment": null, + "job_name": "dd-two-model", + "partition": "batch", + "time_limit": "03:55:00" + } +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json new file mode 100644 index 000000000..8d500c664 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json @@ -0,0 +1,49 @@ +{ + "clusters": { + "lab": { + "container_mounts": [], + "gpu_request_mode": "visible", + "gpus_per_node": "auto", + "host_patterns": [ + "lab-login-*" + ], + "image_build": { + "partition": "cpu" + }, + "scheduler": { + "account": "lab", + "mem_per_gpu": null, + "partition": "gpu" + }, + "schema_version": 1, + "workspace_root": "/workspace/lab" + }, + "primary": { + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "gpu_request_mode": "gres", + "gpus_per_node": 8, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "image_build": { + "partition": "cpu" + }, + "scheduler": { + "account": "research", + "mem_per_gpu": null, + "partition": "batch" + }, + "schema_version": 1, + "workspace_root": "/workspace/primary" + } + }, + "default_cluster": "primary", + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json b/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json new file mode 100644 index 000000000..8cb5348e2 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json @@ -0,0 +1,11 @@ +{ + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.21.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json new file mode 100644 index 000000000..f0794cc75 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -0,0 +1,276 @@ +{ + "array_tasks": { + "count": 1, + "max_concurrent": 1 + }, + "authored_config": { + "path": "/workspace/primary/runs/run-single/authored-config.json", + "sha256": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" + }, + "builder": { + "authored_source": null, + "content_sha256": "b3ef5fc1fe675a8e004633f84842ac60cf82d5ba3dc68b4d50ee4438448b0570", + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + } + ] + } + }, + "model_aliases": [ + "generator" + ], + "referenced_model_aliases": [], + "source": null + }, + "client": { + "authored": { + "cpus": 32, + "dependencies": { + "index_credentials": {}, + "lock_file": null, + "requirements": [] + }, + "image": { + "name": "dd-client-0.9", + "path": null + } + }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-single/dependency-lock.json", + "sha256": "2600c1944897a8e3c9a9410fe912bdee0e169e09769e58180cadebe94f3da52c" + }, + "gpu_count": 0, + "host_node_index": 0, + "image": { + "authored_ref": { + "name": "dd-client-0.9", + "path": null + }, + "inspection": { + "inspection": { + "distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + } + ], + "installer_path": "/usr/bin/pip", + "installer_version": "26.1", + "kind": "client", + "python_abi": "cp312", + "python_implementation": "cpython", + "python_version": "3.12.12" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "path": "/images/dd-client-0.9.sqsh", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "ports": [ + { + "name": "deployment-00000-logical-endpoint", + "node_index": 0, + "port": 17000, + "role": "logical_endpoint" + } + ] + }, + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "deployments": [ + { + "authored": { + "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 1 + }, + "served_model_name": null, + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "image": { + "name": "vllm-0-21", + "path": null + }, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 8 + } + }, + "deployment_id": "deployment-00000", + "gpus_per_node": 8, + "image": { + "authored_ref": { + "name": "vllm-0-21", + "path": null + }, + "inspection": { + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.21.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "node_indices": [ + 0 + ], + "ports": [ + { + "name": "deployment-00000-http-00000", + "node_index": 0, + "port": 18000, + "role": "http" + } + ], + "served_model_name": "example/generator", + "topology": { + "gpus_per_replica": 8, + "node_group_count": 1, + "nodes_per_replica": 1, + "pipeline_parallel": 1, + "replica_count": 1, + "replicas_per_node_group": 1, + "tensor_parallel": 8 + } + } + ], + "invocation": { + "authored": { + "dataset_name": "single-node", + "diagnostics": { + "log_requests": false + }, + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 8 + }, + "num_records": 8, + "resume": "never", + "run_config": {} + }, + "effective_run_config": { + "async_trace": false, + "buffer_size": 16384, + "disable_early_shutdown": true, + "display_tui": false, + "jinja_rendering_engine": "secure", + "max_concurrent_row_groups": 3, + "max_conversation_correction_steps": 0, + "max_conversation_restarts": 0, + "max_in_flight_tasks": 1024, + "non_inference_max_parallel_workers": 4, + "otel_metrics_port": null, + "preserve_dropped_columns": true, + "progress_interval": 5.0, + "request_admission": null, + "shutdown_error_rate": 1.0, + "shutdown_error_window": 10, + "write_scheduler_events": false + } + }, + "output": { + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false, + "root": "/workspace/primary/runs/run-single/output" + }, + "package_version": "0.9.2", + "resolved_gpus_per_node": 8, + "run_id": "run-single", + "runtime_bundle": { + "path": "/workspace/primary/runtime/runtime.tar.gz", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + }, + "schema_version": 1, + "selected_profile": { + "catalog_path": "/workspace/profile.json", + "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "cluster_name": "primary", + "matched_pattern": null, + "profile": { + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "gpu_request_mode": "gres", + "gpus_per_node": 8, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "image_build": { + "partition": "cpu" + }, + "scheduler": { + "account": "research", + "mem_per_gpu": null, + "partition": "batch" + }, + "schema_version": 1, + "workspace_root": "/workspace/primary" + }, + "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "schema_version": 1, + "selection_source": "explicit" + }, + "shards": [ + { + "array_task_index": 0, + "input_partition": null, + "record_range": { + "end_index_exclusive": 8, + "start_index": 0 + }, + "resume_workspace": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/dataset" + }, + "shard_id": "shard-00000", + "shard_index": 0 + } + ], + "submission": { + "account": "research", + "comment": null, + "job_name": "data-designer", + "partition": "batch", + "time_limit": "03:55:00" + } +} diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py new file mode 100644 index 000000000..33959d5c1 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from copy import deepcopy + +import pytest +from pydantic import ValidationError + +from data_designer.config import DataDesignerConfigBuilder, HuggingFaceSeedSource, ModelConfig +from data_designer.slurm.config import ( + ArrayTasksConfig, + BenchmarkBaseRun, + BuilderInput, + ClientDependencies, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageBuildRequest, + ImageInspectionRecord, + ImageRef, + LiteralEnvironmentBinding, + LocalStdioMCPProviderConfig, + QueueBackpressureConfig, + RemoteMCPProviderConfig, + SecretRef, + ServerDeploymentConfig, + SubmissionConfig, + VllmServerConfig, +) + + +@pytest.mark.parametrize("version", [None, 0, 2, "1"]) +def test_run_config_requires_supported_version(authored_run: DataDesignerSlurmConfig, version: object) -> None: + payload = authored_run.model_dump(mode="json") + if version is None: + payload.pop("schema_version") + else: + payload["schema_version"] = version + + with pytest.raises(ValidationError): + DataDesignerSlurmConfig.model_validate(payload) + + +def test_run_config_rejects_unknown_fields(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["placement"] = {"gpu_ids": [0]} + + with pytest.raises(ValidationError, match="placement"): + DataDesignerSlurmConfig.model_validate(payload) + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"name": "image", "path": "/images/image.sqsh"}, + {"path": "relative.sqsh"}, + {"path": "/images/image.tar"}, + ], +) +def test_image_ref_requires_one_registered_alias_or_absolute_sqsh(payload: dict[str, str]) -> None: + with pytest.raises(ValidationError): + ImageRef.model_validate(payload) + + +@pytest.mark.parametrize( + "payload", + [ + {"requirements": None}, + {"requirements": [], "lock_file": "lock.json"}, + {"requirements": ["-e ./plugin"]}, + {"requirements": ["plugin @ git+https://example.test/plugin.git"]}, + {"requirements": ["plugin @ https://example.test/plugin.whl"]}, + {"requirements": ["plugin @ https://user:secret@example.test/plugin.whl#sha256=" + "a" * 64]}, + {"requirements": ["plugin @ https://example.test/plugin.whl?token=secret#sha256=" + "a" * 64]}, + {"requirements": ["plugin @ https://example.test:invalid/plugin.whl#sha256=" + "a" * 64]}, + {"requirements": ["my_pkg==1", "my-pkg==2"]}, + {"requirements": ["not valid !!!"]}, + {"requirements": ["plugin=="]}, + {"requirements": ["plugin==1; python_version >= '3.12'"]}, + {"requirements": None, "lock_file": "../lock.json"}, + ], +) +def test_client_dependencies_reject_mutable_or_ambiguous_sources(payload: dict[str, object]) -> None: + with pytest.raises(ValidationError): + ClientDependencies.model_validate(payload) + + +def test_client_dependencies_accept_digest_bound_wheel() -> None: + dependencies = ClientDependencies(requirements=["plugin @ https://example.test/plugin.whl#sha256=" + "a" * 64]) + + assert dependencies.requirements is not None + + +def test_secret_reference_serializes_only_external_binding() -> None: + secret = SecretRef(type="secret", environment="HF_TOKEN") + + assert secret.model_dump(mode="json") == {"type": "secret", "environment": "HF_TOKEN"} + with pytest.raises(ValidationError): + SecretRef.model_validate({"type": "secret", "environment": "HF_TOKEN", "value": "secret-value"}) + + +def test_literal_environment_rejects_control_characters() -> None: + with pytest.raises(ValidationError, match="control"): + LiteralEnvironmentBinding(type="literal", value="line\nbreak") + + +def test_secret_shaped_environment_requires_external_reference() -> None: + literal = LiteralEnvironmentBinding(type="literal", value="plaintext-secret") + + with pytest.raises(ValidationError, match="external secret references"): + LocalStdioMCPProviderConfig( + provider_type="stdio", + name="provider", + command="provider", + environment={"API_TOKEN": literal}, + ) + with pytest.raises(ValidationError, match="external secret references"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + environment={"MODEL_PASSWORD": literal}, + ) + + +@pytest.mark.parametrize( + "endpoint", + [ + "https://user:password@example.test/mcp", + "https://example.test/mcp?token=plaintext-secret", + "https://example.test/mcp#secret", + ], +) +def test_remote_mcp_endpoint_rejects_embedded_credentials(endpoint: str) -> None: + with pytest.raises(ValidationError, match="must not embed"): + RemoteMCPProviderConfig(provider_type="sse", name="provider", endpoint=endpoint) + + +@pytest.mark.parametrize("endpoint", ["https:///missing-host", "https://example.test:invalid/mcp"]) +def test_remote_mcp_endpoint_requires_valid_host_and_port(endpoint: str) -> None: + with pytest.raises(ValidationError, match=r"HTTP\(S\)"): + RemoteMCPProviderConfig(provider_type="sse", name="provider", endpoint=endpoint) + + +@pytest.mark.parametrize( + "argument", + [ + "--api-key", + "--api-key plaintext-secret", + " --api-key plaintext-secret", + "--access-token=plaintext-secret", + "password", + ], +) +def test_stdio_mcp_rejects_secret_shaped_arguments(argument: str) -> None: + with pytest.raises(ValidationError, match="secret-shaped"): + LocalStdioMCPProviderConfig( + provider_type="stdio", + name="provider", + command="provider", + args=[argument], + ) + + +def test_vllm_defaults_and_backpressure_override() -> None: + server = VllmServerConfig(type="vllm", image=ImageRef(name="vllm")) + overridden = VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + queue_backpressure=QueueBackpressureConfig(max_waiting_requests=0, retry_after_seconds=None), + ) + + assert server.queue_backpressure.model_dump() == {"max_waiting_requests": 128, "retry_after_seconds": 1} + assert overridden.queue_backpressure.model_dump() == {"max_waiting_requests": 0, "retry_after_seconds": None} + + +@pytest.mark.parametrize( + "argument", + [ + "--api-key=plaintext-secret", + "--api-key plaintext-secret", + "--distributed-init-address", + "--host=0.0.0.0", + "--middleware", + "--model", + "--port", + "--port 9000", + " --port 9000", + "--tensor-parallel-size", + ], +) +def test_vllm_rejects_runtime_owned_arguments(argument: str) -> None: + with pytest.raises(ValidationError, match="owned"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=[argument]) + + +@pytest.mark.parametrize( + "extra_args", + [ + ["--hf-token=plaintext-secret"], + ["--hf-token", "plaintext-secret"], + ["--hf-token plaintext-secret"], + ], +) +def test_vllm_rejects_secret_shaped_arguments(extra_args: list[str]) -> None: + with pytest.raises(ValidationError, match="secret-shaped"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=extra_args) + + +def test_vllm_rejects_distributed_timeout_beyond_startup_timeout() -> None: + with pytest.raises(ValidationError, match="must not exceed"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + startup_timeout="10m", + distributed_init_timeout="11m", + ) + + +def test_deployment_rejects_invalid_topology() -> None: + payload = { + "model_alias": "generator", + "model": "example/generator", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + "resources": {"nodes": 3}, + "topology": {"tensor_parallel": 8, "nodes_per_replica": 2}, + } + + with pytest.raises(ValidationError, match="divide"): + ServerDeploymentConfig.model_validate(payload) + + payload["resources"]["nodes"] = 2 + payload["server"]["enable_expert_parallel"] = True + with pytest.raises(ValidationError, match="expert"): + ServerDeploymentConfig.model_validate(payload) + + +def test_model_alias_preserves_public_data_designer_values() -> None: + alias = "judge/v2" + ModelConfig(alias=alias, model="example/judge", provider="openai") + + deployment = ServerDeploymentConfig.model_validate( + { + "model_alias": alias, + "model": "example/judge", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + } + ) + + assert deployment.model_alias == alias + + +def test_run_rejects_duplicate_alias_and_unknown_concurrency(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["deployments"][1]["model_alias"] = "generator" + with pytest.raises(ValidationError, match="aliases"): + DataDesignerSlurmConfig.model_validate(payload) + + payload = authored_run.model_dump(mode="json") + payload["invocation"]["model_concurrency"]["missing"] = 1 + with pytest.raises(ValidationError, match="undeclared"): + DataDesignerSlurmConfig.model_validate(payload) + + +def test_run_rejects_retired_builder_fields(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["server_configs"] = [] + + with pytest.raises(ValidationError, match="retired"): + DataDesignerSlurmConfig.model_validate(payload) + + +@pytest.mark.parametrize( + "secret_key", + [ + "api_key", + "accessToken", + "client-secret", + "consumer_key", + "license_key", + "password", + "private_key", + "signing_key", + "ssh_key", + "subscription_key", + ], +) +def test_builder_input_rejects_secret_values(secret_key: str) -> None: + with pytest.raises(ValidationError, match="secret values"): + BuilderInput(inline={"columns": [], secret_key: "plaintext-secret"}) + + +@pytest.mark.parametrize("plugin_key", ["sort_key", "partition_key", "primary_key", "idempotency_key"]) +def test_builder_input_allows_non_secret_plugin_keys(plugin_key: str) -> None: + inline = {"columns": [{"column_type": "plugin", plugin_key: "value"}]} + + assert BuilderInput(inline=inline).inline == inline + + +def test_builder_input_accepts_exported_and_shorthand_configs() -> None: + exported = DataDesignerConfigBuilder(model_configs=[]).get_builder_config().to_dict() + + assert BuilderInput(inline=exported).inline == exported + assert BuilderInput(inline={"columns": []}).inline == {"columns": []} + + +def test_builder_input_accepts_null_secret_fields_from_canonical_export() -> None: + builder = DataDesignerConfigBuilder(model_configs=[]).with_seed_dataset( + HuggingFaceSeedSource(path="datasets/example/seed/*.parquet") + ) + exported = builder.get_builder_config().to_dict() + + assert exported["data_designer"]["seed_config"]["source"]["token"] is None + assert BuilderInput(inline=exported).inline == exported + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_builder_input_rejects_non_finite_json(value: float) -> None: + with pytest.raises(ValidationError): + BuilderInput(inline={"columns": [], "value": value}) + + +@pytest.mark.parametrize( + "inline", + [ + {"data_designer": {}, "library_version": 1}, + {"data_designer": {}, "unknown": True}, + {"model_configs": []}, + ], +) +def test_builder_input_rejects_invalid_serialized_shapes(inline: dict[str, object]) -> None: + with pytest.raises(ValidationError, match="complete serialized"): + BuilderInput.model_validate({"inline": inline}) + + +def test_run_validates_public_run_config_and_shard_count(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["invocation"]["run_config"]["buffer_size"] = 0 + with pytest.raises(ValidationError, match="buffer_size"): + DataDesignerSlurmConfig.model_validate(payload) + + payload = authored_run.model_dump(mode="json") + payload["array_tasks"]["count"] = 101 + payload["array_tasks"]["max_concurrent"] = 1 + with pytest.raises(ValidationError, match="requested records"): + DataDesignerSlurmConfig.model_validate(payload) + + +def test_small_config_values_validate_at_boundary() -> None: + with pytest.raises(ValidationError, match="concurrency"): + ArrayTasksConfig(count=2, max_concurrent=3) + with pytest.raises(ValidationError, match="minutes"): + SubmissionConfig(time_limit="00:60:00") + with pytest.raises(ValidationError, match="readiness_path"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), readiness_path="health") + + +@pytest.mark.parametrize( + "source", + ["nvcr.io/example/vllm:latest", "relative.sqsh"], +) +def test_image_build_request_rejects_mutable_or_relative_source(source: str) -> None: + with pytest.raises(ValidationError): + ImageBuildRequest(name="vllm", kind="serving", source=source) + + +def test_image_inspection_rejects_duplicate_distribution_names() -> None: + payload = { + "schema_version": 1, + "inspector_version": "v1", + "sqsh_sha256": "a" * 64, + "inspection": { + "kind": "client", + "python_implementation": "cpython", + "python_version": "3.12.1", + "python_abi": "cp312", + "distributions": [ + {"name": "plugin", "version": "1"}, + {"name": "plugin", "version": "2"}, + ], + "installer_path": "/usr/bin/pip", + "installer_version": "1", + }, + } + + with pytest.raises(ValidationError, match="unique"): + ImageInspectionRecord.model_validate_json(json.dumps(payload)) + + +def test_benchmark_config_rejects_duplicate_axes() -> None: + payload = { + "schema_version": 1, + "name": "bench", + "base_run": "./run.yaml", + "model_aliases": ["generator", "generator"], + "concurrency_values": [32, 32], + "deployment_cases": [{"name": "case", "deployments": {"generator": {"nodes": 1, "nodes_per_replica": 1}}}], + "record_policy": {"type": "fixed", "records": 100}, + "analysis": {"target_total_records": 1000, "target_runtime": "1h"}, + } + + with pytest.raises(ValidationError, match="aliases"): + DataDesignerSlurmBenchmarkConfig.model_validate(payload) + + payload["model_aliases"] = ["generator"] + with pytest.raises(ValidationError, match="concurrency"): + DataDesignerSlurmBenchmarkConfig.model_validate(payload) + + +def test_benchmark_base_run_normalizes_local_source() -> None: + base_run = BenchmarkBaseRun.model_validate("./run.yaml") + + assert base_run.source == "run.yaml" + + +def test_config_models_do_not_mutate_input(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + original = deepcopy(payload) + + DataDesignerSlurmConfig.model_validate(payload) + + assert payload == original + + +def test_config_models_are_deeply_immutable(authored_run: DataDesignerSlurmConfig) -> None: + inline = authored_run.builder.inline + assert inline is not None + data_designer = inline["data_designer"] + assert isinstance(data_designer, dict) + columns = data_designer["columns"] + assert isinstance(columns, list) + + with pytest.raises(TypeError, match="frozen list"): + authored_run.deployments.clear() + with pytest.raises(TypeError, match="frozen dictionary"): + authored_run.invocation.run_config["buffer_size"] = 1 + with pytest.raises(TypeError, match="frozen list"): + columns.append({}) diff --git a/packages/data-designer-slurm/tests/contracts/test_golden_records.py b/packages/data-designer-slurm/tests/contracts/test_golden_records.py new file mode 100644 index 000000000..712b158c9 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_golden_records.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageInspectionRecord, + SlurmProfileCatalog, +) +from data_designer.slurm.contracts import AuthoredConfig, ContractRecord +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan + +GOLDEN_DIR = Path(__file__).parent / "golden" + + +@pytest.mark.parametrize( + ("record_type", "filename"), + [ + (DataDesignerSlurmConfig, "authored_run.json"), + (DataDesignerSlurmConfig, "authored_run_single.json"), + (SlurmProfileCatalog, "profile_catalog.json"), + (ImageInspectionRecord, "client_image_inspection.json"), + (ImageInspectionRecord, "serving_image_inspection.json"), + (ResolvedDependencyLock, "dependency_lock.json"), + (ResolvedDependencyLock, "dependency_lock_single.json"), + (ResolvedSlurmRunPlan, "single_node_plan.json"), + (ResolvedSlurmRunPlan, "multi_node_plan.json"), + (ClientResult, "client_result.json"), + (DataDesignerSlurmBenchmarkConfig, "benchmark_config.json"), + (BenchmarkManifest, "benchmark_manifest.json"), + (BenchmarkReport, "benchmark_report.json"), + ], +) +def test_golden_record_round_trip(record_type: type[BaseModel], filename: str) -> None: + fixture = (GOLDEN_DIR / filename).read_text() + record = record_type.model_validate_json(fixture) + + assert record_type.model_validate_json(record.model_dump_json()) == record + if isinstance(record, (AuthoredConfig, ContractRecord)): + assert record_type.model_validate_json(record.serialize_json()) == record + assert record.compute_sha256() == hashlib.sha256(record.serialize_json().encode()).hexdigest() + assert record.serialize_json() == fixture + + +def test_golden_records_are_sanitized() -> None: + contents = "\n".join(path.read_text().casefold() for path in GOLDEN_DIR.glob("*.json")) + + assert "nvidia" not in contents + assert "secret-value" not in contents + assert "lustre" not in contents + + +def test_canonical_serialization_ignores_mapping_order(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["invocation"]["model_concurrency"] = {"judge": 32, "generator": 64} + reordered = DataDesignerSlurmConfig.model_validate(payload) + + assert authored_run.serialize_json() == reordered.serialize_json() + assert authored_run.compute_sha256() == reordered.compute_sha256() diff --git a/packages/data-designer-slurm/tests/contracts/test_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py new file mode 100644 index 000000000..5c64a5c24 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from copy import deepcopy + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig +from data_designer.slurm.contracts import compute_sha256 +from data_designer.slurm.planning import ( + ArtifactReference, + PlanContractError, + ResolvedDependencyLock, + ResolvedDeployment, + ResolvedSlurmRunPlan, + ResolvedSubmission, + validate_resolved_plan, +) + + +def test_multi_node_plan_matches_authored_inputs( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + assert validate_resolved_plan(authored_run, dependency_lock, multi_node_plan) is multi_node_plan + assert multi_node_plan.authored_config.sha256 == authored_run.compute_sha256() + assert [deployment.topology.replica_count for deployment in multi_node_plan.deployments] == [1, 8] + assert multi_node_plan.client.gpu_count == 0 + + +def test_single_node_plan_matches_authored_inputs( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + assert validate_resolved_plan(authored_run_single, dependency_lock_single, single_node_plan) is single_node_plan + assert single_node_plan.authored_config.sha256 == authored_run_single.compute_sha256() + assert [deployment.topology.replica_count for deployment in single_node_plan.deployments] == [1] + assert single_node_plan.client.gpu_count == 0 + + +def test_plan_canonical_json_is_byte_stable(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = json.loads(multi_node_plan.serialize_json()) + payload["invocation"]["authored"]["model_concurrency"] = {"judge": 32, "generator": 64} + reordered = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + assert reordered.serialize_canonical_json() == multi_node_plan.serialize_canonical_json() + + +def test_resolved_plan_is_deeply_immutable(multi_node_plan: ResolvedSlurmRunPlan) -> None: + inline = multi_node_plan.builder.inline + assert inline is not None + data_designer = inline["data_designer"] + assert isinstance(data_designer, dict) + columns = data_designer["columns"] + assert isinstance(columns, list) + + with pytest.raises(TypeError, match="frozen dictionary"): + multi_node_plan.invocation.effective_run_config["buffer_size"] = 1 + with pytest.raises(TypeError, match="frozen list"): + columns.append({}) + + +@pytest.mark.parametrize( + "mutator", + [ + lambda payload: payload.pop("schema_version"), + lambda payload: payload.update(schema_version=2), + lambda payload: payload.update(unknown=True), + lambda payload: payload.update(resolved_gpus_per_node=4), + lambda payload: payload["client"].update(host_node_index=1), + lambda payload: payload["deployments"][1].update(node_indices=[1]), + lambda payload: payload["deployments"][0].update(deployment_id="unrelated-runtime-name"), + lambda payload: payload["deployments"][0]["ports"][1].update(port=18000), + lambda payload: payload["deployments"][1].update(ports=payload["deployments"][1]["ports"][:1]), + lambda payload: payload["client"].update(ports=payload["client"]["ports"][:1]), + lambda payload: payload.update(shards=payload["shards"][:1]), + lambda payload: payload["shards"][1]["record_range"].update(start_index=49), + lambda payload: payload["shards"][1].update(shard_id="shard-00002"), + lambda payload: payload["shards"][1].update(resume_workspace=payload["shards"][0]["resume_workspace"]), + lambda payload: payload.update(run_id="run-other"), + lambda payload: payload["output"].update(root="/outside/output"), + lambda payload: payload.update(container_mounts=[]), + lambda payload: payload["deployments"][0]["topology"].update(replica_count=2), + ], +) +def test_plan_rejects_invalid_boundaries(multi_node_plan: ResolvedSlurmRunPlan, mutator: object) -> None: + payload = deepcopy(multi_node_plan.model_dump(mode="json")) + mutator(payload) + + with pytest.raises(ValidationError): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_rejects_unmaterialized_run_config(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["effective_run_config"] = {} + + with pytest.raises(ValidationError, match="fully materialized"): + ResolvedSlurmRunPlan.model_validate(payload) + + +def test_plan_preserves_default_non_inference_worker_count( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 32 + + with pytest.raises(ValidationError, match="RunConfig default"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_preserves_explicit_non_inference_worker_override( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["authored"]["run_config"]["non_inference_max_parallel_workers"] = 32 + payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 32 + + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + assert plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == 32 + + +def test_plan_rejects_otel_port_collision(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["effective_run_config"]["otel_metrics_port"] = 18000 + + with pytest.raises(ValidationError, match="OTEL metrics port collides"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + "output_root", + [ + "/workspace/primary/runs/run-001/shards/shard-00000/dataset", + "/workspace/primary/runs/run-001", + "/workspace/primary/runs", + ], +) +def test_plan_rejects_output_overlapping_shard_workspace( + multi_node_plan: ResolvedSlurmRunPlan, + output_root: str, +) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["output"]["root"] = output_root + + with pytest.raises(ValidationError, match="must not overlap"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_rejects_deployment_alias_missing_from_builder(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["model_configs"] = payload["builder"]["inline"]["data_designer"][ + "model_configs" + ][:1] + payload["builder"]["model_aliases"] = ["generator"] + payload["builder"]["content_sha256"] = compute_sha256(payload["builder"]["inline"]) + + with pytest.raises(ValidationError, match="deployment alias"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_rejects_referenced_alias_without_deployment(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["columns"] = [{"model_alias": "missing"}] + payload["builder"]["referenced_model_aliases"] = ["missing"] + payload["builder"]["content_sha256"] = compute_sha256(payload["builder"]["inline"]) + + with pytest.raises(ValidationError, match="referenced Data Designer model alias"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_multi_node_tp4_requires_rendezvous_per_replica_lane(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.deployments[0].model_dump(mode="json") + payload["authored"]["topology"]["tensor_parallel"] = 4 + payload["topology"].update( + tensor_parallel=4, + replicas_per_node_group=2, + replica_count=2, + gpus_per_replica=8, + ) + payload["ports"] = [ + { + "name": "deployment-00000-http-00000", + "role": "http", + "node_index": 0, + "port": 18000, + }, + { + "name": "deployment-00000-http-00001", + "role": "http", + "node_index": 0, + "port": 18001, + }, + { + "name": "deployment-00000-rendezvous-00000", + "role": "rendezvous", + "node_index": 0, + "port": 19000, + }, + ] + + with pytest.raises(ValidationError, match="rendezvous"): + ResolvedDeployment.model_validate_json(json.dumps(payload)) + + payload["ports"].append( + { + "name": "deployment-00000-rendezvous-00001", + "role": "rendezvous", + "node_index": 0, + "port": 19001, + } + ) + assert ResolvedDeployment.model_validate_json(json.dumps(payload)).topology.replica_count == 2 + + +def test_sourced_builder_validation_requires_resolved_payload( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + sourced_authored = authored_run.model_copy(update={"builder": BuilderInput(source="builder.json")}) + payload = multi_node_plan.model_dump(mode="json") + payload["authored_config"]["sha256"] = sourced_authored.compute_sha256() + payload["builder"] = { + "authored_source": "builder.json", + "source": {"path": "/workspace/primary/runs/run-001/builder.json", "sha256": "a" * 64}, + "inline": None, + "content_sha256": "a" * 64, + "model_aliases": ["generator", "judge"], + "referenced_model_aliases": [], + } + sourced_plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + with pytest.raises(PlanContractError, match="resolved payload"): + validate_resolved_plan(sourced_authored, dependency_lock, sourced_plan) + + assert ( + validate_resolved_plan( + sourced_authored, + dependency_lock, + sourced_plan, + builder_payload=authored_run.builder.inline, + ) + is sourced_plan + ) + + +@pytest.mark.parametrize( + "update", + [ + {"time_limit": "invalid"}, + {"comment": "bad\ncomment"}, + ], +) +def test_resolved_submission_preserves_authored_validation(update: dict[str, object]) -> None: + payload = {"job_name": "data-designer", "time_limit": "03:55:00", **update} + + with pytest.raises(ValidationError): + ResolvedSubmission.model_validate(payload) + + +def test_resolved_image_rejects_digest_mismatch(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["client"]["image"]["inspection"]["sqsh_sha256"] = "a" * 64 + + with pytest.raises(ValidationError, match="digest"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + "mutator", + [ + lambda payload: payload["overlay_packages"].append( + { + "name": "data-designer", + "version": "0.9.2", + "artifact": {"path": "/wheels/data_designer.whl", "sha256": "a" * 64}, + } + ), + lambda payload: payload.update(image_distributions=list(reversed(payload["image_distributions"]))), + lambda payload: payload["overlay_packages"][0]["artifact"].update(path="/wheels/plugin.tar.gz"), + lambda payload: payload.update(authored_source="lock.json"), + lambda payload: payload.update( + authored_source="lock.yaml", + source={"path": "/workspace/lock.yaml", "sha256": "a" * 64}, + ), + ], +) +def test_dependency_lock_rejects_invalid_boundaries( + dependency_lock: ResolvedDependencyLock, + mutator: object, +) -> None: + payload = deepcopy(dependency_lock.model_dump(mode="json")) + mutator(payload) + + with pytest.raises(ValidationError): + ResolvedDependencyLock.model_validate_json(json.dumps(payload)) + + +def test_dependency_lock_requires_every_authored_package(dependency_lock: ResolvedDependencyLock) -> None: + payload = dependency_lock.model_dump(mode="json") + payload["overlay_packages"] = [] + + with pytest.raises(ValidationError, match="missing from the dependency lock"): + ResolvedDependencyLock.model_validate_json(json.dumps(payload)) + + +def test_dependency_lock_requires_compatible_versions(dependency_lock: ResolvedDependencyLock) -> None: + payload = dependency_lock.model_dump(mode="json") + payload["overlay_packages"][0]["version"] = "0.3.0" + + with pytest.raises(ValidationError, match="does not satisfy"): + ResolvedDependencyLock.model_validate_json(json.dumps(payload)) + + +def test_dependency_lock_binds_direct_wheel_digest(dependency_lock: ResolvedDependencyLock) -> None: + payload = dependency_lock.model_dump(mode="json") + payload["authored_requirements"] = [ + "data-designer-speech @ https://example.test/data_designer_speech.whl#sha256=" + "a" * 64 + ] + + with pytest.raises(ValidationError, match="locked overlay artifact"): + ResolvedDependencyLock.model_validate_json(json.dumps(payload)) + + payload["overlay_packages"][0]["artifact"]["sha256"] = "a" * 64 + assert ResolvedDependencyLock.model_validate_json(json.dumps(payload)).overlay_packages + + +def test_cross_record_validation_rejects_authored_digest( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invalid = multi_node_plan.model_copy( + update={ + "authored_config": ArtifactReference( + path=multi_node_plan.authored_config.path, + sha256="0" * 64, + ) + } + ) + + with pytest.raises(PlanContractError, match="authored config digest"): + validate_resolved_plan(authored_run, dependency_lock, invalid) + + +def test_cross_record_validation_rejects_dependency_lock_digest( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + client = multi_node_plan.client.model_copy( + update={ + "dependency_lock": ArtifactReference( + path=multi_node_plan.client.dependency_lock.path, + sha256="0" * 64, + ) + } + ) + invalid = multi_node_plan.model_copy(update={"client": client}) + + with pytest.raises(PlanContractError, match="dependency lock digest"): + validate_resolved_plan(authored_run, dependency_lock, invalid) + + +def test_cross_record_validation_binds_authored_lock_source( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + dependencies = ClientDependencies(requirements=None, lock_file="locks/user-lock.json") + authored = authored_run.model_copy( + update={"client": authored_run.client.model_copy(update={"dependencies": dependencies})} + ) + plan_payload = multi_node_plan.model_dump(mode="json") + plan_payload["authored_config"]["sha256"] = authored.compute_sha256() + plan_payload["client"]["authored"] = authored.client.model_dump(mode="json") + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(plan_payload)) + + with pytest.raises(PlanContractError, match="authored lock file"): + validate_resolved_plan(authored, dependency_lock, plan) + + lock_payload = dependency_lock.model_dump(mode="json") + lock_payload.update( + authored_source="locks/user-lock.json", + source={ + "path": "/workspace/primary/runs/run-001/inputs/user-lock.json", + "sha256": "a" * 64, + }, + ) + matching_lock = ResolvedDependencyLock.model_validate_json(json.dumps(lock_payload)) + unexpected_source_plan = multi_node_plan.model_copy( + update={ + "client": multi_node_plan.client.model_copy( + update={ + "dependency_lock": multi_node_plan.client.dependency_lock.model_copy( + update={"sha256": matching_lock.compute_sha256()} + ) + } + ) + } + ) + with pytest.raises(PlanContractError, match="present for authored requirements"): + validate_resolved_plan(authored_run, matching_lock, unexpected_source_plan) + + matching_plan = plan.model_copy( + update={ + "client": plan.client.model_copy( + update={ + "dependency_lock": plan.client.dependency_lock.model_copy( + update={"sha256": matching_lock.compute_sha256()} + ) + } + ) + } + ) + + assert validate_resolved_plan(authored, matching_lock, matching_plan) is matching_plan + + +def test_cross_record_validation_rejects_python_abi( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invalid_lock = dependency_lock.model_copy(update={"python_abi": "cp311"}) + client = multi_node_plan.client.model_copy( + update={ + "dependency_lock": multi_node_plan.client.dependency_lock.model_copy( + update={"sha256": invalid_lock.compute_sha256()} + ) + } + ) + invalid_plan = multi_node_plan.model_copy(update={"client": client}) + + with pytest.raises(PlanContractError, match="Python ABI"): + validate_resolved_plan(authored_run, invalid_lock, invalid_plan) + + +def test_cross_record_validation_rejects_image_inventory( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invalid_lock = dependency_lock.model_copy(update={"image_distributions": ()}) + client = multi_node_plan.client.model_copy( + update={ + "dependency_lock": multi_node_plan.client.dependency_lock.model_copy( + update={"sha256": invalid_lock.compute_sha256()} + ) + } + ) + invalid_plan = multi_node_plan.model_copy(update={"client": client}) + + with pytest.raises(PlanContractError, match="image inventory"): + validate_resolved_plan(authored_run, invalid_lock, invalid_plan) + + +def test_cross_record_validation_rejects_changed_invocation( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invocation = multi_node_plan.invocation.model_copy( + update={"authored": authored_run.invocation.model_copy(update={"dataset_name": "other"})} + ) + invalid = multi_node_plan.model_copy(update={"invocation": invocation}) + + with pytest.raises(PlanContractError, match="invocation"): + validate_resolved_plan(authored_run, dependency_lock, invalid) + + +def test_cross_record_validation_preserves_explicit_run_config( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["effective_run_config"]["buffer_size"] = 16384 + invalid = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + with pytest.raises(PlanContractError, match="run_config.buffer_size"): + validate_resolved_plan(authored_run, dependency_lock, invalid) diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py new file mode 100644 index 000000000..1e4086f5e --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from copy import deepcopy + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.config import ( + ProfileSelectionSource, + SlurmProfile, + SlurmProfileCatalog, + injected_profile, + select_profile, + validate_selected_profile, +) + + +def test_profile_selection_precedence(profile_catalog: SlurmProfileCatalog) -> None: + explicit = select_profile(profile_catalog, cluster="lab", hostnames=("primary-login-1",)) + hostname = select_profile(profile_catalog, hostnames=("PRIMARY-LOGIN-1", "host.example")) + default = select_profile(profile_catalog, hostnames=("unmatched",)) + + assert (explicit.cluster_name, explicit.selection_source) == ("lab", ProfileSelectionSource.EXPLICIT) + assert (hostname.cluster_name, hostname.selection_source, hostname.matched_pattern) == ( + "primary", + ProfileSelectionSource.HOSTNAME, + "primary-login-*", + ) + assert (default.cluster_name, default.selection_source) == ("primary", ProfileSelectionSource.DEFAULT) + + +def test_injected_profile_has_no_catalog_provenance(profile_catalog: SlurmProfileCatalog) -> None: + selected = injected_profile(profile_catalog.clusters["lab"]) + + assert selected.selection_source is ProfileSelectionSource.INJECTED + assert selected.cluster_name is None + assert selected.catalog_sha256 is None + with pytest.raises(ValueError, match="no source catalog"): + validate_selected_profile(profile_catalog, selected) + + +def test_catalog_selection_digest_validation(profile_catalog: SlurmProfileCatalog) -> None: + selected = select_profile(profile_catalog, cluster="primary") + + assert validate_selected_profile(profile_catalog, selected) is selected + changed = profile_catalog.model_copy(update={"default_cluster": "lab"}) + with pytest.raises(ValueError, match="catalog digest"): + validate_selected_profile(changed, selected) + + +def test_catalog_selection_revalidates_provenance(profile_catalog: SlurmProfileCatalog) -> None: + forged_default = select_profile(profile_catalog, cluster="lab").model_copy( + update={"selection_source": ProfileSelectionSource.DEFAULT} + ) + with pytest.raises(ValueError, match="catalog default"): + validate_selected_profile(profile_catalog, forged_default) + + forged_pattern = select_profile(profile_catalog, hostnames=("primary-login-1",)).model_copy( + update={"matched_pattern": "lab-*"} + ) + with pytest.raises(ValueError, match="pattern"): + validate_selected_profile(profile_catalog, forged_pattern) + + +def test_unselected_profile_edit_keeps_selected_profile_digest(profile_catalog: SlurmProfileCatalog) -> None: + first = select_profile(profile_catalog, cluster="primary") + payload = profile_catalog.model_dump(mode="json") + payload["clusters"]["lab"]["workspace_root"] = "/workspace/other-lab" + changed = SlurmProfileCatalog.model_validate(payload) + second = select_profile(changed, cluster="primary") + + assert first.profile_sha256 == second.profile_sha256 + assert first.catalog_sha256 != second.catalog_sha256 + + +def test_hostname_selection_rejects_ambiguous_clusters(profile_catalog: SlurmProfileCatalog) -> None: + payload = profile_catalog.model_dump(mode="json") + payload["clusters"]["lab"]["host_patterns"] = ["primary-*"] + catalog = SlurmProfileCatalog.model_validate(payload) + + with pytest.raises(ValueError, match="multiple"): + select_profile(catalog, hostnames=("primary-login-1",)) + + +def test_explicit_selection_rejects_unknown_cluster(profile_catalog: SlurmProfileCatalog) -> None: + with pytest.raises(ValueError, match="unknown cluster"): + select_profile(profile_catalog, cluster="missing") + + +@pytest.mark.parametrize( + "mutator", + [ + lambda payload: payload.pop("schema_version"), + lambda payload: payload.update(default_cluster="missing"), + lambda payload: payload["clusters"]["lab"].update(host_patterns=["primary-login-*"]), + lambda payload: payload["clusters"]["primary"].update(extra="unknown"), + lambda payload: payload["clusters"]["primary"].update(workspace_root="relative"), + lambda payload: payload["clusters"]["primary"].update(host_patterns=["login[broken"]), + lambda payload: payload["clusters"]["primary"].update( + container_mounts=[ + {"source": "/one", "target": "/same"}, + {"source": "/two", "target": "/same"}, + ] + ), + ], +) +def test_profile_catalog_rejects_invalid_boundaries( + profile_catalog: SlurmProfileCatalog, + mutator: object, +) -> None: + payload = deepcopy(profile_catalog.model_dump(mode="json")) + mutator(payload) + + with pytest.raises(ValidationError): + SlurmProfileCatalog.model_validate(payload) + + +def test_profile_requires_explicit_version(profile_catalog: SlurmProfileCatalog) -> None: + payload = profile_catalog.clusters["primary"].model_dump(mode="json") + payload.pop("schema_version") + + with pytest.raises(ValidationError, match="schema_version"): + SlurmProfile.model_validate(payload) diff --git a/packages/data-designer-slurm/tests/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py new file mode 100644 index 000000000..5b511c4c4 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from copy import deepcopy +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.contracts import ( + ArtifactReference as CommonArtifactReference, +) +from data_designer.slurm.contracts import ( + RecordRange as CommonRecordRange, +) +from data_designer.slurm.contracts import ( + ResumeWorkspace as CommonResumeWorkspace, +) +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, +) + +GOLDEN_DIR = Path(__file__).parent / "golden" + + +def test_planning_reexports_common_contract_types() -> None: + assert PlanningArtifactReference is CommonArtifactReference + assert PlanningRecordRange is CommonRecordRange + assert PlanningResumeWorkspace is CommonResumeWorkspace + + +@pytest.mark.parametrize("actual_records", [0, 25]) +def test_client_result_allows_partial_and_failure_facts_to_differ(actual_records: int) -> None: + partial = { + "schema_version": 1, + "run_id": "run-001", + "shard_id": "shard-00000", + "attempt_id": "attempt-0001", + "completed_at": "2026-08-19T12:00:00Z", + "requested_records": 50, + "actual_records": actual_records, + "outcome": "partial", + "dataset_path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset", + "early_shutdown": True, + "requested_resume_mode": "if_possible", + "effective_resume_mode": "never", + "candidate_output_manifest": { + "path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "a" * 64, + }, + } + failed = { + "schema_version": 1, + "run_id": "run-001", + "shard_id": "shard-00000", + "attempt_id": "attempt-0002", + "completed_at": "2026-08-19T12:00:00Z", + "requested_records": 50, + "actual_records": None, + "outcome": "failed", + "early_shutdown": None, + "requested_resume_mode": "if_possible", + "effective_resume_mode": None, + "error_code": "generation_error", + "redacted_message": "generation failed", + } + + assert ClientResult.model_validate_json(json.dumps(partial)).actual_records == actual_records + assert ClientResult.model_validate_json(json.dumps(failed)).dataset_path is None + + +@pytest.mark.parametrize( + "mutation", + [ + {"actual_records": 51}, + {"outcome": "complete", "actual_records": 49}, + {"outcome": "failed", "candidate_output_manifest": {"path": "/x", "sha256": "a" * 64}}, + {"outcome": "failed", "candidate_output_manifest": None, "error_code": None}, + {"effective_resume_mode": "always"}, + {"early_shutdown": None}, + {"early_shutdown": True}, + {"dataset_path": "/workspace/other/dataset"}, + { + "candidate_output_manifest": { + "path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0002/output-manifest.json", + "sha256": "a" * 64, + } + }, + {"completed_at": "2026-08-19T12:00:00+01:00"}, + {"redacted_message": "bad\nmessage"}, + ], +) +def test_client_result_rejects_inconsistent_semantics( + mutation: dict[str, object], + client_result_payload: dict[str, object], +) -> None: + payload = deepcopy(client_result_payload) + payload.update(mutation) + + with pytest.raises(ValidationError): + ClientResult.model_validate_json(json.dumps(payload)) + + +@pytest.fixture +def client_result_payload() -> dict[str, object]: + return { + "schema_version": 1, + "run_id": "run-001", + "shard_id": "shard-00000", + "attempt_id": "attempt-0001", + "completed_at": "2026-08-19T12:00:00Z", + "requested_records": 50, + "actual_records": 50, + "outcome": "complete", + "dataset_path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset", + "early_shutdown": False, + "requested_resume_mode": "never", + "effective_resume_mode": "never", + "candidate_output_manifest": { + "path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "a" * 64, + }, + } + + +@pytest.mark.parametrize( + ("effective_resume_mode", "dataset_path"), + [ + ("never", "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset"), + ("always", "/workspace/runs/run-001/shards/shard-00000/dataset"), + ], +) +def test_if_possible_uses_effective_resume_dataset_location( + effective_resume_mode: str, + dataset_path: str, + client_result_payload: dict[str, object], +) -> None: + payload = deepcopy(client_result_payload) + payload.update( + requested_resume_mode="if_possible", + effective_resume_mode=effective_resume_mode, + dataset_path=dataset_path, + ) + + assert ClientResult.model_validate_json(json.dumps(payload)).dataset_path == dataset_path + + +def test_benchmark_manifest_rejects_duplicate_child_identity() -> None: + payload = { + "schema_version": 1, + "benchmark_id": "bench", + "benchmark_config": {"path": "/workspace/config.json", "sha256": "a" * 64}, + "children": [ + { + "case_id": "case", + "child_run_id": "run", + "child_authored_config": { + "path": "/workspace/runs/run/authored-config.json", + "sha256": "b" * 64, + }, + }, + { + "case_id": "case", + "child_run_id": "run-2", + "child_authored_config": { + "path": "/workspace/runs/run-2/authored-config.json", + "sha256": "c" * 64, + }, + }, + ], + } + + with pytest.raises(ValidationError, match="case IDs"): + BenchmarkManifest.model_validate_json(json.dumps(payload)) + + +def test_benchmark_report_rejects_unknown_or_incomplete_recommendations() -> None: + payload = { + "schema_version": 1, + "benchmark_id": "bench", + "analysis_id": "analysis", + "benchmark_manifest": {"path": "/workspace/benchmark.json", "sha256": "a" * 64}, + "created_at": "2026-08-19T12:00:00Z", + "cases": [ + { + "case_id": "case", + "child_run_id": "run", + "outcome": "pending", + "topology_digest": "b" * 64, + "requested_records": 100, + "gpus_per_job": 8, + "nodes_per_job": 1, + } + ], + "recommendations": [{"kind": "pareto", "case_id": "missing"}], + } + + with pytest.raises(ValidationError, match="unknown cases"): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + payload["recommendations"] = [{"kind": "pareto", "case_id": "case"}] + with pytest.raises(ValidationError, match="successful feasible"): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + +def test_successful_benchmark_case_requires_metrics() -> None: + payload = { + "schema_version": 1, + "benchmark_id": "bench", + "analysis_id": "analysis", + "benchmark_manifest": {"path": "/workspace/benchmark.json", "sha256": "a" * 64}, + "created_at": "2026-08-19T12:00:00Z", + "cases": [ + { + "case_id": "case", + "child_run_id": "run", + "outcome": "succeeded", + "topology_digest": "b" * 64, + "requested_records": 100, + "gpus_per_job": 8, + "nodes_per_job": 1, + } + ], + } + + with pytest.raises(ValidationError, match="complete"): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("actual_records", 999), + ("generation_seconds", 0), + ("wall_seconds", 0), + ("rows_per_second", 0), + ], +) +def test_successful_benchmark_case_requires_complete_positive_output(field: str, value: int) -> None: + payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) + payload["cases"][0][field] = value + + with pytest.raises(ValidationError): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + +def test_benchmark_report_allows_pareto_frontier_but_singleton_minima() -> None: + payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) + second_case = deepcopy(payload["cases"][0]) + second_case.update(case_id="second-case", child_run_id="second-run", topology_digest="f" * 64) + payload["cases"].append(second_case) + payload["recommendations"] = [ + {"kind": "pareto", "case_id": payload["cases"][0]["case_id"]}, + {"kind": "pareto", "case_id": second_case["case_id"]}, + ] + + assert len(BenchmarkReport.model_validate_json(json.dumps(payload)).recommendations) == 2 + + payload["recommendations"] = [ + {"kind": "minimum_jobs", "case_id": payload["cases"][0]["case_id"]}, + {"kind": "minimum_jobs", "case_id": second_case["case_id"]}, + ] + with pytest.raises(ValidationError, match="minimum"): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + +def test_benchmark_report_rejects_duplicate_child_runs() -> None: + payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) + payload["cases"][1]["child_run_id"] = payload["cases"][0]["child_run_id"] + + with pytest.raises(ValidationError, match="child run IDs"): + BenchmarkReport.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/state/test_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py similarity index 95% rename from packages/data-designer-slurm/tests/state/test_golden_records.py rename to packages/data-designer-slurm/tests/state/test_state_golden_records.py index fa31859db..e029b99b4 100644 --- a/packages/data-designer-slurm/tests/state/test_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -37,7 +37,7 @@ @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) -def test_golden_record_round_trip_is_deterministic(filename: str, model: type[StateRecord]) -> None: +def test_state_golden_record_round_trip_is_deterministic(filename: str, model: type[StateRecord]) -> None: serialized = (GOLDEN_DIRECTORY / filename).read_text() record = model.model_validate_json(serialized) diff --git a/packages/data-designer-slurm/tests/test_contracts.py b/packages/data-designer-slurm/tests/test_contracts.py index 3319d16e2..d1983b915 100644 --- a/packages/data-designer-slurm/tests/test_contracts.py +++ b/packages/data-designer-slurm/tests/test_contracts.py @@ -34,6 +34,9 @@ from data_designer.slurm.state import ( RecordRange as StateRecordRange, ) +from data_designer.slurm.state import ( + ResumeWorkspace as StateResumeWorkspace, +) from data_designer.slurm.state import ( StateRecord, StateValue, @@ -43,6 +46,7 @@ def test_state_exports_exact_shared_contract_types() -> None: assert StateArtifactReference is ContractArtifactReference assert StateRecordRange is ContractRecordRange + assert StateResumeWorkspace is ResumeWorkspace assert StateContractValue is ContractValue assert StateContractRecord is ContractRecord assert StateValue is ContractValue diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index d244200d8..bf2cffd28 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -127,12 +127,21 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non assert version("data-designer-slurm") == {version!r} from data_designer.slurm.contracts import ArtifactReference as ContractArtifactReference from data_designer.slurm.contracts import RecordRange as ContractRecordRange +from data_designer.slurm.contracts import ResumeWorkspace as ContractResumeWorkspace +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 from data_designer.slurm.state import ArtifactReference as StateArtifactReference from data_designer.slurm.state import RecordRange as StateRecordRange +from data_designer.slurm.state import ResumeWorkspace as StateResumeWorkspace from data_designer.slurm.state import RunManifest assert RunManifest.__name__ == "RunManifest" +assert PlanningArtifactReference is ContractArtifactReference +assert PlanningRecordRange is ContractRecordRange +assert PlanningResumeWorkspace is ContractResumeWorkspace assert StateArtifactReference is ContractArtifactReference assert StateRecordRange is ContractRecordRange +assert StateResumeWorkspace is ContractResumeWorkspace """ run([str(python), "-c", statement], cwd=cwd) @@ -178,9 +187,11 @@ def main() -> None: base_leaf_requirement = requirement(base_metadata, "data-designer-slurm") leaf_base_requirement = requirement(leaf_metadata, "data-designer") + leaf_packaging_requirement = requirement(leaf_metadata, "packaging") leaf_pydantic_requirement = requirement(leaf_metadata, "pydantic") assert str(base_leaf_requirement.specifier) == f"=={version}" assert str(leaf_base_requirement.specifier) == f"=={version}" + assert leaf_packaging_requirement.specifier == Requirement("packaging>=25,<27").specifier assert leaf_pydantic_requirement.specifier == Requirement("pydantic>=2.9.2,<3").specifier assert base_leaf_requirement.marker is not None assert base_leaf_requirement.marker.evaluate({"extra": "slurm"}) diff --git a/uv.lock b/uv.lock index 6e0db450d..7b6c3394b 100644 --- a/uv.lock +++ b/uv.lock @@ -973,12 +973,14 @@ name = "data-designer-slurm" source = { editable = "packages/data-designer-slurm" } dependencies = [ { name = "data-designer" }, + { name = "packaging" }, { name = "pydantic" }, ] [package.metadata] requires-dist = [ { name = "data-designer", editable = "packages/data-designer" }, + { name = "packaging", specifier = ">=25,<27" }, { name = "pydantic", specifier = ">=2.9.2,<3" }, ]