From b86293c5ecfe511c0e3f71870c9ca183064dd47e Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 19 Aug 2026 09:38:55 -0300 Subject: [PATCH 1/7] feat: define Slurm config and plan contracts Add strict authored configuration, profile, image, client, benchmark, dependency-lock, and resolved-plan records. Validate cross-record identities and digests with sanitized single-node and multi-node golden fixtures. Closes #873 Signed-off-by: Andre Manoel --- .../src/data_designer/slurm/_contracts.py | 128 ++++++ .../data_designer/slurm/benchmark/__init__.py | 26 ++ .../data_designer/slurm/benchmark/records.py | 135 +++++++ .../data_designer/slurm/client/__init__.py | 10 + .../src/data_designer/slurm/client/records.py | 82 ++++ .../data_designer/slurm/config/__init__.py | 106 +++++ .../data_designer/slurm/config/benchmark.py | 103 +++++ .../src/data_designer/slurm/config/images.py | 119 ++++++ .../data_designer/slurm/config/profiles.py | 267 +++++++++++++ .../src/data_designer/slurm/config/run.py | 368 ++++++++++++++++++ .../data_designer/slurm/planning/__init__.py | 43 ++ .../data_designer/slurm/planning/models.py | 325 ++++++++++++++++ .../slurm/planning/validation.py | 97 +++++ .../tests/contracts/conftest.py | 33 ++ .../tests/contracts/golden/authored_run.json | 91 +++++ .../contracts/golden/benchmark_config.json | 43 ++ .../contracts/golden/benchmark_manifest.json | 26 ++ .../contracts/golden/benchmark_report.json | 47 +++ .../golden/client_image_inspection.json | 23 ++ .../tests/contracts/golden/client_result.json | 18 + .../contracts/golden/dependency_lock.json | 29 ++ .../contracts/golden/multi_node_plan.json | 316 +++++++++++++++ .../contracts/golden/profile_catalog.json | 45 +++ .../golden/serving_image_inspection.json | 11 + .../contracts/golden/single_node_plan.json | 194 +++++++++ .../tests/contracts/test_config_records.py | 242 ++++++++++++ .../tests/contracts/test_golden_records.py | 66 ++++ .../tests/contracts/test_planning_records.py | 186 +++++++++ .../tests/contracts/test_profiles.py | 112 ++++++ .../tests/contracts/test_shared_records.py | 170 ++++++++ 30 files changed, 3461 insertions(+) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/_contracts.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/images.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/run.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/models.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py create mode 100644 packages/data-designer-slurm/tests/contracts/conftest.py create mode 100644 packages/data-designer-slurm/tests/contracts/golden/authored_run.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/client_result.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json create mode 100644 packages/data-designer-slurm/tests/contracts/test_config_records.py create mode 100644 packages/data-designer-slurm/tests/contracts/test_golden_records.py create mode 100644 packages/data-designer-slurm/tests/contracts/test_planning_records.py create mode 100644 packages/data-designer-slurm/tests/contracts/test_profiles.py create mode 100644 packages/data-designer-slurm/tests/contracts/test_shared_records.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py new file mode 100644 index 000000000..ee2987ed5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import posixpath +import re +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, StringConstraints + +Identifier = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", + ), +] +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)$")] + + +class AuthoredConfig(BaseModel): + """Base for strict authored configuration values.""" + + model_config = ConfigDict( + extra="forbid", + frozen=True, + protected_namespaces=(), + strict=True, + validate_default=True, + ) + + +class ContractValue(BaseModel): + """Base for strict immutable cross-process values.""" + + model_config = ConfigDict( + extra="forbid", + frozen=True, + protected_namespaces=(), + strict=True, + validate_default=True, + ) + + +class ContractRecord(ContractValue): + """Base for explicitly versioned records with stable serialization.""" + + schema_version: Literal[1] + + def serialize_canonical_json(self) -> bytes: + return canonical_json(self.model_dump(mode="json")) + + def serialize_json(self) -> str: + return ( + json.dumps( + self.model_dump(mode="json"), + allow_nan=False, + ensure_ascii=False, + indent=2, + sort_keys=True, + ) + + "\n" + ) + + def compute_sha256(self) -> Sha256Digest: + return hashlib.sha256(self.serialize_canonical_json()).hexdigest() + + +def canonical_json(value: object) -> bytes: + """Serialize a JSON-compatible value to stable UTF-8 bytes.""" + return json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + +def compute_sha256(value: object) -> Sha256Digest: + """Compute the canonical JSON digest of a JSON-compatible value.""" + return hashlib.sha256(canonical_json(value)).hexdigest() + + +def validate_absolute_path(value: str) -> str: + if not value.startswith("/"): + raise ValueError("path must be absolute") + if value == "/": + raise ValueError("path must not be the filesystem root") + 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: + raise ValueError("path must be normalized") + 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) + if not re.fullmatch(r"https?://[^\s]+", value): + raise ValueError(f"{field_name} must be an HTTP(S) URL") + return value 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..4357646cd --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -0,0 +1,135 @@ +# 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 ContractRecord, ContractValue, Identifier +from data_designer.slurm.planning import ArtifactReference + + +class BenchmarkChildRun(ContractValue): + case_id: Identifier + child_run_id: Identifier + child_config: ArtifactReference + + +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") + 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) + if len(case_ids) != len(set(case_ids)): + raise ValueError("benchmark report case 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))}") + kinds = tuple(recommendation.kind for recommendation in self.recommendations) + if len(kinds) != len(set(kinds)): + raise ValueError("benchmark recommendation kinds must be unique") + 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..0cdb95084 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -0,0 +1,82 @@ +# 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 ContractRecord, Identifier, validate_absolute_path +from data_designer.slurm.planning import ArtifactReference + + +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: Identifier + attempt_id: Identifier + 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.outcome is ClientOutcome.COMPLETE: + if self.actual_records != self.requested_records: + raise ValueError("complete client results require the requested record count") + self._require_success_artifacts() + elif self.outcome is ClientOutcome.PARTIAL: + if self.actual_records is None or not 0 < self.actual_records < self.requested_records: + raise ValueError("partial client results require a positive partial 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") 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..b2450053b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py @@ -0,0 +1,103 @@ +# 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._contracts import AuthoredConfig, Duration, Identifier, validate_local_config_path +from data_designer.slurm.config.run import DataDesignerSlurmConfig + + +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[Identifier, 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: Literal[1] + name: Identifier + base_run: BenchmarkBaseRun + model_aliases: Literal["all"] | list[Identifier] + 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..be14e42c2 --- /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..23c6799d3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -0,0 +1,267 @@ +# 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, + 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: Literal[1] + 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: Literal[1] + 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") + 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..08b7df880 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -0,0 +1,368 @@ +# 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 typing import Annotated, Literal + +from pydantic import ( + Field, + JsonValue, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) + +from data_designer.config import RunConfig +from data_designer.slurm._contracts import ( + AuthoredConfig, + Duration, + EnvironmentName, + Identifier, + validate_absolute_path, + validate_local_config_path, + validate_plain_text, + validate_url, +) +from data_designer.slurm.config.images import ImageRef + +_OWNED_VLLM_FLAGS = { + "--distributed-executor-backend", + "--enable-expert-parallel", + "--headless", + "--host", + "--pipeline-parallel-size", + "--port", + "--served-model-name", + "--tensor-parallel-size", +} + + +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 set(self.inline) != {"data_designer"} or not isinstance(self.inline["data_designer"], dict): + raise ValueError("inline builder input must be one complete serialized Data Designer config") + 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: + return validate_url(value, field_name="MCP endpoint") + + +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") + return values + + +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[Identifier, 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}") + if " @ " in value: + _, source = value.split(" @ ", maxsplit=1) + if not re.fullmatch(r"https://[^\s]+\.whl#sha256=[0-9a-f]{64}", source): + raise ValueError("direct dependency URLs must be HTTPS wheels with a SHA-256 fragment") + elif "://" in value or not re.match(r"^[A-Za-z0-9][A-Za-z0-9._-]*", value): + raise ValueError(f"invalid dependency requirement: {value!r}") + name = re.split(r"\s|\[|[<>=!~@]", value, maxsplit=1)[0].lower().replace("_", "-").replace(".", "-") + names.append(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 = value.split("=", maxsplit=1)[0] + if flag in _OWNED_VLLM_FLAGS: + raise ValueError(f"vLLM argument {flag!r} is owned by the compiler or runtime") + return values + + +class DeploymentResources(AuthoredConfig): + nodes: PositiveInt = 1 + + +class DeploymentTopology(AuthoredConfig): + tensor_parallel: PositiveInt = 1 + nodes_per_replica: PositiveInt = 1 + + +class ServerDeploymentConfig(AuthoredConfig): + model_alias: Identifier + served_model_name: Identifier | 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 + + @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: Literal[1] + 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/planning/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py new file mode 100644 index 000000000..05ae36210 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py @@ -0,0 +1,43 @@ +# 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, + ResolvedBuilderInput, + ResolvedClient, + ResolvedDependencyLock, + ResolvedDeployment, + ResolvedImage, + ResolvedInvocation, + ResolvedOutput, + ResolvedSlurmRunPlan, + ResolvedSubmission, + ResolvedTopology, +) +from data_designer.slurm.planning.validation import PlanContractError, validate_resolved_plan + +__all__ = [ + "ArtifactReference", + "LockedPackage", + "PlanContractError", + "PlannedShard", + "PortClaim", + "ResolvedBuilderInput", + "ResolvedClient", + "ResolvedDependencyLock", + "ResolvedDeployment", + "ResolvedImage", + "ResolvedInvocation", + "ResolvedOutput", + "ResolvedSlurmRunPlan", + "ResolvedSubmission", + "ResolvedTopology", + "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..ad6431bbc --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -0,0 +1,325 @@ +# 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 pydantic import Field, JsonValue, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.config import RunConfig +from data_designer.slurm._contracts import ( + ContractRecord, + ContractValue, + Identifier, + Sha256Digest, + compute_sha256, + validate_absolute_path, +) +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, + InvocationConfig, + ServerDeploymentConfig, +) + + +class ArtifactReference(ContractValue): + """Immutable reference to a persisted artifact and its digest.""" + + path: str + sha256: Sha256Digest + + _path_is_absolute = field_validator("path")(validate_absolute_path) + + +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, ...] + image_distributions: tuple[InstalledDistribution, ...] + overlay_packages: tuple[LockedPackage, ...] + + @model_validator(mode="after") + def validate_packages(self) -> ResolvedDependencyLock: + 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))}") + return self + + +class ResolvedBuilderInput(ContractValue): + authored_source: str | None = None + source: ArtifactReference | None = None + inline: dict[str, JsonValue] | None = None + content_sha256: Sha256Digest + + @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) + else: + if self.authored_source is None: + raise ValueError("resolved builder source requires authored_source") + expected_digest = self.source.sha256 + 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", mode="before") + @classmethod + def materialize_run_config(cls, value: object) -> dict[str, JsonValue]: + return RunConfig.model_validate(value).model_dump(mode="json") + + +class PortClaim(ContractValue): + name: Identifier + 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 + 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: + 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") + return self + + +class ResolvedClient(ContractValue): + authored: ClientConfig + image: ResolvedImage + dependency_lock: ArtifactReference + host_node_index: NonNegativeInt + gpu_count: Literal[0] + + @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") + return self + + +class PlannedShard(ContractValue): + shard_id: Identifier + shard_index: NonNegativeInt + array_task_index: NonNegativeInt + start_index: NonNegativeInt + end_index_exclusive: PositiveInt + requested_records: PositiveInt + resume_workspace: str + + _workspace_is_absolute = field_validator("resume_workspace")(validate_absolute_path) + + @model_validator(mode="after") + def validate_range(self) -> PlannedShard: + if self.end_index_exclusive <= self.start_index: + raise ValueError("shard end_index_exclusive must be greater than start_index") + if self.requested_records != self.end_index_exclusive - self.start_index: + raise ValueError("shard requested_records must match its record range") + return self + + +class ResolvedSubmission(ContractValue): + account: Identifier | None = None + partition: Identifier | None = None + job_name: Identifier + time_limit: str + comment: str | None = None + + +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.""" + + plan_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) + if len(deployment_ids) != len(set(deployment_ids)): + raise ValueError("resolved deployment IDs must be unique") + if len(aliases) != len(set(aliases)): + raise ValueError("resolved deployment aliases must be unique") + + 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") + + port_keys = tuple((port.node_index, port.port) for deployment in self.deployments for port in deployment.ports) + if len(port_keys) != len(set(port_keys)): + raise ValueError("plan port claims must be unique per node") + + self._validate_shards() + if not _is_below(self.output.root, profile.workspace_root): + raise ValueError("resolved output root must be below the selected workspace_root") + return self + + def _validate_shards(self) -> 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 + 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.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_start = shard.end_index_exclusive + if expected_start != requested_records: + raise ValueError("shard record ranges must cover the requested records") + + +def _is_below(path: str, root: str) -> bool: + return path != root and posixpath.commonpath((path, root)) == root 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..3f79d313d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm._contracts import compute_sha256 +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 + + +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, +) -> ResolvedSlurmRunPlan: + """Validate cross-record identities and digests for one resolved plan.""" + _require( + plan.authored_config.sha256 == compute_sha256(authored.model_dump(mode="json")), + "authored config digest does not match the resolved plan", + ) + _require(plan.invocation.authored == authored.invocation, "resolved invocation does not match authored input") + _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", + ) + + 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", + ) + if authored.client.dependencies.requirements is not None: + _require( + dependency_lock.authored_requirements == tuple(authored.client.dependencies.requirements), + "dependency lock requirements do not match authored requirements", + ) + return plan + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise PlanContractError(message) 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..813c24509 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/conftest.py @@ -0,0 +1,33 @@ +# 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 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 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..d8f0470d1 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -0,0 +1,91 @@ +{ + "schema_version": 1, + "name": "two-model-generation", + "builder": { + "inline": { + "data_designer": { + "columns": [], + "model_configs": [] + } + } + }, + "invocation": { + "num_records": 100, + "dataset_name": "training-data", + "resume": "if_possible", + "run_config": { + "buffer_size": 8192 + }, + "model_concurrency": { + "generator": 64, + "judge": 32 + } + }, + "client": { + "image": { + "name": "dd-client-0.9" + }, + "dependencies": { + "requirements": [ + "data-designer-speech==0.2.0" + ], + "index_credentials": { + "private-index": { + "type": "secret", + "environment": "PACKAGE_INDEX_TOKEN" + } + } + } + }, + "deployments": [ + { + "model_alias": "generator", + "model": "example/generator", + "server": { + "type": "vllm", + "image": { + "name": "vllm-0-21" + }, + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "extra_args": [ + "--max-model-len", + "32768" + ], + "environment": { + "HF_TOKEN": { + "type": "secret", + "environment": "HF_TOKEN" + } + } + }, + "resources": { + "nodes": 2 + }, + "topology": { + "tensor_parallel": 8, + "nodes_per_replica": 2 + } + }, + { + "model_alias": "judge", + "served_model_name": "judge-api", + "model": "/models/judge", + "server": { + "type": "vllm", + "image": { + "path": "/images/vllm-0-22.sqsh" + } + } + } + ], + "array_tasks": { + "count": 2, + "max_concurrent": 2 + }, + "submission": { + "job_name": "dd-two-model" + } +} 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..d1ce2232c --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json @@ -0,0 +1,43 @@ +{ + "schema_version": 1, + "name": "generator-scaling", + "base_run": "./run.yaml", + "model_aliases": [ + "generator" + ], + "concurrency_values": [ + 32, + 64, + 128 + ], + "deployment_cases": [ + { + "name": "two-independent-replicas", + "deployments": { + "generator": { + "nodes": 2, + "nodes_per_replica": 1 + } + } + }, + { + "name": "one-two-node-replica", + "deployments": { + "generator": { + "nodes": 2, + "nodes_per_replica": 2 + } + } + } + ], + "record_policy": { + "type": "adaptive", + "base_records": 1000, + "max_records": 5000, + "records_per_concurrency": 1.0 + }, + "analysis": { + "target_total_records": 1000000, + "target_runtime": "4h" + } +} 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..6154a68e6 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json @@ -0,0 +1,26 @@ +{ + "schema_version": 1, + "benchmark_id": "benchmark-001", + "benchmark_config": { + "path": "/workspace/primary/benchmarks/benchmark-001/config.json", + "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + }, + "children": [ + { + "case_id": "two-independent-replicas-c32", + "child_run_id": "run-benchmark-001-c32", + "child_config": { + "path": "/workspace/primary/runs/run-benchmark-001-c32/run.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + } + }, + { + "case_id": "one-two-node-replica-c32", + "child_run_id": "run-benchmark-002-c32", + "child_config": { + "path": "/workspace/primary/runs/run-benchmark-002-c32/run.json", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + } + } + ] +} 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..3e5cdbc3f --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json @@ -0,0 +1,47 @@ +{ + "schema_version": 1, + "benchmark_id": "benchmark-001", + "analysis_id": "analysis-001", + "benchmark_manifest": { + "path": "/workspace/primary/benchmarks/benchmark-001/benchmark.json", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "created_at": "2026-08-19T12:00:00Z", + "cases": [ + { + "case_id": "two-independent-replicas-c32", + "child_run_id": "run-benchmark-001-c32", + "outcome": "succeeded", + "topology_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + "requested_records": 1000, + "actual_records": 1000, + "boot_seconds": 60.0, + "generation_seconds": 120.0, + "wall_seconds": 180.0, + "rows_per_second": 8.333333333333334, + "request_count": 1000, + "token_count": 200000, + "gpus_per_job": 16, + "nodes_per_job": 2, + "gpu_hours_per_job": 0.8, + "total_gpu_hours": 800.0, + "target_jobs": 1000, + "feasible": true + }, + { + "case_id": "one-two-node-replica-c32", + "child_run_id": "run-benchmark-002-c32", + "outcome": "pending", + "topology_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "requested_records": 1000, + "gpus_per_job": 16, + "nodes_per_job": 2 + } + ], + "recommendations": [ + { + "kind": "minimum_gpu_hours", + "case_id": "two-independent-replicas-c32" + } + ] +} 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..066544f57 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json @@ -0,0 +1,23 @@ +{ + "schema_version": 1, + "inspector_version": "inspector-1", + "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "inspection": { + "kind": "client", + "python_implementation": "cpython", + "python_version": "3.12.12", + "python_abi": "cp312", + "distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "installer_path": "/usr/bin/pip", + "installer_version": "26.1" + } +} 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..75caaae9a --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/client_result.json @@ -0,0 +1,18 @@ +{ + "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/primary/runs/run-001/shards/shard-00000/dataset", + "early_shutdown": false, + "requested_resume_mode": "if_possible", + "effective_resume_mode": "never", + "candidate_output_manifest": { + "path": "/workspace/primary/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + } +} 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..4fd9e8e05 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -0,0 +1,29 @@ +{ + "schema_version": 1, + "resolver_version": "resolver-1", + "python_abi": "cp312", + "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "authored_requirements": [ + "data-designer-speech==0.2.0" + ], + "image_distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "overlay_packages": [ + { + "name": "data-designer-speech", + "version": "0.2.0", + "artifact": { + "path": "/workspace/primary/runs/run-001/dependencies/data_designer_speech-0.2.0.whl", + "sha256": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff" + } + } + ] +} 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..2571f88d5 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -0,0 +1,316 @@ +{ + "schema_version": 1, + "plan_id": "plan-001", + "package_version": "0.9.2", + "authored_config": { + "path": "/workspace/primary/runs/run-001/run.json", + "sha256": "a8d4be425bccdd56c4cef65a63b1bda372a28075f4e91e8c9cb724fbea5a3251" + }, + "selected_profile": { + "schema_version": 1, + "cluster_name": "primary", + "selection_source": "explicit", + "catalog_path": "/workspace/profile.json", + "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "profile": { + "schema_version": 1, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "scheduler": { + "account": "research", + "partition": "batch" + }, + "gpus_per_node": 8, + "workspace_root": "/workspace/primary", + "container_mounts": [ + { + "source": "/workspace", + "target": "/workspace", + "read_only": false + } + ], + "image_build": { + "partition": "cpu" + } + } + }, + "resolved_gpus_per_node": 8, + "builder": { + "inline": { + "data_designer": { + "columns": [], + "model_configs": [] + } + }, + "content_sha256": "26631b9b95ed99ac227eff0dfc64fb0c35c8cb9b9b7717d7dc2f2702055e0c4e" + }, + "invocation": { + "authored": { + "num_records": 100, + "dataset_name": "training-data", + "resume": "if_possible", + "run_config": { + "buffer_size": 8192 + }, + "model_concurrency": { + "generator": 64, + "judge": 32 + } + }, + "effective_run_config": { + "disable_early_shutdown": true, + "shutdown_error_rate": 1.0, + "shutdown_error_window": 10, + "buffer_size": 8192, + "max_concurrent_row_groups": 3, + "max_in_flight_tasks": 1024, + "non_inference_max_parallel_workers": 4, + "max_conversation_restarts": 0, + "max_conversation_correction_steps": 0, + "async_trace": false, + "write_scheduler_events": false, + "display_tui": false, + "progress_interval": 5.0, + "otel_metrics_port": null, + "preserve_dropped_columns": true, + "jinja_rendering_engine": "secure", + "request_admission": null + } + }, + "client": { + "authored": { + "image": { + "name": "dd-client-0.9" + }, + "dependencies": { + "requirements": [ + "data-designer-speech==0.2.0" + ], + "index_credentials": { + "private-index": { + "type": "secret", + "environment": "PACKAGE_INDEX_TOKEN" + } + } + } + }, + "image": { + "authored_ref": { + "name": "dd-client-0.9" + }, + "path": "/images/dd-client-0.9.sqsh", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "inspection": { + "schema_version": 1, + "inspector_version": "inspector-1", + "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "inspection": { + "kind": "client", + "python_implementation": "cpython", + "python_version": "3.12.12", + "python_abi": "cp312", + "distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "installer_path": "/usr/bin/pip", + "installer_version": "26.1" + } + } + }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-001/dependency-lock.json", + "sha256": "0db59e7c2db1f1ea4123299a40ca6e42b50d2a27192902721d30fd8bc32a1f8a" + }, + "host_node_index": 0, + "gpu_count": 0 + }, + "deployments": [ + { + "deployment_id": "generator", + "authored": { + "model_alias": "generator", + "model": "example/generator", + "server": { + "type": "vllm", + "image": { + "name": "vllm-0-21" + }, + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "extra_args": [ + "--max-model-len", + "32768" + ], + "environment": { + "HF_TOKEN": { + "type": "secret", + "environment": "HF_TOKEN" + } + } + }, + "resources": { + "nodes": 2 + }, + "topology": { + "tensor_parallel": 8, + "nodes_per_replica": 2 + } + }, + "image": { + "authored_ref": { + "name": "vllm-0-21" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", + "inspection": { + "schema_version": 1, + "inspector_version": "inspector-1", + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", + "inspection": { + "kind": "serving", + "server_type": "vllm", + "runtime_version": "0.21.0", + "executable_path": "/usr/local/bin/vllm" + } + } + }, + "node_indices": [ + 0, + 1 + ], + "gpus_per_node": 8, + "topology": { + "tensor_parallel": 8, + "nodes_per_replica": 2, + "pipeline_parallel": 2, + "node_group_count": 1, + "replicas_per_node_group": 1, + "replica_count": 1, + "gpus_per_replica": 16 + }, + "ports": [ + { + "name": "generator-http", + "node_index": 0, + "port": 18000 + }, + { + "name": "generator-rendezvous", + "node_index": 0, + "port": 19000 + } + ] + }, + { + "deployment_id": "judge", + "authored": { + "model_alias": "judge", + "served_model_name": "judge-api", + "model": "/models/judge", + "server": { + "type": "vllm", + "image": { + "path": "/images/vllm-0-22.sqsh" + } + } + }, + "image": { + "authored_ref": { + "path": "/images/vllm-0-22.sqsh" + }, + "path": "/images/vllm-0-22.sqsh", + "sha256": "9999999999999999999999999999999999999999999999999999999999999999", + "inspection": { + "schema_version": 1, + "inspector_version": "inspector-1", + "sqsh_sha256": "9999999999999999999999999999999999999999999999999999999999999999", + "inspection": { + "kind": "serving", + "server_type": "vllm", + "runtime_version": "0.22.0", + "executable_path": "/usr/local/bin/vllm" + } + } + }, + "node_indices": [ + 2 + ], + "gpus_per_node": 8, + "topology": { + "tensor_parallel": 1, + "nodes_per_replica": 1, + "pipeline_parallel": 1, + "node_group_count": 1, + "replicas_per_node_group": 8, + "replica_count": 8, + "gpus_per_replica": 1 + }, + "ports": [ + { + "name": "judge-http", + "node_index": 2, + "port": 18000 + } + ] + } + ], + "array_tasks": { + "count": 2, + "max_concurrent": 2 + }, + "shards": [ + { + "shard_id": "shard-00000", + "shard_index": 0, + "array_task_index": 0, + "start_index": 0, + "end_index_exclusive": 50, + "requested_records": 50, + "resume_workspace": "/workspace/primary/runs/run-001/shards/shard-00000/dataset" + }, + { + "shard_id": "shard-00001", + "shard_index": 1, + "array_task_index": 1, + "start_index": 50, + "end_index_exclusive": 100, + "requested_records": 50, + "resume_workspace": "/workspace/primary/runs/run-001/shards/shard-00001/dataset" + } + ], + "submission": { + "account": "research", + "partition": "batch", + "job_name": "dd-two-model", + "time_limit": "03:55:00" + }, + "output": { + "root": "/workspace/primary/runs/run-001/output", + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false + }, + "container_mounts": [ + { + "source": "/workspace", + "target": "/workspace", + "read_only": false + } + ], + "runtime_bundle": { + "path": "/workspace/primary/runtime/runtime.tar.gz", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + } +} 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..02771e474 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json @@ -0,0 +1,45 @@ +{ + "schema_version": 1, + "default_cluster": "primary", + "clusters": { + "primary": { + "schema_version": 1, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "scheduler": { + "account": "research", + "partition": "batch" + }, + "gpus_per_node": 8, + "workspace_root": "/workspace/primary", + "container_mounts": [ + { + "source": "/workspace", + "target": "/workspace", + "read_only": false + } + ], + "image_build": { + "partition": "cpu" + } + }, + "lab": { + "schema_version": 1, + "host_patterns": [ + "lab-login-*" + ], + "scheduler": { + "account": "lab", + "partition": "gpu" + }, + "gpus_per_node": "auto", + "workspace_root": "/workspace/lab", + "image_build": { + "partition": "cpu" + }, + "gpu_request_mode": "visible" + } + } +} 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..a2ec95109 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json @@ -0,0 +1,11 @@ +{ + "schema_version": 1, + "inspector_version": "inspector-1", + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", + "inspection": { + "kind": "serving", + "server_type": "vllm", + "runtime_version": "0.21.0", + "executable_path": "/usr/local/bin/vllm" + } +} 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..76b22eac0 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -0,0 +1,194 @@ +{ + "schema_version": 1, + "plan_id": "plan-single", + "package_version": "0.9.2", + "authored_config": { + "path": "/workspace/primary/runs/run-single/run.json", + "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + }, + "selected_profile": { + "schema_version": 1, + "cluster_name": "primary", + "selection_source": "explicit", + "catalog_path": "/workspace/profile.json", + "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "profile": { + "schema_version": 1, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "scheduler": { + "account": "research", + "partition": "batch" + }, + "gpus_per_node": 8, + "workspace_root": "/workspace/primary", + "container_mounts": [ + { + "source": "/workspace", + "target": "/workspace", + "read_only": false + } + ], + "image_build": { + "partition": "cpu" + } + } + }, + "resolved_gpus_per_node": 8, + "builder": { + "inline": { + "data_designer": { + "columns": [], + "model_configs": [] + } + }, + "content_sha256": "26631b9b95ed99ac227eff0dfc64fb0c35c8cb9b9b7717d7dc2f2702055e0c4e" + }, + "invocation": { + "authored": { + "num_records": 8, + "dataset_name": "single-node", + "model_concurrency": { + "generator": 8 + } + }, + "effective_run_config": {} + }, + "client": { + "authored": { + "image": { + "name": "dd-client-0.9" + } + }, + "image": { + "authored_ref": { + "name": "dd-client-0.9" + }, + "path": "/images/dd-client-0.9.sqsh", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "inspection": { + "schema_version": 1, + "inspector_version": "inspector-1", + "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "inspection": { + "kind": "client", + "python_implementation": "cpython", + "python_version": "3.12.12", + "python_abi": "cp312", + "distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + } + ], + "installer_path": "/usr/bin/pip", + "installer_version": "26.1" + } + } + }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-single/dependency-lock.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "host_node_index": 0, + "gpu_count": 0 + }, + "deployments": [ + { + "deployment_id": "generator", + "authored": { + "model_alias": "generator", + "model": "example/generator", + "server": { + "type": "vllm", + "image": { + "name": "vllm-0-21" + } + }, + "topology": { + "tensor_parallel": 8, + "nodes_per_replica": 1 + } + }, + "image": { + "authored_ref": { + "name": "vllm-0-21" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", + "inspection": { + "schema_version": 1, + "inspector_version": "inspector-1", + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", + "inspection": { + "kind": "serving", + "server_type": "vllm", + "runtime_version": "0.21.0", + "executable_path": "/usr/local/bin/vllm" + } + } + }, + "node_indices": [ + 0 + ], + "gpus_per_node": 8, + "topology": { + "tensor_parallel": 8, + "nodes_per_replica": 1, + "pipeline_parallel": 1, + "node_group_count": 1, + "replicas_per_node_group": 1, + "replica_count": 1, + "gpus_per_replica": 8 + }, + "ports": [ + { + "name": "generator-http", + "node_index": 0, + "port": 18000 + } + ] + } + ], + "array_tasks": { + "count": 1, + "max_concurrent": 1 + }, + "shards": [ + { + "shard_id": "shard-00000", + "shard_index": 0, + "array_task_index": 0, + "start_index": 0, + "end_index_exclusive": 8, + "requested_records": 8, + "resume_workspace": "/workspace/primary/runs/run-single/shards/shard-00000/dataset" + } + ], + "submission": { + "account": "research", + "partition": "batch", + "job_name": "data-designer", + "time_limit": "03:55:00" + }, + "output": { + "root": "/workspace/primary/runs/run-single/output", + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false + }, + "container_mounts": [ + { + "source": "/workspace", + "target": "/workspace", + "read_only": false + } + ], + "runtime_bundle": { + "path": "/workspace/primary/runtime/runtime.tar.gz", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + } +} 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..4d0eeb560 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -0,0 +1,242 @@ +# 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 ( + ArrayTasksConfig, + BenchmarkBaseRun, + ClientDependencies, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageBuildRequest, + ImageInspectionRecord, + ImageRef, + LiteralEnvironmentBinding, + QueueBackpressureConfig, + 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": ["my_pkg==1", "my-pkg==2"]}, + {"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_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", ["--port", "--host=0.0.0.0", "--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]) + + +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_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) + + +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 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..9c64f1cdc --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_golden_records.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from data_designer.slurm._contracts import ContractRecord +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.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan + +GOLDEN_DIR = Path(__file__).parent / "golden" + + +@pytest.mark.parametrize( + ("record_type", "filename"), + [ + (DataDesignerSlurmConfig, "authored_run.json"), + (SlurmProfileCatalog, "profile_catalog.json"), + (ImageInspectionRecord, "client_image_inspection.json"), + (ImageInspectionRecord, "serving_image_inspection.json"), + (ResolvedDependencyLock, "dependency_lock.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: + record = record_type.model_validate_json((GOLDEN_DIR / filename).read_text()) + + assert record_type.model_validate_json(record.model_dump_json()) == record + if isinstance(record, ContractRecord): + assert record_type.model_validate_json(record.serialize_json()) == record + assert record.compute_sha256() == record.compute_sha256() + + +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) + + first = json.dumps(authored_run.model_dump(mode="json"), sort_keys=True, separators=(",", ":")) + second = json.dumps(reordered.model_dump(mode="json"), sort_keys=True, separators=(",", ":")) + assert first == second 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..0d2d1e08b --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -0,0 +1,186 @@ +# 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._contracts import compute_sha256 +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.planning import ( + ArtifactReference, + PlanContractError, + ResolvedDependencyLock, + ResolvedSlurmRunPlan, + 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 == compute_sha256(authored_run.model_dump(mode="json")) + assert [deployment.topology.replica_count for deployment in multi_node_plan.deployments] == [1, 8] + assert multi_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() + + +@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]["ports"][1].update(port=18000), + lambda payload: payload.update(shards=payload["shards"][:1]), + lambda payload: payload["shards"][1].update(start_index=49), + 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_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"), + ], +) +def test_dependency_lock_rejects_overlap_order_and_non_wheel( + 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_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_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) 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..1a2a92366 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -0,0 +1,112 @@ +# 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_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..aa5e73fac --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -0,0 +1,170 @@ +# 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.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult + + +def test_client_result_allows_partial_and_failure_facts_to_differ() -> 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": 25, + "outcome": "partial", + "dataset_path": "/workspace/dataset", + "early_shutdown": True, + "requested_resume_mode": "if_possible", + "effective_resume_mode": "never", + "candidate_output_manifest": {"path": "/workspace/output.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 == 25 + 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": "partial", "actual_records": 0}, + {"outcome": "failed", "candidate_output_manifest": {"path": "/x", "sha256": "a" * 64}}, + {"outcome": "failed", "candidate_output_manifest": None, "error_code": None}, + {"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/dataset", + "early_shutdown": False, + "requested_resume_mode": "never", + "effective_resume_mode": "never", + "candidate_output_manifest": {"path": "/workspace/output.json", "sha256": "a" * 64}, + } + + +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_config": {"path": "/workspace/run-1.json", "sha256": "b" * 64}, + }, + { + "case_id": "case", + "child_run_id": "run-2", + "child_config": {"path": "/workspace/run-2.json", "sha256": "c" * 64}, + }, + ], + } + + with pytest.raises(ValidationError, match="case IDs"): + BenchmarkManifest.model_validate_json(json.dumps(payload)) + + +def test_benchmark_report_rejects_unknown_and_duplicate_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"}, + {"kind": "pareto", "case_id": "case"}, + ] + with pytest.raises(ValidationError, match="kinds"): + 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)) From 5ce7b5f75bf86bd51459ffb4b105820f7b68f3de Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 19 Aug 2026 10:27:50 -0300 Subject: [PATCH 2/7] fix: tighten Slurm contract validation Signed-off-by: Andre Manoel --- .../src/data_designer/slurm/config/run.py | 9 +- .../data_designer/slurm/planning/models.py | 30 +- .../tests/contracts/conftest.py | 15 + .../contracts/golden/authored_run_single.json | 40 ++ .../contracts/golden/benchmark_manifest.json | 14 +- .../contracts/golden/benchmark_report.json | 51 +- .../golden/client_image_inspection.json | 18 +- .../tests/contracts/golden/client_result.json | 24 +- .../contracts/golden/dependency_lock.json | 16 +- .../golden/dependency_lock_single.json | 14 + .../contracts/golden/multi_node_plan.json | 477 +++++++++++------- .../golden/serving_image_inspection.json | 12 +- .../contracts/golden/single_node_plan.json | 303 ++++++----- .../tests/contracts/test_config_records.py | 22 + .../tests/contracts/test_golden_records.py | 6 +- .../tests/contracts/test_planning_records.py | 35 ++ 16 files changed, 704 insertions(+), 382 deletions(-) create mode 100644 packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json 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 index 08b7df880..ddde96817 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -82,7 +82,14 @@ def validate_input(self) -> BuilderInput: 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 set(self.inline) != {"data_designer"} or not isinstance(self.inline["data_designer"], dict): + 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") return self 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 index ad6431bbc..cd963e71b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -30,6 +30,7 @@ ClientConfig, InvocationConfig, ServerDeploymentConfig, + SubmissionConfig, ) @@ -134,14 +135,18 @@ class ResolvedInvocation(ContractValue): authored: InvocationConfig effective_run_config: dict[str, JsonValue] - @field_validator("effective_run_config", mode="before") + @field_validator("effective_run_config") @classmethod - def materialize_run_config(cls, value: object) -> dict[str, JsonValue]: - return RunConfig.model_validate(value).model_dump(mode="json") + 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"] node_index: NonNegativeInt port: Annotated[int, Field(ge=1024, le=65535)] @@ -191,6 +196,20 @@ def validate_deployment(self) -> ResolvedDeployment: 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") + + 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_nodes = tuple(port.node_index for port in self.ports if port.role == "http") + if http_nodes != expected_http_nodes: + raise ValueError("deployment requires one ordered HTTP port claim per replica lane") + + expected_rendezvous_nodes = group_heads if self.topology.nodes_per_replica > 1 else () + rendezvous_nodes = tuple(port.node_index for port in self.ports if port.role == "rendezvous") + if rendezvous_nodes != expected_rendezvous_nodes: + raise ValueError("deployment requires one ordered rendezvous port claim per multi-node group") return self @@ -237,6 +256,11 @@ class ResolvedSubmission(ContractValue): 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 diff --git a/packages/data-designer-slurm/tests/contracts/conftest.py b/packages/data-designer-slurm/tests/contracts/conftest.py index 813c24509..44a74d99c 100644 --- a/packages/data-designer-slurm/tests/contracts/conftest.py +++ b/packages/data-designer-slurm/tests/contracts/conftest.py @@ -18,6 +18,11 @@ 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()) @@ -28,6 +33,16 @@ 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_single.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json new file mode 100644 index 000000000..f35b60d8b --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -0,0 +1,40 @@ +{ + "schema_version": 1, + "name": "single-node-generation", + "builder": { + "inline": { + "data_designer": { + "columns": [], + "model_configs": [] + } + } + }, + "invocation": { + "num_records": 8, + "dataset_name": "single-node", + "model_concurrency": { + "generator": 8 + } + }, + "client": { + "image": { + "name": "dd-client-0.9" + } + }, + "deployments": [ + { + "model_alias": "generator", + "model": "example/generator", + "server": { + "type": "vllm", + "image": { + "name": "vllm-0-21" + } + }, + "topology": { + "tensor_parallel": 8, + "nodes_per_replica": 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 index 6154a68e6..4623c2a49 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json @@ -1,26 +1,26 @@ { - "schema_version": 1, - "benchmark_id": "benchmark-001", "benchmark_config": { "path": "/workspace/primary/benchmarks/benchmark-001/config.json", "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" }, + "benchmark_id": "benchmark-001", "children": [ { "case_id": "two-independent-replicas-c32", - "child_run_id": "run-benchmark-001-c32", "child_config": { "path": "/workspace/primary/runs/run-benchmark-001-c32/run.json", "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" - } + }, + "child_run_id": "run-benchmark-001-c32" }, { "case_id": "one-two-node-replica-c32", - "child_run_id": "run-benchmark-002-c32", "child_config": { "path": "/workspace/primary/runs/run-benchmark-002-c32/run.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 index 3e5cdbc3f..99fbf17ac 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json @@ -1,47 +1,58 @@ { - "schema_version": 1, - "benchmark_id": "benchmark-001", "analysis_id": "analysis-001", + "benchmark_id": "benchmark-001", "benchmark_manifest": { "path": "/workspace/primary/benchmarks/benchmark-001/benchmark.json", "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" }, - "created_at": "2026-08-19T12:00:00Z", "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", - "topology_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + "request_count": 1000, "requested_records": 1000, - "actual_records": 1000, - "boot_seconds": 60.0, - "generation_seconds": 120.0, - "wall_seconds": 180.0, "rows_per_second": 8.333333333333334, - "request_count": 1000, + "target_jobs": 1000, "token_count": 200000, - "gpus_per_job": 16, - "nodes_per_job": 2, - "gpu_hours_per_job": 0.8, + "topology_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", "total_gpu_hours": 800.0, - "target_jobs": 1000, - "feasible": true + "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", - "topology_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "request_count": null, "requested_records": 1000, - "gpus_per_job": 16, - "nodes_per_job": 2 + "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": [ { - "kind": "minimum_gpu_hours", - "case_id": "two-independent-replicas-c32" + "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 index 066544f57..6466182d7 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json +++ b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json @@ -1,12 +1,5 @@ { - "schema_version": 1, - "inspector_version": "inspector-1", - "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "inspection": { - "kind": "client", - "python_implementation": "cpython", - "python_version": "3.12.12", - "python_abi": "cp312", "distributions": [ { "name": "data-designer", @@ -18,6 +11,13 @@ } ], "installer_path": "/usr/bin/pip", - "installer_version": "26.1" - } + "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 index 75caaae9a..d4bdc0834 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/client_result.json +++ b/packages/data-designer-slurm/tests/contracts/golden/client_result.json @@ -1,18 +1,20 @@ { - "schema_version": 1, - "run_id": "run-001", - "shard_id": "shard-00000", + "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", - "requested_records": 50, - "actual_records": 50, - "outcome": "complete", "dataset_path": "/workspace/primary/runs/run-001/shards/shard-00000/dataset", "early_shutdown": false, - "requested_resume_mode": "if_possible", "effective_resume_mode": "never", - "candidate_output_manifest": { - "path": "/workspace/primary/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" - } + "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 index 4fd9e8e05..502edcf3d 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -1,11 +1,8 @@ { - "schema_version": 1, - "resolver_version": "resolver-1", - "python_abi": "cp312", - "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "authored_requirements": [ "data-designer-speech==0.2.0" ], + "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "image_distributions": [ { "name": "data-designer", @@ -18,12 +15,15 @@ ], "overlay_packages": [ { - "name": "data-designer-speech", - "version": "0.2.0", "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 } 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..8f7b516c4 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json @@ -0,0 +1,14 @@ +{ + "authored_requirements": [], + "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 +} 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 index 2571f88d5..f82bd5ec3 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -1,117 +1,56 @@ { - "schema_version": 1, - "plan_id": "plan-001", - "package_version": "0.9.2", + "array_tasks": { + "count": 2, + "max_concurrent": 2 + }, "authored_config": { "path": "/workspace/primary/runs/run-001/run.json", "sha256": "a8d4be425bccdd56c4cef65a63b1bda372a28075f4e91e8c9cb724fbea5a3251" }, - "selected_profile": { - "schema_version": 1, - "cluster_name": "primary", - "selection_source": "explicit", - "catalog_path": "/workspace/profile.json", - "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", - "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", - "profile": { - "schema_version": 1, - "host_patterns": [ - "login*.primary.example", - "primary-login-*" - ], - "scheduler": { - "account": "research", - "partition": "batch" - }, - "gpus_per_node": 8, - "workspace_root": "/workspace/primary", - "container_mounts": [ - { - "source": "/workspace", - "target": "/workspace", - "read_only": false - } - ], - "image_build": { - "partition": "cpu" - } - } - }, - "resolved_gpus_per_node": 8, "builder": { + "authored_source": null, + "content_sha256": "26631b9b95ed99ac227eff0dfc64fb0c35c8cb9b9b7717d7dc2f2702055e0c4e", "inline": { "data_designer": { "columns": [], "model_configs": [] } }, - "content_sha256": "26631b9b95ed99ac227eff0dfc64fb0c35c8cb9b9b7717d7dc2f2702055e0c4e" - }, - "invocation": { - "authored": { - "num_records": 100, - "dataset_name": "training-data", - "resume": "if_possible", - "run_config": { - "buffer_size": 8192 - }, - "model_concurrency": { - "generator": 64, - "judge": 32 - } - }, - "effective_run_config": { - "disable_early_shutdown": true, - "shutdown_error_rate": 1.0, - "shutdown_error_window": 10, - "buffer_size": 8192, - "max_concurrent_row_groups": 3, - "max_in_flight_tasks": 1024, - "non_inference_max_parallel_workers": 4, - "max_conversation_restarts": 0, - "max_conversation_correction_steps": 0, - "async_trace": false, - "write_scheduler_events": false, - "display_tui": false, - "progress_interval": 5.0, - "otel_metrics_port": null, - "preserve_dropped_columns": true, - "jinja_rendering_engine": "secure", - "request_admission": null - } + "source": null }, "client": { "authored": { - "image": { - "name": "dd-client-0.9" - }, + "cpus": 32, "dependencies": { - "requirements": [ - "data-designer-speech==0.2.0" - ], "index_credentials": { "private-index": { - "type": "secret", - "environment": "PACKAGE_INDEX_TOKEN" + "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": "0db59e7c2db1f1ea4123299a40ca6e42b50d2a27192902721d30fd8bc32a1f8a" + }, + "gpu_count": 0, + "host_node_index": 0, "image": { "authored_ref": { - "name": "dd-client-0.9" + "name": "dd-client-0.9", + "path": null }, - "path": "/images/dd-client-0.9.sqsh", - "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "inspection": { - "schema_version": 1, - "inspector_version": "inspector-1", - "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "inspection": { - "kind": "client", - "python_implementation": "cpython", - "python_version": "3.12.12", - "python_abi": "cp312", "distributions": [ { "name": "data-designer", @@ -123,194 +62,344 @@ } ], "installer_path": "/usr/bin/pip", - "installer_version": "26.1" - } - } - }, - "dependency_lock": { - "path": "/workspace/primary/runs/run-001/dependency-lock.json", - "sha256": "0db59e7c2db1f1ea4123299a40ca6e42b50d2a27192902721d30fd8bc32a1f8a" - }, - "host_node_index": 0, - "gpu_count": 0 + "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" + } }, + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], "deployments": [ { - "deployment_id": "generator", "authored": { - "model_alias": "generator", "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 2 + }, + "served_model_name": null, "server": { - "type": "vllm", + "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" + "name": "vllm-0-21", + "path": null }, "queue_backpressure": { "max_waiting_requests": 1024, "retry_after_seconds": 2 }, - "extra_args": [ - "--max-model-len", - "32768" - ], - "environment": { - "HF_TOKEN": { - "type": "secret", - "environment": "HF_TOKEN" - } - } - }, - "resources": { - "nodes": 2 + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" }, "topology": { - "tensor_parallel": 8, - "nodes_per_replica": 2 + "nodes_per_replica": 2, + "tensor_parallel": 8 } }, + "deployment_id": "generator", + "gpus_per_node": 8, "image": { "authored_ref": { - "name": "vllm-0-21" + "name": "vllm-0-21", + "path": null }, - "path": "/images/vllm-0-21.sqsh", - "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", "inspection": { - "schema_version": 1, - "inspector_version": "inspector-1", - "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", "inspection": { + "executable_path": "/usr/local/bin/vllm", "kind": "serving", - "server_type": "vllm", "runtime_version": "0.21.0", - "executable_path": "/usr/local/bin/vllm" - } - } + "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 ], - "gpus_per_node": 8, - "topology": { - "tensor_parallel": 8, - "nodes_per_replica": 2, - "pipeline_parallel": 2, - "node_group_count": 1, - "replicas_per_node_group": 1, - "replica_count": 1, - "gpus_per_replica": 16 - }, "ports": [ { "name": "generator-http", "node_index": 0, - "port": 18000 + "port": 18000, + "role": "http" }, { "name": "generator-rendezvous", "node_index": 0, - "port": 19000 + "port": 19000, + "role": "rendezvous" } - ] + ], + "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 + } }, { - "deployment_id": "judge", "authored": { + "model": "/models/judge", "model_alias": "judge", + "resources": { + "nodes": 1 + }, "served_model_name": "judge-api", - "model": "/models/judge", "server": { - "type": "vllm", + "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": "judge", + "gpus_per_node": 8, "image": { "authored_ref": { + "name": null, "path": "/images/vllm-0-22.sqsh" }, - "path": "/images/vllm-0-22.sqsh", - "sha256": "9999999999999999999999999999999999999999999999999999999999999999", "inspection": { - "schema_version": 1, - "inspector_version": "inspector-1", - "sqsh_sha256": "9999999999999999999999999999999999999999999999999999999999999999", "inspection": { + "executable_path": "/usr/local/bin/vllm", "kind": "serving", - "server_type": "vllm", "runtime_version": "0.22.0", - "executable_path": "/usr/local/bin/vllm" - } - } + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "9999999999999999999999999999999999999999999999999999999999999999" + }, + "path": "/images/vllm-0-22.sqsh", + "sha256": "9999999999999999999999999999999999999999999999999999999999999999" }, "node_indices": [ 2 ], - "gpus_per_node": 8, + "ports": [ + { + "name": "judge-http-0", + "node_index": 2, + "port": 18000, + "role": "http" + }, + { + "name": "judge-http-1", + "node_index": 2, + "port": 18001, + "role": "http" + }, + { + "name": "judge-http-2", + "node_index": 2, + "port": 18002, + "role": "http" + }, + { + "name": "judge-http-3", + "node_index": 2, + "port": 18003, + "role": "http" + }, + { + "name": "judge-http-4", + "node_index": 2, + "port": 18004, + "role": "http" + }, + { + "name": "judge-http-5", + "node_index": 2, + "port": 18005, + "role": "http" + }, + { + "name": "judge-http-6", + "node_index": 2, + "port": 18006, + "role": "http" + }, + { + "name": "judge-http-7", + "node_index": 2, + "port": 18007, + "role": "http" + } + ], "topology": { - "tensor_parallel": 1, + "gpus_per_replica": 1, + "node_group_count": 1, "nodes_per_replica": 1, "pipeline_parallel": 1, - "node_group_count": 1, - "replicas_per_node_group": 8, "replica_count": 8, - "gpus_per_replica": 1 + "replicas_per_node_group": 8, + "tensor_parallel": 1 + } + } + ], + "invocation": { + "authored": { + "dataset_name": "training-data", + "diagnostics": { + "log_requests": false }, - "ports": [ + "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", + "plan_id": "plan-001", + "resolved_gpus_per_node": 8, + "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": [ { - "name": "judge-http", - "node_index": 2, - "port": 18000 + "read_only": false, + "source": "/workspace", + "target": "/workspace" } - ] - } - ], - "array_tasks": { - "count": 2, - "max_concurrent": 2 + ], + "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": [ { - "shard_id": "shard-00000", - "shard_index": 0, "array_task_index": 0, - "start_index": 0, "end_index_exclusive": 50, "requested_records": 50, - "resume_workspace": "/workspace/primary/runs/run-001/shards/shard-00000/dataset" + "resume_workspace": "/workspace/primary/runs/run-001/shards/shard-00000/dataset", + "shard_id": "shard-00000", + "shard_index": 0, + "start_index": 0 }, { - "shard_id": "shard-00001", - "shard_index": 1, "array_task_index": 1, - "start_index": 50, "end_index_exclusive": 100, "requested_records": 50, - "resume_workspace": "/workspace/primary/runs/run-001/shards/shard-00001/dataset" + "resume_workspace": "/workspace/primary/runs/run-001/shards/shard-00001/dataset", + "shard_id": "shard-00001", + "shard_index": 1, + "start_index": 50 } ], "submission": { "account": "research", - "partition": "batch", + "comment": null, "job_name": "dd-two-model", + "partition": "batch", "time_limit": "03:55:00" - }, - "output": { - "root": "/workspace/primary/runs/run-001/output", - "format": "parquet", - "partitions": 1, - "require_exact_record_count": false - }, - "container_mounts": [ - { - "source": "/workspace", - "target": "/workspace", - "read_only": false - } - ], - "runtime_bundle": { - "path": "/workspace/primary/runtime/runtime.tar.gz", - "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" } } 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 index a2ec95109..8cb5348e2 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json +++ b/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json @@ -1,11 +1,11 @@ { - "schema_version": 1, - "inspector_version": "inspector-1", - "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", "inspection": { + "executable_path": "/usr/local/bin/vllm", "kind": "serving", - "server_type": "vllm", "runtime_version": "0.21.0", - "executable_path": "/usr/local/bin/vllm" - } + "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 index 76b22eac0..feec12025 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -1,83 +1,49 @@ { - "schema_version": 1, - "plan_id": "plan-single", - "package_version": "0.9.2", + "array_tasks": { + "count": 1, + "max_concurrent": 1 + }, "authored_config": { "path": "/workspace/primary/runs/run-single/run.json", - "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + "sha256": "1ca141405769834f7b1e53461861d9ae349417008306e5fea80402d2a368bfa8" }, - "selected_profile": { - "schema_version": 1, - "cluster_name": "primary", - "selection_source": "explicit", - "catalog_path": "/workspace/profile.json", - "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", - "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", - "profile": { - "schema_version": 1, - "host_patterns": [ - "login*.primary.example", - "primary-login-*" - ], - "scheduler": { - "account": "research", - "partition": "batch" - }, - "gpus_per_node": 8, - "workspace_root": "/workspace/primary", - "container_mounts": [ - { - "source": "/workspace", - "target": "/workspace", - "read_only": false - } - ], - "image_build": { - "partition": "cpu" - } - } - }, - "resolved_gpus_per_node": 8, "builder": { + "authored_source": null, + "content_sha256": "26631b9b95ed99ac227eff0dfc64fb0c35c8cb9b9b7717d7dc2f2702055e0c4e", "inline": { "data_designer": { "columns": [], "model_configs": [] } }, - "content_sha256": "26631b9b95ed99ac227eff0dfc64fb0c35c8cb9b9b7717d7dc2f2702055e0c4e" - }, - "invocation": { - "authored": { - "num_records": 8, - "dataset_name": "single-node", - "model_concurrency": { - "generator": 8 - } - }, - "effective_run_config": {} + "source": null }, "client": { "authored": { + "cpus": 32, + "dependencies": { + "index_credentials": {}, + "lock_file": null, + "requirements": [] + }, "image": { - "name": "dd-client-0.9" + "name": "dd-client-0.9", + "path": null } }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-single/dependency-lock.json", + "sha256": "9b4e1db6dbcc62df9ee90e329578cfeb727e4060d7423c79c52558a69a42c156" + }, + "gpu_count": 0, + "host_node_index": 0, "image": { "authored_ref": { - "name": "dd-client-0.9" + "name": "dd-client-0.9", + "path": null }, - "path": "/images/dd-client-0.9.sqsh", - "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "inspection": { - "schema_version": 1, - "inspector_version": "inspector-1", - "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "inspection": { - "kind": "client", - "python_implementation": "cpython", - "python_version": "3.12.12", - "python_abi": "cp312", "distributions": [ { "name": "data-designer", @@ -85,110 +51,203 @@ } ], "installer_path": "/usr/bin/pip", - "installer_version": "26.1" - } - } - }, - "dependency_lock": { - "path": "/workspace/primary/runs/run-single/dependency-lock.json", - "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" - }, - "host_node_index": 0, - "gpu_count": 0 + "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" + } }, + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], "deployments": [ { - "deployment_id": "generator", "authored": { - "model_alias": "generator", "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 1 + }, + "served_model_name": null, "server": { - "type": "vllm", + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], "image": { - "name": "vllm-0-21" - } + "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": { - "tensor_parallel": 8, - "nodes_per_replica": 1 + "nodes_per_replica": 1, + "tensor_parallel": 8 } }, + "deployment_id": "generator", + "gpus_per_node": 8, "image": { "authored_ref": { - "name": "vllm-0-21" + "name": "vllm-0-21", + "path": null }, - "path": "/images/vllm-0-21.sqsh", - "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", "inspection": { - "schema_version": 1, - "inspector_version": "inspector-1", - "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd", "inspection": { + "executable_path": "/usr/local/bin/vllm", "kind": "serving", - "server_type": "vllm", "runtime_version": "0.21.0", - "executable_path": "/usr/local/bin/vllm" - } - } + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" }, "node_indices": [ 0 ], - "gpus_per_node": 8, - "topology": { - "tensor_parallel": 8, - "nodes_per_replica": 1, - "pipeline_parallel": 1, - "node_group_count": 1, - "replicas_per_node_group": 1, - "replica_count": 1, - "gpus_per_replica": 8 - }, "ports": [ { "name": "generator-http", "node_index": 0, - "port": 18000 + "port": 18000, + "role": "http" } - ] + ], + "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 + } } ], - "array_tasks": { - "count": 1, - "max_concurrent": 1 + "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", + "plan_id": "plan-single", + "resolved_gpus_per_node": 8, + "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": [ { - "shard_id": "shard-00000", - "shard_index": 0, "array_task_index": 0, - "start_index": 0, "end_index_exclusive": 8, "requested_records": 8, - "resume_workspace": "/workspace/primary/runs/run-single/shards/shard-00000/dataset" + "resume_workspace": "/workspace/primary/runs/run-single/shards/shard-00000/dataset", + "shard_id": "shard-00000", + "shard_index": 0, + "start_index": 0 } ], "submission": { "account": "research", - "partition": "batch", + "comment": null, "job_name": "data-designer", + "partition": "batch", "time_limit": "03:55:00" - }, - "output": { - "root": "/workspace/primary/runs/run-single/output", - "format": "parquet", - "partitions": 1, - "require_exact_record_count": false - }, - "container_mounts": [ - { - "source": "/workspace", - "target": "/workspace", - "read_only": false - } - ], - "runtime_bundle": { - "path": "/workspace/primary/runtime/runtime.tar.gz", - "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" } } diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index 4d0eeb560..e5af7b501 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -9,9 +9,11 @@ import pytest from pydantic import ValidationError +from data_designer.config import DataDesignerConfigBuilder from data_designer.slurm.config import ( ArrayTasksConfig, BenchmarkBaseRun, + BuilderInput, ClientDependencies, DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, @@ -153,6 +155,26 @@ def test_run_rejects_retired_builder_fields(authored_run: DataDesignerSlurmConfi DataDesignerSlurmConfig.model_validate(payload) +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": []} + + +@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 diff --git a/packages/data-designer-slurm/tests/contracts/test_golden_records.py b/packages/data-designer-slurm/tests/contracts/test_golden_records.py index 9c64f1cdc..c6c963b6e 100644 --- a/packages/data-designer-slurm/tests/contracts/test_golden_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_golden_records.py @@ -27,10 +27,12 @@ ("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"), @@ -40,10 +42,12 @@ ], ) def test_golden_record_round_trip(record_type: type[BaseModel], filename: str) -> None: - record = record_type.model_validate_json((GOLDEN_DIR / filename).read_text()) + 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, ContractRecord): + assert record.serialize_json() == fixture assert record_type.model_validate_json(record.serialize_json()) == record assert record.compute_sha256() == record.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 index 0d2d1e08b..a27ee60b5 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -16,6 +16,7 @@ PlanContractError, ResolvedDependencyLock, ResolvedSlurmRunPlan, + ResolvedSubmission, validate_resolved_plan, ) @@ -31,6 +32,17 @@ def test_multi_node_plan_matches_authored_inputs( 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 == compute_sha256(authored_run_single.model_dump(mode="json")) + 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} @@ -49,6 +61,7 @@ def test_plan_canonical_json_is_byte_stable(multi_node_plan: ResolvedSlurmRunPla lambda payload: payload["client"].update(host_node_index=1), lambda payload: payload["deployments"][1].update(node_indices=[1]), 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.update(shards=payload["shards"][:1]), lambda payload: payload["shards"][1].update(start_index=49), lambda payload: payload["output"].update(root="/outside/output"), @@ -64,6 +77,28 @@ def test_plan_rejects_invalid_boundaries(multi_node_plan: ResolvedSlurmRunPlan, 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) + + +@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 From 1d37add2f8c2054ea77849137ca239eee2ef680e Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 19 Aug 2026 11:37:05 -0300 Subject: [PATCH 3/7] fix: align Slurm F3 contract seam --- .../src/data_designer/slurm/_contracts.py | 88 ++++++-- .../data_designer/slurm/benchmark/records.py | 12 +- .../src/data_designer/slurm/client/records.py | 32 ++- .../data_designer/slurm/config/benchmark.py | 15 +- .../data_designer/slurm/config/profiles.py | 5 +- .../src/data_designer/slurm/config/run.py | 15 +- .../data_designer/slurm/planning/__init__.py | 4 + .../data_designer/slurm/planning/models.py | 191 ++++++++++++++---- .../slurm/planning/validation.py | 21 +- .../tests/contracts/golden/authored_run.json | 13 +- .../contracts/golden/authored_run_single.json | 8 +- .../contracts/golden/benchmark_manifest.json | 8 +- .../contracts/golden/multi_node_plan.json | 97 ++++++--- .../contracts/golden/single_node_plan.json | 50 +++-- .../tests/contracts/test_config_records.py | 23 ++- .../tests/contracts/test_golden_records.py | 14 +- .../tests/contracts/test_planning_records.py | 112 +++++++++- .../tests/contracts/test_shared_records.py | 39 +++- 18 files changed, 612 insertions(+), 135 deletions(-) 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 ee2987ed5..b0b199686 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py @@ -9,7 +9,15 @@ import re from typing import Annotated, Literal -from pydantic import BaseModel, ConfigDict, StringConstraints +from pydantic import ( + BaseModel, + ConfigDict, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) Identifier = Annotated[ str, @@ -19,6 +27,10 @@ pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", ), ] +ModelAlias = str +ShardId = Annotated[str, StringConstraints(pattern=r"^shard-[0-9]{5,}$")] +AttemptId = Annotated[str, StringConstraints(pattern=r"^attempt-[0-9]{4,}$")] +SchemaVersion = Literal[1] 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)$")] @@ -30,11 +42,21 @@ class AuthoredConfig(BaseModel): model_config = ConfigDict( extra="forbid", frozen=True, + allow_inf_nan=False, protected_namespaces=(), strict=True, validate_default=True, ) + 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 ContractValue(BaseModel): """Base for strict immutable cross-process values.""" @@ -42,6 +64,7 @@ class ContractValue(BaseModel): model_config = ConfigDict( extra="forbid", frozen=True, + allow_inf_nan=False, protected_namespaces=(), strict=True, validate_default=True, @@ -51,25 +74,16 @@ class ContractValue(BaseModel): class ContractRecord(ContractValue): """Base for explicitly versioned records with stable serialization.""" - schema_version: Literal[1] + schema_version: SchemaVersion def serialize_canonical_json(self) -> bytes: return canonical_json(self.model_dump(mode="json")) def serialize_json(self) -> str: - return ( - json.dumps( - self.model_dump(mode="json"), - allow_nan=False, - ensure_ascii=False, - indent=2, - sort_keys=True, - ) - + "\n" - ) + return pretty_json(self.model_dump(mode="json")) def compute_sha256(self) -> Sha256Digest: - return hashlib.sha256(self.serialize_canonical_json()).hexdigest() + return hashlib.sha256(self.serialize_json().encode("utf-8")).hexdigest() def canonical_json(value: object) -> bytes: @@ -83,6 +97,20 @@ def canonical_json(value: object) -> bytes: ).encode("utf-8") +def pretty_json(value: object) -> str: + """Serialize a JSON-compatible value to deterministic persisted text.""" + return ( + json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + indent=2, + sort_keys=True, + ) + + "\n" + ) + + def compute_sha256(value: object) -> Sha256Digest: """Compute the canonical JSON digest of a JSON-compatible value.""" return hashlib.sha256(canonical_json(value)).hexdigest() @@ -126,3 +154,37 @@ def validate_url(value: str, *, field_name: str) -> str: if not re.fullmatch(r"https?://[^\s]+", 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.""" + + path: str + sha256: Sha256Digest + + _path_is_absolute = field_validator("path")(validate_absolute_path) + + +class RecordRange(ContractValue): + """Half-open global record range assigned to one shard.""" + + start_index: NonNegativeInt + end_index_exclusive: PositiveInt + + @property + def record_count(self) -> int: + return self.end_index_exclusive - self.start_index + + @model_validator(mode="after") + def validate_bounds(self) -> RecordRange: + if self.end_index_exclusive <= self.start_index: + raise ValueError("end_index_exclusive must be greater than start_index") + return self + + +class ResumeWorkspace(ContractValue): + """Canonical shard-owned dataset workspace.""" + + path: str + + _path_is_absolute = field_validator("path")(validate_absolute_path) 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 index 4357646cd..396015c24 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -17,14 +17,20 @@ model_validator, ) -from data_designer.slurm._contracts import ContractRecord, ContractValue, Identifier -from data_designer.slurm.planning import ArtifactReference +from data_designer.slurm._contracts import ArtifactReference, ContractRecord, ContractValue, Identifier class BenchmarkChildRun(ContractValue): case_id: Identifier child_run_id: Identifier - child_config: ArtifactReference + 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): 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 index 0cdb95084..9fb2ef51f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -9,8 +9,14 @@ from pydantic import NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator -from data_designer.slurm._contracts import ContractRecord, Identifier, validate_absolute_path -from data_designer.slurm.planning import ArtifactReference +from data_designer.slurm._contracts import ( + ArtifactReference, + AttemptId, + ContractRecord, + Identifier, + ShardId, + validate_absolute_path, +) class ClientOutcome(str, Enum): @@ -23,8 +29,8 @@ class ClientResult(ContractRecord): """Semantic Data Designer outcome independent of engine-internal result types.""" run_id: Identifier - shard_id: Identifier - attempt_id: Identifier + shard_id: ShardId + attempt_id: AttemptId completed_at: datetime requested_records: PositiveInt actual_records: NonNegativeInt | None @@ -60,6 +66,11 @@ def validate_message(cls, value: str | None) -> str | None: 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 ClientOutcome.COMPLETE: if self.actual_records != self.requested_records: raise ValueError("complete client results require the requested record count") @@ -73,6 +84,9 @@ def validate_outcome(self) -> ClientResult: 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") + 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") return self def _require_success_artifacts(self) -> None: @@ -80,3 +94,13 @@ def _require_success_artifacts(self) -> 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.requested_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/benchmark.py b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py index b2450053b..88f711227 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py @@ -7,7 +7,14 @@ from pydantic import Field, PositiveFloat, PositiveInt, field_validator, model_validator -from data_designer.slurm._contracts import AuthoredConfig, Duration, Identifier, validate_local_config_path +from data_designer.slurm._contracts import ( + AuthoredConfig, + Duration, + Identifier, + ModelAlias, + SchemaVersion, + validate_local_config_path, +) from data_designer.slurm.config.run import DataDesignerSlurmConfig @@ -47,7 +54,7 @@ def validate_topology(self) -> BenchmarkDeploymentOverride: class BenchmarkDeploymentCase(AuthoredConfig): name: Identifier - deployments: dict[Identifier, BenchmarkDeploymentOverride] = Field(min_length=1) + deployments: dict[ModelAlias, BenchmarkDeploymentOverride] = Field(min_length=1) class FixedRecordPolicy(AuthoredConfig): @@ -79,10 +86,10 @@ class BenchmarkAnalysisTargets(AuthoredConfig): class DataDesignerSlurmBenchmarkConfig(AuthoredConfig): """Authored benchmark intent expanded into ordinary Slurm run configs.""" - schema_version: Literal[1] + schema_version: SchemaVersion name: Identifier base_run: BenchmarkBaseRun - model_aliases: Literal["all"] | list[Identifier] + 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 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 index 23c6799d3..de163faed 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -13,6 +13,7 @@ AuthoredConfig, ContractRecord, Identifier, + SchemaVersion, Sha256Digest, compute_sha256, validate_absolute_path, @@ -46,7 +47,7 @@ class ContainerMount(AuthoredConfig): class SlurmProfile(AuthoredConfig): """Strict facts for one Slurm cluster.""" - schema_version: Literal[1] + schema_version: SchemaVersion host_patterns: list[str] = Field(default_factory=list) scheduler: SchedulerProfile = Field(default_factory=SchedulerProfile) gpus_per_node: PositiveInt | Literal["auto"] @@ -81,7 +82,7 @@ def validate_mounts(self) -> SlurmProfile: class SlurmProfileCatalog(AuthoredConfig): """Versioned catalog of independently complete cluster profiles.""" - schema_version: Literal[1] + schema_version: SchemaVersion default_cluster: Identifier clusters: dict[Identifier, SlurmProfile] = Field(min_length=1) 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 index ddde96817..a94a0cf26 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -23,6 +23,8 @@ Duration, EnvironmentName, Identifier, + ModelAlias, + SchemaVersion, validate_absolute_path, validate_local_config_path, validate_plain_text, @@ -156,7 +158,7 @@ class InvocationConfig(AuthoredConfig): 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[Identifier, PositiveInt] = Field(default_factory=dict) + model_concurrency: dict[ModelAlias, PositiveInt] = Field(default_factory=dict) diagnostics: InvocationDiagnostics = Field(default_factory=InvocationDiagnostics) @field_validator("run_config") @@ -276,8 +278,8 @@ class DeploymentTopology(AuthoredConfig): class ServerDeploymentConfig(AuthoredConfig): - model_alias: Identifier - served_model_name: Identifier | None = None + model_alias: ModelAlias + served_model_name: str | None = None model: str server: VllmServerConfig resources: DeploymentResources = Field(default_factory=DeploymentResources) @@ -293,6 +295,11 @@ def validate_model(cls, value: str) -> str: 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: @@ -350,7 +357,7 @@ def validate_root(cls, value: str | None) -> str | None: class DataDesignerSlurmConfig(AuthoredConfig): """Complete portable intent for one Data Designer Slurm run.""" - schema_version: Literal[1] + schema_version: SchemaVersion name: Identifier builder: BuilderInput invocation: InvocationConfig 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 index 05ae36210..160b7ef77 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py @@ -10,6 +10,7 @@ LockedPackage, PlannedShard, PortClaim, + RecordRange, ResolvedBuilderInput, ResolvedClient, ResolvedDependencyLock, @@ -20,6 +21,7 @@ ResolvedSlurmRunPlan, ResolvedSubmission, ResolvedTopology, + ResumeWorkspace, ) from data_designer.slurm.planning.validation import PlanContractError, validate_resolved_plan @@ -29,6 +31,7 @@ "PlanContractError", "PlannedShard", "PortClaim", + "RecordRange", "ResolvedBuilderInput", "ResolvedClient", "ResolvedDependencyLock", @@ -39,5 +42,6 @@ "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 index cd963e71b..bc4259b44 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -10,12 +10,18 @@ from data_designer.config import RunConfig from data_designer.slurm._contracts import ( + ArtifactReference, ContractRecord, ContractValue, Identifier, + ModelAlias, + RecordRange, + ResumeWorkspace, Sha256Digest, + ShardId, compute_sha256, validate_absolute_path, + validate_plain_text, ) from data_designer.slurm.config.images import ( DistributionName, @@ -34,15 +40,6 @@ ) -class ArtifactReference(ContractValue): - """Immutable reference to a persisted artifact and its digest.""" - - path: str - sha256: Sha256Digest - - _path_is_absolute = field_validator("path")(validate_absolute_path) - - class ResolvedImage(ContractValue): """Immutable SQSH path and the digest-bound inspection that approved it.""" @@ -113,6 +110,8 @@ class ResolvedBuilderInput(ContractValue): 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: @@ -122,10 +121,19 @@ def validate_input(self) -> ResolvedBuilderInput: 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 @@ -146,7 +154,7 @@ def validate_run_config_is_materialized(cls, value: dict[str, JsonValue]) -> dic class PortClaim(ContractValue): name: Identifier - role: Literal["http", "rendezvous"] + role: Literal["http", "rendezvous", "logical_endpoint"] node_index: NonNegativeInt port: Annotated[int, Field(ge=1024, le=65535)] @@ -164,6 +172,7 @@ class ResolvedTopology(ContractValue): 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 @@ -172,6 +181,9 @@ class ResolvedDeployment(ContractValue): @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: @@ -199,17 +211,35 @@ def validate_deployment(self) -> ResolvedDeployment: 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_nodes = tuple(port.node_index for port in self.ports if port.role == "http") + 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_rendezvous_nodes = group_heads if self.topology.nodes_per_replica > 1 else () - rendezvous_nodes = tuple(port.node_index for port in self.ports if port.role == "rendezvous") + 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 group") + 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 @@ -219,6 +249,7 @@ class ResolvedClient(ContractValue): dependency_lock: ArtifactReference host_node_index: NonNegativeInt gpu_count: Literal[0] + ports: tuple[PortClaim, ...] = () @model_validator(mode="after") def validate_client(self) -> ResolvedClient: @@ -226,27 +257,27 @@ def validate_client(self) -> ResolvedClient: 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: Identifier + shard_id: ShardId shard_index: NonNegativeInt array_task_index: NonNegativeInt - start_index: NonNegativeInt - end_index_exclusive: PositiveInt - requested_records: PositiveInt - resume_workspace: str - - _workspace_is_absolute = field_validator("resume_workspace")(validate_absolute_path) + record_range: RecordRange + input_partition: ArtifactReference | None = None + resume_workspace: ResumeWorkspace - @model_validator(mode="after") - def validate_range(self) -> PlannedShard: - if self.end_index_exclusive <= self.start_index: - raise ValueError("shard end_index_exclusive must be greater than start_index") - if self.requested_records != self.end_index_exclusive - self.start_index: - raise ValueError("shard requested_records must match its record range") - return self + @property + def requested_records(self) -> int: + return self.record_range.record_count class ResolvedSubmission(ContractValue): @@ -274,6 +305,7 @@ class ResolvedOutput(ContractValue): class ResolvedSlurmRunPlan(ContractRecord): """Immutable allocation input consumed without ambient configuration.""" + run_id: Identifier plan_id: Identifier package_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] authored_config: ArtifactReference @@ -302,10 +334,15 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: deployment_ids = tuple(deployment.deployment_id for deployment in self.deployments) aliases = tuple(deployment.authored.model_alias for deployment in self.deployments) - if len(deployment_ids) != len(set(deployment_ids)): - raise ValueError("resolved deployment IDs must be unique") + 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))): @@ -313,25 +350,47 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: 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") - port_keys = tuple((port.node_index, port.port) for deployment in self.deployments for port in deployment.ports) + 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") - - self._validate_shards() + 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") + + 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") return self - def _validate_shards(self) -> None: + 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.start_index != expected_start: + 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) @@ -340,10 +399,68 @@ def _validate_shards(self) -> None: ) if shard.requested_records != expected_count: raise ValueError("shards must use deterministic floor/remainder record counts") - expected_start = shard.end_index_exclusive + 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 index 3f79d313d..4ab7d514b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -3,10 +3,15 @@ from __future__ import annotations -from data_designer.slurm._contracts import compute_sha256 +from pydantic import JsonValue + 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 +from data_designer.slurm.planning.models import ( + ResolvedDependencyLock, + ResolvedSlurmRunPlan, + _extract_builder_aliases, +) class PlanContractError(ValueError): @@ -17,10 +22,12 @@ 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 == compute_sha256(authored.model_dump(mode="json")), + 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") @@ -40,6 +47,14 @@ def validate_resolved_plan( 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 diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json index d8f0470d1..ba58722ab 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -5,7 +5,18 @@ "inline": { "data_designer": { "columns": [], - "model_configs": [] + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + }, + { + "alias": "judge", + "model": "/models/judge", + "provider": "openai" + } + ] } } }, 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 index f35b60d8b..5fcf9e524 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -5,7 +5,13 @@ "inline": { "data_designer": { "columns": [], - "model_configs": [] + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + } + ] } } }, diff --git a/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json b/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json index 4623c2a49..d72acb06f 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json @@ -7,16 +7,16 @@ "children": [ { "case_id": "two-independent-replicas-c32", - "child_config": { - "path": "/workspace/primary/runs/run-benchmark-001-c32/run.json", + "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_config": { - "path": "/workspace/primary/runs/run-benchmark-002-c32/run.json", + "child_authored_config": { + "path": "/workspace/primary/runs/run-benchmark-002-c32/authored-config.json", "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" }, "child_run_id": "run-benchmark-002-c32" 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 index f82bd5ec3..271367b73 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -4,18 +4,34 @@ "max_concurrent": 2 }, "authored_config": { - "path": "/workspace/primary/runs/run-001/run.json", - "sha256": "a8d4be425bccdd56c4cef65a63b1bda372a28075f4e91e8c9cb724fbea5a3251" + "path": "/workspace/primary/runs/run-001/authored-config.json", + "sha256": "9bfdbcb1ddd374a5d43eea44890d13c08798ea9e18c30ff5d5258fa0d27f539d" }, "builder": { "authored_source": null, - "content_sha256": "26631b9b95ed99ac227eff0dfc64fb0c35c8cb9b9b7717d7dc2f2702055e0c4e", + "content_sha256": "f37227ca7c67e203abe881c0228b4308a8e741364296d293159a1201949732f2", "inline": { "data_designer": { "columns": [], - "model_configs": [] + "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": { @@ -40,7 +56,7 @@ }, "dependency_lock": { "path": "/workspace/primary/runs/run-001/dependency-lock.json", - "sha256": "0db59e7c2db1f1ea4123299a40ca6e42b50d2a27192902721d30fd8bc32a1f8a" + "sha256": "fd80a86f035fd1d7cab2e4988722eb7441fe07922f809d4a07e73a0c75190a75" }, "gpu_count": 0, "host_node_index": 0, @@ -74,7 +90,21 @@ }, "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": [ { @@ -122,7 +152,7 @@ "tensor_parallel": 8 } }, - "deployment_id": "generator", + "deployment_id": "deployment-00000", "gpus_per_node": 8, "image": { "authored_ref": { @@ -149,18 +179,19 @@ ], "ports": [ { - "name": "generator-http", + "name": "deployment-00000-http-00000", "node_index": 0, "port": 18000, "role": "http" }, { - "name": "generator-rendezvous", + "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, @@ -201,7 +232,7 @@ "tensor_parallel": 1 } }, - "deployment_id": "judge", + "deployment_id": "deployment-00001", "gpus_per_node": 8, "image": { "authored_ref": { @@ -227,54 +258,55 @@ ], "ports": [ { - "name": "judge-http-0", + "name": "deployment-00001-http-00000", "node_index": 2, "port": 18000, "role": "http" }, { - "name": "judge-http-1", + "name": "deployment-00001-http-00001", "node_index": 2, "port": 18001, "role": "http" }, { - "name": "judge-http-2", + "name": "deployment-00001-http-00002", "node_index": 2, "port": 18002, "role": "http" }, { - "name": "judge-http-3", + "name": "deployment-00001-http-00003", "node_index": 2, "port": 18003, "role": "http" }, { - "name": "judge-http-4", + "name": "deployment-00001-http-00004", "node_index": 2, "port": 18004, "role": "http" }, { - "name": "judge-http-5", + "name": "deployment-00001-http-00005", "node_index": 2, "port": 18005, "role": "http" }, { - "name": "judge-http-6", + "name": "deployment-00001-http-00006", "node_index": 2, "port": 18006, "role": "http" }, { - "name": "judge-http-7", + "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, @@ -336,6 +368,7 @@ "package_version": "0.9.2", "plan_id": "plan-001", "resolved_gpus_per_node": 8, + "run_id": "run-001", "runtime_bundle": { "path": "/workspace/primary/runtime/runtime.tar.gz", "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" @@ -378,21 +411,29 @@ "shards": [ { "array_task_index": 0, - "end_index_exclusive": 50, - "requested_records": 50, - "resume_workspace": "/workspace/primary/runs/run-001/shards/shard-00000/dataset", + "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, - "start_index": 0 + "shard_index": 0 }, { "array_task_index": 1, - "end_index_exclusive": 100, - "requested_records": 50, - "resume_workspace": "/workspace/primary/runs/run-001/shards/shard-00001/dataset", + "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, - "start_index": 50 + "shard_index": 1 } ], "submission": { 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 index feec12025..d39595143 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -4,18 +4,28 @@ "max_concurrent": 1 }, "authored_config": { - "path": "/workspace/primary/runs/run-single/run.json", - "sha256": "1ca141405769834f7b1e53461861d9ae349417008306e5fea80402d2a368bfa8" + "path": "/workspace/primary/runs/run-single/authored-config.json", + "sha256": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" }, "builder": { "authored_source": null, - "content_sha256": "26631b9b95ed99ac227eff0dfc64fb0c35c8cb9b9b7717d7dc2f2702055e0c4e", + "content_sha256": "b3ef5fc1fe675a8e004633f84842ac60cf82d5ba3dc68b4d50ee4438448b0570", "inline": { "data_designer": { "columns": [], - "model_configs": [] + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + } + ] } }, + "model_aliases": [ + "generator" + ], + "referenced_model_aliases": [], "source": null }, "client": { @@ -33,7 +43,7 @@ }, "dependency_lock": { "path": "/workspace/primary/runs/run-single/dependency-lock.json", - "sha256": "9b4e1db6dbcc62df9ee90e329578cfeb727e4060d7423c79c52558a69a42c156" + "sha256": "4349a233e932c09effba025d7983bfdc0b1c3f7ddbd264676942fee9ec9c3cb1" }, "gpu_count": 0, "host_node_index": 0, @@ -63,7 +73,15 @@ }, "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": [ { @@ -103,7 +121,7 @@ "tensor_parallel": 8 } }, - "deployment_id": "generator", + "deployment_id": "deployment-00000", "gpus_per_node": 8, "image": { "authored_ref": { @@ -129,12 +147,13 @@ ], "ports": [ { - "name": "generator-http", + "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, @@ -193,6 +212,7 @@ "package_version": "0.9.2", "plan_id": "plan-single", "resolved_gpus_per_node": 8, + "run_id": "run-single", "runtime_bundle": { "path": "/workspace/primary/runtime/runtime.tar.gz", "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" @@ -235,12 +255,16 @@ "shards": [ { "array_task_index": 0, - "end_index_exclusive": 8, - "requested_records": 8, - "resume_workspace": "/workspace/primary/runs/run-single/shards/shard-00000/dataset", + "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, - "start_index": 0 + "shard_index": 0 } ], "submission": { diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index e5af7b501..bc5c9ed61 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -9,7 +9,7 @@ import pytest from pydantic import ValidationError -from data_designer.config import DataDesignerConfigBuilder +from data_designer.config import DataDesignerConfigBuilder, ModelConfig from data_designer.slurm.config import ( ArrayTasksConfig, BenchmarkBaseRun, @@ -135,6 +135,21 @@ def test_deployment_rejects_invalid_topology() -> None: 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" @@ -162,6 +177,12 @@ def test_builder_input_accepts_exported_and_shorthand_configs() -> None: assert BuilderInput(inline={"columns": []}).inline == {"columns": []} +@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", [ diff --git a/packages/data-designer-slurm/tests/contracts/test_golden_records.py b/packages/data-designer-slurm/tests/contracts/test_golden_records.py index c6c963b6e..3a377de1a 100644 --- a/packages/data-designer-slurm/tests/contracts/test_golden_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_golden_records.py @@ -3,13 +3,13 @@ from __future__ import annotations -import json +import hashlib from pathlib import Path import pytest from pydantic import BaseModel -from data_designer.slurm._contracts import ContractRecord +from data_designer.slurm._contracts import AuthoredConfig, ContractRecord from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport from data_designer.slurm.client import ClientResult from data_designer.slurm.config import ( @@ -46,10 +46,11 @@ def test_golden_record_round_trip(record_type: type[BaseModel], filename: str) - 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() if isinstance(record, ContractRecord): assert record.serialize_json() == fixture - assert record_type.model_validate_json(record.serialize_json()) == record - assert record.compute_sha256() == record.compute_sha256() def test_golden_records_are_sanitized() -> None: @@ -65,6 +66,5 @@ def test_canonical_serialization_ignores_mapping_order(authored_run: DataDesigne payload["invocation"]["model_concurrency"] = {"judge": 32, "generator": 64} reordered = DataDesignerSlurmConfig.model_validate(payload) - first = json.dumps(authored_run.model_dump(mode="json"), sort_keys=True, separators=(",", ":")) - second = json.dumps(reordered.model_dump(mode="json"), sort_keys=True, separators=(",", ":")) - assert first == second + 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 index a27ee60b5..d6466e05f 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -10,11 +10,12 @@ from pydantic import ValidationError from data_designer.slurm._contracts import compute_sha256 -from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.config import BuilderInput, DataDesignerSlurmConfig from data_designer.slurm.planning import ( ArtifactReference, PlanContractError, ResolvedDependencyLock, + ResolvedDeployment, ResolvedSlurmRunPlan, ResolvedSubmission, validate_resolved_plan, @@ -27,7 +28,7 @@ def test_multi_node_plan_matches_authored_inputs( 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 == compute_sha256(authored_run.model_dump(mode="json")) + 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 @@ -38,7 +39,7 @@ def test_single_node_plan_matches_authored_inputs( 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 == compute_sha256(authored_run_single.model_dump(mode="json")) + 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 @@ -60,10 +61,15 @@ def test_plan_canonical_json_is_byte_stable(multi_node_plan: ResolvedSlurmRunPla 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].update(start_index=49), + 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), @@ -85,6 +91,104 @@ def test_plan_rejects_unmaterialized_run_config(multi_node_plan: ResolvedSlurmRu ResolvedSlurmRunPlan.model_validate(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", [ diff --git a/packages/data-designer-slurm/tests/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py index aa5e73fac..23cebfa71 100644 --- a/packages/data-designer-slurm/tests/contracts/test_shared_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -9,8 +9,14 @@ import pytest from pydantic import ValidationError +from data_designer.slurm._contracts import ArtifactReference as CommonArtifactReference from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport from data_designer.slurm.client import ClientResult +from data_designer.slurm.planning import ArtifactReference as PlanningArtifactReference + + +def test_planning_reexports_common_artifact_reference() -> None: + assert PlanningArtifactReference is CommonArtifactReference def test_client_result_allows_partial_and_failure_facts_to_differ() -> None: @@ -23,11 +29,14 @@ def test_client_result_allows_partial_and_failure_facts_to_differ() -> None: "requested_records": 50, "actual_records": 25, "outcome": "partial", - "dataset_path": "/workspace/dataset", + "dataset_path": "/workspace/runs/run-001/shards/shard-00000/dataset", "early_shutdown": True, "requested_resume_mode": "if_possible", "effective_resume_mode": "never", - "candidate_output_manifest": {"path": "/workspace/output.json", "sha256": "a" * 64}, + "candidate_output_manifest": { + "path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "a" * 64, + }, } failed = { "schema_version": 1, @@ -57,6 +66,15 @@ def test_client_result_allows_partial_and_failure_facts_to_differ() -> None: {"outcome": "partial", "actual_records": 0}, {"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}, + {"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"}, ], @@ -83,11 +101,14 @@ def client_result_payload() -> dict[str, object]: "requested_records": 50, "actual_records": 50, "outcome": "complete", - "dataset_path": "/workspace/dataset", + "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/output.json", "sha256": "a" * 64}, + "candidate_output_manifest": { + "path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "a" * 64, + }, } @@ -100,12 +121,18 @@ def test_benchmark_manifest_rejects_duplicate_child_identity() -> None: { "case_id": "case", "child_run_id": "run", - "child_config": {"path": "/workspace/run-1.json", "sha256": "b" * 64}, + "child_authored_config": { + "path": "/workspace/runs/run/authored-config.json", + "sha256": "b" * 64, + }, }, { "case_id": "case", "child_run_id": "run-2", - "child_config": {"path": "/workspace/run-2.json", "sha256": "c" * 64}, + "child_authored_config": { + "path": "/workspace/runs/run-2/authored-config.json", + "sha256": "c" * 64, + }, }, ], } From 472cf09bf8faa3d9ba5bc4b6a20a5f37f9feb06d Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 19 Aug 2026 11:59:01 -0300 Subject: [PATCH 4/7] fix: bind client dataset to effective resume mode --- .../src/data_designer/slurm/client/records.py | 8 +++---- .../tests/contracts/golden/client_result.json | 2 +- .../tests/contracts/test_shared_records.py | 24 ++++++++++++++++++- 3 files changed, 28 insertions(+), 6 deletions(-) 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 index 9fb2ef51f..b61ac4e79 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -71,6 +71,9 @@ def validate_outcome(self) -> ClientResult: 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") @@ -84,9 +87,6 @@ def validate_outcome(self) -> ClientResult: 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") - 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") return self def _require_success_artifacts(self) -> None: @@ -95,7 +95,7 @@ def _require_success_artifacts(self) -> None: 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.requested_resume_mode == "never": + if self.effective_resume_mode == "never": expected_dataset = f"{shard_root}/attempts/{self.attempt_id}/dataset" else: expected_dataset = f"{shard_root}/dataset" diff --git a/packages/data-designer-slurm/tests/contracts/golden/client_result.json b/packages/data-designer-slurm/tests/contracts/golden/client_result.json index d4bdc0834..f3b5edac6 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/client_result.json +++ b/packages/data-designer-slurm/tests/contracts/golden/client_result.json @@ -6,7 +6,7 @@ "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" }, "completed_at": "2026-08-19T12:00:00Z", - "dataset_path": "/workspace/primary/runs/run-001/shards/shard-00000/dataset", + "dataset_path": "/workspace/primary/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset", "early_shutdown": false, "effective_resume_mode": "never", "error_code": null, diff --git a/packages/data-designer-slurm/tests/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py index 23cebfa71..e81e551ad 100644 --- a/packages/data-designer-slurm/tests/contracts/test_shared_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -29,7 +29,7 @@ def test_client_result_allows_partial_and_failure_facts_to_differ() -> None: "requested_records": 50, "actual_records": 25, "outcome": "partial", - "dataset_path": "/workspace/runs/run-001/shards/shard-00000/dataset", + "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", @@ -112,6 +112,28 @@ def client_result_payload() -> dict[str, object]: } +@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, From 4896ea9187de844870d4ae5f1dddc45ae0417b24 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 19 Aug 2026 12:01:24 -0600 Subject: [PATCH 5/7] fix: harden Slurm shared contracts Prevent secret material from entering persisted configuration and make nested contract collections immutable. Tighten client and benchmark semantic invariants with focused negative coverage. Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/_contracts.py | 64 ++++++++++++- .../data_designer/slurm/benchmark/records.py | 27 +++++- .../src/data_designer/slurm/client/records.py | 2 + .../src/data_designer/slurm/config/run.py | 77 +++++++++++++++- .../tests/contracts/test_config_records.py | 91 ++++++++++++++++++- .../tests/contracts/test_planning_records.py | 14 +++ .../tests/contracts/test_shared_records.py | 58 ++++++++++-- 7 files changed, 320 insertions(+), 13 deletions(-) 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 b0b199686..ae4f3f654 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py @@ -7,7 +7,8 @@ import json import posixpath import re -from typing import Annotated, Literal +from collections.abc import Mapping +from typing import Annotated, Literal, TypeVar from pydantic import ( BaseModel, @@ -35,6 +36,57 @@ 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 AuthoredConfig(BaseModel): """Base for strict authored configuration values.""" @@ -48,6 +100,11 @@ class AuthoredConfig(BaseModel): validate_default=True, ) + @field_validator("*", mode="after") + @classmethod + def freeze_collections(cls, value: object) -> object: + return _freeze_collections(value) + def serialize_canonical_json(self) -> bytes: return canonical_json(self.model_dump(mode="json")) @@ -70,6 +127,11 @@ class ContractValue(BaseModel): validate_default=True, ) + @field_validator("*", mode="after") + @classmethod + def freeze_collections(cls, value: object) -> object: + return _freeze_collections(value) + class ContractRecord(ContractValue): """Base for explicitly versioned records with stable serialization.""" 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 index 396015c24..418bef7b4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -96,6 +96,11 @@ def validate_metrics(self) -> BenchmarkCaseResult: ) 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 @@ -130,12 +135,28 @@ def validate_created_at(cls, value: datetime) -> datetime: @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))}") - kinds = tuple(recommendation.kind for recommendation in self.recommendations) - if len(kinds) != len(set(kinds)): - raise ValueError("benchmark recommendation kinds must be unique") + 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/records.py b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py index b61ac4e79..4eb746810 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -77,6 +77,8 @@ def validate_outcome(self) -> ClientResult: 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 not 0 < self.actual_records < self.requested_records: 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 index a94a0cf26..bc8ca8198 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -5,7 +5,9 @@ import posixpath import re +from collections.abc import Mapping from typing import Annotated, Literal +from urllib.parse import urlsplit from pydantic import ( Field, @@ -33,15 +35,57 @@ from data_designer.slurm.config.images import ImageRef _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} +_SECRET_NAME_PARTS = frozenset({"credential", "credentials", "password", "secret", "token"}) + + +def _duration_seconds(value: Duration) -> int: + return int(value[:-1]) * _DURATION_FACTORS[value[-1]] + + +def _is_secret_name(value: str) -> bool: + snake_case = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", value) + normalized = re.sub(r"[^a-z0-9]+", "_", snake_case.casefold()).strip("_") + segments = normalized.split("_") + return bool( + _SECRET_NAME_PARTS.intersection(segments) + or {"access", "key"}.issubset(segments) + or segments[-1] in {"auth", "key"} + ) + + +def _contains_secret_key(value: object) -> bool: + if isinstance(value, Mapping): + return any(_is_secret_name(str(key)) 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): @@ -93,6 +137,8 @@ def validate_input(self) -> BuilderInput: 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 @@ -115,7 +161,11 @@ class RemoteMCPProviderConfig(AuthoredConfig): @field_validator("endpoint") @classmethod def validate_endpoint(cls, value: str) -> str: - return validate_url(value, field_name="MCP endpoint") + 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): @@ -138,8 +188,13 @@ def validate_command(cls, value: str) -> str: def validate_args(cls, values: list[str]) -> list[str]: for value in values: validate_plain_text(value, field_name="MCP argument") + option = value.partition("=")[0].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, @@ -195,7 +250,17 @@ def validate_requirements(cls, values: list[str] | None) -> list[str] | None: raise ValueError(f"dependency requirement must identify a package or immutable wheel: {value!r}") if " @ " in value: _, source = value.split(" @ ", maxsplit=1) - if not re.fullmatch(r"https://[^\s]+\.whl#sha256=[0-9a-f]{64}", source): + parsed = urlsplit(source) + 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") elif "://" in value or not re.match(r"^[A-Za-z0-9][A-Za-z0-9._-]*", value): raise ValueError(f"invalid dependency requirement: {value!r}") @@ -267,6 +332,14 @@ def validate_extra_args(cls, values: list[str]) -> list[str]: raise ValueError(f"vLLM argument {flag!r} is owned by the compiler or runtime") 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 diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index bc5c9ed61..e0ff01c4e 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -21,7 +21,9 @@ ImageInspectionRecord, ImageRef, LiteralEnvironmentBinding, + LocalStdioMCPProviderConfig, QueueBackpressureConfig, + RemoteMCPProviderConfig, SecretRef, ServerDeploymentConfig, SubmissionConfig, @@ -71,6 +73,8 @@ def test_image_ref_requires_one_registered_alias_or_absolute_sqsh(payload: dict[ {"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": ["my_pkg==1", "my-pkg==2"]}, {"requirements": None, "lock_file": "../lock.json"}, ], @@ -99,6 +103,48 @@ def test_literal_environment_rejects_control_characters() -> None: 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("argument", ["--api-key", "--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( @@ -111,12 +157,33 @@ def test_vllm_defaults_and_backpressure_override() -> None: assert overridden.queue_backpressure.model_dump() == {"max_waiting_requests": 0, "retry_after_seconds": None} -@pytest.mark.parametrize("argument", ["--port", "--host=0.0.0.0", "--tensor-parallel-size"]) +@pytest.mark.parametrize( + "argument", + [ + "--api-key=plaintext-secret", + "--distributed-init-address", + "--host=0.0.0.0", + "--middleware", + "--model", + "--port", + "--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]) +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", @@ -170,6 +237,12 @@ def test_run_rejects_retired_builder_fields(authored_run: DataDesignerSlurmConfi DataDesignerSlurmConfig.model_validate(payload) +@pytest.mark.parametrize("secret_key", ["api_key", "accessToken", "client-secret", "password"]) +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"}) + + def test_builder_input_accepts_exported_and_shorthand_configs() -> None: exported = DataDesignerConfigBuilder(model_configs=[]).get_builder_config().to_dict() @@ -283,3 +356,19 @@ def test_config_models_do_not_mutate_input(authored_run: DataDesignerSlurmConfig 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_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py index d6466e05f..3236dd6fa 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -52,6 +52,20 @@ def test_plan_canonical_json_is_byte_stable(multi_node_plan: ResolvedSlurmRunPla 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", [ diff --git a/packages/data-designer-slurm/tests/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py index e81e551ad..0bea20c40 100644 --- a/packages/data-designer-slurm/tests/contracts/test_shared_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -5,6 +5,7 @@ import json from copy import deepcopy +from pathlib import Path import pytest from pydantic import ValidationError @@ -14,6 +15,8 @@ from data_designer.slurm.client import ClientResult from data_designer.slurm.planning import ArtifactReference as PlanningArtifactReference +GOLDEN_DIR = Path(__file__).parent / "golden" + def test_planning_reexports_common_artifact_reference() -> None: assert PlanningArtifactReference is CommonArtifactReference @@ -68,6 +71,7 @@ def test_client_result_allows_partial_and_failure_facts_to_differ() -> None: {"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": { @@ -163,7 +167,7 @@ def test_benchmark_manifest_rejects_duplicate_child_identity() -> None: BenchmarkManifest.model_validate_json(json.dumps(payload)) -def test_benchmark_report_rejects_unknown_and_duplicate_recommendations() -> None: +def test_benchmark_report_rejects_unknown_or_incomplete_recommendations() -> None: payload = { "schema_version": 1, "benchmark_id": "bench", @@ -187,11 +191,8 @@ def test_benchmark_report_rejects_unknown_and_duplicate_recommendations() -> Non with pytest.raises(ValidationError, match="unknown cases"): BenchmarkReport.model_validate_json(json.dumps(payload)) - payload["recommendations"] = [ - {"kind": "pareto", "case_id": "case"}, - {"kind": "pareto", "case_id": "case"}, - ] - with pytest.raises(ValidationError, match="kinds"): + payload["recommendations"] = [{"kind": "pareto", "case_id": "case"}] + with pytest.raises(ValidationError, match="successful feasible"): BenchmarkReport.model_validate_json(json.dumps(payload)) @@ -217,3 +218,48 @@ def test_successful_benchmark_case_requires_metrics() -> None: 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)) From ad5a915275bfb80ee742be8a409812ac04ae1b27 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 20 Aug 2026 16:38:44 -0300 Subject: [PATCH 6/7] fix: address Slurm contract review findings Signed-off-by: Andre Manoel --- .../src/data_designer/slurm/_contracts.py | 13 +++- .../data_designer/slurm/config/profiles.py | 7 ++ .../src/data_designer/slurm/config/run.py | 23 ++++--- .../data_designer/slurm/planning/models.py | 15 +++++ .../slurm/planning/validation.py | 15 ++++- .../contracts/golden/dependency_lock.json | 4 +- .../golden/dependency_lock_single.json | 4 +- .../contracts/golden/multi_node_plan.json | 2 +- .../contracts/golden/single_node_plan.json | 2 +- .../tests/contracts/test_config_records.py | 32 +++++++++- .../tests/contracts/test_planning_records.py | 64 ++++++++++++++++++- .../tests/contracts/test_profiles.py | 14 ++++ 12 files changed, 176 insertions(+), 19 deletions(-) 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 ae4f3f654..dea1540ce 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/_contracts.py @@ -6,9 +6,9 @@ import hashlib import json import posixpath -import re from collections.abc import Mapping from typing import Annotated, Literal, TypeVar +from urllib.parse import urlsplit from pydantic import ( BaseModel, @@ -213,7 +213,16 @@ def validate_plain_text(value: str, *, field_name: str) -> str: def validate_url(value: str, *, field_name: str) -> str: validate_plain_text(value, field_name=field_name) - if not re.fullmatch(r"https?://[^\s]+", value): + 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 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 index de163faed..76200e18f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -222,6 +222,13 @@ def validate_selected_profile( 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 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 index bc8ca8198..8e8ee4ec3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -9,6 +9,8 @@ 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, @@ -69,7 +71,10 @@ def _is_secret_name(value: str) -> bool: def _contains_secret_key(value: object) -> bool: if isinstance(value, Mapping): - return any(_is_secret_name(str(key)) or _contains_secret_key(item) for key, item in value.items()) + return any( + (_is_secret_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 @@ -248,9 +253,12 @@ def validate_requirements(cls, values: list[str] | None) -> list[str] | None: 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}") - if " @ " in value: - _, source = value.split(" @ ", maxsplit=1) - parsed = urlsplit(source) + try: + requirement = Requirement(value) + except InvalidRequirement as error: + raise ValueError(f"invalid dependency requirement: {value!r}") from error + if requirement.url is not None: + parsed = urlsplit(requirement.url) valid_wheel = ( parsed.scheme == "https" and parsed.hostname is not None @@ -262,10 +270,7 @@ def validate_requirements(cls, values: list[str] | None) -> list[str] | None: ) if not valid_wheel: raise ValueError("direct dependency URLs must be HTTPS wheels with a SHA-256 fragment") - elif "://" in value or not re.match(r"^[A-Za-z0-9][A-Za-z0-9._-]*", value): - raise ValueError(f"invalid dependency requirement: {value!r}") - name = re.split(r"\s|\[|[<>=!~@]", value, maxsplit=1)[0].lower().replace("_", "-").replace(".", "-") - names.append(name) + names.append(canonicalize_name(requirement.name)) if len(names) != len(set(names)): raise ValueError("dependency requirements must have unique normalized names") return values @@ -330,6 +335,8 @@ def validate_extra_args(cls, values: list[str]) -> list[str]: flag = value.split("=", maxsplit=1)[0] 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) 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 index bc4259b44..885b9b7a9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -21,6 +21,7 @@ ShardId, compute_sha256, validate_absolute_path, + validate_local_config_path, validate_plain_text, ) from data_designer.slurm.config.images import ( @@ -88,11 +89,25 @@ class ResolvedDependencyLock(ContractRecord): 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") 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)): 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 index 4ab7d514b..fbc015047 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -99,11 +99,22 @@ def validate_resolved_plan( dependency_lock.image_distributions == inspection.distributions, "dependency lock image inventory does not match client image inspection", ) - if authored.client.dependencies.requirements is not None: + authored_requirements = authored.client.dependencies.requirements + if authored_requirements is not None: _require( - dependency_lock.authored_requirements == tuple(authored.client.dependencies.requirements), + 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 diff --git a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json index 502edcf3d..87190b04c 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -2,6 +2,7 @@ "authored_requirements": [ "data-designer-speech==0.2.0" ], + "authored_source": null, "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "image_distributions": [ { @@ -25,5 +26,6 @@ ], "python_abi": "cp312", "resolver_version": "resolver-1", - "schema_version": 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 index 8f7b516c4..b86f29d6c 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json @@ -1,5 +1,6 @@ { "authored_requirements": [], + "authored_source": null, "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "image_distributions": [ { @@ -10,5 +11,6 @@ "overlay_packages": [], "python_abi": "cp312", "resolver_version": "resolver-1", - "schema_version": 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 index 271367b73..07e461352 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -56,7 +56,7 @@ }, "dependency_lock": { "path": "/workspace/primary/runs/run-001/dependency-lock.json", - "sha256": "fd80a86f035fd1d7cab2e4988722eb7441fe07922f809d4a07e73a0c75190a75" + "sha256": "a86032b310aa6bdb95fc7f35ef606fec2e56b977ba60e36b8b9293341ac43e00" }, "gpu_count": 0, "host_node_index": 0, 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 index d39595143..aef26ae9e 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -43,7 +43,7 @@ }, "dependency_lock": { "path": "/workspace/primary/runs/run-single/dependency-lock.json", - "sha256": "4349a233e932c09effba025d7983bfdc0b1c3f7ddbd264676942fee9ec9c3cb1" + "sha256": "2600c1944897a8e3c9a9410fe912bdee0e169e09769e58180cadebe94f3da52c" }, "gpu_count": 0, "host_node_index": 0, diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index e0ff01c4e..08bc3e3ff 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -9,7 +9,7 @@ import pytest from pydantic import ValidationError -from data_designer.config import DataDesignerConfigBuilder, ModelConfig +from data_designer.config import DataDesignerConfigBuilder, HuggingFaceSeedSource, ModelConfig from data_designer.slurm.config import ( ArrayTasksConfig, BenchmarkBaseRun, @@ -76,6 +76,8 @@ def test_image_ref_requires_one_registered_alias_or_absolute_sqsh(payload: dict[ {"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": ["my_pkg==1", "my-pkg==2"]}, + {"requirements": ["not valid !!!"]}, + {"requirements": ["plugin=="]}, {"requirements": None, "lock_file": "../lock.json"}, ], ) @@ -134,6 +136,12 @@ def test_remote_mcp_endpoint_rejects_embedded_credentials(endpoint: str) -> None 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", "--access-token=plaintext-secret", "password"]) def test_stdio_mcp_rejects_secret_shaped_arguments(argument: str) -> None: with pytest.raises(ValidationError, match="secret-shaped"): @@ -174,6 +182,18 @@ def test_vllm_rejects_runtime_owned_arguments(argument: str) -> None: VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=[argument]) +@pytest.mark.parametrize( + "extra_args", + [ + ["--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( @@ -250,6 +270,16 @@ def test_builder_input_accepts_exported_and_shorthand_configs() -> None: 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): diff --git a/packages/data-designer-slurm/tests/contracts/test_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py index 3236dd6fa..bc583842e 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -10,7 +10,7 @@ from pydantic import ValidationError from data_designer.slurm._contracts import compute_sha256 -from data_designer.slurm.config import BuilderInput, DataDesignerSlurmConfig +from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig from data_designer.slurm.planning import ( ArtifactReference, PlanContractError, @@ -237,9 +237,14 @@ def test_resolved_image_rejects_digest_mismatch(multi_node_plan: ResolvedSlurmRu ), 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_overlap_order_and_non_wheel( +def test_dependency_lock_rejects_invalid_boundaries( dependency_lock: ResolvedDependencyLock, mutator: object, ) -> None: @@ -287,6 +292,61 @@ def test_cross_record_validation_rejects_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, diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py index 1a2a92366..1e4086f5e 100644 --- a/packages/data-designer-slurm/tests/contracts/test_profiles.py +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -51,6 +51,20 @@ def test_catalog_selection_digest_validation(profile_catalog: SlurmProfileCatalo 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") From 6888df95d630de7a9b8ee6519088172708914af2 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 20 Aug 2026 18:41:20 -0300 Subject: [PATCH 7/7] fix: close Slurm argument validation bypasses Signed-off-by: Andre Manoel --- .../src/data_designer/slurm/config/run.py | 9 +++++++-- .../tests/contracts/test_config_records.py | 16 +++++++++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) 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 index 8e8ee4ec3..511754e10 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -69,6 +69,10 @@ def _is_secret_name(value: str) -> bool: ) +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( @@ -193,7 +197,7 @@ def validate_command(cls, value: str) -> str: def validate_args(cls, values: list[str]) -> list[str]: for value in values: validate_plain_text(value, field_name="MCP argument") - option = value.partition("=")[0].lstrip("-") + option = _option_flag(value).lstrip("-") if _is_secret_name(option): raise ValueError("secret-shaped MCP arguments must use an environment secret reference") return values @@ -258,6 +262,7 @@ def validate_requirements(cls, values: list[str] | None) -> list[str] | None: except InvalidRequirement as error: raise ValueError(f"invalid dependency requirement: {value!r}") from error if requirement.url is not None: + validate_url(requirement.url, field_name="direct dependency URL") parsed = urlsplit(requirement.url) valid_wheel = ( parsed.scheme == "https" @@ -332,7 +337,7 @@ def validate_readiness_path(cls, value: str) -> str: def validate_extra_args(cls, values: list[str]) -> list[str]: for value in values: validate_plain_text(value, field_name="vLLM argument") - flag = value.split("=", maxsplit=1)[0] + 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("-")): diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index 08bc3e3ff..e252ef987 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -75,6 +75,7 @@ def test_image_ref_requires_one_registered_alias_or_absolute_sqsh(payload: dict[ {"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=="]}, @@ -142,7 +143,16 @@ def test_remote_mcp_endpoint_requires_valid_host_and_port(endpoint: str) -> None RemoteMCPProviderConfig(provider_type="sse", name="provider", endpoint=endpoint) -@pytest.mark.parametrize("argument", ["--api-key", "--access-token=plaintext-secret", "password"]) +@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( @@ -169,11 +179,14 @@ def test_vllm_defaults_and_backpressure_override() -> None: "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", ], ) @@ -187,6 +200,7 @@ def test_vllm_rejects_runtime_owned_arguments(argument: str) -> None: [ ["--hf-token=plaintext-secret"], ["--hf-token", "plaintext-secret"], + ["--hf-token plaintext-secret"], ], ) def test_vllm_rejects_secret_shaped_arguments(extra_args: list[str]) -> None: