From 51c604cd959a0580ed102842e8c1ea7e893a3047 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 26 Aug 2026 13:36:35 -0300 Subject: [PATCH 1/9] feat: compile deterministic Slurm run plans Signed-off-by: Andre Manoel --- packages/data-designer-slurm/pyproject.toml | 1 + .../data_designer/slurm/config/__init__.py | 17 + .../src/data_designer/slurm/config/builder.py | 174 +++++++ .../src/data_designer/slurm/config/loading.py | 197 ++++++++ .../data_designer/slurm/planning/__init__.py | 14 + .../data_designer/slurm/planning/compiler.py | 471 ++++++++++++++++++ .../tests/config/test_loading_builder.py | 183 +++++++ .../tests/planning/test_compiler.py | 308 ++++++++++++ uv.lock | 2 + 9 files changed, 1367 insertions(+) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/builder.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/loading.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py create mode 100644 packages/data-designer-slurm/tests/config/test_loading_builder.py create mode 100644 packages/data-designer-slurm/tests/planning/test_compiler.py diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 340c8f289..9e4610eff 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "data-designer=={{ version }}", "packaging>=25,<27", "pydantic>=2.9.2,<3", + "pyyaml>=6.0.1,<7", ] [tool.hatch.build.targets.wheel] 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 index 5cda08684..66e2734cf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -14,6 +14,7 @@ DataDesignerSlurmBenchmarkConfig, FixedRecordPolicy, ) +from data_designer.slurm.config.builder import ConfigBuilderError, DataDesignerSlurmConfigBuilder from data_designer.slurm.config.images import ( ClientImageInspection, ImageBuildRequest, @@ -23,6 +24,14 @@ InstalledDistribution, ServingImageInspection, ) +from data_designer.slurm.config.loading import ( + DEFAULT_PROFILE_FILE_NAME, + PROFILE_FILE_ENVIRONMENT, + ConfigLoadError, + load_profile_catalog, + load_run_config, + resolve_profile, +) from data_designer.slurm.config.profiles import ( ContainerMount, GpuRequestMode, @@ -69,9 +78,13 @@ "ClientConfig", "ClientDependencies", "ClientImageInspection", + "ConfigBuilderError", + "ConfigLoadError", "ContainerMount", "DataDesignerSlurmBenchmarkConfig", "DataDesignerSlurmConfig", + "DataDesignerSlurmConfigBuilder", + "DEFAULT_PROFILE_FILE_NAME", "DeploymentResources", "DeploymentTopology", "FixedRecordPolicy", @@ -89,6 +102,7 @@ "LocalStdioMCPProviderConfig", "OutputConfig", "ProfileSelectionSource", + "PROFILE_FILE_ENVIRONMENT", "QueueBackpressureConfig", "RemoteMCPProviderConfig", "SchedulerProfile", @@ -101,6 +115,9 @@ "SubmissionConfig", "VllmServerConfig", "injected_profile", + "load_profile_catalog", + "load_run_config", + "resolve_profile", "select_profile", "validate_selected_profile", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py new file mode 100644 index 000000000..accb18bd4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure convenience builder for authored Slurm run declarations.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from pathlib import Path + +import yaml + +from data_designer.config import DataDesignerConfigBuilder +from data_designer.slurm.config.images import ImageRef +from data_designer.slurm.config.run import ( + ArrayTasksConfig, + BuilderInput, + ClientConfig, + ClientDependencies, + DataDesignerSlurmConfig, + InputBindings, + InvocationConfig, + InvocationDiagnostics, + MCPProviderConfig, + OutputConfig, + ServerDeploymentConfig, + SubmissionConfig, +) + + +class ConfigBuilderError(ValueError): + """Raised when the Slurm config builder is incomplete or cannot serialize.""" + + +class DataDesignerSlurmConfigBuilder: + """Build one strict authored Slurm run without resolving or submitting it.""" + + def __init__(self, builder: BuilderInput, *, name: str = "data-designer") -> None: + self._name = name + self._builder = builder + self._invocation: InvocationConfig | None = None + self._client: ClientConfig | None = None + self._deployments: list[ServerDeploymentConfig] = [] + self._array_tasks = ArrayTasksConfig() + self._submission = SubmissionConfig() + self._output = OutputConfig() + + @classmethod + def from_config_builder( + cls, + builder: DataDesignerConfigBuilder, + *, + name: str = "data-designer", + ) -> DataDesignerSlurmConfigBuilder: + """Start from one public Data Designer configuration builder.""" + return cls(BuilderInput(inline=builder.get_builder_config().to_dict()), name=name) + + @classmethod + def from_builder_source( + cls, + source: str, + *, + name: str = "data-designer", + ) -> DataDesignerSlurmConfigBuilder: + """Start from one local serialized Data Designer builder path.""" + return cls(BuilderInput(source=source), name=name) + + def with_invocation( + self, + *, + num_records: int, + dataset_name: str, + resume: str = "never", + run_config: Mapping[str, object] | None = None, + input_bindings: InputBindings | Mapping[str, object] | None = None, + mcp_providers: Sequence[MCPProviderConfig | Mapping[str, object]] = (), + model_concurrency: Mapping[str, int] | None = None, + diagnostics: InvocationDiagnostics | Mapping[str, object] | None = None, + ) -> DataDesignerSlurmConfigBuilder: + """Set typed Data Designer invocation intent.""" + self._invocation = InvocationConfig.model_validate( + { + "num_records": num_records, + "dataset_name": dataset_name, + "resume": resume, + "run_config": {} if run_config is None else dict(run_config), + "input_bindings": {} if input_bindings is None else input_bindings, + "mcp_providers": list(mcp_providers), + "model_concurrency": {} if model_concurrency is None else dict(model_concurrency), + "diagnostics": {} if diagnostics is None else diagnostics, + } + ) + return self + + def with_client( + self, + *, + image: ImageRef | Mapping[str, object], + cpus: int = 32, + dependencies: ClientDependencies | Mapping[str, object] | None = None, + ) -> DataDesignerSlurmConfigBuilder: + """Set the separate zero-GPU Data Designer client declaration.""" + self._client = ClientConfig.model_validate( + { + "cpus": cpus, + "image": image, + "dependencies": {} if dependencies is None else dependencies, + } + ) + return self + + def with_deployment( + self, + deployment: ServerDeploymentConfig | Mapping[str, object], + ) -> DataDesignerSlurmConfigBuilder: + """Append one deployment while preserving authored order.""" + self._deployments.append(ServerDeploymentConfig.model_validate(deployment)) + return self + + def with_array_tasks(self, *, count: int, max_concurrent: int = 1) -> DataDesignerSlurmConfigBuilder: + """Set deterministic horizontal sharding.""" + self._array_tasks = ArrayTasksConfig(count=count, max_concurrent=max_concurrent) + return self + + def with_submission(self, **values: object) -> DataDesignerSlurmConfigBuilder: + """Set typed Slurm submission intent.""" + self._submission = SubmissionConfig.model_validate(values) + return self + + def with_output(self, **values: object) -> DataDesignerSlurmConfigBuilder: + """Set typed dataset output intent.""" + self._output = OutputConfig.model_validate(values) + return self + + def build(self) -> DataDesignerSlurmConfig: + """Return the complete authored declaration without resolving ambient state.""" + missing = [] + if self._invocation is None: + missing.append("invocation") + if self._client is None: + missing.append("client") + if not self._deployments: + missing.append("deployment") + if missing: + raise ConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") + assert self._invocation is not None + assert self._client is not None + return DataDesignerSlurmConfig( + schema_version=1, + name=self._name, + builder=self._builder, + invocation=self._invocation, + client=self._client, + deployments=self._deployments, + array_tasks=self._array_tasks, + submission=self._submission, + output=self._output, + ) + + def write_config(self, path: str | Path) -> None: + """Serialize the authored declaration as deterministic JSON or YAML.""" + output_path = Path(path) + config = self.build() + if output_path.suffix == ".json": + contents = config.serialize_json() + elif output_path.suffix in {".yaml", ".yml"}: + contents = yaml.safe_dump( + config.model_dump(mode="json"), + default_flow_style=False, + sort_keys=True, + ) + else: + raise ConfigBuilderError("config path must end in .json, .yaml, or .yml") + output_path.write_text(contents, encoding="utf-8") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py new file mode 100644 index 000000000..9e127af67 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict local loading and cluster-profile selection.""" + +from __future__ import annotations + +import json +import os +import socket +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import TypeVar, cast + +import yaml +from pydantic import ValidationError +from yaml.nodes import MappingNode + +from data_designer.slurm.config.profiles import ( + SelectedSlurmProfile, + SlurmProfile, + SlurmProfileCatalog, + injected_profile, + select_profile, +) +from data_designer.slurm.config.run import DataDesignerSlurmConfig + +PROFILE_FILE_ENVIRONMENT = "DATA_DESIGNER_SLURM_PROFILE_FILE" +DEFAULT_PROFILE_FILE_NAME = ".data-designer-slurm-profile.yml" + +_Config = TypeVar("_Config", DataDesignerSlurmConfig, SlurmProfileCatalog) +_HostnameResolver = Callable[[], tuple[str, ...]] + + +class ConfigLoadError(ValueError): + """Raised when a local Slurm configuration file is not strict and valid.""" + + +class _StrictYamlLoader(yaml.SafeLoader): + pass + + +def _construct_unique_mapping( + loader: _StrictYamlLoader, + node: MappingNode, + deep: bool = False, +) -> dict[object, object]: + mapping: dict[object, object] = {} + for key_node, value_node in node.value: + if key_node.tag == "tag:yaml.org,2002:merge": + raise ConfigLoadError("YAML merge keys are not supported") + key = loader.construct_object(key_node, deep=deep) + try: + duplicate = key in mapping + except TypeError as error: + raise ConfigLoadError("configuration mapping keys must be scalar values") from error + if duplicate: + raise ConfigLoadError(f"duplicate configuration key: {key!r}") + mapping[key] = loader.construct_object(value_node, deep=deep) + return mapping + + +_StrictYamlLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, + _construct_unique_mapping, +) + + +def load_run_config(path: str | Path) -> DataDesignerSlurmConfig: + """Load one strict local YAML or JSON run declaration.""" + return _load_config(path, DataDesignerSlurmConfig) + + +def load_profile_catalog(path: str | Path) -> SlurmProfileCatalog: + """Load one strict local YAML or JSON cluster-profile catalog.""" + return _load_config(path, SlurmProfileCatalog) + + +def resolve_profile( + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + hostnames: tuple[str, ...] | None = None, + hostname_resolver: _HostnameResolver | None = None, + environ: Mapping[str, str] | None = None, + home_directory: str | Path | None = None, +) -> SelectedSlurmProfile: + """Resolve an injected profile or select one catalog entry.""" + sources = sum(source is not None for source in (profile, catalog, profile_file)) + if sources > 1: + raise ConfigLoadError("profile, catalog, and profile_file are mutually exclusive") + if profile is not None: + if cluster is not None: + raise ConfigLoadError("an injected profile cannot be combined with cluster selection") + return injected_profile(profile) + + catalog_path: str | None = None + if catalog is None: + path = _resolve_profile_path( + profile_file, + environ=os.environ if environ is None else environ, + home_directory=home_directory, + ) + catalog = load_profile_catalog(path) + catalog_path = path.as_posix() + + if hostnames is None: + resolver = hostname_resolver or _local_hostnames + hostnames = resolver() + normalized_hostnames = tuple(dict.fromkeys(hostname.strip().casefold() for hostname in hostnames if hostname)) + return select_profile( + catalog, + cluster=cluster, + hostnames=normalized_hostnames, + catalog_path=catalog_path, + ) + + +def _load_config(path: str | Path, config_type: type[_Config]) -> _Config: + resolved_path = _normalize_file_path(path) + try: + contents = resolved_path.read_text(encoding="utf-8") + except OSError as error: + raise ConfigLoadError(f"cannot read configuration file {resolved_path}") from error + try: + payload = _parse_mapping(contents, suffix=resolved_path.suffix) + _reject_environment_interpolation(payload) + return config_type.model_validate(payload) + except ConfigLoadError: + raise + except (ValidationError, json.JSONDecodeError, yaml.YAMLError) as error: + raise ConfigLoadError(f"invalid configuration file {resolved_path}: {error}") from error + + +def _parse_mapping(contents: str, *, suffix: str) -> dict[str, object]: + if suffix == ".json": + payload = json.loads(contents, object_pairs_hook=_unique_json_object) + else: + events = yaml.parse(contents, Loader=yaml.SafeLoader) + if any(getattr(event, "anchor", None) is not None for event in events): + raise ConfigLoadError("YAML anchors and aliases are not supported") + payload = yaml.load(contents, Loader=_StrictYamlLoader) + if not isinstance(payload, dict) or any(not isinstance(key, str) for key in payload): + raise ConfigLoadError("configuration root must be an object with string keys") + return cast(dict[str, object], payload) + + +def _unique_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise ConfigLoadError(f"duplicate configuration key: {key!r}") + result[key] = value + return result + + +def _reject_environment_interpolation(value: object) -> None: + if isinstance(value, str) and "${" in value: + raise ConfigLoadError("environment interpolation is not supported") + if isinstance(value, Mapping): + for key, item in value.items(): + _reject_environment_interpolation(key) + _reject_environment_interpolation(item) + elif isinstance(value, list | tuple): + for item in value: + _reject_environment_interpolation(item) + + +def _resolve_profile_path( + explicit_path: str | Path | None, + *, + environ: Mapping[str, str], + home_directory: str | Path | None, +) -> Path: + source = explicit_path + if source is None: + environment_path = environ.get(PROFILE_FILE_ENVIRONMENT) + if environment_path is not None and not environment_path: + raise ConfigLoadError(f"{PROFILE_FILE_ENVIRONMENT} must not be empty") + source = environment_path + if source is None: + home = Path.home() if home_directory is None else Path(home_directory) + source = home / DEFAULT_PROFILE_FILE_NAME + return _normalize_file_path(source) + + +def _normalize_file_path(path: str | Path) -> Path: + resolved = Path(path).expanduser().resolve() + if resolved.suffix not in {".json", ".yaml", ".yml"}: + raise ConfigLoadError("configuration path must end in .json, .yaml, or .yml") + return resolved + + +def _local_hostnames() -> tuple[str, ...]: + return socket.gethostname(), socket.getfqdn() 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 160b7ef77..a1c46d7e6 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 @@ -5,6 +5,14 @@ from __future__ import annotations +from data_designer.slurm.planning.compiler import ( + ConfigurationResolutionError, + EffectiveDataDesignerSlurmConfig, + PlanCompilationError, + SlurmRunCompiler, + compile_slurm_run_plan, + resolve_slurm_config, +) from data_designer.slurm.planning.models import ( ArtifactReference, LockedPackage, @@ -27,8 +35,11 @@ __all__ = [ "ArtifactReference", + "ConfigurationResolutionError", + "EffectiveDataDesignerSlurmConfig", "LockedPackage", "PlanContractError", + "PlanCompilationError", "PlannedShard", "PortClaim", "RecordRange", @@ -43,5 +54,8 @@ "ResolvedSubmission", "ResolvedTopology", "ResumeWorkspace", + "SlurmRunCompiler", + "compile_slurm_run_plan", + "resolve_slurm_config", "validate_resolved_plan", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py new file mode 100644 index 000000000..0f1f89caf --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -0,0 +1,471 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure authored-configuration resolution and deterministic plan compilation.""" + +from __future__ import annotations + +import hashlib +import posixpath +from typing import Annotated + +from pydantic import JsonValue, PositiveInt, StringConstraints, model_validator + +from data_designer.config import RunConfig +from data_designer.slurm.config.images import ClientImageInspection, ImageKind +from data_designer.slurm.config.profiles import SelectedSlurmProfile +from data_designer.slurm.config.run import BuilderInput, DataDesignerSlurmConfig +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractValue, + Identifier, + RecordRange, + ResumeWorkspace, + compute_sha256, + pretty_json, +) +from data_designer.slurm.planning.models import ( + PlannedShard, + PortClaim, + ResolvedBuilderInput, + ResolvedClient, + ResolvedDependencyLock, + ResolvedDeployment, + ResolvedImage, + ResolvedInvocation, + ResolvedOutput, + ResolvedSlurmRunPlan, + ResolvedSubmission, + ResolvedTopology, + _extract_builder_aliases, +) +from data_designer.slurm.planning.validation import validate_resolved_plan + +_LOGICAL_ENDPOINT_PORT = 17000 +_HTTP_PORT = 18000 +_RENDEZVOUS_PORT = 19000 +_PORT_RANGE_SIZE = 1000 +_COMPATIBILITY_RUN_DEFAULTS: dict[str, JsonValue] = { + "buffer_size": 16384, + "disable_early_shutdown": True, + "display_tui": False, + "max_conversation_correction_steps": 0, + "max_conversation_restarts": 0, + "otel_metrics_port": None, + "shutdown_error_rate": 1.0, +} + + +class ConfigurationResolutionError(ValueError): + """Raised when resolved inputs do not match one authored declaration.""" + + +class PlanCompilationError(ValueError): + """Raised when one effective configuration cannot produce a valid plan.""" + + +class EffectiveDataDesignerSlurmConfig(ContractValue): + """Fully materialized, side-effect-free input to the plan compiler.""" + + run_id: Identifier + package_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + authored: DataDesignerSlurmConfig + selected_profile: SelectedSlurmProfile + resolved_gpus_per_node: PositiveInt + builder: ResolvedBuilderInput + builder_payload: dict[str, JsonValue] | None = None + invocation: ResolvedInvocation + client_image: ResolvedImage + deployment_images: tuple[ResolvedImage, ...] + dependency_lock: ResolvedDependencyLock + submission: ResolvedSubmission + output: ResolvedOutput + runtime_bundle: ArtifactReference + + @model_validator(mode="after") + def validate_resolution(self) -> EffectiveDataDesignerSlurmConfig: + profile_gpus = self.selected_profile.profile.gpus_per_node + if profile_gpus != "auto" and profile_gpus != self.resolved_gpus_per_node: + raise ValueError("resolved GPU count does not match the selected profile") + if self.client_image.kind is not ImageKind.CLIENT: + raise ValueError("resolved client image must contain client inspection facts") + if self.client_image.authored_ref != self.authored.client.image: + raise ValueError("resolved client image does not match the authored reference") + if len(self.deployment_images) != len(self.authored.deployments): + raise ValueError("resolved serving images must match the authored deployment count") + for deployment, image in zip(self.authored.deployments, self.deployment_images, strict=True): + if image.kind is not ImageKind.SERVING: + raise ValueError("resolved deployment image must contain serving inspection facts") + if image.authored_ref != deployment.server.image: + raise ValueError("resolved deployment image does not match the authored reference") + if self.builder_payload is not None and self.authored.builder.source is None: + raise ValueError("only sourced builder input may retain a resolved payload") + if self.authored.builder.source is not None and self.builder_payload is None: + raise ValueError("sourced builder input requires its resolved payload") + runtime_root = posixpath.join(self.selected_profile.profile.workspace_root, "runtime") + if not _is_below(self.runtime_bundle.path, runtime_root) or not self.runtime_bundle.path.endswith(".tar.gz"): + raise ValueError("runtime bundle must be a tar archive below the selected workspace runtime root") + return self + + +def resolve_slurm_config( + authored: DataDesignerSlurmConfig, + *, + selected_profile: SelectedSlurmProfile, + client_image: ResolvedImage, + deployment_images: tuple[ResolvedImage, ...], + dependency_lock: ResolvedDependencyLock, + runtime_bundle: ArtifactReference, + run_id: str, + package_version: str, + resolved_gpus_per_node: int | None = None, + builder_payload: dict[str, JsonValue] | None = None, +) -> EffectiveDataDesignerSlurmConfig: + """Materialize all non-secret defaults from explicitly supplied resolved inputs.""" + try: + gpus_per_node = _resolve_gpu_count(selected_profile, resolved_gpus_per_node) + run_root = posixpath.join(selected_profile.profile.workspace_root, "runs", run_id) + builder = _resolve_builder(authored, run_root=run_root, builder_payload=builder_payload) + invocation = ResolvedInvocation( + authored=authored.invocation, + effective_run_config=_materialize_run_config(authored), + ) + submission = ResolvedSubmission( + account=authored.submission.account or selected_profile.profile.scheduler.account, + partition=authored.submission.partition or selected_profile.profile.scheduler.partition, + job_name=authored.submission.job_name, + time_limit=authored.submission.time_limit, + comment=authored.submission.comment, + ) + output_root = authored.output.root or posixpath.join(run_root, "output") + _validate_output_destination(output_root, selected_profile.profile.workspace_root) + if authored.output.partitions > authored.invocation.num_records: + raise ConfigurationResolutionError("output partitions must not exceed requested records") + output = ResolvedOutput( + root=output_root, + format=authored.output.format, + partitions=authored.output.partitions, + require_exact_record_count=authored.output.require_exact_record_count, + ) + _validate_dependency_resolution(authored, client_image, dependency_lock) + return EffectiveDataDesignerSlurmConfig( + run_id=run_id, + package_version=package_version, + authored=authored, + selected_profile=selected_profile, + resolved_gpus_per_node=gpus_per_node, + builder=builder, + builder_payload=builder_payload, + invocation=invocation, + client_image=client_image, + deployment_images=deployment_images, + dependency_lock=dependency_lock, + submission=submission, + output=output, + runtime_bundle=runtime_bundle, + ) + except ConfigurationResolutionError: + raise + except ValueError as error: + raise ConfigurationResolutionError(str(error)) from error + + +class SlurmRunCompiler: + """Compile one fully resolved configuration without ambient state or I/O.""" + + @staticmethod + def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: + """Return one immutable deterministic execution plan.""" + try: + deployments = _compile_deployments(effective) + client = _compile_client(effective, deployments) + _validate_port_claims(effective, client, deployments) + plan = ResolvedSlurmRunPlan( + schema_version=1, + run_id=effective.run_id, + package_version=effective.package_version, + authored_config=ArtifactReference( + path=posixpath.join(_run_root(effective), "authored-config.json"), + sha256=effective.authored.compute_sha256(), + ), + selected_profile=effective.selected_profile, + resolved_gpus_per_node=effective.resolved_gpus_per_node, + builder=effective.builder, + invocation=effective.invocation, + client=client, + deployments=deployments, + array_tasks=effective.authored.array_tasks, + shards=_compile_shards(effective), + submission=effective.submission, + output=effective.output, + container_mounts=tuple(effective.selected_profile.profile.container_mounts), + runtime_bundle=effective.runtime_bundle, + ) + return validate_resolved_plan( + effective.authored, + effective.dependency_lock, + plan, + builder_payload=effective.builder_payload, + ) + except PlanCompilationError: + raise + except ValueError as error: + raise PlanCompilationError(str(error)) from error + + +def compile_slurm_run_plan(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: + """Compile one effective configuration with the package-owned compiler.""" + return SlurmRunCompiler.compile(effective) + + +def _resolve_gpu_count(selected: SelectedSlurmProfile, resolved: int | None) -> int: + configured = selected.profile.gpus_per_node + if configured == "auto": + if type(resolved) is not int or resolved <= 0: + raise ConfigurationResolutionError("auto gpus_per_node requires one resolved positive integer") + return resolved + if resolved is not None and resolved != configured: + raise ConfigurationResolutionError("resolved GPU count conflicts with the selected profile") + return configured + + +def _resolve_builder( + authored: DataDesignerSlurmConfig, + *, + run_root: str, + builder_payload: dict[str, JsonValue] | None, +) -> ResolvedBuilderInput: + if authored.builder.inline is not None: + if builder_payload is not None: + raise ConfigurationResolutionError("inline builder input must not provide a separate payload") + aliases, referenced_aliases = _extract_builder_aliases(authored.builder.inline) + return ResolvedBuilderInput( + inline=authored.builder.inline, + content_sha256=compute_sha256(authored.builder.inline), + model_aliases=aliases, + referenced_model_aliases=referenced_aliases, + ) + if builder_payload is None: + raise ConfigurationResolutionError("sourced builder input requires its resolved payload") + validated_payload = BuilderInput(inline=builder_payload).inline + assert validated_payload is not None + aliases, referenced_aliases = _extract_builder_aliases(validated_payload) + serialized = pretty_json(validated_payload).encode("utf-8") + source = ArtifactReference( + path=posixpath.join(run_root, "builder-config.json"), + sha256=hashlib.sha256(serialized).hexdigest(), + ) + return ResolvedBuilderInput( + authored_source=authored.builder.source, + source=source, + content_sha256=source.sha256, + model_aliases=aliases, + referenced_model_aliases=referenced_aliases, + ) + + +def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: + values = dict(authored.invocation.run_config) + for name, value in _COMPATIBILITY_RUN_DEFAULTS.items(): + values.setdefault(name, value) + return RunConfig.model_validate(values).model_dump(mode="json") + + +def _validate_dependency_resolution( + authored: DataDesignerSlurmConfig, + client_image: ResolvedImage, + dependency_lock: ResolvedDependencyLock, +) -> None: + inspection = client_image.inspection.inspection + if not isinstance(inspection, ClientImageInspection): + raise ConfigurationResolutionError("resolved client image lacks dependency inspection facts") + if dependency_lock.client_image_sha256 != client_image.sha256: + raise ConfigurationResolutionError("dependency lock does not match the resolved client image") + if dependency_lock.python_abi != inspection.python_abi: + raise ConfigurationResolutionError("dependency lock Python ABI does not match the client image") + if dependency_lock.image_distributions != inspection.distributions: + raise ConfigurationResolutionError("dependency lock inventory does not match the client image") + requirements = authored.client.dependencies.requirements + if requirements is not None: + if dependency_lock.authored_source is not None or dependency_lock.source is not None: + raise ConfigurationResolutionError("inline requirements cannot resolve from an authored lock file") + if dependency_lock.authored_requirements != tuple(requirements): + raise ConfigurationResolutionError("dependency lock requirements do not match authored requirements") + elif dependency_lock.authored_source != authored.client.dependencies.lock_file or dependency_lock.source is None: + raise ConfigurationResolutionError("dependency lock source does not match the authored lock file") + + +def _compile_deployments( + effective: EffectiveDataDesignerSlurmConfig, +) -> tuple[ResolvedDeployment, ...]: + if len(effective.authored.deployments) > _PORT_RANGE_SIZE: + raise PlanCompilationError("deployment count exceeds the compiler-owned logical endpoint port range") + resolved: list[ResolvedDeployment] = [] + next_node_index = 0 + for index, (authored, image) in enumerate( + zip(effective.authored.deployments, effective.deployment_images, strict=True) + ): + tensor_parallel = authored.topology.tensor_parallel + if effective.resolved_gpus_per_node % tensor_parallel: + raise PlanCompilationError("tensor_parallel must divide resolved GPUs per node") + replicas_per_group = effective.resolved_gpus_per_node // tensor_parallel + if replicas_per_group > _PORT_RANGE_SIZE: + raise PlanCompilationError("replica lanes exceed the compiler-owned deployment port range") + node_group_count = authored.resources.nodes // authored.topology.nodes_per_replica + replica_count = node_group_count * replicas_per_group + topology = ResolvedTopology( + tensor_parallel=tensor_parallel, + nodes_per_replica=authored.topology.nodes_per_replica, + pipeline_parallel=authored.topology.nodes_per_replica, + node_group_count=node_group_count, + replicas_per_node_group=replicas_per_group, + replica_count=replica_count, + gpus_per_replica=tensor_parallel * authored.topology.nodes_per_replica, + ) + deployment_id = f"deployment-{index:05d}" + node_indices = tuple(range(next_node_index, next_node_index + authored.resources.nodes)) + ports = _compile_deployment_ports(deployment_id, node_indices, topology) + resolved.append( + ResolvedDeployment( + deployment_id=deployment_id, + authored=authored, + served_model_name=authored.served_model_name or authored.model, + image=image, + node_indices=node_indices, + gpus_per_node=effective.resolved_gpus_per_node, + topology=topology, + ports=ports, + ) + ) + next_node_index += authored.resources.nodes + return tuple(resolved) + + +def _compile_deployment_ports( + deployment_id: str, + node_indices: tuple[int, ...], + topology: ResolvedTopology, +) -> tuple[PortClaim, ...]: + http: list[PortClaim] = [] + rendezvous: list[PortClaim] = [] + for group_index in range(topology.node_group_count): + head = node_indices[group_index * topology.nodes_per_replica] + for lane_index in range(topology.replicas_per_node_group): + replica_index = group_index * topology.replicas_per_node_group + lane_index + http.append( + PortClaim( + name=f"{deployment_id}-http-{replica_index:05d}", + role="http", + node_index=head, + port=_HTTP_PORT + lane_index, + ) + ) + if topology.nodes_per_replica > 1: + rendezvous.append( + PortClaim( + name=f"{deployment_id}-rendezvous-{replica_index:05d}", + role="rendezvous", + node_index=head, + port=_RENDEZVOUS_PORT + lane_index, + ) + ) + return tuple(http + rendezvous) + + +def _compile_client( + effective: EffectiveDataDesignerSlurmConfig, + deployments: tuple[ResolvedDeployment, ...], +) -> ResolvedClient: + host_node_index = deployments[0].node_indices[0] + ports = tuple( + PortClaim( + name=f"{deployment.deployment_id}-logical-endpoint", + role="logical_endpoint", + node_index=host_node_index, + port=_LOGICAL_ENDPOINT_PORT + index, + ) + for index, deployment in enumerate(deployments) + ) + return ResolvedClient( + authored=effective.authored.client, + image=effective.client_image, + dependency_lock=ArtifactReference( + path=posixpath.join(_run_root(effective), "dependency-lock.json"), + sha256=effective.dependency_lock.compute_sha256(), + ), + host_node_index=host_node_index, + gpu_count=0, + ports=ports, + ) + + +def _validate_port_claims( + effective: EffectiveDataDesignerSlurmConfig, + client: ResolvedClient, + deployments: tuple[ResolvedDeployment, ...], +) -> None: + ports = client.ports + tuple(port for deployment in deployments for port in deployment.ports) + addresses = tuple((port.node_index, port.port) for port in ports) + if len(addresses) != len(set(addresses)): + raise PlanCompilationError("compiler-owned port claims collide on one node") + otel_port = effective.invocation.effective_run_config.get("otel_metrics_port") + if type(otel_port) is int and (client.host_node_index, otel_port) in addresses: + raise PlanCompilationError("client OTEL metrics port collides with a compiler-owned port") + + +def _compile_shards(effective: EffectiveDataDesignerSlurmConfig) -> tuple[PlannedShard, ...]: + count = effective.authored.array_tasks.count + requested = effective.authored.invocation.num_records + floor_count = requested // count + start = 0 + shards: list[PlannedShard] = [] + for index in range(count): + record_count = requested - floor_count * (count - 1) if index == count - 1 else floor_count + end = start + record_count + shard_id = f"shard-{index:05d}" + shard_root = posixpath.join(_run_root(effective), "shards", shard_id) + record_range = RecordRange(start_index=start, end_index_exclusive=end) + partition = None + seed_path = effective.authored.invocation.input_bindings.seed_path + if seed_path is not None: + partition = ArtifactReference( + path=posixpath.join(shard_root, "input-partition.json"), + sha256=compute_sha256( + { + "record_range": record_range.model_dump(mode="json"), + "seed_path": seed_path, + } + ), + ) + shards.append( + PlannedShard( + shard_id=shard_id, + shard_index=index, + array_task_index=index, + record_range=record_range, + input_partition=partition, + resume_workspace=ResumeWorkspace(path=posixpath.join(shard_root, "dataset")), + ) + ) + start = end + return tuple(shards) + + +def _validate_output_destination(output_root: str, workspace_root: str) -> None: + if not _is_below(output_root, workspace_root): + raise ConfigurationResolutionError("output root must be below the selected workspace_root") + reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) + if any(_paths_overlap(output_root, path) for path in reserved): + raise ConfigurationResolutionError("output root must not overlap package-managed workspace state") + + +def _run_root(effective: EffectiveDataDesignerSlurmConfig) -> str: + return posixpath.join(effective.selected_profile.profile.workspace_root, "runs", effective.run_id) + + +def _is_below(path: str, root: str) -> bool: + return path != root and posixpath.commonpath((path, root)) == root + + +def _paths_overlap(left: str, right: str) -> bool: + return left == right or _is_below(left, right) or _is_below(right, left) diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py new file mode 100644 index 000000000..98fb8b610 --- /dev/null +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -0,0 +1,183 @@ +# 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 +import yaml + +from data_designer.config import DataDesignerConfigBuilder, ModelConfig +from data_designer.slurm.config import ( + DEFAULT_PROFILE_FILE_NAME, + PROFILE_FILE_ENVIRONMENT, + ConfigBuilderError, + ConfigLoadError, + DataDesignerSlurmConfig, + DataDesignerSlurmConfigBuilder, + ProfileSelectionSource, + SlurmProfileCatalog, + load_profile_catalog, + load_run_config, + resolve_profile, +) + + +def _config_builder() -> DataDesignerSlurmConfigBuilder: + data_designer = DataDesignerConfigBuilder( + model_configs=[ModelConfig(alias="generator", model="example/generator", provider="openai")] + ) + return ( + DataDesignerSlurmConfigBuilder.from_config_builder(data_designer, name="generated-run") + .with_invocation( + num_records=8, + dataset_name="generated", + model_concurrency={"generator": 8}, + ) + .with_client(image={"name": "dd-client"}) + .with_deployment( + { + "model_alias": "generator", + "model": "example/generator", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + "topology": {"tensor_parallel": 8}, + } + ) + ) + + +def test_builder_builds_without_file_discovery_or_serialization(tmp_path: Path) -> None: + builder = _config_builder() + + config = builder.build() + + assert config.name == "generated-run" + assert config.builder.inline is not None + assert config.deployments[0].model_alias == "generator" + assert not tuple(tmp_path.iterdir()) + + +def test_builder_requires_complete_authored_intent() -> None: + builder = DataDesignerSlurmConfigBuilder.from_builder_source("builder.json") + + with pytest.raises(ConfigBuilderError, match="invocation, client, deployment"): + builder.build() + + +@pytest.mark.parametrize("suffix", [".json", ".yaml", ".yml"]) +def test_builder_write_config_round_trips_supported_formats(tmp_path: Path, suffix: str) -> None: + builder = _config_builder() + path = tmp_path / f"run{suffix}" + + builder.write_config(path) + + assert load_run_config(path) == builder.build() + + +def test_builder_rejects_unsupported_output_format(tmp_path: Path) -> None: + with pytest.raises(ConfigBuilderError, match="must end"): + _config_builder().write_config(tmp_path / "run.txt") + + +@pytest.mark.parametrize( + ("suffix", "contents", "message"), + [ + (".json", '{"schema_version": 1, "schema_version": 1}', "duplicate"), + (".yaml", "schema_version: 1\nschema_version: 1\n", "duplicate"), + (".yaml", "defaults: &defaults\n schema_version: 1\nrun: *defaults\n", "anchors"), + (".yaml", "schema_version: 1\nname: ${RUN_NAME}\n", "interpolation"), + ], +) +def test_strict_loader_rejects_ambiguous_yaml_and_json( + tmp_path: Path, + suffix: str, + contents: str, + message: str, +) -> None: + path = tmp_path / f"run{suffix}" + path.write_text(contents) + + with pytest.raises(ConfigLoadError, match=message): + load_run_config(path) + + +def test_strict_loader_rejects_non_object_and_unknown_extension(tmp_path: Path) -> None: + json_path = tmp_path / "run.json" + json_path.write_text("[]") + + with pytest.raises(ConfigLoadError, match="root must be an object"): + load_run_config(json_path) + with pytest.raises(ConfigLoadError, match="must end"): + load_run_config(tmp_path / "run.toml") + + +def test_profile_source_and_selection_precedence( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, +) -> None: + explicit_path = tmp_path / "explicit.json" + environment_path = tmp_path / "environment.yaml" + default_path = tmp_path / DEFAULT_PROFILE_FILE_NAME + explicit_path.write_text(profile_catalog.serialize_json()) + environment_path.write_text(yaml.safe_dump(profile_catalog.model_dump(mode="json"), sort_keys=True)) + default_path.write_text(profile_catalog.serialize_json()) + + explicit = resolve_profile( + profile_file=explicit_path, + cluster="lab", + hostname_resolver=lambda: ("primary-login-1",), + environ={PROFILE_FILE_ENVIRONMENT: str(environment_path)}, + ) + environment = resolve_profile( + hostname_resolver=lambda: ("LAB-LOGIN-1",), + environ={PROFILE_FILE_ENVIRONMENT: str(environment_path)}, + ) + default = resolve_profile( + hostname_resolver=lambda: ("unmatched",), + environ={}, + home_directory=tmp_path, + ) + + assert (explicit.cluster_name, explicit.selection_source) == ("lab", ProfileSelectionSource.EXPLICIT) + assert explicit.catalog_path == explicit_path.as_posix() + assert (environment.cluster_name, environment.selection_source) == ( + "lab", + ProfileSelectionSource.HOSTNAME, + ) + assert (default.cluster_name, default.selection_source) == ("primary", ProfileSelectionSource.DEFAULT) + assert load_profile_catalog(default_path) == profile_catalog + + +def test_injected_profile_bypasses_catalog_lookup(profile_catalog: SlurmProfileCatalog) -> None: + selected = resolve_profile( + profile=profile_catalog.clusters["primary"], + hostname_resolver=lambda: pytest.fail("hostname lookup must not run"), + environ={PROFILE_FILE_ENVIRONMENT: "/missing/profile.json"}, + ) + + assert selected.selection_source is ProfileSelectionSource.INJECTED + assert selected.catalog_path is None + + +def test_profile_resolution_rejects_conflicting_or_empty_sources( + profile_catalog: SlurmProfileCatalog, +) -> None: + with pytest.raises(ConfigLoadError, match="mutually exclusive"): + resolve_profile(catalog=profile_catalog, profile_file="profile.json") + with pytest.raises(ConfigLoadError, match="must not be empty"): + resolve_profile(environ={PROFILE_FILE_ENVIRONMENT: ""}) + with pytest.raises(ConfigLoadError, match="cluster selection"): + resolve_profile(profile=profile_catalog.clusters["primary"], cluster="primary") + + +def test_json_builder_output_is_stable(tmp_path: Path) -> None: + path = tmp_path / "run.json" + builder = _config_builder() + + builder.write_config(path) + + assert path.read_text() == builder.build().serialize_json() + assert json.loads(path.read_text())["schema_version"] == 1 + assert DataDesignerSlurmConfig.model_validate_json(path.read_text()) == builder.build() diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py new file mode 100644 index 000000000..bc26326bb --- /dev/null +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -0,0 +1,308 @@ +# 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 ( + BuilderInput, + DataDesignerSlurmConfig, + InputBindings, + SlurmProfileCatalog, + select_profile, +) +from data_designer.slurm.planning import ( + ArtifactReference, + ConfigurationResolutionError, + EffectiveDataDesignerSlurmConfig, + PlanCompilationError, + ResolvedDependencyLock, + ResolvedSlurmRunPlan, + compile_slurm_run_plan, + resolve_slurm_config, +) + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "contracts" / "golden" + + +def _resolve_fixture( + authored: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + expected: ResolvedSlurmRunPlan, + **updates: object, +) -> EffectiveDataDesignerSlurmConfig: + values = { + "selected_profile": expected.selected_profile, + "client_image": expected.client.image, + "deployment_images": tuple(deployment.image for deployment in expected.deployments), + "dependency_lock": dependency_lock, + "runtime_bundle": expected.runtime_bundle, + "run_id": expected.run_id, + "package_version": expected.package_version, + "resolved_gpus_per_node": expected.resolved_gpus_per_node, + } + values.update(updates) + return resolve_slurm_config(authored, **values) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("authored_fixture", "lock_fixture", "plan_fixture", "golden_name"), + [ + ("authored_run_single", "dependency_lock_single", "single_node_plan", "single_node_plan.json"), + ("authored_run", "dependency_lock", "multi_node_plan", "multi_node_plan.json"), + ], +) +def test_compiler_reproduces_plan_goldens_byte_for_byte( + request: pytest.FixtureRequest, + authored_fixture: str, + lock_fixture: str, + plan_fixture: str, + golden_name: str, +) -> None: + authored = request.getfixturevalue(authored_fixture) + dependency_lock = request.getfixturevalue(lock_fixture) + expected = request.getfixturevalue(plan_fixture) + effective = _resolve_fixture(authored, dependency_lock, expected) + + first = compile_slurm_run_plan(effective) + second = compile_slurm_run_plan(effective) + + assert first.serialize_json() == (GOLDEN_DIRECTORY / golden_name).read_text() + assert first.serialize_canonical_json() == second.serialize_canonical_json() + assert first.compute_sha256() == second.compute_sha256() == expected.compute_sha256() + + +def test_compiler_resolves_explicit_hostname_and_default_profile_selection( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + profile_catalog: SlurmProfileCatalog, +) -> None: + selections = ( + select_profile(profile_catalog, cluster="primary", hostnames=("lab-login-1",)), + select_profile(profile_catalog, hostnames=("PRIMARY-LOGIN-1",)), + select_profile(profile_catalog, hostnames=("unmatched",)), + ) + + plans = tuple( + compile_slurm_run_plan( + _resolve_fixture( + authored_run_single, + dependency_lock_single, + single_node_plan, + selected_profile=selected, + ) + ) + for selected in selections + ) + + assert {plan.selected_profile.selection_source.value for plan in plans} == { + "explicit", + "hostname", + "default", + } + assert {plan.output.root for plan in plans} == {"/workspace/primary/runs/run-single/output"} + assert all(plan.deployments[0].node_indices == (0,) for plan in plans) + + +def test_auto_gpu_resolution_is_explicit_and_scheduler_free( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + profile_catalog: SlurmProfileCatalog, +) -> None: + selected = select_profile(profile_catalog, cluster="lab") + runtime_bundle = single_node_plan.runtime_bundle.model_copy( + update={"path": "/workspace/lab/runtime/runtime.tar.gz"} + ) + + with pytest.raises(ConfigurationResolutionError, match="auto gpus_per_node"): + _resolve_fixture( + authored_run_single, + dependency_lock_single, + single_node_plan, + selected_profile=selected, + runtime_bundle=runtime_bundle, + resolved_gpus_per_node=None, + ) + + plan = compile_slurm_run_plan( + _resolve_fixture( + authored_run_single, + dependency_lock_single, + single_node_plan, + selected_profile=selected, + runtime_bundle=runtime_bundle, + resolved_gpus_per_node=8, + ) + ) + + assert plan.resolved_gpus_per_node == 8 + assert plan.output.root == "/workspace/lab/runs/run-single/output" + assert plan.submission.account == "lab" + + +def test_compiler_preserves_explicit_compatibility_run_values( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["run_config"] = { + "buffer_size": 2048, + "disable_early_shutdown": False, + "max_conversation_restarts": 3, + "otel_metrics_port": 24000, + "shutdown_error_rate": 0.25, + } + authored = DataDesignerSlurmConfig.model_validate(payload) + + effective = _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + assert effective.invocation.effective_run_config["buffer_size"] == 2048 + assert effective.invocation.effective_run_config["disable_early_shutdown"] is False + assert effective.invocation.effective_run_config["max_conversation_restarts"] == 3 + assert effective.invocation.effective_run_config["otel_metrics_port"] == 24000 + assert effective.invocation.effective_run_config["shutdown_error_rate"] == 0.25 + + +def test_compiler_rejects_tensor_parallelism_that_does_not_divide_gpu_shape( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["deployments"][0]["topology"]["tensor_parallel"] = 3 + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(PlanCompilationError, match="tensor_parallel"): + compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + +@pytest.mark.parametrize( + "output_update", + [ + {"root": "/outside/output"}, + {"root": "/workspace/primary/images/output"}, + {"root": "/workspace/primary/runtime/output"}, + {"partitions": 9}, + ], +) +def test_resolution_rejects_invalid_output_destinations( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + output_update: dict[str, object], +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["output"].update(output_update) + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(ConfigurationResolutionError, match="output"): + _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + +def test_compiler_rejects_model_alias_missing_from_builder( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["model_configs"][0]["alias"] = "other" + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(PlanCompilationError, match="deployment alias"): + compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + +def test_compiler_rejects_otel_collision_before_runtime( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["run_config"] = {"otel_metrics_port": 17000} + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(PlanCompilationError, match="OTEL"): + compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + +def test_sharded_seed_inputs_have_stable_ranges_and_partition_digests( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invocation = authored_run.invocation.model_copy( + update={"input_bindings": InputBindings(seed_path="/datasets/seed.parquet")} + ) + authored = authored_run.model_copy(update={"invocation": invocation}) + + first = compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + second = compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + + assert [(shard.record_range.start_index, shard.record_range.end_index_exclusive) for shard in first.shards] == [ + (0, 50), + (50, 100), + ] + assert all(shard.input_partition is not None for shard in first.shards) + assert first.shards == second.shards + + +def test_sourced_builder_is_resolved_to_one_digest_bound_run_input( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + builder_payload = authored_run_single.model_dump(mode="json")["builder"]["inline"] + assert isinstance(builder_payload, dict) + authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) + + plan = compile_slurm_run_plan( + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + ) + + assert plan.builder.authored_source == "builder.json" + assert plan.builder.source is not None + assert plan.builder.source.path == "/workspace/primary/runs/run-single/builder-config.json" + assert plan.builder.content_sha256 == plan.builder.source.sha256 + + +def test_resolution_rejects_artifact_identity_mismatches( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + wrong_lock = dependency_lock_single.model_copy(update={"client_image_sha256": "a" * 64}) + wrong_runtime = ArtifactReference(path="/tmp/runtime.tar.gz", sha256="e" * 64) + + with pytest.raises(ConfigurationResolutionError, match="client image"): + _resolve_fixture(authored_run_single, wrong_lock, single_node_plan) + with pytest.raises(ConfigurationResolutionError, match="runtime bundle"): + _resolve_fixture( + authored_run_single, + dependency_lock_single, + single_node_plan, + runtime_bundle=wrong_runtime, + ) + + +def test_plan_contains_secret_references_without_credentials( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = compile_slurm_run_plan(_resolve_fixture(authored_run, dependency_lock, multi_node_plan)) + serialized = plan.serialize_json() + + assert "PACKAGE_INDEX_TOKEN" in serialized + assert "HF_TOKEN" in serialized + assert "secret-value" not in serialized diff --git a/uv.lock b/uv.lock index 7b6c3394b..794b09e35 100644 --- a/uv.lock +++ b/uv.lock @@ -975,6 +975,7 @@ dependencies = [ { name = "data-designer" }, { name = "packaging" }, { name = "pydantic" }, + { name = "pyyaml" }, ] [package.metadata] @@ -982,6 +983,7 @@ requires-dist = [ { name = "data-designer", editable = "packages/data-designer" }, { name = "packaging", specifier = ">=25,<27" }, { name = "pydantic", specifier = ">=2.9.2,<3" }, + { name = "pyyaml", specifier = ">=6.0.1,<7" }, ] [[package]] From fe796f67ab966f72bc06f26527c5853ae0b9b0f6 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 26 Aug 2026 14:35:53 -0300 Subject: [PATCH 2/9] fix: enforce Slurm planning invariants Reject unsupported multi-shard semantics and managed output collisions at the effective configuration boundary. Normalize config errors and cover the corrected Big Iron field disposition. Signed-off-by: Andre Manoel --- .../src/data_designer/slurm/config/builder.py | 64 ++++-- .../src/data_designer/slurm/config/loading.py | 78 ++++--- .../data_designer/slurm/planning/compiler.py | 63 +++++- .../tests/config/test_loading_builder.py | 57 ++++- .../tests/planning/test_compiler.py | 207 ++++++++++++++++++ 5 files changed, 412 insertions(+), 57 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py index accb18bd4..9cbb9fb21 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -7,8 +7,10 @@ from collections.abc import Mapping, Sequence from pathlib import Path +from typing import TypeVar import yaml +from pydantic import BaseModel, ValidationError from data_designer.config import DataDesignerConfigBuilder from data_designer.slurm.config.images import ImageRef @@ -27,6 +29,8 @@ SubmissionConfig, ) +_ConfigValue = TypeVar("_ConfigValue", bound=BaseModel) + class ConfigBuilderError(ValueError): """Raised when the Slurm config builder is incomplete or cannot serialize.""" @@ -53,7 +57,7 @@ def from_config_builder( name: str = "data-designer", ) -> DataDesignerSlurmConfigBuilder: """Start from one public Data Designer configuration builder.""" - return cls(BuilderInput(inline=builder.get_builder_config().to_dict()), name=name) + return cls(_validate_model(BuilderInput, {"inline": builder.get_builder_config().to_dict()}), name=name) @classmethod def from_builder_source( @@ -63,7 +67,7 @@ def from_builder_source( name: str = "data-designer", ) -> DataDesignerSlurmConfigBuilder: """Start from one local serialized Data Designer builder path.""" - return cls(BuilderInput(source=source), name=name) + return cls(_validate_model(BuilderInput, {"source": source}), name=name) def with_invocation( self, @@ -78,7 +82,8 @@ def with_invocation( diagnostics: InvocationDiagnostics | Mapping[str, object] | None = None, ) -> DataDesignerSlurmConfigBuilder: """Set typed Data Designer invocation intent.""" - self._invocation = InvocationConfig.model_validate( + self._invocation = _validate_model( + InvocationConfig, { "num_records": num_records, "dataset_name": dataset_name, @@ -88,7 +93,7 @@ def with_invocation( "mcp_providers": list(mcp_providers), "model_concurrency": {} if model_concurrency is None else dict(model_concurrency), "diagnostics": {} if diagnostics is None else diagnostics, - } + }, ) return self @@ -100,12 +105,13 @@ def with_client( dependencies: ClientDependencies | Mapping[str, object] | None = None, ) -> DataDesignerSlurmConfigBuilder: """Set the separate zero-GPU Data Designer client declaration.""" - self._client = ClientConfig.model_validate( + self._client = _validate_model( + ClientConfig, { "cpus": cpus, "image": image, "dependencies": {} if dependencies is None else dependencies, - } + }, ) return self @@ -114,22 +120,25 @@ def with_deployment( deployment: ServerDeploymentConfig | Mapping[str, object], ) -> DataDesignerSlurmConfigBuilder: """Append one deployment while preserving authored order.""" - self._deployments.append(ServerDeploymentConfig.model_validate(deployment)) + self._deployments.append(_validate_model(ServerDeploymentConfig, deployment)) return self def with_array_tasks(self, *, count: int, max_concurrent: int = 1) -> DataDesignerSlurmConfigBuilder: """Set deterministic horizontal sharding.""" - self._array_tasks = ArrayTasksConfig(count=count, max_concurrent=max_concurrent) + self._array_tasks = _validate_model( + ArrayTasksConfig, + {"count": count, "max_concurrent": max_concurrent}, + ) return self def with_submission(self, **values: object) -> DataDesignerSlurmConfigBuilder: """Set typed Slurm submission intent.""" - self._submission = SubmissionConfig.model_validate(values) + self._submission = _validate_model(SubmissionConfig, values) return self def with_output(self, **values: object) -> DataDesignerSlurmConfigBuilder: """Set typed dataset output intent.""" - self._output = OutputConfig.model_validate(values) + self._output = _validate_model(OutputConfig, values) return self def build(self) -> DataDesignerSlurmConfig: @@ -145,16 +154,19 @@ def build(self) -> DataDesignerSlurmConfig: raise ConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") assert self._invocation is not None assert self._client is not None - return DataDesignerSlurmConfig( - schema_version=1, - name=self._name, - builder=self._builder, - invocation=self._invocation, - client=self._client, - deployments=self._deployments, - array_tasks=self._array_tasks, - submission=self._submission, - output=self._output, + return _validate_model( + DataDesignerSlurmConfig, + { + "schema_version": 1, + "name": self._name, + "builder": self._builder, + "invocation": self._invocation, + "client": self._client, + "deployments": self._deployments, + "array_tasks": self._array_tasks, + "submission": self._submission, + "output": self._output, + }, ) def write_config(self, path: str | Path) -> None: @@ -171,4 +183,14 @@ def write_config(self, path: str | Path) -> None: ) else: raise ConfigBuilderError("config path must end in .json, .yaml, or .yml") - output_path.write_text(contents, encoding="utf-8") + try: + output_path.write_text(contents, encoding="utf-8") + except OSError as error: + raise ConfigBuilderError(f"cannot write Slurm config {output_path}") from error + + +def _validate_model(config_type: type[_ConfigValue], value: object) -> _ConfigValue: + try: + return config_type.model_validate(value) + except ValidationError as error: + raise ConfigBuilderError(str(error)) from error diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py index 9e127af67..647510697 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -88,34 +88,41 @@ def resolve_profile( home_directory: str | Path | None = None, ) -> SelectedSlurmProfile: """Resolve an injected profile or select one catalog entry.""" - sources = sum(source is not None for source in (profile, catalog, profile_file)) - if sources > 1: - raise ConfigLoadError("profile, catalog, and profile_file are mutually exclusive") - if profile is not None: - if cluster is not None: - raise ConfigLoadError("an injected profile cannot be combined with cluster selection") - return injected_profile(profile) - - catalog_path: str | None = None - if catalog is None: - path = _resolve_profile_path( - profile_file, - environ=os.environ if environ is None else environ, - home_directory=home_directory, + try: + sources = sum(source is not None for source in (profile, catalog, profile_file)) + if sources > 1: + raise ConfigLoadError("profile, catalog, and profile_file are mutually exclusive") + if profile is not None: + if cluster is not None: + raise ConfigLoadError("an injected profile cannot be combined with cluster selection") + return injected_profile(profile) + + catalog_path: str | None = None + if catalog is None: + path = _resolve_profile_path( + profile_file, + environ=os.environ if environ is None else environ, + home_directory=home_directory, + ) + catalog = load_profile_catalog(path) + catalog_path = path.as_posix() + + if cluster is None and hostnames is None: + resolver = hostname_resolver or _local_hostnames + hostnames = resolver() + normalized_hostnames = tuple( + dict.fromkeys(hostname.strip().casefold() for hostname in (hostnames or ()) if hostname) ) - catalog = load_profile_catalog(path) - catalog_path = path.as_posix() - - if hostnames is None: - resolver = hostname_resolver or _local_hostnames - hostnames = resolver() - normalized_hostnames = tuple(dict.fromkeys(hostname.strip().casefold() for hostname in hostnames if hostname)) - return select_profile( - catalog, - cluster=cluster, - hostnames=normalized_hostnames, - catalog_path=catalog_path, - ) + return select_profile( + catalog, + cluster=cluster, + hostnames=normalized_hostnames, + catalog_path=catalog_path, + ) + except ConfigLoadError: + raise + except ValueError as error: + raise ConfigLoadError(str(error)) from error def _load_config(path: str | Path, config_type: type[_Config]) -> _Config: @@ -126,7 +133,10 @@ def _load_config(path: str | Path, config_type: type[_Config]) -> _Config: raise ConfigLoadError(f"cannot read configuration file {resolved_path}") from error try: payload = _parse_mapping(contents, suffix=resolved_path.suffix) - _reject_environment_interpolation(payload) + if config_type is DataDesignerSlurmConfig: + _reject_run_environment_interpolation(payload) + else: + _reject_environment_interpolation(payload) return config_type.model_validate(payload) except ConfigLoadError: raise @@ -168,6 +178,18 @@ def _reject_environment_interpolation(value: object) -> None: _reject_environment_interpolation(item) +def _reject_run_environment_interpolation(payload: Mapping[str, object]) -> None: + for key, value in payload.items(): + _reject_environment_interpolation(key) + if key != "builder" or not isinstance(value, Mapping): + _reject_environment_interpolation(value) + continue + for builder_key, builder_value in value.items(): + _reject_environment_interpolation(builder_key) + if builder_key != "inline": + _reject_environment_interpolation(builder_value) + + def _resolve_profile_path( explicit_path: str | Path | None, *, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index 0f1f89caf..c2e4b18d6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -84,6 +84,8 @@ class EffectiveDataDesignerSlurmConfig(ContractValue): @model_validator(mode="after") def validate_resolution(self) -> EffectiveDataDesignerSlurmConfig: + workspace_root = self.selected_profile.profile.workspace_root + run_root = posixpath.join(workspace_root, "runs", self.run_id) profile_gpus = self.selected_profile.profile.gpus_per_node if profile_gpus != "auto" and profile_gpus != self.resolved_gpus_per_node: raise ValueError("resolved GPU count does not match the selected profile") @@ -102,7 +104,19 @@ def validate_resolution(self) -> EffectiveDataDesignerSlurmConfig: raise ValueError("only sourced builder input may retain a resolved payload") if self.authored.builder.source is not None and self.builder_payload is None: raise ValueError("sourced builder input requires its resolved payload") - runtime_root = posixpath.join(self.selected_profile.profile.workspace_root, "runtime") + _validate_sharding_constraints(self.authored, builder_payload=self.builder_payload) + expected_output = ResolvedOutput( + root=self.authored.output.root or posixpath.join(run_root, "output"), + format=self.authored.output.format, + partitions=self.authored.output.partitions, + require_exact_record_count=self.authored.output.require_exact_record_count, + ) + if self.output != expected_output: + raise ValueError("resolved output does not match the authored output") + _validate_output_destination(self.output.root, workspace_root, run_root) + if self.output.partitions > self.authored.invocation.num_records: + raise ValueError("output partitions must not exceed requested records") + runtime_root = posixpath.join(workspace_root, "runtime") if not _is_below(self.runtime_bundle.path, runtime_root) or not self.runtime_bundle.path.endswith(".tar.gz"): raise ValueError("runtime bundle must be a tar archive below the selected workspace runtime root") return self @@ -138,9 +152,6 @@ def resolve_slurm_config( comment=authored.submission.comment, ) output_root = authored.output.root or posixpath.join(run_root, "output") - _validate_output_destination(output_root, selected_profile.profile.workspace_root) - if authored.output.partitions > authored.invocation.num_records: - raise ConfigurationResolutionError("output partitions must not exceed requested records") output = ResolvedOutput( root=output_root, format=authored.output.format, @@ -271,6 +282,42 @@ def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, Json return RunConfig.model_validate(values).model_dump(mode="json") +def _validate_sharding_constraints( + authored: DataDesignerSlurmConfig, + *, + builder_payload: dict[str, JsonValue] | None, +) -> None: + if authored.array_tasks.count == 1: + return + if authored.output.format != "parquet": + raise ConfigurationResolutionError("multi-shard runs require parquet output") + + payload = authored.builder.inline if authored.builder.inline is not None else builder_payload + assert payload is not None + data_designer = payload.get("data_designer", payload) + if not isinstance(data_designer, dict): + raise ConfigurationResolutionError("builder data_designer value must be an object") + if data_designer.get("processors"): + raise ConfigurationResolutionError("multi-shard runs do not support global processors") + if data_designer.get("profilers"): + raise ConfigurationResolutionError("multi-shard runs do not support global profilers") + + seed_config = data_designer.get("seed_config") + if isinstance(seed_config, dict): + if seed_config.get("sampling_strategy") == "shuffle": + raise ConfigurationResolutionError("multi-shard runs do not support shuffled seed input") + if seed_config.get("selection_strategy") is not None: + raise ConfigurationResolutionError("multi-shard runs do not support authored seed selection strategies") + if authored.invocation.input_bindings.seed_path is None: + raise ConfigurationResolutionError("multi-shard seed input requires a typed seed_path binding") + + columns = data_designer.get("columns") + if isinstance(columns, list) and any( + isinstance(column, dict) and column.get("column_type") == "image" for column in columns + ): + raise ConfigurationResolutionError("multi-shard runs do not support media output columns") + + def _validate_dependency_resolution( authored: DataDesignerSlurmConfig, client_image: ResolvedImage, @@ -451,12 +498,18 @@ def _compile_shards(effective: EffectiveDataDesignerSlurmConfig) -> tuple[Planne return tuple(shards) -def _validate_output_destination(output_root: str, workspace_root: str) -> None: +def _validate_output_destination(output_root: str, workspace_root: str, run_root: str) -> None: if not _is_below(output_root, workspace_root): raise ConfigurationResolutionError("output root must be below the selected workspace_root") reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) if any(_paths_overlap(output_root, path) for path in reserved): raise ConfigurationResolutionError("output root must not overlap package-managed workspace state") + runs_root = posixpath.join(workspace_root, "runs") + run_output_root = posixpath.join(run_root, "output") + if _paths_overlap(output_root, runs_root) and not ( + output_root == run_output_root or _is_below(output_root, run_output_root) + ): + raise ConfigurationResolutionError("output root must not overlap another package-managed run") def _run_root(effective: EffectiveDataDesignerSlurmConfig) -> str: diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 98fb8b610..5e75c90df 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -9,7 +9,7 @@ import pytest import yaml -from data_designer.config import DataDesignerConfigBuilder, ModelConfig +from data_designer.config import DataDesignerConfigBuilder, LLMTextColumnConfig, ModelConfig from data_designer.slurm.config import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, @@ -25,10 +25,12 @@ ) -def _config_builder() -> DataDesignerSlurmConfigBuilder: +def _config_builder(*, prompt: str | None = None) -> DataDesignerSlurmConfigBuilder: data_designer = DataDesignerConfigBuilder( model_configs=[ModelConfig(alias="generator", model="example/generator", provider="openai")] ) + if prompt is not None: + data_designer.add_column(LLMTextColumnConfig(name="generated", prompt=prompt, model_alias="generator")) return ( DataDesignerSlurmConfigBuilder.from_config_builder(data_designer, name="generated-run") .with_invocation( @@ -81,6 +83,31 @@ def test_builder_rejects_unsupported_output_format(tmp_path: Path) -> None: _config_builder().write_config(tmp_path / "run.txt") +@pytest.mark.parametrize( + ("method", "values"), + [ + ("with_invocation", {"num_records": 0, "dataset_name": "generated"}), + ("with_array_tasks", {"count": 1, "max_concurrent": 2}), + ("with_submission", {"time_limit": "invalid"}), + ("with_output", {"partitions": 0}), + ], +) +def test_builder_normalizes_invalid_authored_values( + method: str, + values: dict[str, object], +) -> None: + with pytest.raises(ConfigBuilderError): + getattr(_config_builder(), method)(**values) + + +def test_builder_normalizes_write_failures(tmp_path: Path) -> None: + path = tmp_path / "run.json" + path.mkdir() + + with pytest.raises(ConfigBuilderError, match="cannot write"): + _config_builder().write_config(path) + + @pytest.mark.parametrize( ("suffix", "contents", "message"), [ @@ -88,6 +115,7 @@ def test_builder_rejects_unsupported_output_format(tmp_path: Path) -> None: (".yaml", "schema_version: 1\nschema_version: 1\n", "duplicate"), (".yaml", "defaults: &defaults\n schema_version: 1\nrun: *defaults\n", "anchors"), (".yaml", "schema_version: 1\nname: ${RUN_NAME}\n", "interpolation"), + (".yaml", "schema_version: 1\nbuilder:\n source: ${HOME}/builder.json\n", "interpolation"), ], ) def test_strict_loader_rejects_ambiguous_yaml_and_json( @@ -113,6 +141,16 @@ def test_strict_loader_rejects_non_object_and_unknown_extension(tmp_path: Path) load_run_config(tmp_path / "run.toml") +@pytest.mark.parametrize("suffix", [".json", ".yaml", ".yml"]) +def test_loader_preserves_literal_interpolation_inside_builder_payload(tmp_path: Path, suffix: str) -> None: + path = tmp_path / f"run{suffix}" + builder = _config_builder(prompt="Use the literal ${HOME} value") + + builder.write_config(path) + + assert load_run_config(path) == builder.build() + + def test_profile_source_and_selection_precedence( tmp_path: Path, profile_catalog: SlurmProfileCatalog, @@ -127,7 +165,7 @@ def test_profile_source_and_selection_precedence( explicit = resolve_profile( profile_file=explicit_path, cluster="lab", - hostname_resolver=lambda: ("primary-login-1",), + hostname_resolver=lambda: pytest.fail("explicit cluster selection must not resolve hostnames"), environ={PROFILE_FILE_ENVIRONMENT: str(environment_path)}, ) environment = resolve_profile( @@ -170,6 +208,19 @@ def test_profile_resolution_rejects_conflicting_or_empty_sources( resolve_profile(environ={PROFILE_FILE_ENVIRONMENT: ""}) with pytest.raises(ConfigLoadError, match="cluster selection"): resolve_profile(profile=profile_catalog.clusters["primary"], cluster="primary") + with pytest.raises(ConfigLoadError, match="unknown cluster"): + resolve_profile(catalog=profile_catalog, cluster="missing") + + +def test_profile_resolution_normalizes_ambiguous_hostname_errors( + profile_catalog: SlurmProfileCatalog, +) -> None: + payload = profile_catalog.model_dump(mode="json") + payload["clusters"]["lab"]["host_patterns"] = ["*-login-*"] + catalog = SlurmProfileCatalog.model_validate(payload) + + with pytest.raises(ConfigLoadError, match="multiple clusters"): + resolve_profile(catalog=catalog, hostnames=("primary-login-1",)) def test_json_builder_output_is_stable(tmp_path: Path) -> None: diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index bc26326bb..889a5a399 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -7,6 +7,13 @@ import pytest +from data_designer.config import ( + DataDesignerConfigBuilder, + DropColumnsProcessorConfig, + JudgeScoreProfilerConfig, + LLMTextColumnConfig, + ModelConfig, +) from data_designer.slurm.config import ( BuilderInput, DataDesignerSlurmConfig, @@ -188,6 +195,8 @@ def test_compiler_rejects_tensor_parallelism_that_does_not_divide_gpu_shape( {"root": "/outside/output"}, {"root": "/workspace/primary/images/output"}, {"root": "/workspace/primary/runtime/output"}, + {"root": "/workspace/primary/runs/other-run/output"}, + {"root": "/workspace/primary/runs/run-single/shards"}, {"partitions": 9}, ], ) @@ -205,6 +214,45 @@ def test_resolution_rejects_invalid_output_destinations( _resolve_fixture(authored, dependency_lock_single, single_node_plan) +def test_effective_config_rejects_invalid_direct_construction( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run, dependency_lock, multi_node_plan) + values = {name: getattr(effective, name) for name in EffectiveDataDesignerSlurmConfig.model_fields} + values["authored"] = authored_run.model_copy( + update={"output": authored_run.output.model_copy(update={"format": "jsonl"})} + ) + values["output"] = effective.output.model_copy(update={"format": "jsonl"}) + + with pytest.raises(ValueError, match="parquet output"): + EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] + + other_output = "/workspace/primary/runs/other-run/output" + values["authored"] = authored_run.model_copy( + update={"output": authored_run.output.model_copy(update={"root": other_output})} + ) + values["output"] = effective.output.model_copy(update={"root": other_output}) + + with pytest.raises(ValueError, match="another package-managed run"): + EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] + + values["authored"] = authored_run.model_copy( + update={"output": authored_run.output.model_copy(update={"partitions": 101})} + ) + values["output"] = effective.output.model_copy(update={"partitions": 101}) + + with pytest.raises(ValueError, match="requested records"): + EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] + + values["authored"] = authored_run + values["output"] = effective.output.model_copy(update={"format": "jsonl"}) + + with pytest.raises(ValueError, match="resolved output"): + EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] + + def test_compiler_rejects_model_alias_missing_from_builder( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -252,6 +300,147 @@ def test_sharded_seed_inputs_have_stable_ranges_and_partition_digests( assert first.shards == second.shards +@pytest.mark.parametrize( + ("builder_update", "output_update", "message"), + [ + ( + { + "seed_config": { + "sampling_strategy": "shuffle", + "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, + } + }, + {}, + "shuffled seed", + ), + ( + { + "seed_config": { + "sampling_strategy": "ordered", + "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, + } + }, + {}, + "seed_path", + ), + ( + { + "seed_config": { + "sampling_strategy": "ordered", + "selection_strategy": {"start": 0, "end": 9}, + "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, + } + }, + {}, + "selection strategies", + ), + ( + { + "columns": [ + { + "column_type": "image", + "model_alias": "generator", + "name": "picture", + "prompt": "an image", + } + ] + }, + {}, + "media output", + ), + ({}, {"format": "jsonl"}, "parquet output"), + ], +) +def test_resolution_rejects_unshardable_big_iron_fields( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + builder_update: dict[str, object], + output_update: dict[str, object], + message: str, +) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"].update(builder_update) + payload["output"].update(output_update) + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(ConfigurationResolutionError, match=message): + _resolve_fixture(authored, dependency_lock, multi_node_plan) + + +@pytest.mark.parametrize("field", ["processors", "profilers"]) +def test_resolution_rejects_real_global_builder_configs( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + field: str, +) -> None: + builder = DataDesignerConfigBuilder( + model_configs=[ModelConfig(alias="generator", model="example/generator", provider="openai")] + ) + builder.add_column(LLMTextColumnConfig(name="generated", prompt="hello", model_alias="generator")) + if field == "processors": + builder.add_processor(DropColumnsProcessorConfig(name="global", column_names=["generated"])) + else: + builder.add_profiler(JudgeScoreProfilerConfig(model_alias="generator")) + data_designer = builder.get_builder_config().to_dict()["data_designer"] + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"][field] = data_designer[field] + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(ConfigurationResolutionError, match=field): + _resolve_fixture(authored, dependency_lock, multi_node_plan) + + +def test_sharded_seed_binding_may_override_authored_source( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["seed_config"] = { + "sampling_strategy": "ordered", + "source": {"path": "/datasets/original.parquet", "seed_type": "local"}, + } + payload["invocation"]["input_bindings"]["seed_path"] = "/datasets/override.parquet" + authored = DataDesignerSlurmConfig.model_validate(payload) + + plan = compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + + assert all(shard.input_partition is not None for shard in plan.shards) + + +def test_single_shard_allows_non_collectable_big_iron_fields( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"].update( + { + "columns": [ + { + "column_type": "image", + "model_alias": "generator", + "name": "picture", + "prompt": "an image", + } + ], + "seed_config": { + "sampling_strategy": "shuffle", + "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, + }, + } + ) + payload["output"]["format"] = "jsonl" + authored = DataDesignerSlurmConfig.model_validate(payload) + + plan = compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + assert plan.array_tasks.count == 1 + assert plan.output.format == "jsonl" + + def test_sourced_builder_is_resolved_to_one_digest_bound_run_input( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -276,6 +465,24 @@ def test_sourced_builder_is_resolved_to_one_digest_bound_run_input( assert plan.builder.content_sha256 == plan.builder.source.sha256 +def test_resolution_rejects_secret_values_in_sourced_builder_payload( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + builder_payload = authored_run_single.model_dump(mode="json")["builder"]["inline"] + builder_payload["data_designer"]["api_key"] = "secret-value" + authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) + + with pytest.raises(ConfigurationResolutionError, match="secret values"): + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + + def test_resolution_rejects_artifact_identity_mismatches( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, From d3264a9285ab930d8c48a2ee9b04ccd236cd3dc2 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 27 Aug 2026 11:46:19 -0300 Subject: [PATCH 3/9] fix: harden Slurm plan resolution Preserve partial early-shutdown intent and hide secret values in validation errors. Reject unsafe sharding semantics, require exact deployment coverage, and bind sourced payload identity. Signed-off-by: Andre Manoel --- .../src/data_designer/slurm/contracts.py | 1 + .../data_designer/slurm/planning/compiler.py | 62 ++++++-- .../data_designer/slurm/planning/models.py | 14 +- .../slurm/planning/validation.py | 5 +- .../tests/config/test_loading_builder.py | 26 ++++ .../tests/contracts/test_planning_records.py | 8 +- .../tests/planning/test_compiler.py | 145 +++++++++++++++++- 7 files changed, 242 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 fc47926b1..df6222d74 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -96,6 +96,7 @@ class ContractValue(BaseModel): model_config = ConfigDict( extra="forbid", frozen=True, + hide_input_in_errors=True, allow_inf_nan=False, protected_namespaces=(), strict=True, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index c2e4b18d6..4c6a5b93c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -5,7 +5,6 @@ from __future__ import annotations -import hashlib import posixpath from typing import Annotated @@ -22,7 +21,6 @@ RecordRange, ResumeWorkspace, compute_sha256, - pretty_json, ) from data_designer.slurm.planning.models import ( PlannedShard, @@ -38,6 +36,7 @@ ResolvedSubmission, ResolvedTopology, _extract_builder_aliases, + _extract_builder_identity, ) from data_designer.slurm.planning.validation import validate_resolved_plan @@ -45,6 +44,19 @@ _HTTP_PORT = 18000 _RENDEZVOUS_PORT = 19000 _PORT_RANGE_SIZE = 1000 +_SHARDABLE_COLUMN_TYPES = frozenset( + { + "embedding", + "expression", + "llm-code", + "llm-judge", + "llm-structured", + "llm-text", + "sampler", + "seed-dataset", + "validation", + } +) _COMPATIBILITY_RUN_DEFAULTS: dict[str, JsonValue] = { "buffer_size": 16384, "disable_early_shutdown": True, @@ -104,6 +116,22 @@ def validate_resolution(self) -> EffectiveDataDesignerSlurmConfig: raise ValueError("only sourced builder input may retain a resolved payload") if self.authored.builder.source is not None and self.builder_payload is None: raise ValueError("sourced builder input requires its resolved payload") + if self.authored.builder.source is not None: + assert self.builder_payload is not None + validated_payload = BuilderInput(inline=self.builder_payload).inline + assert validated_payload is not None + aliases, referenced_aliases, digest = _extract_builder_identity(validated_payload) + if self.builder.authored_source != self.authored.builder.source or self.builder.source is None: + raise ValueError("resolved builder source does not match the authored input") + expected_path = posixpath.join(run_root, "builder-config.json") + if self.builder.source.path != expected_path: + raise ValueError("resolved builder artifact path does not match the package-managed run") + if self.builder.model_aliases != aliases: + raise ValueError("resolved model aliases do not match the sourced builder payload") + if self.builder.referenced_model_aliases != referenced_aliases: + raise ValueError("resolved referenced aliases do not match the sourced builder payload") + if self.builder.content_sha256 != digest: + raise ValueError("resolved builder digest does not match the sourced builder payload") _validate_sharding_constraints(self.authored, builder_payload=self.builder_payload) expected_output = ResolvedOutput( root=self.authored.output.root or posixpath.join(run_root, "output"), @@ -260,11 +288,10 @@ def _resolve_builder( raise ConfigurationResolutionError("sourced builder input requires its resolved payload") validated_payload = BuilderInput(inline=builder_payload).inline assert validated_payload is not None - aliases, referenced_aliases = _extract_builder_aliases(validated_payload) - serialized = pretty_json(validated_payload).encode("utf-8") + aliases, referenced_aliases, digest = _extract_builder_identity(validated_payload) source = ArtifactReference( path=posixpath.join(run_root, "builder-config.json"), - sha256=hashlib.sha256(serialized).hexdigest(), + sha256=digest, ) return ResolvedBuilderInput( authored_source=authored.builder.source, @@ -277,7 +304,12 @@ def _resolve_builder( def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: values = dict(authored.invocation.run_config) + preserve_authored_early_shutdown = "disable_early_shutdown" not in values and ( + "shutdown_error_rate" in values or "shutdown_error_window" in values + ) for name, value in _COMPATIBILITY_RUN_DEFAULTS.items(): + if preserve_authored_early_shutdown and name in {"disable_early_shutdown", "shutdown_error_rate"}: + continue values.setdefault(name, value) return RunConfig.model_validate(values).model_dump(mode="json") @@ -311,11 +343,21 @@ def _validate_sharding_constraints( if authored.invocation.input_bindings.seed_path is None: raise ConfigurationResolutionError("multi-shard seed input requires a typed seed_path binding") - columns = data_designer.get("columns") - if isinstance(columns, list) and any( - isinstance(column, dict) and column.get("column_type") == "image" for column in columns - ): - raise ConfigurationResolutionError("multi-shard runs do not support media output columns") + columns = data_designer.get("columns", []) + if not isinstance(columns, list): + raise ConfigurationResolutionError("builder columns must be a list") + for column in columns: + if not isinstance(column, dict) or not isinstance(column.get("column_type"), str): + raise ConfigurationResolutionError("multi-shard runs require known column semantics") + column_type = column["column_type"] + if column_type == "image": + raise ConfigurationResolutionError("multi-shard runs do not support media output columns") + if column_type not in _SHARDABLE_COLUMN_TYPES: + raise ConfigurationResolutionError( + "multi-shard runs do not support custom, plugin, or unknown column semantics" + ) + if column_type == "validation" and column.get("validator_type") == "local_callable": + raise ConfigurationResolutionError("multi-shard runs do not support local callable validators") def _validate_dependency_resolution( 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 d1eff7ab6..2871ff1ac 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 @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import posixpath from typing import Annotated, Literal from urllib.parse import urlsplit @@ -40,6 +41,7 @@ Sha256Digest, ShardId, compute_sha256, + pretty_json, validate_absolute_path, validate_local_config_path, validate_plain_text, @@ -379,8 +381,8 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: 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 set(aliases) != set(self.builder.model_aliases): + raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") if not set(self.builder.referenced_model_aliases).issubset(aliases): raise ValueError("each referenced Data Designer model alias requires a deployment") @@ -520,3 +522,11 @@ def collect(value: JsonValue, *, key: str | None = None) -> None: collect(data_designer) return tuple(model_aliases), tuple(dict.fromkeys(referenced_aliases)) + + +def _extract_builder_identity( + builder: dict[str, JsonValue], +) -> tuple[tuple[ModelAlias, ...], tuple[ModelAlias, ...], Sha256Digest]: + model_aliases, referenced_aliases = _extract_builder_aliases(builder) + digest = hashlib.sha256(pretty_json(builder).encode("utf-8")).hexdigest() + return model_aliases, referenced_aliases, digest 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 c47ebc512..e1d17fd8d 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 @@ -11,7 +11,7 @@ from data_designer.slurm.planning.models import ( ResolvedDependencyLock, ResolvedSlurmRunPlan, - _extract_builder_aliases, + _extract_builder_identity, ) @@ -55,12 +55,13 @@ def validate_resolved_plan( ) if builder_payload is None: raise PlanContractError("sourced builder validation requires its resolved payload") - model_aliases, referenced_aliases = _extract_builder_aliases(builder_payload) + model_aliases, referenced_aliases, digest = _extract_builder_identity(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", ) + _require(plan.builder.content_sha256 == digest, "resolved builder digest does 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/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 5e75c90df..2e24a3ba7 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -100,6 +100,17 @@ def test_builder_normalizes_invalid_authored_values( getattr(_config_builder(), method)(**values) +def test_builder_validation_errors_hide_secret_inputs() -> None: + secret = "super-secret-token" + + with pytest.raises(ConfigBuilderError) as error: + _config_builder().with_client(image={"name": "dd-client", "api_key": secret}) + + assert secret not in str(error.value) + assert error.value.__cause__ is not None + assert secret not in str(error.value.__cause__) + + def test_builder_normalizes_write_failures(tmp_path: Path) -> None: path = tmp_path / "run.json" path.mkdir() @@ -141,6 +152,21 @@ def test_strict_loader_rejects_non_object_and_unknown_extension(tmp_path: Path) load_run_config(tmp_path / "run.toml") +def test_loader_validation_errors_hide_secret_inputs(tmp_path: Path) -> None: + secret = "super-secret-token" + payload = _config_builder().build().model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["api_key"] = secret + path = tmp_path / "run.json" + path.write_text(json.dumps(payload)) + + with pytest.raises(ConfigLoadError) as error: + load_run_config(path) + + assert secret not in str(error.value) + assert error.value.__cause__ is not None + assert secret not in str(error.value.__cause__) + + @pytest.mark.parametrize("suffix", [".json", ".yaml", ".yml"]) def test_loader_preserves_literal_interpolation_inside_builder_payload(tmp_path: Path, suffix: str) -> None: path = tmp_path / f"run{suffix}" 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 5c64a5c24..da15830be 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import json from copy import deepcopy @@ -10,7 +11,7 @@ from pydantic import ValidationError from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig -from data_designer.slurm.contracts import compute_sha256 +from data_designer.slurm.contracts import compute_sha256, pretty_json from data_designer.slurm.planning import ( ArtifactReference, PlanContractError, @@ -226,13 +227,14 @@ def test_sourced_builder_validation_requires_resolved_payload( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: sourced_authored = authored_run.model_copy(update={"builder": BuilderInput(source="builder.json")}) + builder_digest = hashlib.sha256(pretty_json(authored_run.builder.inline).encode("utf-8")).hexdigest() 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}, + "source": {"path": "/workspace/primary/runs/run-001/builder.json", "sha256": builder_digest}, "inline": None, - "content_sha256": "a" * 64, + "content_sha256": builder_digest, "model_aliases": ["generator", "judge"], "referenced_model_aliases": [], } diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index 889a5a399..2f2fd8737 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -3,6 +3,7 @@ from __future__ import annotations +from copy import deepcopy from pathlib import Path import pytest @@ -176,6 +177,32 @@ def test_compiler_preserves_explicit_compatibility_run_values( assert effective.invocation.effective_run_config["shutdown_error_rate"] == 0.25 +@pytest.mark.parametrize( + ("run_config", "expected_rate", "expected_window"), + [ + ({"shutdown_error_rate": 0.25}, 0.25, 10), + ({"shutdown_error_window": 25}, 0.5, 25), + ], +) +def test_compiler_preserves_partial_early_shutdown_configuration( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + run_config: dict[str, object], + expected_rate: float, + expected_window: int, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["run_config"] = run_config + authored = DataDesignerSlurmConfig.model_validate(payload) + + effective = _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + assert effective.invocation.effective_run_config["disable_early_shutdown"] is False + assert effective.invocation.effective_run_config["shutdown_error_rate"] == expected_rate + assert effective.invocation.effective_run_config["shutdown_error_window"] == expected_window + + def test_compiler_rejects_tensor_parallelism_that_does_not_divide_gpu_shape( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -266,6 +293,34 @@ def test_compiler_rejects_model_alias_missing_from_builder( compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) +@pytest.mark.parametrize("sourced", [False, True], ids=["inline", "sourced"]) +def test_compiler_rejects_builder_model_alias_without_deployment( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + sourced: bool, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["model_configs"].append( + {"alias": "undeployed", "model": "example/undeployed", "provider": "openai"} + ) + builder_payload = None + if sourced: + builder_payload = payload["builder"]["inline"] + payload["builder"] = {"source": "builder.json"} + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(PlanCompilationError, match="exactly cover"): + compile_slurm_run_plan( + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + ) + + def test_compiler_rejects_otel_collision_before_runtime( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -392,6 +447,53 @@ def test_resolution_rejects_real_global_builder_configs( _resolve_fixture(authored, dependency_lock, multi_node_plan) +@pytest.mark.parametrize("sourced", [False, True], ids=["inline", "sourced"]) +@pytest.mark.parametrize( + ("column", "message"), + [ + ({"column_type": "future-column", "name": "future"}, "unknown column semantics"), + ({"column_type": "fake-slurm-column", "name": "plugin"}, "plugin"), + ({"column_type": "custom", "generator_function": "generate", "name": "custom"}, "custom"), + ( + { + "column_type": "validation", + "name": "validate", + "target_columns": ["generated"], + "validator_params": { + "validation_function": "validate", + "validator_type": "local_callable", + }, + "validator_type": "local_callable", + }, + "local callable", + ), + ], +) +def test_resolution_rejects_unportable_multi_shard_columns( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + sourced: bool, + column: dict[str, object], + message: str, +) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["columns"] = [column] + builder_payload = None + if sourced: + builder_payload = payload["builder"]["inline"] + payload["builder"] = {"source": "builder.json"} + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(ConfigurationResolutionError, match=message): + _resolve_fixture( + authored, + dependency_lock, + multi_node_plan, + builder_payload=builder_payload, + ) + + def test_sharded_seed_binding_may_override_authored_source( authored_run: DataDesignerSlurmConfig, dependency_lock: ResolvedDependencyLock, @@ -465,16 +567,51 @@ def test_sourced_builder_is_resolved_to_one_digest_bound_run_input( assert plan.builder.content_sha256 == plan.builder.source.sha256 +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("bytes", "builder digest"), + ("aliases", "model aliases"), + ], +) +def test_effective_config_rejects_sourced_builder_payload_identity_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + mutation: str, + message: str, +) -> None: + builder_payload = authored_run_single.model_dump(mode="json")["builder"]["inline"] + authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) + effective = _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + drifted_payload = deepcopy(builder_payload) + if mutation == "bytes": + drifted_payload["library_version"] = "drifted" + else: + drifted_payload["data_designer"]["model_configs"][0]["alias"] = "drifted" + values = {name: getattr(effective, name) for name in EffectiveDataDesignerSlurmConfig.model_fields} + values["builder_payload"] = drifted_payload + + with pytest.raises(ValueError, match=message): + EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] + + def test_resolution_rejects_secret_values_in_sourced_builder_payload( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, single_node_plan: ResolvedSlurmRunPlan, ) -> None: builder_payload = authored_run_single.model_dump(mode="json")["builder"]["inline"] - builder_payload["data_designer"]["api_key"] = "secret-value" + secret = "super-secret-token" + builder_payload["data_designer"]["api_key"] = secret authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) - with pytest.raises(ConfigurationResolutionError, match="secret values"): + with pytest.raises(ConfigurationResolutionError, match="secret values") as error: _resolve_fixture( authored, dependency_lock_single, @@ -482,6 +619,10 @@ def test_resolution_rejects_secret_values_in_sourced_builder_payload( builder_payload=builder_payload, ) + assert secret not in str(error.value) + assert error.value.__cause__ is not None + assert secret not in str(error.value.__cause__) + def test_resolution_rejects_artifact_identity_mismatches( authored_run_single: DataDesignerSlurmConfig, From 8026ffe475f42c7b9d59de920dd6f60245dcccbc Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 27 Aug 2026 12:03:45 -0300 Subject: [PATCH 4/9] fix: bind Slurm effective plan inputs Preserve all authored early-shutdown controls and reject direct effective invocation drift. Bind shard partition digests to deterministic persisted JSON bytes. Signed-off-by: Andre Manoel --- .../src/data_designer/slurm/contracts.py | 6 +++++ .../data_designer/slurm/planning/compiler.py | 15 ++++++++--- .../data_designer/slurm/planning/models.py | 6 ++--- .../tests/contracts/test_planning_records.py | 5 ++-- .../tests/planning/test_compiler.py | 26 +++++++++++++++++++ 5 files changed, 48 insertions(+), 10 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 df6222d74..cad317e73 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -170,6 +170,11 @@ def compute_sha256(value: object) -> Sha256Digest: return hashlib.sha256(canonical_json(value)).hexdigest() +def compute_pretty_sha256(value: object) -> Sha256Digest: + """Compute the digest of deterministic persisted JSON bytes.""" + return hashlib.sha256(pretty_json(value).encode("utf-8")).hexdigest() + + def validate_absolute_path(value: str) -> str: """Validate a normalized, absolute POSIX path below the filesystem root.""" if not value.startswith("/"): @@ -288,6 +293,7 @@ class ResumeWorkspace(ContractValue): "Sha256Digest", "ShardId", "canonical_json", + "compute_pretty_sha256", "compute_sha256", "pretty_json", "validate_absolute_path", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index 4c6a5b93c..03a566865 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -20,6 +20,7 @@ Identifier, RecordRange, ResumeWorkspace, + compute_pretty_sha256, compute_sha256, ) from data_designer.slurm.planning.models import ( @@ -133,6 +134,12 @@ def validate_resolution(self) -> EffectiveDataDesignerSlurmConfig: if self.builder.content_sha256 != digest: raise ValueError("resolved builder digest does not match the sourced builder payload") _validate_sharding_constraints(self.authored, builder_payload=self.builder_payload) + expected_invocation = ResolvedInvocation( + authored=self.authored.invocation, + effective_run_config=_materialize_run_config(self.authored), + ) + if self.invocation != expected_invocation: + raise ValueError("resolved invocation does not match the authored invocation") expected_output = ResolvedOutput( root=self.authored.output.root or posixpath.join(run_root, "output"), format=self.authored.output.format, @@ -304,11 +311,11 @@ def _resolve_builder( def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: values = dict(authored.invocation.run_config) - preserve_authored_early_shutdown = "disable_early_shutdown" not in values and ( - "shutdown_error_rate" in values or "shutdown_error_window" in values + authored_early_shutdown = {"disable_early_shutdown", "shutdown_error_rate", "shutdown_error_window"}.intersection( + values ) for name, value in _COMPATIBILITY_RUN_DEFAULTS.items(): - if preserve_authored_early_shutdown and name in {"disable_early_shutdown", "shutdown_error_rate"}: + if authored_early_shutdown and name in {"disable_early_shutdown", "shutdown_error_rate"}: continue values.setdefault(name, value) return RunConfig.model_validate(values).model_dump(mode="json") @@ -519,7 +526,7 @@ def _compile_shards(effective: EffectiveDataDesignerSlurmConfig) -> tuple[Planne if seed_path is not None: partition = ArtifactReference( path=posixpath.join(shard_root, "input-partition.json"), - sha256=compute_sha256( + sha256=compute_pretty_sha256( { "record_range": record_range.model_dump(mode="json"), "seed_path": seed_path, 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 2871ff1ac..1e56f9580 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 @@ -3,7 +3,6 @@ from __future__ import annotations -import hashlib import posixpath from typing import Annotated, Literal from urllib.parse import urlsplit @@ -40,8 +39,8 @@ ResumeWorkspace, Sha256Digest, ShardId, + compute_pretty_sha256, compute_sha256, - pretty_json, validate_absolute_path, validate_local_config_path, validate_plain_text, @@ -152,6 +151,7 @@ class ResolvedBuilderInput(ContractValue): authored_source: str | None = None source: ArtifactReference | None = None inline: dict[str, JsonValue] | None = None + # Inline input uses canonical JSON; sourced input uses its persisted artifact bytes. content_sha256: Sha256Digest model_aliases: tuple[ModelAlias, ...] referenced_model_aliases: tuple[ModelAlias, ...] = () @@ -528,5 +528,5 @@ def _extract_builder_identity( builder: dict[str, JsonValue], ) -> tuple[tuple[ModelAlias, ...], tuple[ModelAlias, ...], Sha256Digest]: model_aliases, referenced_aliases = _extract_builder_aliases(builder) - digest = hashlib.sha256(pretty_json(builder).encode("utf-8")).hexdigest() + digest = compute_pretty_sha256(builder) return model_aliases, referenced_aliases, digest 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 da15830be..102b07769 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -3,7 +3,6 @@ from __future__ import annotations -import hashlib import json from copy import deepcopy @@ -11,7 +10,7 @@ from pydantic import ValidationError from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig -from data_designer.slurm.contracts import compute_sha256, pretty_json +from data_designer.slurm.contracts import compute_pretty_sha256, compute_sha256 from data_designer.slurm.planning import ( ArtifactReference, PlanContractError, @@ -227,7 +226,7 @@ def test_sourced_builder_validation_requires_resolved_payload( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: sourced_authored = authored_run.model_copy(update={"builder": BuilderInput(source="builder.json")}) - builder_digest = hashlib.sha256(pretty_json(authored_run.builder.inline).encode("utf-8")).hexdigest() + builder_digest = compute_pretty_sha256(authored_run.builder.inline) payload = multi_node_plan.model_dump(mode="json") payload["authored_config"]["sha256"] = sourced_authored.compute_sha256() payload["builder"] = { diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index 2f2fd8737..0f8c67cc5 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -22,6 +22,7 @@ SlurmProfileCatalog, select_profile, ) +from data_designer.slurm.contracts import compute_pretty_sha256 from data_designer.slurm.planning import ( ArtifactReference, ConfigurationResolutionError, @@ -182,6 +183,8 @@ def test_compiler_preserves_explicit_compatibility_run_values( [ ({"shutdown_error_rate": 0.25}, 0.25, 10), ({"shutdown_error_window": 25}, 0.5, 25), + ({"disable_early_shutdown": False}, 0.5, 10), + ({"disable_early_shutdown": False, "shutdown_error_window": 25}, 0.5, 25), ], ) def test_compiler_preserves_partial_early_shutdown_configuration( @@ -280,6 +283,21 @@ def test_effective_config_rejects_invalid_direct_construction( EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] +def test_effective_config_rejects_invocation_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + run_config = dict(effective.invocation.effective_run_config) + run_config["jinja_rendering_engine"] = "native" + values = {name: getattr(effective, name) for name in EffectiveDataDesignerSlurmConfig.model_fields} + values["invocation"] = effective.invocation.model_copy(update={"effective_run_config": run_config}) + + with pytest.raises(ValueError, match="resolved invocation"): + EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] + + def test_compiler_rejects_model_alias_missing_from_builder( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -352,6 +370,14 @@ def test_sharded_seed_inputs_have_stable_ranges_and_partition_digests( (50, 100), ] assert all(shard.input_partition is not None for shard in first.shards) + for shard in first.shards: + assert shard.input_partition is not None + assert shard.input_partition.sha256 == compute_pretty_sha256( + { + "record_range": shard.record_range.model_dump(mode="json"), + "seed_path": "/datasets/seed.parquet", + } + ) assert first.shards == second.shards From 879995c772f9b7c2802aa437f26d0b2f806d0abb Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 27 Aug 2026 16:48:45 -0300 Subject: [PATCH 5/9] fix: align Slurm planning boundaries --- .../src/data_designer/slurm/_errors.py | 32 ++ .../data_designer/slurm/config/__init__.py | 8 +- .../src/data_designer/slurm/config/builder.py | 21 +- .../src/data_designer/slurm/config/errors.py | 18 + .../src/data_designer/slurm/config/loading.py | 51 +-- .../src/data_designer/slurm/contracts.py | 11 +- .../data_designer/slurm/planning/__init__.py | 17 - .../slurm/planning/builder_identity.py | 34 ++ .../data_designer/slurm/planning/compiler.py | 367 +----------------- .../data_designer/slurm/planning/errors.py | 22 ++ .../data_designer/slurm/planning/models.py | 60 +-- .../slurm/planning/resolution.py | 318 +++++++++++++++ .../slurm/planning/validation.py | 17 +- .../tests/config/test_loading_builder.py | 63 ++- .../contracts/golden/multi_node_plan.json | 1 - .../contracts/golden/single_node_plan.json | 1 - .../tests/contracts/test_planning_records.py | 37 +- .../golden/finalization_chain.json | 4 +- .../tests/planning/test_compiler.py | 163 ++++---- .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- 22 files changed, 624 insertions(+), 629 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/_errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py new file mode 100644 index 000000000..80234fafa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Safe formatting for normalized Slurm boundary errors.""" + +from __future__ import annotations + +import json + +import yaml +from pydantic import ValidationError + + +def format_validation_error(error: ValidationError, *, subject: str) -> str: + """Summarize validation without rendering user-controlled values.""" + error_types = sorted( + {str(detail["type"]) for detail in error.errors(include_url=False, include_context=False, include_input=False)} + ) + count = error.error_count() + noun = "error" if count == 1 else "errors" + kinds = f": {', '.join(error_types)}" if error_types else "" + return f"{subject} failed validation ({count} {noun}{kinds})" + + +def format_parse_error(error: json.JSONDecodeError | yaml.YAMLError) -> str: + """Summarize a parse failure without rendering source text.""" + if isinstance(error, json.JSONDecodeError): + return f"invalid JSON at line {error.lineno}, column {error.colno}" + mark = getattr(error, "problem_mark", None) + if mark is None: + return "invalid YAML" + return f"invalid YAML at line {mark.line + 1}, column {mark.column + 1}" 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 index 66e2734cf..0355bd93b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -14,7 +14,8 @@ DataDesignerSlurmBenchmarkConfig, FixedRecordPolicy, ) -from data_designer.slurm.config.builder import ConfigBuilderError, DataDesignerSlurmConfigBuilder +from data_designer.slurm.config.builder import DataDesignerSlurmConfigBuilder +from data_designer.slurm.config.errors import SlurmConfigBuilderError, SlurmConfigLoadError from data_designer.slurm.config.images import ( ClientImageInspection, ImageBuildRequest, @@ -27,7 +28,6 @@ from data_designer.slurm.config.loading import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, - ConfigLoadError, load_profile_catalog, load_run_config, resolve_profile, @@ -78,8 +78,6 @@ "ClientConfig", "ClientDependencies", "ClientImageInspection", - "ConfigBuilderError", - "ConfigLoadError", "ContainerMount", "DataDesignerSlurmBenchmarkConfig", "DataDesignerSlurmConfig", @@ -110,6 +108,8 @@ "SelectedSlurmProfile", "ServerDeploymentConfig", "ServingImageInspection", + "SlurmConfigBuilderError", + "SlurmConfigLoadError", "SlurmProfile", "SlurmProfileCatalog", "SubmissionConfig", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py index 9cbb9fb21..062fa54dc 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -13,6 +13,8 @@ from pydantic import BaseModel, ValidationError from data_designer.config import DataDesignerConfigBuilder +from data_designer.slurm._errors import format_validation_error +from data_designer.slurm.config.errors import SlurmConfigBuilderError from data_designer.slurm.config.images import ImageRef from data_designer.slurm.config.run import ( ArrayTasksConfig, @@ -29,11 +31,7 @@ SubmissionConfig, ) -_ConfigValue = TypeVar("_ConfigValue", bound=BaseModel) - - -class ConfigBuilderError(ValueError): - """Raised when the Slurm config builder is incomplete or cannot serialize.""" +_ConfigValueT = TypeVar("_ConfigValueT", bound=BaseModel) class DataDesignerSlurmConfigBuilder: @@ -151,7 +149,7 @@ def build(self) -> DataDesignerSlurmConfig: if not self._deployments: missing.append("deployment") if missing: - raise ConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") + raise SlurmConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") assert self._invocation is not None assert self._client is not None return _validate_model( @@ -182,15 +180,16 @@ def write_config(self, path: str | Path) -> None: sort_keys=True, ) else: - raise ConfigBuilderError("config path must end in .json, .yaml, or .yml") + raise SlurmConfigBuilderError("config path must end in .json, .yaml, or .yml") try: output_path.write_text(contents, encoding="utf-8") - except OSError as error: - raise ConfigBuilderError(f"cannot write Slurm config {output_path}") from error + except OSError: + raise SlurmConfigBuilderError(f"cannot write Slurm config {output_path}") from None -def _validate_model(config_type: type[_ConfigValue], value: object) -> _ConfigValue: +def _validate_model(config_type: type[_ConfigValueT], value: object) -> _ConfigValueT: try: return config_type.model_validate(value) except ValidationError as error: - raise ConfigBuilderError(str(error)) from error + message = format_validation_error(error, subject=config_type.__name__) + raise SlurmConfigBuilderError(message) from None diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/config/errors.py new file mode 100644 index 000000000..e45a33b33 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/errors.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalized authored Slurm configuration errors.""" + +from __future__ import annotations + + +class SlurmConfigError(ValueError): + """Base error for authored Slurm configuration boundaries.""" + + +class SlurmConfigBuilderError(SlurmConfigError): + """Raised when the Slurm config builder is incomplete or cannot serialize.""" + + +class SlurmConfigLoadError(SlurmConfigError): + """Raised when a local Slurm configuration file is not strict and valid.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py index 647510697..2d9472286 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -16,6 +16,8 @@ from pydantic import ValidationError from yaml.nodes import MappingNode +from data_designer.slurm._errors import format_parse_error, format_validation_error +from data_designer.slurm.config.errors import SlurmConfigLoadError from data_designer.slurm.config.profiles import ( SelectedSlurmProfile, SlurmProfile, @@ -28,14 +30,10 @@ PROFILE_FILE_ENVIRONMENT = "DATA_DESIGNER_SLURM_PROFILE_FILE" DEFAULT_PROFILE_FILE_NAME = ".data-designer-slurm-profile.yml" -_Config = TypeVar("_Config", DataDesignerSlurmConfig, SlurmProfileCatalog) +_ConfigT = TypeVar("_ConfigT", DataDesignerSlurmConfig, SlurmProfileCatalog) _HostnameResolver = Callable[[], tuple[str, ...]] -class ConfigLoadError(ValueError): - """Raised when a local Slurm configuration file is not strict and valid.""" - - class _StrictYamlLoader(yaml.SafeLoader): pass @@ -48,14 +46,14 @@ def _construct_unique_mapping( mapping: dict[object, object] = {} for key_node, value_node in node.value: if key_node.tag == "tag:yaml.org,2002:merge": - raise ConfigLoadError("YAML merge keys are not supported") + raise SlurmConfigLoadError("YAML merge keys are not supported") key = loader.construct_object(key_node, deep=deep) try: duplicate = key in mapping - except TypeError as error: - raise ConfigLoadError("configuration mapping keys must be scalar values") from error + except TypeError: + raise SlurmConfigLoadError("configuration mapping keys must be scalar values") from None if duplicate: - raise ConfigLoadError(f"duplicate configuration key: {key!r}") + raise SlurmConfigLoadError("duplicate configuration key") mapping[key] = loader.construct_object(value_node, deep=deep) return mapping @@ -91,10 +89,10 @@ def resolve_profile( try: sources = sum(source is not None for source in (profile, catalog, profile_file)) if sources > 1: - raise ConfigLoadError("profile, catalog, and profile_file are mutually exclusive") + raise SlurmConfigLoadError("profile, catalog, and profile_file are mutually exclusive") if profile is not None: if cluster is not None: - raise ConfigLoadError("an injected profile cannot be combined with cluster selection") + raise SlurmConfigLoadError("an injected profile cannot be combined with cluster selection") return injected_profile(profile) catalog_path: str | None = None @@ -119,18 +117,18 @@ def resolve_profile( hostnames=normalized_hostnames, catalog_path=catalog_path, ) - except ConfigLoadError: + except SlurmConfigLoadError: raise except ValueError as error: - raise ConfigLoadError(str(error)) from error + raise SlurmConfigLoadError(str(error)) from None -def _load_config(path: str | Path, config_type: type[_Config]) -> _Config: +def _load_config(path: str | Path, config_type: type[_ConfigT]) -> _ConfigT: resolved_path = _normalize_file_path(path) try: contents = resolved_path.read_text(encoding="utf-8") - except OSError as error: - raise ConfigLoadError(f"cannot read configuration file {resolved_path}") from error + except OSError: + raise SlurmConfigLoadError(f"cannot read configuration file {resolved_path}") from None try: payload = _parse_mapping(contents, suffix=resolved_path.suffix) if config_type is DataDesignerSlurmConfig: @@ -138,10 +136,13 @@ def _load_config(path: str | Path, config_type: type[_Config]) -> _Config: else: _reject_environment_interpolation(payload) return config_type.model_validate(payload) - except ConfigLoadError: + except SlurmConfigLoadError: raise - except (ValidationError, json.JSONDecodeError, yaml.YAMLError) as error: - raise ConfigLoadError(f"invalid configuration file {resolved_path}: {error}") from error + except ValidationError as error: + message = format_validation_error(error, subject=f"configuration file {resolved_path}") + raise SlurmConfigLoadError(message) from None + except (json.JSONDecodeError, yaml.YAMLError) as error: + raise SlurmConfigLoadError(f"invalid configuration file {resolved_path}: {format_parse_error(error)}") from None def _parse_mapping(contents: str, *, suffix: str) -> dict[str, object]: @@ -150,10 +151,10 @@ def _parse_mapping(contents: str, *, suffix: str) -> dict[str, object]: else: events = yaml.parse(contents, Loader=yaml.SafeLoader) if any(getattr(event, "anchor", None) is not None for event in events): - raise ConfigLoadError("YAML anchors and aliases are not supported") + raise SlurmConfigLoadError("YAML anchors and aliases are not supported") payload = yaml.load(contents, Loader=_StrictYamlLoader) if not isinstance(payload, dict) or any(not isinstance(key, str) for key in payload): - raise ConfigLoadError("configuration root must be an object with string keys") + raise SlurmConfigLoadError("configuration root must be an object with string keys") return cast(dict[str, object], payload) @@ -161,14 +162,14 @@ def _unique_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]: result: dict[str, object] = {} for key, value in pairs: if key in result: - raise ConfigLoadError(f"duplicate configuration key: {key!r}") + raise SlurmConfigLoadError("duplicate configuration key") result[key] = value return result def _reject_environment_interpolation(value: object) -> None: if isinstance(value, str) and "${" in value: - raise ConfigLoadError("environment interpolation is not supported") + raise SlurmConfigLoadError("environment interpolation is not supported") if isinstance(value, Mapping): for key, item in value.items(): _reject_environment_interpolation(key) @@ -200,7 +201,7 @@ def _resolve_profile_path( if source is None: environment_path = environ.get(PROFILE_FILE_ENVIRONMENT) if environment_path is not None and not environment_path: - raise ConfigLoadError(f"{PROFILE_FILE_ENVIRONMENT} must not be empty") + raise SlurmConfigLoadError(f"{PROFILE_FILE_ENVIRONMENT} must not be empty") source = environment_path if source is None: home = Path.home() if home_directory is None else Path(home_directory) @@ -211,7 +212,7 @@ def _resolve_profile_path( def _normalize_file_path(path: str | Path) -> Path: resolved = Path(path).expanduser().resolve() if resolved.suffix not in {".json", ".yaml", ".yml"}: - raise ConfigLoadError("configuration path must end in .json, .yaml, or .yml") + raise SlurmConfigLoadError("configuration path must end in .json, .yaml, or .yml") return resolved 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 cad317e73..04203616a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -38,11 +38,11 @@ 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") +_KeyT = TypeVar("_KeyT") +_ValueT = TypeVar("_ValueT") -class _FrozenList(list[_Value]): +class _FrozenList(list[_ValueT]): """List that retains JSON compatibility without exposing mutation.""" def _immutable(self, *args: object, **kwargs: object) -> None: @@ -63,7 +63,7 @@ def _immutable(self, *args: object, **kwargs: object) -> None: sort = _immutable -class _FrozenDict(dict[_Key, _Value]): +class _FrozenDict(dict[_KeyT, _ValueT]): """Dictionary that retains JSON compatibility without exposing mutation.""" def _immutable(self, *args: object, **kwargs: object) -> None: @@ -170,7 +170,7 @@ def compute_sha256(value: object) -> Sha256Digest: return hashlib.sha256(canonical_json(value)).hexdigest() -def compute_pretty_sha256(value: object) -> Sha256Digest: +def compute_serialized_json_sha256(value: object) -> Sha256Digest: """Compute the digest of deterministic persisted JSON bytes.""" return hashlib.sha256(pretty_json(value).encode("utf-8")).hexdigest() @@ -293,7 +293,6 @@ class ResumeWorkspace(ContractValue): "Sha256Digest", "ShardId", "canonical_json", - "compute_pretty_sha256", "compute_sha256", "pretty_json", "validate_absolute_path", 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 a1c46d7e6..46dfcbcea 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 @@ -5,14 +5,6 @@ from __future__ import annotations -from data_designer.slurm.planning.compiler import ( - ConfigurationResolutionError, - EffectiveDataDesignerSlurmConfig, - PlanCompilationError, - SlurmRunCompiler, - compile_slurm_run_plan, - resolve_slurm_config, -) from data_designer.slurm.planning.models import ( ArtifactReference, LockedPackage, @@ -31,15 +23,10 @@ ResolvedTopology, ResumeWorkspace, ) -from data_designer.slurm.planning.validation import PlanContractError, validate_resolved_plan __all__ = [ "ArtifactReference", - "ConfigurationResolutionError", - "EffectiveDataDesignerSlurmConfig", "LockedPackage", - "PlanContractError", - "PlanCompilationError", "PlannedShard", "PortClaim", "RecordRange", @@ -54,8 +41,4 @@ "ResolvedSubmission", "ResolvedTopology", "ResumeWorkspace", - "SlurmRunCompiler", - "compile_slurm_run_plan", - "resolve_slurm_config", - "validate_resolved_plan", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py new file mode 100644 index 000000000..21f59e381 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Submission-safe identity extraction for Data Designer builder payloads.""" + +from __future__ import annotations + +from pydantic import JsonValue + +from data_designer.slurm.contracts import ModelAlias, Sha256Digest, compute_serialized_json_sha256 + + +def get_declared_model_aliases(builder: dict[str, JsonValue]) -> tuple[ModelAlias, ...]: + """Return aliases declared by the known ``model_configs`` envelope.""" + 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"]) + return tuple(model_aliases) + + +def get_persisted_builder_identity( + builder: dict[str, JsonValue], +) -> tuple[tuple[ModelAlias, ...], Sha256Digest]: + """Return declared aliases and the digest of persisted builder JSON.""" + return get_declared_model_aliases(builder), compute_serialized_json_sha256(builder) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index 03a566865..b4008d60f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -1,219 +1,37 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Pure authored-configuration resolution and deterministic plan compilation.""" +"""Pure deterministic Slurm plan compilation.""" from __future__ import annotations import posixpath -from typing import Annotated -from pydantic import JsonValue, PositiveInt, StringConstraints, model_validator +from pydantic import ValidationError -from data_designer.config import RunConfig -from data_designer.slurm.config.images import ClientImageInspection, ImageKind -from data_designer.slurm.config.profiles import SelectedSlurmProfile -from data_designer.slurm.config.run import BuilderInput, DataDesignerSlurmConfig +from data_designer.slurm._errors import format_validation_error from data_designer.slurm.contracts import ( ArtifactReference, - ContractValue, - Identifier, RecordRange, ResumeWorkspace, - compute_pretty_sha256, - compute_sha256, + compute_serialized_json_sha256, ) +from data_designer.slurm.planning.errors import SlurmPlanCompilationError from data_designer.slurm.planning.models import ( PlannedShard, PortClaim, - ResolvedBuilderInput, ResolvedClient, - ResolvedDependencyLock, ResolvedDeployment, - ResolvedImage, - ResolvedInvocation, - ResolvedOutput, ResolvedSlurmRunPlan, - ResolvedSubmission, ResolvedTopology, - _extract_builder_aliases, - _extract_builder_identity, ) +from data_designer.slurm.planning.resolution import EffectiveDataDesignerSlurmConfig from data_designer.slurm.planning.validation import validate_resolved_plan _LOGICAL_ENDPOINT_PORT = 17000 _HTTP_PORT = 18000 _RENDEZVOUS_PORT = 19000 _PORT_RANGE_SIZE = 1000 -_SHARDABLE_COLUMN_TYPES = frozenset( - { - "embedding", - "expression", - "llm-code", - "llm-judge", - "llm-structured", - "llm-text", - "sampler", - "seed-dataset", - "validation", - } -) -_COMPATIBILITY_RUN_DEFAULTS: dict[str, JsonValue] = { - "buffer_size": 16384, - "disable_early_shutdown": True, - "display_tui": False, - "max_conversation_correction_steps": 0, - "max_conversation_restarts": 0, - "otel_metrics_port": None, - "shutdown_error_rate": 1.0, -} - - -class ConfigurationResolutionError(ValueError): - """Raised when resolved inputs do not match one authored declaration.""" - - -class PlanCompilationError(ValueError): - """Raised when one effective configuration cannot produce a valid plan.""" - - -class EffectiveDataDesignerSlurmConfig(ContractValue): - """Fully materialized, side-effect-free input to the plan compiler.""" - - run_id: Identifier - package_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] - authored: DataDesignerSlurmConfig - selected_profile: SelectedSlurmProfile - resolved_gpus_per_node: PositiveInt - builder: ResolvedBuilderInput - builder_payload: dict[str, JsonValue] | None = None - invocation: ResolvedInvocation - client_image: ResolvedImage - deployment_images: tuple[ResolvedImage, ...] - dependency_lock: ResolvedDependencyLock - submission: ResolvedSubmission - output: ResolvedOutput - runtime_bundle: ArtifactReference - - @model_validator(mode="after") - def validate_resolution(self) -> EffectiveDataDesignerSlurmConfig: - workspace_root = self.selected_profile.profile.workspace_root - run_root = posixpath.join(workspace_root, "runs", self.run_id) - profile_gpus = self.selected_profile.profile.gpus_per_node - if profile_gpus != "auto" and profile_gpus != self.resolved_gpus_per_node: - raise ValueError("resolved GPU count does not match the selected profile") - if self.client_image.kind is not ImageKind.CLIENT: - raise ValueError("resolved client image must contain client inspection facts") - if self.client_image.authored_ref != self.authored.client.image: - raise ValueError("resolved client image does not match the authored reference") - if len(self.deployment_images) != len(self.authored.deployments): - raise ValueError("resolved serving images must match the authored deployment count") - for deployment, image in zip(self.authored.deployments, self.deployment_images, strict=True): - if image.kind is not ImageKind.SERVING: - raise ValueError("resolved deployment image must contain serving inspection facts") - if image.authored_ref != deployment.server.image: - raise ValueError("resolved deployment image does not match the authored reference") - if self.builder_payload is not None and self.authored.builder.source is None: - raise ValueError("only sourced builder input may retain a resolved payload") - if self.authored.builder.source is not None and self.builder_payload is None: - raise ValueError("sourced builder input requires its resolved payload") - if self.authored.builder.source is not None: - assert self.builder_payload is not None - validated_payload = BuilderInput(inline=self.builder_payload).inline - assert validated_payload is not None - aliases, referenced_aliases, digest = _extract_builder_identity(validated_payload) - if self.builder.authored_source != self.authored.builder.source or self.builder.source is None: - raise ValueError("resolved builder source does not match the authored input") - expected_path = posixpath.join(run_root, "builder-config.json") - if self.builder.source.path != expected_path: - raise ValueError("resolved builder artifact path does not match the package-managed run") - if self.builder.model_aliases != aliases: - raise ValueError("resolved model aliases do not match the sourced builder payload") - if self.builder.referenced_model_aliases != referenced_aliases: - raise ValueError("resolved referenced aliases do not match the sourced builder payload") - if self.builder.content_sha256 != digest: - raise ValueError("resolved builder digest does not match the sourced builder payload") - _validate_sharding_constraints(self.authored, builder_payload=self.builder_payload) - expected_invocation = ResolvedInvocation( - authored=self.authored.invocation, - effective_run_config=_materialize_run_config(self.authored), - ) - if self.invocation != expected_invocation: - raise ValueError("resolved invocation does not match the authored invocation") - expected_output = ResolvedOutput( - root=self.authored.output.root or posixpath.join(run_root, "output"), - format=self.authored.output.format, - partitions=self.authored.output.partitions, - require_exact_record_count=self.authored.output.require_exact_record_count, - ) - if self.output != expected_output: - raise ValueError("resolved output does not match the authored output") - _validate_output_destination(self.output.root, workspace_root, run_root) - if self.output.partitions > self.authored.invocation.num_records: - raise ValueError("output partitions must not exceed requested records") - runtime_root = posixpath.join(workspace_root, "runtime") - if not _is_below(self.runtime_bundle.path, runtime_root) or not self.runtime_bundle.path.endswith(".tar.gz"): - raise ValueError("runtime bundle must be a tar archive below the selected workspace runtime root") - return self - - -def resolve_slurm_config( - authored: DataDesignerSlurmConfig, - *, - selected_profile: SelectedSlurmProfile, - client_image: ResolvedImage, - deployment_images: tuple[ResolvedImage, ...], - dependency_lock: ResolvedDependencyLock, - runtime_bundle: ArtifactReference, - run_id: str, - package_version: str, - resolved_gpus_per_node: int | None = None, - builder_payload: dict[str, JsonValue] | None = None, -) -> EffectiveDataDesignerSlurmConfig: - """Materialize all non-secret defaults from explicitly supplied resolved inputs.""" - try: - gpus_per_node = _resolve_gpu_count(selected_profile, resolved_gpus_per_node) - run_root = posixpath.join(selected_profile.profile.workspace_root, "runs", run_id) - builder = _resolve_builder(authored, run_root=run_root, builder_payload=builder_payload) - invocation = ResolvedInvocation( - authored=authored.invocation, - effective_run_config=_materialize_run_config(authored), - ) - submission = ResolvedSubmission( - account=authored.submission.account or selected_profile.profile.scheduler.account, - partition=authored.submission.partition or selected_profile.profile.scheduler.partition, - job_name=authored.submission.job_name, - time_limit=authored.submission.time_limit, - comment=authored.submission.comment, - ) - output_root = authored.output.root or posixpath.join(run_root, "output") - output = ResolvedOutput( - root=output_root, - format=authored.output.format, - partitions=authored.output.partitions, - require_exact_record_count=authored.output.require_exact_record_count, - ) - _validate_dependency_resolution(authored, client_image, dependency_lock) - return EffectiveDataDesignerSlurmConfig( - run_id=run_id, - package_version=package_version, - authored=authored, - selected_profile=selected_profile, - resolved_gpus_per_node=gpus_per_node, - builder=builder, - builder_payload=builder_payload, - invocation=invocation, - client_image=client_image, - deployment_images=deployment_images, - dependency_lock=dependency_lock, - submission=submission, - output=output, - runtime_bundle=runtime_bundle, - ) - except ConfigurationResolutionError: - raise - except ValueError as error: - raise ConfigurationResolutionError(str(error)) from error class SlurmRunCompiler: @@ -253,149 +71,20 @@ def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan plan, builder_payload=effective.builder_payload, ) - except PlanCompilationError: + except SlurmPlanCompilationError: raise + except ValidationError as error: + message = format_validation_error(error, subject="Slurm plan compilation") + raise SlurmPlanCompilationError(message) from None except ValueError as error: - raise PlanCompilationError(str(error)) from error - - -def compile_slurm_run_plan(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: - """Compile one effective configuration with the package-owned compiler.""" - return SlurmRunCompiler.compile(effective) - - -def _resolve_gpu_count(selected: SelectedSlurmProfile, resolved: int | None) -> int: - configured = selected.profile.gpus_per_node - if configured == "auto": - if type(resolved) is not int or resolved <= 0: - raise ConfigurationResolutionError("auto gpus_per_node requires one resolved positive integer") - return resolved - if resolved is not None and resolved != configured: - raise ConfigurationResolutionError("resolved GPU count conflicts with the selected profile") - return configured - - -def _resolve_builder( - authored: DataDesignerSlurmConfig, - *, - run_root: str, - builder_payload: dict[str, JsonValue] | None, -) -> ResolvedBuilderInput: - if authored.builder.inline is not None: - if builder_payload is not None: - raise ConfigurationResolutionError("inline builder input must not provide a separate payload") - aliases, referenced_aliases = _extract_builder_aliases(authored.builder.inline) - return ResolvedBuilderInput( - inline=authored.builder.inline, - content_sha256=compute_sha256(authored.builder.inline), - model_aliases=aliases, - referenced_model_aliases=referenced_aliases, - ) - if builder_payload is None: - raise ConfigurationResolutionError("sourced builder input requires its resolved payload") - validated_payload = BuilderInput(inline=builder_payload).inline - assert validated_payload is not None - aliases, referenced_aliases, digest = _extract_builder_identity(validated_payload) - source = ArtifactReference( - path=posixpath.join(run_root, "builder-config.json"), - sha256=digest, - ) - return ResolvedBuilderInput( - authored_source=authored.builder.source, - source=source, - content_sha256=source.sha256, - model_aliases=aliases, - referenced_model_aliases=referenced_aliases, - ) - - -def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: - values = dict(authored.invocation.run_config) - authored_early_shutdown = {"disable_early_shutdown", "shutdown_error_rate", "shutdown_error_window"}.intersection( - values - ) - for name, value in _COMPATIBILITY_RUN_DEFAULTS.items(): - if authored_early_shutdown and name in {"disable_early_shutdown", "shutdown_error_rate"}: - continue - values.setdefault(name, value) - return RunConfig.model_validate(values).model_dump(mode="json") - - -def _validate_sharding_constraints( - authored: DataDesignerSlurmConfig, - *, - builder_payload: dict[str, JsonValue] | None, -) -> None: - if authored.array_tasks.count == 1: - return - if authored.output.format != "parquet": - raise ConfigurationResolutionError("multi-shard runs require parquet output") - - payload = authored.builder.inline if authored.builder.inline is not None else builder_payload - assert payload is not None - data_designer = payload.get("data_designer", payload) - if not isinstance(data_designer, dict): - raise ConfigurationResolutionError("builder data_designer value must be an object") - if data_designer.get("processors"): - raise ConfigurationResolutionError("multi-shard runs do not support global processors") - if data_designer.get("profilers"): - raise ConfigurationResolutionError("multi-shard runs do not support global profilers") - - seed_config = data_designer.get("seed_config") - if isinstance(seed_config, dict): - if seed_config.get("sampling_strategy") == "shuffle": - raise ConfigurationResolutionError("multi-shard runs do not support shuffled seed input") - if seed_config.get("selection_strategy") is not None: - raise ConfigurationResolutionError("multi-shard runs do not support authored seed selection strategies") - if authored.invocation.input_bindings.seed_path is None: - raise ConfigurationResolutionError("multi-shard seed input requires a typed seed_path binding") - - columns = data_designer.get("columns", []) - if not isinstance(columns, list): - raise ConfigurationResolutionError("builder columns must be a list") - for column in columns: - if not isinstance(column, dict) or not isinstance(column.get("column_type"), str): - raise ConfigurationResolutionError("multi-shard runs require known column semantics") - column_type = column["column_type"] - if column_type == "image": - raise ConfigurationResolutionError("multi-shard runs do not support media output columns") - if column_type not in _SHARDABLE_COLUMN_TYPES: - raise ConfigurationResolutionError( - "multi-shard runs do not support custom, plugin, or unknown column semantics" - ) - if column_type == "validation" and column.get("validator_type") == "local_callable": - raise ConfigurationResolutionError("multi-shard runs do not support local callable validators") - - -def _validate_dependency_resolution( - authored: DataDesignerSlurmConfig, - client_image: ResolvedImage, - dependency_lock: ResolvedDependencyLock, -) -> None: - inspection = client_image.inspection.inspection - if not isinstance(inspection, ClientImageInspection): - raise ConfigurationResolutionError("resolved client image lacks dependency inspection facts") - if dependency_lock.client_image_sha256 != client_image.sha256: - raise ConfigurationResolutionError("dependency lock does not match the resolved client image") - if dependency_lock.python_abi != inspection.python_abi: - raise ConfigurationResolutionError("dependency lock Python ABI does not match the client image") - if dependency_lock.image_distributions != inspection.distributions: - raise ConfigurationResolutionError("dependency lock inventory does not match the client image") - requirements = authored.client.dependencies.requirements - if requirements is not None: - if dependency_lock.authored_source is not None or dependency_lock.source is not None: - raise ConfigurationResolutionError("inline requirements cannot resolve from an authored lock file") - if dependency_lock.authored_requirements != tuple(requirements): - raise ConfigurationResolutionError("dependency lock requirements do not match authored requirements") - elif dependency_lock.authored_source != authored.client.dependencies.lock_file or dependency_lock.source is None: - raise ConfigurationResolutionError("dependency lock source does not match the authored lock file") + raise SlurmPlanCompilationError(str(error)) from None def _compile_deployments( effective: EffectiveDataDesignerSlurmConfig, ) -> tuple[ResolvedDeployment, ...]: if len(effective.authored.deployments) > _PORT_RANGE_SIZE: - raise PlanCompilationError("deployment count exceeds the compiler-owned logical endpoint port range") + raise SlurmPlanCompilationError("deployment count exceeds the compiler-owned logical endpoint port range") resolved: list[ResolvedDeployment] = [] next_node_index = 0 for index, (authored, image) in enumerate( @@ -403,10 +92,10 @@ def _compile_deployments( ): tensor_parallel = authored.topology.tensor_parallel if effective.resolved_gpus_per_node % tensor_parallel: - raise PlanCompilationError("tensor_parallel must divide resolved GPUs per node") + raise SlurmPlanCompilationError("tensor_parallel must divide resolved GPUs per node") replicas_per_group = effective.resolved_gpus_per_node // tensor_parallel if replicas_per_group > _PORT_RANGE_SIZE: - raise PlanCompilationError("replica lanes exceed the compiler-owned deployment port range") + raise SlurmPlanCompilationError("replica lanes exceed the compiler-owned deployment port range") node_group_count = authored.resources.nodes // authored.topology.nodes_per_replica replica_count = node_group_count * replicas_per_group topology = ResolvedTopology( @@ -503,10 +192,10 @@ def _validate_port_claims( ports = client.ports + tuple(port for deployment in deployments for port in deployment.ports) addresses = tuple((port.node_index, port.port) for port in ports) if len(addresses) != len(set(addresses)): - raise PlanCompilationError("compiler-owned port claims collide on one node") + raise SlurmPlanCompilationError("compiler-owned port claims collide on one node") otel_port = effective.invocation.effective_run_config.get("otel_metrics_port") if type(otel_port) is int and (client.host_node_index, otel_port) in addresses: - raise PlanCompilationError("client OTEL metrics port collides with a compiler-owned port") + raise SlurmPlanCompilationError("client OTEL metrics port collides with a compiler-owned port") def _compile_shards(effective: EffectiveDataDesignerSlurmConfig) -> tuple[PlannedShard, ...]: @@ -526,7 +215,7 @@ def _compile_shards(effective: EffectiveDataDesignerSlurmConfig) -> tuple[Planne if seed_path is not None: partition = ArtifactReference( path=posixpath.join(shard_root, "input-partition.json"), - sha256=compute_pretty_sha256( + sha256=compute_serialized_json_sha256( { "record_range": record_range.model_dump(mode="json"), "seed_path": seed_path, @@ -547,27 +236,5 @@ def _compile_shards(effective: EffectiveDataDesignerSlurmConfig) -> tuple[Planne return tuple(shards) -def _validate_output_destination(output_root: str, workspace_root: str, run_root: str) -> None: - if not _is_below(output_root, workspace_root): - raise ConfigurationResolutionError("output root must be below the selected workspace_root") - reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) - if any(_paths_overlap(output_root, path) for path in reserved): - raise ConfigurationResolutionError("output root must not overlap package-managed workspace state") - runs_root = posixpath.join(workspace_root, "runs") - run_output_root = posixpath.join(run_root, "output") - if _paths_overlap(output_root, runs_root) and not ( - output_root == run_output_root or _is_below(output_root, run_output_root) - ): - raise ConfigurationResolutionError("output root must not overlap another package-managed run") - - def _run_root(effective: EffectiveDataDesignerSlurmConfig) -> str: return posixpath.join(effective.selected_profile.profile.workspace_root, "runs", effective.run_id) - - -def _is_below(path: str, root: str) -> bool: - return path != root and posixpath.commonpath((path, root)) == root - - -def _paths_overlap(left: str, right: str) -> bool: - return left == right or _is_below(left, right) or _is_below(right, left) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/errors.py new file mode 100644 index 000000000..63a6471a3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/errors.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalized Slurm planning errors.""" + +from __future__ import annotations + + +class SlurmPlanningError(ValueError): + """Base error for Slurm configuration resolution and plan validation.""" + + +class SlurmConfigResolutionError(SlurmPlanningError): + """Raised when resolved inputs do not match one authored declaration.""" + + +class SlurmPlanCompilationError(SlurmPlanningError): + """Raised when one effective configuration cannot produce a valid plan.""" + + +class SlurmPlanContractError(SlurmPlanningError): + """Raised when a resolved plan does not match its authored inputs.""" 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 1e56f9580..f49c580db 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 @@ -39,12 +39,12 @@ ResumeWorkspace, Sha256Digest, ShardId, - compute_pretty_sha256, compute_sha256, validate_absolute_path, validate_local_config_path, validate_plain_text, ) +from data_designer.slurm.planning.builder_identity import get_declared_model_aliases class ResolvedImage(ContractValue): @@ -154,7 +154,6 @@ class ResolvedBuilderInput(ContractValue): # Inline input uses canonical JSON; sourced input uses its persisted artifact bytes. content_sha256: Sha256Digest model_aliases: tuple[ModelAlias, ...] - referenced_model_aliases: tuple[ModelAlias, ...] = () @model_validator(mode="after") def validate_input(self) -> ResolvedBuilderInput: @@ -164,19 +163,14 @@ 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: + if self.model_aliases != get_declared_model_aliases(self.inline): 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 @@ -383,8 +377,6 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: raise ValueError("resolved deployment aliases must be unique") if set(aliases) != set(self.builder.model_aliases): raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") - 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))): @@ -482,51 +474,3 @@ def _validate_shards(self, run_root: str) -> None: 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)) - - -def _extract_builder_identity( - builder: dict[str, JsonValue], -) -> tuple[tuple[ModelAlias, ...], tuple[ModelAlias, ...], Sha256Digest]: - model_aliases, referenced_aliases = _extract_builder_aliases(builder) - digest = compute_pretty_sha256(builder) - return model_aliases, referenced_aliases, digest diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py new file mode 100644 index 000000000..7668af072 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure authored-configuration resolution for Slurm planning.""" + +from __future__ import annotations + +import posixpath +from typing import Annotated + +from pydantic import JsonValue, PositiveInt, StringConstraints, TypeAdapter, ValidationError + +from data_designer.config import RunConfig +from data_designer.slurm._errors import format_validation_error +from data_designer.slurm.config.images import ClientImageInspection, ImageKind +from data_designer.slurm.config.profiles import SelectedSlurmProfile +from data_designer.slurm.config.run import BuilderInput, DataDesignerSlurmConfig +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractValue, + Identifier, + compute_sha256, +) +from data_designer.slurm.planning.builder_identity import ( + get_declared_model_aliases, + get_persisted_builder_identity, +) +from data_designer.slurm.planning.errors import SlurmConfigResolutionError +from data_designer.slurm.planning.models import ( + ResolvedBuilderInput, + ResolvedDependencyLock, + ResolvedImage, + ResolvedInvocation, + ResolvedOutput, + ResolvedSubmission, +) + +_SHARDABLE_COLUMN_TYPES = frozenset( + { + "embedding", + "expression", + "llm-code", + "llm-judge", + "llm-structured", + "llm-text", + "sampler", + "seed-dataset", + "validation", + } +) +_COMPATIBILITY_RUN_DEFAULTS: dict[str, JsonValue] = { + "buffer_size": 16384, + "disable_early_shutdown": True, + "display_tui": False, + "max_conversation_correction_steps": 0, + "max_conversation_restarts": 0, + "otel_metrics_port": None, + "shutdown_error_rate": 1.0, +} +_RUN_ID_ADAPTER = TypeAdapter(Identifier) +_PACKAGE_VERSION_ADAPTER = TypeAdapter(Annotated[str, StringConstraints(min_length=1, max_length=128)]) + + +class EffectiveDataDesignerSlurmConfig(ContractValue): + """Fully materialized, side-effect-free input to the plan compiler.""" + + run_id: Identifier + package_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + authored: DataDesignerSlurmConfig + selected_profile: SelectedSlurmProfile + resolved_gpus_per_node: PositiveInt + builder: ResolvedBuilderInput + builder_payload: dict[str, JsonValue] | None = None + invocation: ResolvedInvocation + client_image: ResolvedImage + deployment_images: tuple[ResolvedImage, ...] + dependency_lock: ResolvedDependencyLock + submission: ResolvedSubmission + output: ResolvedOutput + runtime_bundle: ArtifactReference + + +def resolve_slurm_config( + authored: DataDesignerSlurmConfig, + *, + selected_profile: SelectedSlurmProfile, + client_image: ResolvedImage, + deployment_images: tuple[ResolvedImage, ...], + dependency_lock: ResolvedDependencyLock, + runtime_bundle: ArtifactReference, + run_id: str, + package_version: str, + resolved_gpus_per_node: int | None = None, + builder_payload: dict[str, JsonValue] | None = None, +) -> EffectiveDataDesignerSlurmConfig: + """Materialize all non-secret defaults from explicitly supplied resolved inputs.""" + try: + run_id = _RUN_ID_ADAPTER.validate_python(run_id, strict=True) + package_version = _PACKAGE_VERSION_ADAPTER.validate_python(package_version, strict=True) + gpus_per_node = _resolve_gpu_count(selected_profile, resolved_gpus_per_node) + workspace_root = selected_profile.profile.workspace_root + run_root = posixpath.join(workspace_root, "runs", run_id) + builder = _resolve_builder(authored, run_root=run_root, builder_payload=builder_payload) + invocation = ResolvedInvocation( + authored=authored.invocation, + effective_run_config=_materialize_run_config(authored), + ) + submission = ResolvedSubmission( + account=authored.submission.account or selected_profile.profile.scheduler.account, + partition=authored.submission.partition or selected_profile.profile.scheduler.partition, + job_name=authored.submission.job_name, + time_limit=authored.submission.time_limit, + comment=authored.submission.comment, + ) + output = ResolvedOutput( + root=authored.output.root or posixpath.join(run_root, "output"), + format=authored.output.format, + partitions=authored.output.partitions, + require_exact_record_count=authored.output.require_exact_record_count, + ) + _validate_dependency_resolution(authored, client_image, dependency_lock) + _validate_resolved_images(authored, client_image, deployment_images) + _validate_sharding_constraints(authored, builder_payload=builder_payload) + _validate_output_destination(output.root, workspace_root, run_root) + if output.partitions > authored.invocation.num_records: + raise SlurmConfigResolutionError("output partitions must not exceed requested records") + runtime_root = posixpath.join(workspace_root, "runtime") + if not _is_below(runtime_bundle.path, runtime_root) or not runtime_bundle.path.endswith(".tar.gz"): + raise SlurmConfigResolutionError( + "runtime bundle must be a tar archive below the selected workspace runtime root" + ) + return EffectiveDataDesignerSlurmConfig( + run_id=run_id, + package_version=package_version, + authored=authored, + selected_profile=selected_profile, + resolved_gpus_per_node=gpus_per_node, + builder=builder, + builder_payload=builder_payload, + invocation=invocation, + client_image=client_image, + deployment_images=deployment_images, + dependency_lock=dependency_lock, + submission=submission, + output=output, + runtime_bundle=runtime_bundle, + ) + except SlurmConfigResolutionError: + raise + except ValidationError as error: + message = format_validation_error(error, subject="Slurm configuration resolution") + raise SlurmConfigResolutionError(message) from None + except ValueError as error: + raise SlurmConfigResolutionError(str(error)) from None + + +def _resolve_gpu_count(selected: SelectedSlurmProfile, resolved: int | None) -> int: + configured = selected.profile.gpus_per_node + if configured == "auto": + if type(resolved) is not int or resolved <= 0: + raise SlurmConfigResolutionError("auto gpus_per_node requires one resolved positive integer") + return resolved + if resolved is not None and resolved != configured: + raise SlurmConfigResolutionError("resolved GPU count conflicts with the selected profile") + return configured + + +def _resolve_builder( + authored: DataDesignerSlurmConfig, + *, + run_root: str, + builder_payload: dict[str, JsonValue] | None, +) -> ResolvedBuilderInput: + if authored.builder.inline is not None: + if builder_payload is not None: + raise SlurmConfigResolutionError("inline builder input must not provide a separate payload") + return ResolvedBuilderInput( + inline=authored.builder.inline, + content_sha256=compute_sha256(authored.builder.inline), + model_aliases=get_declared_model_aliases(authored.builder.inline), + ) + if builder_payload is None: + raise SlurmConfigResolutionError("sourced builder input requires its resolved payload") + validated_payload = BuilderInput(inline=builder_payload).inline + assert validated_payload is not None + aliases, digest = get_persisted_builder_identity(validated_payload) + source = ArtifactReference( + path=posixpath.join(run_root, "builder-config.json"), + sha256=digest, + ) + return ResolvedBuilderInput( + authored_source=authored.builder.source, + source=source, + content_sha256=source.sha256, + model_aliases=aliases, + ) + + +def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: + values = dict(authored.invocation.run_config) + authored_early_shutdown = {"disable_early_shutdown", "shutdown_error_rate", "shutdown_error_window"}.intersection( + values + ) + for name, value in _COMPATIBILITY_RUN_DEFAULTS.items(): + if authored_early_shutdown and name in {"disable_early_shutdown", "shutdown_error_rate"}: + continue + values.setdefault(name, value) + return RunConfig.model_validate(values).model_dump(mode="json") + + +def _validate_sharding_constraints( + authored: DataDesignerSlurmConfig, + *, + builder_payload: dict[str, JsonValue] | None, +) -> None: + if authored.array_tasks.count == 1: + return + if authored.output.format != "parquet": + raise SlurmConfigResolutionError("multi-shard runs require parquet output") + + payload = authored.builder.inline if authored.builder.inline is not None else builder_payload + assert payload is not None + data_designer = payload.get("data_designer", payload) + if not isinstance(data_designer, dict): + raise SlurmConfigResolutionError("builder data_designer value must be an object") + if data_designer.get("processors"): + raise SlurmConfigResolutionError("multi-shard runs do not support global processors") + if data_designer.get("profilers"): + raise SlurmConfigResolutionError("multi-shard runs do not support global profilers") + + seed_config = data_designer.get("seed_config") + if isinstance(seed_config, dict): + if seed_config.get("sampling_strategy") == "shuffle": + raise SlurmConfigResolutionError("multi-shard runs do not support shuffled seed input") + if seed_config.get("selection_strategy") is not None: + raise SlurmConfigResolutionError("multi-shard runs do not support authored seed selection strategies") + if authored.invocation.input_bindings.seed_path is None: + raise SlurmConfigResolutionError("multi-shard seed input requires a typed seed_path binding") + + columns = data_designer.get("columns", []) + if not isinstance(columns, list): + raise SlurmConfigResolutionError("builder columns must be a list") + for column in columns: + if not isinstance(column, dict) or not isinstance(column.get("column_type"), str): + raise SlurmConfigResolutionError("multi-shard runs require known column semantics") + column_type = column["column_type"] + if column_type == "image": + raise SlurmConfigResolutionError("multi-shard runs do not support media output columns") + if column_type not in _SHARDABLE_COLUMN_TYPES: + raise SlurmConfigResolutionError( + "multi-shard runs do not support custom, plugin, or unknown column semantics" + ) + if column_type == "validation" and column.get("validator_type") == "local_callable": + raise SlurmConfigResolutionError("multi-shard runs do not support local callable validators") + + +def _validate_dependency_resolution( + authored: DataDesignerSlurmConfig, + client_image: ResolvedImage, + dependency_lock: ResolvedDependencyLock, +) -> None: + inspection = client_image.inspection.inspection + if not isinstance(inspection, ClientImageInspection): + raise SlurmConfigResolutionError("resolved client image lacks dependency inspection facts") + if dependency_lock.client_image_sha256 != client_image.sha256: + raise SlurmConfigResolutionError("dependency lock does not match the resolved client image") + if dependency_lock.python_abi != inspection.python_abi: + raise SlurmConfigResolutionError("dependency lock Python ABI does not match the client image") + if dependency_lock.image_distributions != inspection.distributions: + raise SlurmConfigResolutionError("dependency lock inventory does not match the client image") + requirements = authored.client.dependencies.requirements + if requirements is not None: + if dependency_lock.authored_source is not None or dependency_lock.source is not None: + raise SlurmConfigResolutionError("inline requirements cannot resolve from an authored lock file") + if dependency_lock.authored_requirements != tuple(requirements): + raise SlurmConfigResolutionError("dependency lock requirements do not match authored requirements") + elif dependency_lock.authored_source != authored.client.dependencies.lock_file or dependency_lock.source is None: + raise SlurmConfigResolutionError("dependency lock source does not match the authored lock file") + + +def _validate_resolved_images( + authored: DataDesignerSlurmConfig, + client_image: ResolvedImage, + deployment_images: tuple[ResolvedImage, ...], +) -> None: + if client_image.kind is not ImageKind.CLIENT: + raise SlurmConfigResolutionError("resolved client image must contain client inspection facts") + if client_image.authored_ref != authored.client.image: + raise SlurmConfigResolutionError("resolved client image does not match the authored reference") + if len(deployment_images) != len(authored.deployments): + raise SlurmConfigResolutionError("resolved serving images must match the authored deployment count") + for deployment, image in zip(authored.deployments, deployment_images, strict=True): + if image.kind is not ImageKind.SERVING: + raise SlurmConfigResolutionError("resolved deployment image must contain serving inspection facts") + if image.authored_ref != deployment.server.image: + raise SlurmConfigResolutionError("resolved deployment image does not match the authored reference") + + +def _validate_output_destination(output_root: str, workspace_root: str, run_root: str) -> None: + if not _is_below(output_root, workspace_root): + raise SlurmConfigResolutionError("output root must be below the selected workspace_root") + reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) + if any(_paths_overlap(output_root, path) for path in reserved): + raise SlurmConfigResolutionError("output root must not overlap package-managed workspace state") + runs_root = posixpath.join(workspace_root, "runs") + run_output_root = posixpath.join(run_root, "output") + if _paths_overlap(output_root, runs_root) and not ( + output_root == run_output_root or _is_below(output_root, run_output_root) + ): + raise SlurmConfigResolutionError("output root must not overlap another package-managed run") + + +def _is_below(path: str, root: str) -> bool: + return path != root and posixpath.commonpath((path, root)) == root + + +def _paths_overlap(left: str, right: str) -> bool: + return left == right or _is_below(left, right) or _is_below(right, left) 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 e1d17fd8d..0351c21c9 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 @@ -8,17 +8,14 @@ from data_designer.config import RunConfig from data_designer.slurm.config.images import ClientImageInspection from data_designer.slurm.config.run import DataDesignerSlurmConfig +from data_designer.slurm.planning.builder_identity import get_persisted_builder_identity +from data_designer.slurm.planning.errors import SlurmPlanContractError from data_designer.slurm.planning.models import ( ResolvedDependencyLock, ResolvedSlurmRunPlan, - _extract_builder_identity, ) -class PlanContractError(ValueError): - """Raised when a resolved plan does not match its authored inputs.""" - - def validate_resolved_plan( authored: DataDesignerSlurmConfig, dependency_lock: ResolvedDependencyLock, @@ -54,13 +51,9 @@ def validate_resolved_plan( "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, digest = _extract_builder_identity(builder_payload) + raise SlurmPlanContractError("sourced builder validation requires its resolved payload") + model_aliases, digest = get_persisted_builder_identity(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", - ) _require(plan.builder.content_sha256 == digest, "resolved builder digest does not match builder source") expected_account = authored.submission.account or plan.selected_profile.profile.scheduler.account @@ -127,7 +120,7 @@ def validate_resolved_plan( def _require(condition: bool, message: str) -> None: if not condition: - raise PlanContractError(message) + raise SlurmPlanContractError(message) def _require_json_subset(actual: JsonValue, expected: JsonValue, *, path: str) -> None: diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 2e24a3ba7..d5794e33b 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -13,11 +13,11 @@ from data_designer.slurm.config import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, - ConfigBuilderError, - ConfigLoadError, DataDesignerSlurmConfig, DataDesignerSlurmConfigBuilder, ProfileSelectionSource, + SlurmConfigBuilderError, + SlurmConfigLoadError, SlurmProfileCatalog, load_profile_catalog, load_run_config, @@ -64,7 +64,7 @@ def test_builder_builds_without_file_discovery_or_serialization(tmp_path: Path) def test_builder_requires_complete_authored_intent() -> None: builder = DataDesignerSlurmConfigBuilder.from_builder_source("builder.json") - with pytest.raises(ConfigBuilderError, match="invocation, client, deployment"): + with pytest.raises(SlurmConfigBuilderError, match="invocation, client, deployment"): builder.build() @@ -79,7 +79,7 @@ def test_builder_write_config_round_trips_supported_formats(tmp_path: Path, suff def test_builder_rejects_unsupported_output_format(tmp_path: Path) -> None: - with pytest.raises(ConfigBuilderError, match="must end"): + with pytest.raises(SlurmConfigBuilderError, match="must end"): _config_builder().write_config(tmp_path / "run.txt") @@ -96,26 +96,38 @@ def test_builder_normalizes_invalid_authored_values( method: str, values: dict[str, object], ) -> None: - with pytest.raises(ConfigBuilderError): + with pytest.raises(SlurmConfigBuilderError): getattr(_config_builder(), method)(**values) def test_builder_validation_errors_hide_secret_inputs() -> None: secret = "super-secret-token" - with pytest.raises(ConfigBuilderError) as error: + with pytest.raises(SlurmConfigBuilderError) as error: _config_builder().with_client(image={"name": "dd-client", "api_key": secret}) assert secret not in str(error.value) - assert error.value.__cause__ is not None - assert secret not in str(error.value.__cause__) + assert error.value.__cause__ is None + + +def test_builder_custom_validation_errors_hide_secret_values() -> None: + secret = "super-secret-token" + + with pytest.raises(SlurmConfigBuilderError) as error: + _config_builder().with_client( + image={"name": "dd-client"}, + dependencies={"requirements": [f"pkg=={secret}"]}, + ) + + assert secret not in str(error.value) + assert error.value.__cause__ is None def test_builder_normalizes_write_failures(tmp_path: Path) -> None: path = tmp_path / "run.json" path.mkdir() - with pytest.raises(ConfigBuilderError, match="cannot write"): + with pytest.raises(SlurmConfigBuilderError, match="cannot write"): _config_builder().write_config(path) @@ -138,7 +150,7 @@ def test_strict_loader_rejects_ambiguous_yaml_and_json( path = tmp_path / f"run{suffix}" path.write_text(contents) - with pytest.raises(ConfigLoadError, match=message): + with pytest.raises(SlurmConfigLoadError, match=message): load_run_config(path) @@ -146,9 +158,9 @@ def test_strict_loader_rejects_non_object_and_unknown_extension(tmp_path: Path) json_path = tmp_path / "run.json" json_path.write_text("[]") - with pytest.raises(ConfigLoadError, match="root must be an object"): + with pytest.raises(SlurmConfigLoadError, match="root must be an object"): load_run_config(json_path) - with pytest.raises(ConfigLoadError, match="must end"): + with pytest.raises(SlurmConfigLoadError, match="must end"): load_run_config(tmp_path / "run.toml") @@ -159,12 +171,23 @@ def test_loader_validation_errors_hide_secret_inputs(tmp_path: Path) -> None: path = tmp_path / "run.json" path.write_text(json.dumps(payload)) - with pytest.raises(ConfigLoadError) as error: + with pytest.raises(SlurmConfigLoadError) as error: + load_run_config(path) + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + +def test_loader_parse_errors_hide_source_values(tmp_path: Path) -> None: + secret = "super-secret-token" + path = tmp_path / "run.yaml" + path.write_text(f"api_key: {secret}: x\n") + + with pytest.raises(SlurmConfigLoadError, match="invalid YAML at line 1") as error: load_run_config(path) assert secret not in str(error.value) - assert error.value.__cause__ is not None - assert secret not in str(error.value.__cause__) + assert error.value.__cause__ is None @pytest.mark.parametrize("suffix", [".json", ".yaml", ".yml"]) @@ -228,13 +251,13 @@ def test_injected_profile_bypasses_catalog_lookup(profile_catalog: SlurmProfileC def test_profile_resolution_rejects_conflicting_or_empty_sources( profile_catalog: SlurmProfileCatalog, ) -> None: - with pytest.raises(ConfigLoadError, match="mutually exclusive"): + with pytest.raises(SlurmConfigLoadError, match="mutually exclusive"): resolve_profile(catalog=profile_catalog, profile_file="profile.json") - with pytest.raises(ConfigLoadError, match="must not be empty"): + with pytest.raises(SlurmConfigLoadError, match="must not be empty"): resolve_profile(environ={PROFILE_FILE_ENVIRONMENT: ""}) - with pytest.raises(ConfigLoadError, match="cluster selection"): + with pytest.raises(SlurmConfigLoadError, match="cluster selection"): resolve_profile(profile=profile_catalog.clusters["primary"], cluster="primary") - with pytest.raises(ConfigLoadError, match="unknown cluster"): + with pytest.raises(SlurmConfigLoadError, match="unknown cluster"): resolve_profile(catalog=profile_catalog, cluster="missing") @@ -245,7 +268,7 @@ def test_profile_resolution_normalizes_ambiguous_hostname_errors( payload["clusters"]["lab"]["host_patterns"] = ["*-login-*"] catalog = SlurmProfileCatalog.model_validate(payload) - with pytest.raises(ConfigLoadError, match="multiple clusters"): + with pytest.raises(SlurmConfigLoadError, match="multiple clusters"): resolve_profile(catalog=catalog, hostnames=("primary-login-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 0f2fd69b1..6d369e061 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 @@ -31,7 +31,6 @@ "generator", "judge" ], - "referenced_model_aliases": [], "source": null }, "client": { 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 f0794cc75..d6e35f670 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 @@ -25,7 +25,6 @@ "model_aliases": [ "generator" ], - "referenced_model_aliases": [], "source": null }, "client": { 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 102b07769..64e825e0f 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -10,16 +10,16 @@ from pydantic import ValidationError from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig -from data_designer.slurm.contracts import compute_pretty_sha256, compute_sha256 +from data_designer.slurm.contracts import compute_serialized_json_sha256, compute_sha256 from data_designer.slurm.planning import ( ArtifactReference, - PlanContractError, ResolvedDependencyLock, ResolvedDeployment, ResolvedSlurmRunPlan, ResolvedSubmission, - validate_resolved_plan, ) +from data_designer.slurm.planning.errors import SlurmPlanContractError +from data_designer.slurm.planning.validation import validate_resolved_plan def test_multi_node_plan_matches_authored_inputs( @@ -166,16 +166,6 @@ def test_plan_rejects_deployment_alias_missing_from_builder(multi_node_plan: Res 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 @@ -226,7 +216,7 @@ def test_sourced_builder_validation_requires_resolved_payload( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: sourced_authored = authored_run.model_copy(update={"builder": BuilderInput(source="builder.json")}) - builder_digest = compute_pretty_sha256(authored_run.builder.inline) + builder_digest = compute_serialized_json_sha256(authored_run.builder.inline) payload = multi_node_plan.model_dump(mode="json") payload["authored_config"]["sha256"] = sourced_authored.compute_sha256() payload["builder"] = { @@ -235,11 +225,10 @@ def test_sourced_builder_validation_requires_resolved_payload( "inline": None, "content_sha256": builder_digest, "model_aliases": ["generator", "judge"], - "referenced_model_aliases": [], } sourced_plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) - with pytest.raises(PlanContractError, match="resolved payload"): + with pytest.raises(SlurmPlanContractError, match="resolved payload"): validate_resolved_plan(sourced_authored, dependency_lock, sourced_plan) assert ( @@ -348,7 +337,7 @@ def test_cross_record_validation_rejects_authored_digest( } ) - with pytest.raises(PlanContractError, match="authored config digest"): + with pytest.raises(SlurmPlanContractError, match="authored config digest"): validate_resolved_plan(authored_run, dependency_lock, invalid) @@ -367,7 +356,7 @@ def test_cross_record_validation_rejects_dependency_lock_digest( ) invalid = multi_node_plan.model_copy(update={"client": client}) - with pytest.raises(PlanContractError, match="dependency lock digest"): + with pytest.raises(SlurmPlanContractError, match="dependency lock digest"): validate_resolved_plan(authored_run, dependency_lock, invalid) @@ -385,7 +374,7 @@ def test_cross_record_validation_binds_authored_lock_source( 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"): + with pytest.raises(SlurmPlanContractError, match="authored lock file"): validate_resolved_plan(authored, dependency_lock, plan) lock_payload = dependency_lock.model_dump(mode="json") @@ -408,7 +397,7 @@ def test_cross_record_validation_binds_authored_lock_source( ) } ) - with pytest.raises(PlanContractError, match="present for authored requirements"): + with pytest.raises(SlurmPlanContractError, match="present for authored requirements"): validate_resolved_plan(authored_run, matching_lock, unexpected_source_plan) matching_plan = plan.model_copy( @@ -441,7 +430,7 @@ def test_cross_record_validation_rejects_python_abi( ) invalid_plan = multi_node_plan.model_copy(update={"client": client}) - with pytest.raises(PlanContractError, match="Python ABI"): + with pytest.raises(SlurmPlanContractError, match="Python ABI"): validate_resolved_plan(authored_run, invalid_lock, invalid_plan) @@ -460,7 +449,7 @@ def test_cross_record_validation_rejects_image_inventory( ) invalid_plan = multi_node_plan.model_copy(update={"client": client}) - with pytest.raises(PlanContractError, match="image inventory"): + with pytest.raises(SlurmPlanContractError, match="image inventory"): validate_resolved_plan(authored_run, invalid_lock, invalid_plan) @@ -474,7 +463,7 @@ def test_cross_record_validation_rejects_changed_invocation( ) invalid = multi_node_plan.model_copy(update={"invocation": invocation}) - with pytest.raises(PlanContractError, match="invocation"): + with pytest.raises(SlurmPlanContractError, match="invocation"): validate_resolved_plan(authored_run, dependency_lock, invalid) @@ -487,5 +476,5 @@ def test_cross_record_validation_preserves_explicit_run_config( payload["invocation"]["effective_run_config"]["buffer_size"] = 16384 invalid = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) - with pytest.raises(PlanContractError, match="run_config.buffer_size"): + with pytest.raises(SlurmPlanContractError, match="run_config.buffer_size"): validate_resolved_plan(authored_run, dependency_lock, invalid) diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index 517a45ac5..674560f42 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -9,7 +9,7 @@ "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" + "sha256": "9bd1694bbe32fedec605de7e835a4b8506dc2b74be6bd640d4e2b8577951ed9c" }, "run_id": "run-single", "scheduler": { @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" + "sha256": "9bd1694bbe32fedec605de7e835a4b8506dc2b74be6bd640d4e2b8577951ed9c" }, "run_id": "run-single", "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index 0f8c67cc5..d041f3de4 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -22,17 +22,15 @@ SlurmProfileCatalog, select_profile, ) -from data_designer.slurm.contracts import compute_pretty_sha256 +from data_designer.slurm.contracts import compute_serialized_json_sha256 from data_designer.slurm.planning import ( ArtifactReference, - ConfigurationResolutionError, - EffectiveDataDesignerSlurmConfig, - PlanCompilationError, ResolvedDependencyLock, ResolvedSlurmRunPlan, - compile_slurm_run_plan, - resolve_slurm_config, ) +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.errors import SlurmConfigResolutionError, SlurmPlanCompilationError +from data_designer.slurm.planning.resolution import EffectiveDataDesignerSlurmConfig, resolve_slurm_config GOLDEN_DIRECTORY = Path(__file__).parents[1] / "contracts" / "golden" @@ -76,8 +74,8 @@ def test_compiler_reproduces_plan_goldens_byte_for_byte( expected = request.getfixturevalue(plan_fixture) effective = _resolve_fixture(authored, dependency_lock, expected) - first = compile_slurm_run_plan(effective) - second = compile_slurm_run_plan(effective) + first = SlurmRunCompiler.compile(effective) + second = SlurmRunCompiler.compile(effective) assert first.serialize_json() == (GOLDEN_DIRECTORY / golden_name).read_text() assert first.serialize_canonical_json() == second.serialize_canonical_json() @@ -97,7 +95,7 @@ def test_compiler_resolves_explicit_hostname_and_default_profile_selection( ) plans = tuple( - compile_slurm_run_plan( + SlurmRunCompiler.compile( _resolve_fixture( authored_run_single, dependency_lock_single, @@ -128,7 +126,7 @@ def test_auto_gpu_resolution_is_explicit_and_scheduler_free( update={"path": "/workspace/lab/runtime/runtime.tar.gz"} ) - with pytest.raises(ConfigurationResolutionError, match="auto gpus_per_node"): + with pytest.raises(SlurmConfigResolutionError, match="auto gpus_per_node"): _resolve_fixture( authored_run_single, dependency_lock_single, @@ -138,7 +136,7 @@ def test_auto_gpu_resolution_is_explicit_and_scheduler_free( resolved_gpus_per_node=None, ) - plan = compile_slurm_run_plan( + plan = SlurmRunCompiler.compile( _resolve_fixture( authored_run_single, dependency_lock_single, @@ -215,8 +213,8 @@ def test_compiler_rejects_tensor_parallelism_that_does_not_divide_gpu_shape( payload["deployments"][0]["topology"]["tensor_parallel"] = 3 authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(PlanCompilationError, match="tensor_parallel"): - compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + with pytest.raises(SlurmPlanCompilationError, match="tensor_parallel"): + SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) @pytest.mark.parametrize( @@ -240,64 +238,10 @@ def test_resolution_rejects_invalid_output_destinations( payload["output"].update(output_update) authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(ConfigurationResolutionError, match="output"): + with pytest.raises(SlurmConfigResolutionError, match="output"): _resolve_fixture(authored, dependency_lock_single, single_node_plan) -def test_effective_config_rejects_invalid_direct_construction( - authored_run: DataDesignerSlurmConfig, - dependency_lock: ResolvedDependencyLock, - multi_node_plan: ResolvedSlurmRunPlan, -) -> None: - effective = _resolve_fixture(authored_run, dependency_lock, multi_node_plan) - values = {name: getattr(effective, name) for name in EffectiveDataDesignerSlurmConfig.model_fields} - values["authored"] = authored_run.model_copy( - update={"output": authored_run.output.model_copy(update={"format": "jsonl"})} - ) - values["output"] = effective.output.model_copy(update={"format": "jsonl"}) - - with pytest.raises(ValueError, match="parquet output"): - EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] - - other_output = "/workspace/primary/runs/other-run/output" - values["authored"] = authored_run.model_copy( - update={"output": authored_run.output.model_copy(update={"root": other_output})} - ) - values["output"] = effective.output.model_copy(update={"root": other_output}) - - with pytest.raises(ValueError, match="another package-managed run"): - EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] - - values["authored"] = authored_run.model_copy( - update={"output": authored_run.output.model_copy(update={"partitions": 101})} - ) - values["output"] = effective.output.model_copy(update={"partitions": 101}) - - with pytest.raises(ValueError, match="requested records"): - EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] - - values["authored"] = authored_run - values["output"] = effective.output.model_copy(update={"format": "jsonl"}) - - with pytest.raises(ValueError, match="resolved output"): - EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] - - -def test_effective_config_rejects_invocation_drift( - authored_run_single: DataDesignerSlurmConfig, - dependency_lock_single: ResolvedDependencyLock, - single_node_plan: ResolvedSlurmRunPlan, -) -> None: - effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) - run_config = dict(effective.invocation.effective_run_config) - run_config["jinja_rendering_engine"] = "native" - values = {name: getattr(effective, name) for name in EffectiveDataDesignerSlurmConfig.model_fields} - values["invocation"] = effective.invocation.model_copy(update={"effective_run_config": run_config}) - - with pytest.raises(ValueError, match="resolved invocation"): - EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] - - def test_compiler_rejects_model_alias_missing_from_builder( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -307,8 +251,8 @@ def test_compiler_rejects_model_alias_missing_from_builder( payload["builder"]["inline"]["data_designer"]["model_configs"][0]["alias"] = "other" authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(PlanCompilationError, match="deployment alias"): - compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + with pytest.raises(SlurmPlanCompilationError, match="failed validation"): + SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) @pytest.mark.parametrize("sourced", [False, True], ids=["inline", "sourced"]) @@ -328,8 +272,8 @@ def test_compiler_rejects_builder_model_alias_without_deployment( payload["builder"] = {"source": "builder.json"} authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(PlanCompilationError, match="exactly cover"): - compile_slurm_run_plan( + with pytest.raises(SlurmPlanCompilationError, match="failed validation"): + SlurmRunCompiler.compile( _resolve_fixture( authored, dependency_lock_single, @@ -348,8 +292,8 @@ def test_compiler_rejects_otel_collision_before_runtime( payload["invocation"]["run_config"] = {"otel_metrics_port": 17000} authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(PlanCompilationError, match="OTEL"): - compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + with pytest.raises(SlurmPlanCompilationError, match="OTEL"): + SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) def test_sharded_seed_inputs_have_stable_ranges_and_partition_digests( @@ -362,8 +306,8 @@ def test_sharded_seed_inputs_have_stable_ranges_and_partition_digests( ) authored = authored_run.model_copy(update={"invocation": invocation}) - first = compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock, multi_node_plan)) - second = compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + first = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + second = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock, multi_node_plan)) assert [(shard.record_range.start_index, shard.record_range.end_index_exclusive) for shard in first.shards] == [ (0, 50), @@ -372,7 +316,7 @@ def test_sharded_seed_inputs_have_stable_ranges_and_partition_digests( assert all(shard.input_partition is not None for shard in first.shards) for shard in first.shards: assert shard.input_partition is not None - assert shard.input_partition.sha256 == compute_pretty_sha256( + assert shard.input_partition.sha256 == compute_serialized_json_sha256( { "record_range": shard.record_range.model_dump(mode="json"), "seed_path": "/datasets/seed.parquet", @@ -445,7 +389,7 @@ def test_resolution_rejects_unshardable_big_iron_fields( payload["output"].update(output_update) authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(ConfigurationResolutionError, match=message): + with pytest.raises(SlurmConfigResolutionError, match=message): _resolve_fixture(authored, dependency_lock, multi_node_plan) @@ -469,7 +413,7 @@ def test_resolution_rejects_real_global_builder_configs( payload["builder"]["inline"]["data_designer"][field] = data_designer[field] authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(ConfigurationResolutionError, match=field): + with pytest.raises(SlurmConfigResolutionError, match=field): _resolve_fixture(authored, dependency_lock, multi_node_plan) @@ -511,7 +455,7 @@ def test_resolution_rejects_unportable_multi_shard_columns( payload["builder"] = {"source": "builder.json"} authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(ConfigurationResolutionError, match=message): + with pytest.raises(SlurmConfigResolutionError, match=message): _resolve_fixture( authored, dependency_lock, @@ -533,7 +477,7 @@ def test_sharded_seed_binding_may_override_authored_source( payload["invocation"]["input_bindings"]["seed_path"] = "/datasets/override.parquet" authored = DataDesignerSlurmConfig.model_validate(payload) - plan = compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + plan = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock, multi_node_plan)) assert all(shard.input_partition is not None for shard in plan.shards) @@ -563,12 +507,45 @@ def test_single_shard_allows_non_collectable_big_iron_fields( payload["output"]["format"] = "jsonl" authored = DataDesignerSlurmConfig.model_validate(payload) - plan = compile_slurm_run_plan(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + plan = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) assert plan.array_tasks.count == 1 assert plan.output.format == "jsonl" +@pytest.mark.parametrize("sourced", [False, True], ids=["inline", "sourced"]) +def test_single_shard_defers_opaque_plugin_aliases_to_client_preflight( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + sourced: bool, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["columns"] = [ + { + "column_type": "plugin-column", + "fallback_model_alias": None, + "name": "plugin", + } + ] + builder_payload = None + if sourced: + builder_payload = payload["builder"]["inline"] + payload["builder"] = {"source": "builder.json"} + authored = DataDesignerSlurmConfig.model_validate(payload) + + plan = SlurmRunCompiler.compile( + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + ) + + assert plan.builder.model_aliases == ("generator",) + + def test_sourced_builder_is_resolved_to_one_digest_bound_run_input( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -578,7 +555,7 @@ def test_sourced_builder_is_resolved_to_one_digest_bound_run_input( assert isinstance(builder_payload, dict) authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) - plan = compile_slurm_run_plan( + plan = SlurmRunCompiler.compile( _resolve_fixture( authored, dependency_lock_single, @@ -600,7 +577,7 @@ def test_sourced_builder_is_resolved_to_one_digest_bound_run_input( ("aliases", "model aliases"), ], ) -def test_effective_config_rejects_sourced_builder_payload_identity_drift( +def test_compiler_rejects_sourced_builder_payload_identity_drift( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, single_node_plan: ResolvedSlurmRunPlan, @@ -620,11 +597,10 @@ def test_effective_config_rejects_sourced_builder_payload_identity_drift( drifted_payload["library_version"] = "drifted" else: drifted_payload["data_designer"]["model_configs"][0]["alias"] = "drifted" - values = {name: getattr(effective, name) for name in EffectiveDataDesignerSlurmConfig.model_fields} - values["builder_payload"] = drifted_payload + drifted = effective.model_copy(update={"builder_payload": drifted_payload}) - with pytest.raises(ValueError, match=message): - EffectiveDataDesignerSlurmConfig(**values) # type: ignore[arg-type] + with pytest.raises(SlurmPlanCompilationError, match=message): + SlurmRunCompiler.compile(drifted) def test_resolution_rejects_secret_values_in_sourced_builder_payload( @@ -637,7 +613,7 @@ def test_resolution_rejects_secret_values_in_sourced_builder_payload( builder_payload["data_designer"]["api_key"] = secret authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) - with pytest.raises(ConfigurationResolutionError, match="secret values") as error: + with pytest.raises(SlurmConfigResolutionError, match="failed validation") as error: _resolve_fixture( authored, dependency_lock_single, @@ -646,8 +622,7 @@ def test_resolution_rejects_secret_values_in_sourced_builder_payload( ) assert secret not in str(error.value) - assert error.value.__cause__ is not None - assert secret not in str(error.value.__cause__) + assert error.value.__cause__ is None def test_resolution_rejects_artifact_identity_mismatches( @@ -658,9 +633,9 @@ def test_resolution_rejects_artifact_identity_mismatches( wrong_lock = dependency_lock_single.model_copy(update={"client_image_sha256": "a" * 64}) wrong_runtime = ArtifactReference(path="/tmp/runtime.tar.gz", sha256="e" * 64) - with pytest.raises(ConfigurationResolutionError, match="client image"): + with pytest.raises(SlurmConfigResolutionError, match="client image"): _resolve_fixture(authored_run_single, wrong_lock, single_node_plan) - with pytest.raises(ConfigurationResolutionError, match="runtime bundle"): + with pytest.raises(SlurmConfigResolutionError, match="runtime bundle"): _resolve_fixture( authored_run_single, dependency_lock_single, @@ -674,7 +649,7 @@ def test_plan_contains_secret_references_without_credentials( dependency_lock: ResolvedDependencyLock, multi_node_plan: ResolvedSlurmRunPlan, ) -> None: - plan = compile_slurm_run_plan(_resolve_fixture(authored_run, dependency_lock, multi_node_plan)) + plan = SlurmRunCompiler.compile(_resolve_fixture(authored_run, dependency_lock, multi_node_plan)) serialized = plan.serialize_json() assert "PACKAGE_INDEX_TOKEN" in serialized diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 5a78c97b7..93ee2b213 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -11,7 +11,7 @@ set -Eeuo pipefail readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="dbcae1da6ce8ef6799faf2add8e08b1b093390ae08a37dd13e336da522b9a3fd" +readonly DD_PLAN_SHA256="9efc57d4ef15dc32ac5276cb97afb053155c5d43df03ed28db5fa09546bfa3eb" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index e10478f7d..3fc811602 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -11,7 +11,7 @@ set -Eeuo pipefail readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" +readonly DD_PLAN_SHA256="9bd1694bbe32fedec605de7e835a4b8506dc2b74be6bd640d4e2b8577951ed9c" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index c7d119389..9660aa7eb 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="faa1dac9b9b0423423c9d06a13d71bee339932ffa45d1e02a8a95012a7934520", + expected_fixture_sha256="d1875aa6705dd862ae005e10e346d0fd38752377a020d33d46351e6806bdf9f9", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="a4542e56b124c2346d0a92ebadc4e89b45d94c9355d7b86a4a1b05180d331a48", + expected_fixture_sha256="864188c9631f69073f93a2e6c75ce62be7412a9c941de48ba35a090e5880fc6a", ) From 079da69064e6db50130824344f3ffd95b732615c Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 27 Aug 2026 17:05:48 -0300 Subject: [PATCH 6/9] fix: unify Slurm builder artifact identity --- .../src/data_designer/slurm/config/loading.py | 3 ++ .../data_designer/slurm/planning/compiler.py | 4 +- .../data_designer/slurm/planning/models.py | 6 +-- .../slurm/planning/resolution.py | 38 ++++++++++++------- .../tests/config/test_loading_builder.py | 16 ++++++++ .../contracts/golden/multi_node_plan.json | 2 +- .../contracts/golden/single_node_plan.json | 2 +- .../tests/contracts/test_planning_records.py | 4 +- .../golden/finalization_chain.json | 4 +- .../tests/planning/test_compiler.py | 24 +++++++++++- .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- 13 files changed, 80 insertions(+), 31 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py index 2d9472286..031a5b9c8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -119,6 +119,9 @@ def resolve_profile( ) except SlurmConfigLoadError: raise + except ValidationError as error: + message = format_validation_error(error, subject="profile selection") + raise SlurmConfigLoadError(message) from None except ValueError as error: raise SlurmConfigLoadError(str(error)) from None diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index b4008d60f..afeb8cfae 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -16,7 +16,7 @@ ResumeWorkspace, compute_serialized_json_sha256, ) -from data_designer.slurm.planning.errors import SlurmPlanCompilationError +from data_designer.slurm.planning.errors import SlurmPlanCompilationError, SlurmPlanContractError from data_designer.slurm.planning.models import ( PlannedShard, PortClaim, @@ -71,7 +71,7 @@ def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan plan, builder_payload=effective.builder_payload, ) - except SlurmPlanCompilationError: + except (SlurmPlanCompilationError, SlurmPlanContractError): raise except ValidationError as error: message = format_validation_error(error, subject="Slurm plan compilation") 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 f49c580db..9af44a1a8 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 @@ -39,7 +39,7 @@ ResumeWorkspace, Sha256Digest, ShardId, - compute_sha256, + compute_serialized_json_sha256, validate_absolute_path, validate_local_config_path, validate_plain_text, @@ -151,7 +151,7 @@ class ResolvedBuilderInput(ContractValue): authored_source: str | None = None source: ArtifactReference | None = None inline: dict[str, JsonValue] | None = None - # Inline input uses canonical JSON; sourced input uses its persisted artifact bytes. + # Both forms use deterministic persisted builder JSON bytes. content_sha256: Sha256Digest model_aliases: tuple[ModelAlias, ...] @@ -162,7 +162,7 @@ def validate_input(self) -> ResolvedBuilderInput: 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) + expected_digest = compute_serialized_json_sha256(self.inline) if self.model_aliases != get_declared_model_aliases(self.inline): raise ValueError("resolved model aliases do not match the inline builder") else: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py index 7668af072..2508860cd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -19,7 +19,7 @@ ArtifactReference, ContractValue, Identifier, - compute_sha256, + compute_serialized_json_sha256, ) from data_designer.slurm.planning.builder_identity import ( get_declared_model_aliases, @@ -100,7 +100,11 @@ def resolve_slurm_config( gpus_per_node = _resolve_gpu_count(selected_profile, resolved_gpus_per_node) workspace_root = selected_profile.profile.workspace_root run_root = posixpath.join(workspace_root, "runs", run_id) - builder = _resolve_builder(authored, run_root=run_root, builder_payload=builder_payload) + builder, resolved_builder_payload = _resolve_builder( + authored, + run_root=run_root, + builder_payload=builder_payload, + ) invocation = ResolvedInvocation( authored=authored.invocation, effective_run_config=_materialize_run_config(authored), @@ -120,7 +124,7 @@ def resolve_slurm_config( ) _validate_dependency_resolution(authored, client_image, dependency_lock) _validate_resolved_images(authored, client_image, deployment_images) - _validate_sharding_constraints(authored, builder_payload=builder_payload) + _validate_sharding_constraints(authored, builder_payload=resolved_builder_payload) _validate_output_destination(output.root, workspace_root, run_root) if output.partitions > authored.invocation.num_records: raise SlurmConfigResolutionError("output partitions must not exceed requested records") @@ -136,7 +140,7 @@ def resolve_slurm_config( selected_profile=selected_profile, resolved_gpus_per_node=gpus_per_node, builder=builder, - builder_payload=builder_payload, + builder_payload=resolved_builder_payload, invocation=invocation, client_image=client_image, deployment_images=deployment_images, @@ -170,14 +174,17 @@ def _resolve_builder( *, run_root: str, builder_payload: dict[str, JsonValue] | None, -) -> ResolvedBuilderInput: +) -> tuple[ResolvedBuilderInput, dict[str, JsonValue] | None]: if authored.builder.inline is not None: if builder_payload is not None: raise SlurmConfigResolutionError("inline builder input must not provide a separate payload") - return ResolvedBuilderInput( - inline=authored.builder.inline, - content_sha256=compute_sha256(authored.builder.inline), - model_aliases=get_declared_model_aliases(authored.builder.inline), + return ( + ResolvedBuilderInput( + inline=authored.builder.inline, + content_sha256=compute_serialized_json_sha256(authored.builder.inline), + model_aliases=get_declared_model_aliases(authored.builder.inline), + ), + None, ) if builder_payload is None: raise SlurmConfigResolutionError("sourced builder input requires its resolved payload") @@ -188,11 +195,14 @@ def _resolve_builder( path=posixpath.join(run_root, "builder-config.json"), sha256=digest, ) - return ResolvedBuilderInput( - authored_source=authored.builder.source, - source=source, - content_sha256=source.sha256, - model_aliases=aliases, + return ( + ResolvedBuilderInput( + authored_source=authored.builder.source, + source=source, + content_sha256=source.sha256, + model_aliases=aliases, + ), + validated_payload, ) diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index d5794e33b..8c21d8c3c 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -272,6 +272,22 @@ def test_profile_resolution_normalizes_ambiguous_hostname_errors( resolve_profile(catalog=catalog, hostnames=("primary-login-1",)) +def test_profile_resolution_hides_validation_inputs(profile_catalog: SlurmProfileCatalog) -> None: + secret = "super-secret-token: x" + catalog = profile_catalog.model_copy( + update={ + "clusters": {secret: profile_catalog.clusters["primary"]}, + "default_cluster": secret, + } + ) + + with pytest.raises(SlurmConfigLoadError, match="profile selection failed validation") as error: + resolve_profile(catalog=catalog, cluster=secret) + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + def test_json_builder_output_is_stable(tmp_path: Path) -> None: path = tmp_path / "run.json" builder = _config_builder() 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 6d369e061..d294085fc 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 @@ -9,7 +9,7 @@ }, "builder": { "authored_source": null, - "content_sha256": "f37227ca7c67e203abe881c0228b4308a8e741364296d293159a1201949732f2", + "content_sha256": "2d4536fc8ff5a32e06269769bb563cb354c53cee95f45da31c9c6781130b4222", "inline": { "data_designer": { "columns": [], 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 d6e35f670..f4d8396be 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 @@ -9,7 +9,7 @@ }, "builder": { "authored_source": null, - "content_sha256": "b3ef5fc1fe675a8e004633f84842ac60cf82d5ba3dc68b4d50ee4438448b0570", + "content_sha256": "0bac3a88498774a30b64f0a62617511c8e75601ca54e0b775d86484bb13788d8", "inline": { "data_designer": { "columns": [], 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 64e825e0f..63fc2f061 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.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig -from data_designer.slurm.contracts import compute_serialized_json_sha256, compute_sha256 +from data_designer.slurm.contracts import compute_serialized_json_sha256 from data_designer.slurm.planning import ( ArtifactReference, ResolvedDependencyLock, @@ -160,7 +160,7 @@ def test_plan_rejects_deployment_alias_missing_from_builder(multi_node_plan: Res "model_configs" ][:1] payload["builder"]["model_aliases"] = ["generator"] - payload["builder"]["content_sha256"] = compute_sha256(payload["builder"]["inline"]) + payload["builder"]["content_sha256"] = compute_serialized_json_sha256(payload["builder"]["inline"]) with pytest.raises(ValidationError, match="deployment alias"): ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index 674560f42..0c647119b 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -9,7 +9,7 @@ "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "9bd1694bbe32fedec605de7e835a4b8506dc2b74be6bd640d4e2b8577951ed9c" + "sha256": "623a3d81c02f86a712d27a8ef2ce93d6018d534ccd255d924a00781ce7f448cc" }, "run_id": "run-single", "scheduler": { @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "9bd1694bbe32fedec605de7e835a4b8506dc2b74be6bd640d4e2b8577951ed9c" + "sha256": "623a3d81c02f86a712d27a8ef2ce93d6018d534ccd255d924a00781ce7f448cc" }, "run_id": "run-single", "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index d041f3de4..265cc1b1d 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -29,7 +29,11 @@ ResolvedSlurmRunPlan, ) from data_designer.slurm.planning.compiler import SlurmRunCompiler -from data_designer.slurm.planning.errors import SlurmConfigResolutionError, SlurmPlanCompilationError +from data_designer.slurm.planning.errors import ( + SlurmConfigResolutionError, + SlurmPlanCompilationError, + SlurmPlanContractError, +) from data_designer.slurm.planning.resolution import EffectiveDataDesignerSlurmConfig, resolve_slurm_config GOLDEN_DIRECTORY = Path(__file__).parents[1] / "contracts" / "golden" @@ -325,6 +329,22 @@ def test_sharded_seed_inputs_have_stable_ranges_and_partition_digests( assert first.shards == second.shards +def test_uneven_shards_assign_remainder_to_the_final_task( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run.model_copy(update={"array_tasks": authored_run.array_tasks.model_copy(update={"count": 3})}) + + plan = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + + assert tuple((shard.record_range.start_index, shard.record_range.end_index_exclusive) for shard in plan.shards) == ( + (0, 33), + (33, 66), + (66, 100), + ) + + @pytest.mark.parametrize( ("builder_update", "output_update", "message"), [ @@ -599,7 +619,7 @@ def test_compiler_rejects_sourced_builder_payload_identity_drift( drifted_payload["data_designer"]["model_configs"][0]["alias"] = "drifted" drifted = effective.model_copy(update={"builder_payload": drifted_payload}) - with pytest.raises(SlurmPlanCompilationError, match=message): + with pytest.raises(SlurmPlanContractError, match=message): SlurmRunCompiler.compile(drifted) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 93ee2b213..c25c4324d 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -11,7 +11,7 @@ set -Eeuo pipefail readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="9efc57d4ef15dc32ac5276cb97afb053155c5d43df03ed28db5fa09546bfa3eb" +readonly DD_PLAN_SHA256="e61513fe9dde370755619ffaaf128649daa939426c88226f619aa2cf82046a71" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 3fc811602..93a512e1e 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -11,7 +11,7 @@ set -Eeuo pipefail readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="9bd1694bbe32fedec605de7e835a4b8506dc2b74be6bd640d4e2b8577951ed9c" +readonly DD_PLAN_SHA256="623a3d81c02f86a712d27a8ef2ce93d6018d534ccd255d924a00781ce7f448cc" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 9660aa7eb..4451aa001 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="d1875aa6705dd862ae005e10e346d0fd38752377a020d33d46351e6806bdf9f9", + expected_fixture_sha256="b8905b808fc0d3161b43c2d6bdabe5ec1c884ee7ea3a342700118ae6a0789430", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="864188c9631f69073f93a2e6c75ce62be7412a9c941de48ba35a090e5880fc6a", + expected_fixture_sha256="16cb158afd0d999b962331dd88450251a1d68b5bb817b9e93d5df0a3a0c809ec", ) From 98f6712f4c2c50d62ad4783d69fbc53909f74bff Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 28 Aug 2026 10:10:22 -0300 Subject: [PATCH 7/9] fix: validate Slurm compiler inputs --- .../src/data_designer/slurm/_errors.py | 64 ++++++++- .../data_designer/slurm/planning/compiler.py | 10 +- .../slurm/planning/resolution.py | 123 ++++++++++++----- .../slurm/planning/validation.py | 11 ++ .../tests/config/test_loading_builder.py | 15 +++ .../tests/contracts/test_planning_records.py | 22 ++- .../tests/planning/test_compiler.py | 127 +++++++++++++++++- 7 files changed, 329 insertions(+), 43 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py index 80234fafa..5d2f2b7df 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py @@ -6,20 +6,74 @@ from __future__ import annotations import json +import re +from collections.abc import Iterable +from typing import Any import yaml from pydantic import ValidationError +_LOCATION_SEGMENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_SENSITIVE_LOCATION = re.compile(r"(?:api_?key|credential|password|secret|token)", re.IGNORECASE) +_ERROR_DESCRIPTIONS = { + "extra_forbidden": "field is not permitted", + "greater_than": "must be greater than the allowed minimum", + "greater_than_equal": "must be at least the allowed minimum", + "less_than": "must be less than the allowed maximum", + "less_than_equal": "must not exceed the allowed maximum", + "literal_error": "must use an allowed value", + "missing": "field is required", + "string_pattern_mismatch": "does not match the required pattern", + "string_too_long": "is longer than allowed", + "string_too_short": "is shorter than allowed", + "value_error": "value is invalid", +} +_SAFE_VALUE_ERROR_MESSAGES = { + "Value error, builder content digest does not match the resolved input": ( + "resolved builder digest does not match its input" + ), + "Value error, resolved model aliases do not match the inline builder": ( + "resolved model aliases do not match the inline builder" + ), + "Value error, resolved deployment aliases must exactly cover Data Designer model aliases": ( + "resolved deployment aliases must exactly cover Data Designer model aliases" + ), +} + def format_validation_error(error: ValidationError, *, subject: str) -> str: """Summarize validation without rendering user-controlled values.""" - error_types = sorted( - {str(detail["type"]) for detail in error.errors(include_url=False, include_context=False, include_input=False)} - ) + details = error.errors(include_url=False, include_context=False, include_input=False) + summaries = sorted({_format_error_detail(detail) for detail in details}) count = error.error_count() noun = "error" if count == 1 else "errors" - kinds = f": {', '.join(error_types)}" if error_types else "" - return f"{subject} failed validation ({count} {noun}{kinds})" + summary = f": {'; '.join(summaries)}" if summaries else "" + return f"{subject} failed validation ({count} {noun}{summary})" + + +def _format_error_detail(detail: dict[str, Any]) -> str: + error_type = str(detail["type"]) + message = _SAFE_VALUE_ERROR_MESSAGES.get(str(detail.get("msg"))) + description = message or _ERROR_DESCRIPTIONS.get(error_type, error_type.replace("_", " ")) + location = _format_location(detail.get("loc", ())) + return f"{location}: {description}" if location else description + + +def _format_location(location: Iterable[object]) -> str: + parts: list[str] = [] + for segment in location: + if isinstance(segment, int): + if parts: + parts[-1] = f"{parts[-1]}[{segment}]" + continue + if ( + not isinstance(segment, str) + or _LOCATION_SEGMENT.fullmatch(segment) is None + or _SENSITIVE_LOCATION.search(segment) is not None + ): + break + parts.append(segment) + return ".".join(parts) def format_parse_error(error: json.JSONDecodeError | yaml.YAMLError) -> str: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index afeb8cfae..b725a396c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Pure deterministic Slurm plan compilation.""" +"""Internal pure deterministic Slurm plan compilation.""" from __future__ import annotations @@ -25,9 +25,14 @@ ResolvedSlurmRunPlan, ResolvedTopology, ) -from data_designer.slurm.planning.resolution import EffectiveDataDesignerSlurmConfig +from data_designer.slurm.planning.resolution import ( + EffectiveDataDesignerSlurmConfig, + validate_effective_slurm_config, +) from data_designer.slurm.planning.validation import validate_resolved_plan +__all__: list[str] = [] + _LOGICAL_ENDPOINT_PORT = 17000 _HTTP_PORT = 18000 _RENDEZVOUS_PORT = 19000 @@ -41,6 +46,7 @@ class SlurmRunCompiler: def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: """Return one immutable deterministic execution plan.""" try: + validate_effective_slurm_config(effective) deployments = _compile_deployments(effective) client = _compile_client(effective, deployments) _validate_port_claims(effective, client, deployments) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py index 2508860cd..79201d72d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Pure authored-configuration resolution for Slurm planning.""" +"""Internal pure authored-configuration resolution for Slurm planning.""" from __future__ import annotations @@ -35,6 +35,8 @@ ResolvedSubmission, ) +__all__: list[str] = [] + _SHARDABLE_COLUMN_TYPES = frozenset( { "embedding", @@ -105,35 +107,7 @@ def resolve_slurm_config( run_root=run_root, builder_payload=builder_payload, ) - invocation = ResolvedInvocation( - authored=authored.invocation, - effective_run_config=_materialize_run_config(authored), - ) - submission = ResolvedSubmission( - account=authored.submission.account or selected_profile.profile.scheduler.account, - partition=authored.submission.partition or selected_profile.profile.scheduler.partition, - job_name=authored.submission.job_name, - time_limit=authored.submission.time_limit, - comment=authored.submission.comment, - ) - output = ResolvedOutput( - root=authored.output.root or posixpath.join(run_root, "output"), - format=authored.output.format, - partitions=authored.output.partitions, - require_exact_record_count=authored.output.require_exact_record_count, - ) - _validate_dependency_resolution(authored, client_image, dependency_lock) - _validate_resolved_images(authored, client_image, deployment_images) - _validate_sharding_constraints(authored, builder_payload=resolved_builder_payload) - _validate_output_destination(output.root, workspace_root, run_root) - if output.partitions > authored.invocation.num_records: - raise SlurmConfigResolutionError("output partitions must not exceed requested records") - runtime_root = posixpath.join(workspace_root, "runtime") - if not _is_below(runtime_bundle.path, runtime_root) or not runtime_bundle.path.endswith(".tar.gz"): - raise SlurmConfigResolutionError( - "runtime bundle must be a tar archive below the selected workspace runtime root" - ) - return EffectiveDataDesignerSlurmConfig( + effective = EffectiveDataDesignerSlurmConfig( run_id=run_id, package_version=package_version, authored=authored, @@ -141,14 +115,15 @@ def resolve_slurm_config( resolved_gpus_per_node=gpus_per_node, builder=builder, builder_payload=resolved_builder_payload, - invocation=invocation, + invocation=_materialize_invocation(authored), client_image=client_image, deployment_images=deployment_images, dependency_lock=dependency_lock, - submission=submission, - output=output, + submission=_materialize_submission(authored, selected_profile), + output=_materialize_output(authored, run_root), runtime_bundle=runtime_bundle, ) + return validate_effective_slurm_config(effective) except SlurmConfigResolutionError: raise except ValidationError as error: @@ -158,6 +133,50 @@ def resolve_slurm_config( raise SlurmConfigResolutionError(str(error)) from None +def validate_effective_slurm_config( + effective: EffectiveDataDesignerSlurmConfig, +) -> EffectiveDataDesignerSlurmConfig: + """Validate one fully materialized compiler input.""" + authored = effective.authored + profile = effective.selected_profile.profile + workspace_root = profile.workspace_root + run_root = posixpath.join(workspace_root, "runs", effective.run_id) + if profile.gpus_per_node != "auto" and profile.gpus_per_node != effective.resolved_gpus_per_node: + raise SlurmConfigResolutionError("resolved GPU count does not match the selected profile") + + if authored.builder.inline is not None: + if effective.builder_payload is not None: + raise SlurmConfigResolutionError("inline builder input must not provide a separate payload") + else: + if effective.builder_payload is None: + raise SlurmConfigResolutionError("sourced builder input requires its resolved payload") + BuilderInput(inline=effective.builder_payload) + expected_invocation = _materialize_invocation(authored) + if effective.invocation != expected_invocation: + raise SlurmConfigResolutionError("resolved invocation does not match the authored invocation") + expected_submission = _materialize_submission(authored, effective.selected_profile) + if effective.submission != expected_submission: + raise SlurmConfigResolutionError("resolved submission does not match the authored and profile input") + expected_output = _materialize_output(authored, run_root) + if effective.output != expected_output: + raise SlurmConfigResolutionError("resolved output does not match the authored output") + + _validate_dependency_resolution(authored, effective.client_image, effective.dependency_lock) + _validate_resolved_images(authored, effective.client_image, effective.deployment_images) + _validate_sharding_constraints(authored, builder_payload=effective.builder_payload) + _validate_output_destination(effective.output.root, workspace_root, run_root) + if effective.output.partitions > authored.invocation.num_records: + raise SlurmConfigResolutionError("output partitions must not exceed requested records") + runtime_root = posixpath.join(workspace_root, "runtime") + if not _is_below(effective.runtime_bundle.path, runtime_root) or not effective.runtime_bundle.path.endswith( + ".tar.gz" + ): + raise SlurmConfigResolutionError( + "runtime bundle must be a tar archive below the selected workspace runtime root" + ) + return effective + + def _resolve_gpu_count(selected: SelectedSlurmProfile, resolved: int | None) -> int: configured = selected.profile.gpus_per_node if configured == "auto": @@ -206,6 +225,35 @@ def _resolve_builder( ) +def _materialize_invocation(authored: DataDesignerSlurmConfig) -> ResolvedInvocation: + return ResolvedInvocation( + authored=authored.invocation, + effective_run_config=_materialize_run_config(authored), + ) + + +def _materialize_submission( + authored: DataDesignerSlurmConfig, + selected_profile: SelectedSlurmProfile, +) -> ResolvedSubmission: + return ResolvedSubmission( + account=authored.submission.account or selected_profile.profile.scheduler.account, + partition=authored.submission.partition or selected_profile.profile.scheduler.partition, + job_name=authored.submission.job_name, + time_limit=authored.submission.time_limit, + comment=authored.submission.comment, + ) + + +def _materialize_output(authored: DataDesignerSlurmConfig, run_root: str) -> ResolvedOutput: + return ResolvedOutput( + root=authored.output.root or posixpath.join(run_root, "output"), + format=authored.output.format, + partitions=authored.output.partitions, + require_exact_record_count=authored.output.require_exact_record_count, + ) + + def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: values = dict(authored.invocation.run_config) authored_early_shutdown = {"disable_early_shutdown", "shutdown_error_rate", "shutdown_error_window"}.intersection( @@ -293,6 +341,7 @@ def _validate_resolved_images( client_image: ResolvedImage, deployment_images: tuple[ResolvedImage, ...], ) -> None: + _validate_resolved_image_identity(client_image) if client_image.kind is not ImageKind.CLIENT: raise SlurmConfigResolutionError("resolved client image must contain client inspection facts") if client_image.authored_ref != authored.client.image: @@ -300,12 +349,20 @@ def _validate_resolved_images( if len(deployment_images) != len(authored.deployments): raise SlurmConfigResolutionError("resolved serving images must match the authored deployment count") for deployment, image in zip(authored.deployments, deployment_images, strict=True): + _validate_resolved_image_identity(image) if image.kind is not ImageKind.SERVING: raise SlurmConfigResolutionError("resolved deployment image must contain serving inspection facts") if image.authored_ref != deployment.server.image: raise SlurmConfigResolutionError("resolved deployment image does not match the authored reference") +def _validate_resolved_image_identity(image: ResolvedImage) -> None: + if image.sha256 != image.inspection.sqsh_sha256: + raise SlurmConfigResolutionError("resolved image digest does not match its inspection record") + if image.authored_ref.path is not None and image.path != image.authored_ref.path: + raise SlurmConfigResolutionError("resolved image path does not match the authored path") + + def _validate_output_destination(output_root: str, workspace_root: str, run_root: str) -> None: if not _is_below(output_root, workspace_root): raise SlurmConfigResolutionError("output root must be below the selected workspace_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 index 0351c21c9..b23b715e2 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,6 +3,8 @@ from __future__ import annotations +import posixpath + from pydantic import JsonValue from data_designer.config import RunConfig @@ -42,9 +44,13 @@ def validate_resolved_plan( _require(plan.array_tasks == authored.array_tasks, "resolved array task policy does not match authored input") if authored.builder.inline is not None: + _require(builder_payload is None, "inline builder input must not retain a separate payload") _require( plan.builder.inline == authored.builder.inline, "resolved inline builder does not match authored input" ) + model_aliases, digest = get_persisted_builder_identity(authored.builder.inline) + _require(plan.builder.model_aliases == model_aliases, "resolved model aliases do not match inline builder") + _require(plan.builder.content_sha256 == digest, "resolved builder digest does not match inline builder") else: _require( plan.builder.authored_source == authored.builder.source, @@ -53,7 +59,12 @@ def validate_resolved_plan( if builder_payload is None: raise SlurmPlanContractError("sourced builder validation requires its resolved payload") model_aliases, digest = get_persisted_builder_identity(builder_payload) + _require(plan.builder.source is not None, "resolved builder source artifact is missing") + assert plan.builder.source is not None + expected_path = posixpath.join(posixpath.dirname(plan.authored_config.path), "builder-config.json") + _require(plan.builder.source.path == expected_path, "resolved builder artifact path does not match the run") _require(plan.builder.model_aliases == model_aliases, "resolved model aliases do not match builder source") + _require(plan.builder.source.sha256 == digest, "resolved builder digest does not match builder source") _require(plan.builder.content_sha256 == digest, "resolved builder digest does not match builder source") expected_account = authored.submission.account or plan.selected_profile.profile.scheduler.account diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 8c21d8c3c..964cfd835 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -100,6 +100,11 @@ def test_builder_normalizes_invalid_authored_values( getattr(_config_builder(), method)(**values) +def test_builder_validation_errors_identify_the_invalid_field() -> None: + with pytest.raises(SlurmConfigBuilderError, match="num_records: must be greater"): + _config_builder().with_invocation(num_records=0, dataset_name="generated") + + def test_builder_validation_errors_hide_secret_inputs() -> None: secret = "super-secret-token" @@ -178,6 +183,16 @@ def test_loader_validation_errors_hide_secret_inputs(tmp_path: Path) -> None: assert error.value.__cause__ is None +def test_loader_validation_errors_identify_the_invalid_field(tmp_path: Path) -> None: + payload = _config_builder().build().model_dump(mode="json") + payload["invocation"]["num_records"] = 0 + path = tmp_path / "run.json" + path.write_text(json.dumps(payload)) + + with pytest.raises(SlurmConfigLoadError, match=r"invocation\.num_records: must be greater"): + load_run_config(path) + + def test_loader_parse_errors_hide_source_values(tmp_path: Path) -> None: secret = "super-secret-token" path = tmp_path / "run.yaml" 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 63fc2f061..cd0f1a126 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -221,7 +221,7 @@ def test_sourced_builder_validation_requires_resolved_payload( payload["authored_config"]["sha256"] = sourced_authored.compute_sha256() payload["builder"] = { "authored_source": "builder.json", - "source": {"path": "/workspace/primary/runs/run-001/builder.json", "sha256": builder_digest}, + "source": {"path": "/workspace/primary/runs/run-001/builder-config.json", "sha256": builder_digest}, "inline": None, "content_sha256": builder_digest, "model_aliases": ["generator", "judge"], @@ -242,6 +242,26 @@ def test_sourced_builder_validation_requires_resolved_payload( ) +@pytest.mark.parametrize( + ("builder_update", "message"), + [ + ({"content_sha256": "a" * 64}, "builder digest"), + ({"model_aliases": ("other",)}, "model aliases"), + ], +) +def test_inline_builder_validation_rederives_identity( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + builder_update: dict[str, object], + message: str, +) -> None: + invalid = multi_node_plan.model_copy(update={"builder": multi_node_plan.builder.model_copy(update=builder_update)}) + + with pytest.raises(SlurmPlanContractError, match=message): + validate_resolved_plan(authored_run, dependency_lock, invalid) + + @pytest.mark.parametrize( "update", [ diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index 265cc1b1d..307ecf162 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -246,6 +246,83 @@ def test_resolution_rejects_invalid_output_destinations( _resolve_fixture(authored, dependency_lock_single, single_node_plan) +def test_compiler_rejects_direct_effective_input_resolution_bypass( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run, dependency_lock, multi_node_plan) + + jsonl_output = authored_run.output.model_copy(update={"format": "jsonl"}) + jsonl_authored = authored_run.model_copy(update={"output": jsonl_output}) + with pytest.raises(SlurmPlanCompilationError, match="parquet output"): + SlurmRunCompiler.compile( + effective.model_copy( + update={ + "authored": jsonl_authored, + "output": effective.output.model_copy(update={"format": "jsonl"}), + } + ) + ) + + other_output = "/workspace/primary/runs/other-run/output" + other_authored = authored_run.model_copy( + update={"output": authored_run.output.model_copy(update={"root": other_output})} + ) + with pytest.raises(SlurmPlanCompilationError, match="another package-managed run"): + SlurmRunCompiler.compile( + effective.model_copy( + update={ + "authored": other_authored, + "output": effective.output.model_copy(update={"root": other_output}), + } + ) + ) + + partitioned_authored = authored_run.model_copy( + update={"output": authored_run.output.model_copy(update={"partitions": 101})} + ) + with pytest.raises(SlurmPlanCompilationError, match="requested records"): + SlurmRunCompiler.compile( + effective.model_copy( + update={ + "authored": partitioned_authored, + "output": effective.output.model_copy(update={"partitions": 101}), + } + ) + ) + + with pytest.raises(SlurmPlanCompilationError, match="runtime bundle"): + SlurmRunCompiler.compile( + effective.model_copy( + update={ + "runtime_bundle": ArtifactReference(path="/tmp/runtime.tar.gz", sha256="e" * 64), + } + ) + ) + + +def test_compiler_rejects_direct_effective_record_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + run_config = dict(effective.invocation.effective_run_config) + run_config["jinja_rendering_engine"] = "native" + + with pytest.raises(SlurmPlanCompilationError, match="resolved invocation"): + SlurmRunCompiler.compile( + effective.model_copy( + update={"invocation": effective.invocation.model_copy(update={"effective_run_config": run_config})} + ) + ) + with pytest.raises(SlurmPlanCompilationError, match="resolved output"): + SlurmRunCompiler.compile( + effective.model_copy(update={"output": effective.output.model_copy(update={"format": "jsonl"})}) + ) + + def test_compiler_rejects_model_alias_missing_from_builder( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -255,7 +332,7 @@ def test_compiler_rejects_model_alias_missing_from_builder( payload["builder"]["inline"]["data_designer"]["model_configs"][0]["alias"] = "other" authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(SlurmPlanCompilationError, match="failed validation"): + with pytest.raises(SlurmPlanCompilationError, match="exactly cover"): SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) @@ -276,7 +353,7 @@ def test_compiler_rejects_builder_model_alias_without_deployment( payload["builder"] = {"source": "builder.json"} authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(SlurmPlanCompilationError, match="failed validation"): + with pytest.raises(SlurmPlanCompilationError, match="exactly cover"): SlurmRunCompiler.compile( _resolve_fixture( authored, @@ -287,6 +364,52 @@ def test_compiler_rejects_builder_model_alias_without_deployment( ) +def test_compiler_rejects_inline_builder_identity_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + + with pytest.raises(SlurmPlanCompilationError, match="builder digest"): + SlurmRunCompiler.compile( + effective.model_copy(update={"builder": effective.builder.model_copy(update={"content_sha256": "a" * 64})}) + ) + + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["model_configs"].append( + {"alias": "undeployed", "model": "example/undeployed", "provider": "openai"} + ) + authored = DataDesignerSlurmConfig.model_validate(payload) + effective = _resolve_fixture(authored, dependency_lock_single, single_node_plan) + forged_builder = effective.builder.model_copy(update={"model_aliases": ("generator",)}) + + with pytest.raises(SlurmPlanCompilationError, match="model aliases"): + SlurmRunCompiler.compile(effective.model_copy(update={"builder": forged_builder})) + + +def test_compiler_rejects_resolved_image_identity_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + client_image = effective.client_image.model_copy(update={"sha256": "a" * 64}) + dependency_lock = effective.dependency_lock.model_copy(update={"client_image_sha256": "a" * 64}) + + with pytest.raises(SlurmPlanCompilationError, match="inspection record"): + SlurmRunCompiler.compile( + effective.model_copy(update={"client_image": client_image, "dependency_lock": dependency_lock}) + ) + + deployment_images = ( + effective.deployment_images[0].model_copy(update={"sha256": "a" * 64}), + *effective.deployment_images[1:], + ) + with pytest.raises(SlurmPlanCompilationError, match="inspection record"): + SlurmRunCompiler.compile(effective.model_copy(update={"deployment_images": deployment_images})) + + def test_compiler_rejects_otel_collision_before_runtime( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, From 555d7f7ae2a7f85a22a1fb75372a20076045e71a Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 28 Aug 2026 10:53:54 -0300 Subject: [PATCH 8/9] fix: complete Slurm plan resolution defaults --- .../src/data_designer/slurm/config/builder.py | 7 +- .../data_designer/slurm/config/profiles.py | 2 + .../src/data_designer/slurm/config/run.py | 4 +- .../src/data_designer/slurm/contracts.py | 15 +++ .../data_designer/slurm/launcher/renderer.py | 5 +- .../data_designer/slurm/planning/compiler.py | 2 +- .../data_designer/slurm/planning/models.py | 49 +++++++--- .../slurm/planning/resolution.py | 67 +++++++++---- .../slurm/planning/validation.py | 9 ++ .../contracts/golden/authored_run_single.json | 2 +- .../contracts/golden/multi_node_plan.json | 8 +- .../contracts/golden/single_node_plan.json | 12 ++- .../tests/contracts/test_config_records.py | 1 + .../tests/contracts/test_planning_records.py | 37 ++++++-- .../tests/contracts/test_profiles.py | 9 ++ .../golden/finalization_chain.json | 6 +- .../tests/launcher/test_renderer.py | 13 ++- .../tests/planning/test_compiler.py | 94 ++++++++++++++++++- .../golden/rendered/multi_node.sbatch | 4 +- .../golden/rendered/single_node.sbatch | 4 +- .../slurm_test_fakes/test_rendered_scripts.py | 6 +- 21 files changed, 291 insertions(+), 65 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py index 062fa54dc..5c51edc23 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -121,7 +121,12 @@ def with_deployment( self._deployments.append(_validate_model(ServerDeploymentConfig, deployment)) return self - def with_array_tasks(self, *, count: int, max_concurrent: int = 1) -> DataDesignerSlurmConfigBuilder: + def with_array_tasks( + self, + *, + count: int, + max_concurrent: int | None = None, + ) -> DataDesignerSlurmConfigBuilder: """Set deterministic horizontal sharding.""" self._array_tasks = _validate_model( ArrayTasksConfig, 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 87412d84d..732794584 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 @@ -76,6 +76,8 @@ 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") + if self.gpu_request_mode == "visible" and self.scheduler.mem_per_gpu is not None: + raise ValueError("mem_per_gpu requires GRES GPU request mode") return self 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 c6f7d69d4..afdb4d58c 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 @@ -409,11 +409,11 @@ def validate_topology(self) -> ServerDeploymentConfig: class ArrayTasksConfig(AuthoredConfig): count: PositiveInt = 1 - max_concurrent: PositiveInt = 1 + max_concurrent: PositiveInt | None = None @model_validator(mode="after") def validate_concurrency(self) -> ArrayTasksConfig: - if self.max_concurrent > self.count: + if self.max_concurrent is not None and self.max_concurrent > self.count: raise ValueError("array task concurrency must not exceed task count") return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py index 04203616a..72b4f667c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -175,6 +175,21 @@ def compute_serialized_json_sha256(value: object) -> Sha256Digest: return hashlib.sha256(pretty_json(value).encode("utf-8")).hexdigest() +def derive_managed_assets_path(workspace_root: str) -> str: + """Derive the default managed-assets path from a workspace root.""" + return posixpath.join(workspace_root, "managed-assets") + + +def is_path_below(path: str, root: str) -> bool: + """Return whether a path is strictly below a root.""" + return path != root and posixpath.commonpath((path, root)) == root + + +def paths_overlap(left: str, right: str) -> bool: + """Return whether either path contains the other.""" + return left == right or is_path_below(left, right) or is_path_below(right, left) + + def validate_absolute_path(value: str) -> str: """Validate a normalized, absolute POSIX path below the filesystem root.""" if not value.startswith("/"): diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index 203cfe14a..e8941df39 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -91,7 +91,9 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire node_count = max(node_indices) + 1 array = "0" if plan.array_tasks.count > 1: - array = f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + array = f"0-{plan.array_tasks.count - 1}" + if plan.array_tasks.max_concurrent is not None: + array = f"{array}%{plan.array_tasks.max_concurrent}" values: list[tuple[str, str | None]] = [ ("job-name", plan.submission.job_name), @@ -106,7 +108,6 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire if profile.gpu_request_mode == "gres": values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) elif profile.scheduler.mem_per_gpu is not None: - # TODO(#875): Remove this defense once config and plan validation make this state unrepresentable. raise SlurmBatchRenderError("mem_per_gpu requires GRES GPU request mode") if profile.scheduler.mem_per_gpu is not None: values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index b725a396c..c93ec99a1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -217,7 +217,7 @@ def _compile_shards(effective: EffectiveDataDesignerSlurmConfig) -> tuple[Planne shard_root = posixpath.join(_run_root(effective), "shards", shard_id) record_range = RecordRange(start_index=start, end_index_exclusive=end) partition = None - seed_path = effective.authored.invocation.input_bindings.seed_path + seed_path = effective.invocation.effective_input_bindings.seed_path if seed_path is not None: partition = ArtifactReference( path=posixpath.join(shard_root, "input-partition.json"), 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 9af44a1a8..98378fe19 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 @@ -25,6 +25,7 @@ ArrayTasksConfig, ClientConfig, ClientDependencies, + InputBindings, InvocationConfig, ServerDeploymentConfig, SubmissionConfig, @@ -40,6 +41,9 @@ Sha256Digest, ShardId, compute_serialized_json_sha256, + derive_managed_assets_path, + is_path_below, + paths_overlap, validate_absolute_path, validate_local_config_path, validate_plain_text, @@ -178,6 +182,7 @@ def validate_input(self) -> ResolvedBuilderInput: class ResolvedInvocation(ContractValue): authored: InvocationConfig + effective_input_bindings: InputBindings effective_run_config: dict[str, JsonValue] @field_validator("effective_run_config") @@ -363,6 +368,8 @@ 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 profile.gpu_request_mode == "visible" and profile.scheduler.mem_per_gpu is not None: + raise ValueError("mem_per_gpu requires GRES GPU request mode") 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: @@ -383,10 +390,25 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: 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") + authored_bindings = self.invocation.authored.input_bindings + expected_bindings = InputBindings( + seed_path=authored_bindings.seed_path, + managed_assets_path=authored_bindings.managed_assets_path + or derive_managed_assets_path(profile.workspace_root), + ) + if self.invocation.effective_input_bindings != expected_bindings: + raise ValueError("effective input bindings must match the authored and selected profile input") + managed_assets_path = self.invocation.effective_input_bindings.managed_assets_path + assert managed_assets_path is not None + workspace_state = tuple( + posixpath.join(profile.workspace_root, name) for name in ("images", "runtime", "benchmarks", "runs") + ) + if any(paths_overlap(managed_assets_path, path) for path in workspace_state): + raise ValueError("managed_assets_path must not overlap package-managed workspace state") if "non_inference_max_parallel_workers" not in self.invocation.authored.run_config: workers = self.invocation.effective_run_config["non_inference_max_parallel_workers"] - if workers != RunConfig().non_inference_max_parallel_workers: - raise ValueError("default non-inference worker count must match the Data Designer RunConfig default") + if workers != self.client.authored.cpus: + raise ValueError("default non-inference worker count must match the client CPU count") expected_logical_names = tuple( f"{deployment.deployment_id}-logical-endpoint" for deployment in self.deployments @@ -413,15 +435,20 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: 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") + runtime_root = posixpath.join(profile.workspace_root, "runtime") + runtime_name = posixpath.basename(self.runtime_bundle.path) + if ( + not is_path_below(self.runtime_bundle.path, runtime_root) + or runtime_name != f"{self.runtime_bundle.sha256}.tar.gz" + ): + raise ValueError("runtime bundle must be a content-addressed tar archive below the workspace runtime root") self._validate_shards(run_root) - if not _is_below(self.output.root, profile.workspace_root): + if not is_path_below(self.output.root, profile.workspace_root): raise ValueError("resolved output root must be below the selected workspace_root") + if paths_overlap(self.output.root, managed_assets_path): + raise ValueError("resolved output root must not overlap managed assets") shards_root = posixpath.join(run_root, "shards") - if ( - self.output.root == shards_root - or _is_below(self.output.root, shards_root) - or _is_below(shards_root, self.output.root) - ): + if paths_overlap(self.output.root, shards_root): raise ValueError("resolved output root must not overlap the run shard workspace") return self @@ -434,7 +461,7 @@ def _validate_shards(self, run_root: str) -> None: shard_ids: list[ShardId] = [] workspace_paths: list[str] = [] partition_paths: list[str] = [] - requires_partition = self.invocation.authored.input_bindings.seed_path is not None + requires_partition = self.invocation.effective_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") @@ -470,7 +497,3 @@ def _validate_shards(self, run_root: str) -> None: 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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py index 79201d72d..9cadf4a3e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -14,12 +14,15 @@ from data_designer.slurm._errors import format_validation_error from data_designer.slurm.config.images import ClientImageInspection, ImageKind from data_designer.slurm.config.profiles import SelectedSlurmProfile -from data_designer.slurm.config.run import BuilderInput, DataDesignerSlurmConfig +from data_designer.slurm.config.run import BuilderInput, DataDesignerSlurmConfig, InputBindings from data_designer.slurm.contracts import ( ArtifactReference, ContractValue, Identifier, compute_serialized_json_sha256, + derive_managed_assets_path, + is_path_below, + paths_overlap, ) from data_designer.slurm.planning.builder_identity import ( get_declared_model_aliases, @@ -115,7 +118,7 @@ def resolve_slurm_config( resolved_gpus_per_node=gpus_per_node, builder=builder, builder_payload=resolved_builder_payload, - invocation=_materialize_invocation(authored), + invocation=_materialize_invocation(authored, workspace_root), client_image=client_image, deployment_images=deployment_images, dependency_lock=dependency_lock, @@ -143,6 +146,8 @@ def validate_effective_slurm_config( run_root = posixpath.join(workspace_root, "runs", effective.run_id) if profile.gpus_per_node != "auto" and profile.gpus_per_node != effective.resolved_gpus_per_node: raise SlurmConfigResolutionError("resolved GPU count does not match the selected profile") + if profile.gpu_request_mode == "visible" and profile.scheduler.mem_per_gpu is not None: + raise SlurmConfigResolutionError("mem_per_gpu requires GRES GPU request mode") if authored.builder.inline is not None: if effective.builder_payload is not None: @@ -151,7 +156,7 @@ def validate_effective_slurm_config( if effective.builder_payload is None: raise SlurmConfigResolutionError("sourced builder input requires its resolved payload") BuilderInput(inline=effective.builder_payload) - expected_invocation = _materialize_invocation(authored) + expected_invocation = _materialize_invocation(authored, workspace_root) if effective.invocation != expected_invocation: raise SlurmConfigResolutionError("resolved invocation does not match the authored invocation") expected_submission = _materialize_submission(authored, effective.selected_profile) @@ -164,15 +169,25 @@ def validate_effective_slurm_config( _validate_dependency_resolution(authored, effective.client_image, effective.dependency_lock) _validate_resolved_images(authored, effective.client_image, effective.deployment_images) _validate_sharding_constraints(authored, builder_payload=effective.builder_payload) - _validate_output_destination(effective.output.root, workspace_root, run_root) + managed_assets_path = effective.invocation.effective_input_bindings.managed_assets_path + assert managed_assets_path is not None + _validate_managed_assets_path(managed_assets_path, workspace_root) + _validate_output_destination( + effective.output.root, + workspace_root, + run_root, + managed_assets_path=managed_assets_path, + ) if effective.output.partitions > authored.invocation.num_records: raise SlurmConfigResolutionError("output partitions must not exceed requested records") runtime_root = posixpath.join(workspace_root, "runtime") - if not _is_below(effective.runtime_bundle.path, runtime_root) or not effective.runtime_bundle.path.endswith( - ".tar.gz" + runtime_name = posixpath.basename(effective.runtime_bundle.path) + if ( + not is_path_below(effective.runtime_bundle.path, runtime_root) + or runtime_name != f"{effective.runtime_bundle.sha256}.tar.gz" ): raise SlurmConfigResolutionError( - "runtime bundle must be a tar archive below the selected workspace runtime root" + "runtime bundle must be a content-addressed tar archive below the selected workspace runtime root" ) return effective @@ -225,9 +240,14 @@ def _resolve_builder( ) -def _materialize_invocation(authored: DataDesignerSlurmConfig) -> ResolvedInvocation: +def _materialize_invocation(authored: DataDesignerSlurmConfig, workspace_root: str) -> ResolvedInvocation: + input_bindings = authored.invocation.input_bindings return ResolvedInvocation( authored=authored.invocation, + effective_input_bindings=InputBindings( + seed_path=input_bindings.seed_path, + managed_assets_path=input_bindings.managed_assets_path or derive_managed_assets_path(workspace_root), + ), effective_run_config=_materialize_run_config(authored), ) @@ -256,6 +276,7 @@ def _materialize_output(authored: DataDesignerSlurmConfig, run_root: str) -> Res def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: values = dict(authored.invocation.run_config) + values.setdefault("non_inference_max_parallel_workers", authored.client.cpus) authored_early_shutdown = {"disable_early_shutdown", "shutdown_error_rate", "shutdown_error_window"}.intersection( values ) @@ -363,23 +384,31 @@ def _validate_resolved_image_identity(image: ResolvedImage) -> None: raise SlurmConfigResolutionError("resolved image path does not match the authored path") -def _validate_output_destination(output_root: str, workspace_root: str, run_root: str) -> None: - if not _is_below(output_root, workspace_root): +def _validate_output_destination( + output_root: str, + workspace_root: str, + run_root: str, + *, + managed_assets_path: str, +) -> None: + if not is_path_below(output_root, workspace_root): raise SlurmConfigResolutionError("output root must be below the selected workspace_root") reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) - if any(_paths_overlap(output_root, path) for path in reserved): + if any(paths_overlap(output_root, path) for path in reserved): raise SlurmConfigResolutionError("output root must not overlap package-managed workspace state") + if paths_overlap(output_root, managed_assets_path): + raise SlurmConfigResolutionError("output root must not overlap managed assets") runs_root = posixpath.join(workspace_root, "runs") run_output_root = posixpath.join(run_root, "output") - if _paths_overlap(output_root, runs_root) and not ( - output_root == run_output_root or _is_below(output_root, run_output_root) + if paths_overlap(output_root, runs_root) and not ( + output_root == run_output_root or is_path_below(output_root, run_output_root) ): raise SlurmConfigResolutionError("output root must not overlap another package-managed run") -def _is_below(path: str, root: str) -> bool: - return path != root and posixpath.commonpath((path, root)) == root - - -def _paths_overlap(left: str, right: str) -> bool: - return left == right or _is_below(left, right) or _is_below(right, left) +def _validate_managed_assets_path(managed_assets_path: str, workspace_root: str) -> None: + workspace_state = tuple( + posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks", "runs") + ) + if any(paths_overlap(managed_assets_path, path) for path in workspace_state): + raise SlurmConfigResolutionError("managed_assets_path must not overlap package-managed workspace state") 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 b23b715e2..5cf7ee9ba 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 @@ -10,6 +10,7 @@ from data_designer.config import RunConfig from data_designer.slurm.config.images import ClientImageInspection from data_designer.slurm.config.run import DataDesignerSlurmConfig +from data_designer.slurm.contracts import derive_managed_assets_path from data_designer.slurm.planning.builder_identity import get_persisted_builder_identity from data_designer.slurm.planning.errors import SlurmPlanContractError from data_designer.slurm.planning.models import ( @@ -31,6 +32,14 @@ def validate_resolved_plan( "authored config digest does not match the resolved plan", ) _require(plan.invocation.authored == authored.invocation, "resolved invocation does not match authored input") + expected_managed_assets = authored.invocation.input_bindings.managed_assets_path or derive_managed_assets_path( + plan.selected_profile.profile.workspace_root + ) + _require( + plan.invocation.effective_input_bindings.seed_path == authored.invocation.input_bindings.seed_path + and plan.invocation.effective_input_bindings.managed_assets_path == expected_managed_assets, + "resolved input bindings do not match authored/profile input", + ) explicit_run_config = RunConfig.model_validate(authored.invocation.run_config).model_dump( mode="json", exclude_unset=True, 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 c6a46897a..3a518814f 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 @@ -1,7 +1,7 @@ { "array_tasks": { "count": 1, - "max_concurrent": 1 + "max_concurrent": null }, "builder": { "inline": { 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 175fd7b47..1d1d0fd96 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 @@ -350,6 +350,10 @@ "buffer_size": 8192 } }, + "effective_input_bindings": { + "managed_assets_path": "/workspace/primary/managed-assets", + "seed_path": null + }, "effective_run_config": { "async_trace": false, "buffer_size": 8192, @@ -360,7 +364,7 @@ "max_conversation_correction_steps": 0, "max_conversation_restarts": 0, "max_in_flight_tasks": 1024, - "non_inference_max_parallel_workers": 4, + "non_inference_max_parallel_workers": 32, "otel_metrics_port": null, "preserve_dropped_columns": true, "progress_interval": 5.0, @@ -380,7 +384,7 @@ "resolved_gpus_per_node": 8, "run_id": "run-001", "runtime_bundle": { - "path": "/workspace/primary/runtime/runtime.tar.gz", + "path": "/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz", "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" }, "schema_version": 1, 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 70129468d..1ce8e8a5e 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,11 +1,11 @@ { "array_tasks": { "count": 1, - "max_concurrent": 1 + "max_concurrent": null }, "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" + "sha256": "3ca8b4c53642796435ae782cb7710904681df6eb6750d414ad6701861eb85b01" }, "builder": { "authored_source": null, @@ -198,6 +198,10 @@ "resume": "never", "run_config": {} }, + "effective_input_bindings": { + "managed_assets_path": "/workspace/primary/managed-assets", + "seed_path": null + }, "effective_run_config": { "async_trace": false, "buffer_size": 16384, @@ -208,7 +212,7 @@ "max_conversation_correction_steps": 0, "max_conversation_restarts": 0, "max_in_flight_tasks": 1024, - "non_inference_max_parallel_workers": 4, + "non_inference_max_parallel_workers": 32, "otel_metrics_port": null, "preserve_dropped_columns": true, "progress_interval": 5.0, @@ -228,7 +232,7 @@ "resolved_gpus_per_node": 8, "run_id": "run-single", "runtime_bundle": { - "path": "/workspace/primary/runtime/runtime.tar.gz", + "path": "/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz", "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" }, "schema_version": 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 33959d5c1..770998539 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -349,6 +349,7 @@ def test_run_validates_public_run_config_and_shard_count(authored_run: DataDesig def test_small_config_values_validate_at_boundary() -> None: + assert ArrayTasksConfig(count=2).max_concurrent is None with pytest.raises(ValidationError, match="concurrency"): ArrayTasksConfig(count=2, max_concurrent=3) with pytest.raises(ValidationError, match="minutes"): 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 cd0f1a126..63975bed8 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -87,6 +87,12 @@ def test_resolved_plan_is_deeply_immutable(multi_node_plan: ResolvedSlurmRunPlan lambda payload: payload["output"].update(root="/outside/output"), lambda payload: payload.update(container_mounts=[]), lambda payload: payload["deployments"][0]["topology"].update(replica_count=2), + lambda payload: payload["invocation"]["effective_input_bindings"].update( + managed_assets_path="/workspace/other-assets" + ), + lambda payload: payload["runtime_bundle"].update(path="/tmp/" + "e" * 64 + ".tar.gz"), + lambda payload: payload["runtime_bundle"].update(path="/workspace/primary/runtime/runtime.tar.gz"), + lambda payload: payload["output"].update(root="/workspace/primary/managed-assets"), ], ) def test_plan_rejects_invalid_boundaries(multi_node_plan: ResolvedSlurmRunPlan, mutator: object) -> None: @@ -105,13 +111,13 @@ def test_plan_rejects_unmaterialized_run_config(multi_node_plan: ResolvedSlurmRu ResolvedSlurmRunPlan.model_validate(payload) -def test_plan_preserves_default_non_inference_worker_count( +def test_plan_derives_default_non_inference_worker_count_from_client_cpus( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: payload = multi_node_plan.model_dump(mode="json") - payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 32 + payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 4 - with pytest.raises(ValidationError, match="RunConfig default"): + with pytest.raises(ValidationError, match="client CPU count"): ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) @@ -119,12 +125,31 @@ def test_plan_preserves_explicit_non_inference_worker_override( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: payload = multi_node_plan.model_dump(mode="json") - payload["invocation"]["authored"]["run_config"]["non_inference_max_parallel_workers"] = 32 - payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 32 + payload["invocation"]["authored"]["run_config"]["non_inference_max_parallel_workers"] = 16 + payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 16 plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) - assert plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == 32 + assert plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == 16 + + +def test_plan_materializes_default_managed_assets_path(single_node_plan: ResolvedSlurmRunPlan) -> None: + assert single_node_plan.invocation.authored.input_bindings.managed_assets_path is None + assert single_node_plan.invocation.effective_input_bindings.managed_assets_path == ( + "/workspace/primary/managed-assets" + ) + + +def test_plan_rejects_managed_assets_overlapping_workspace_state( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = single_node_plan.model_dump(mode="json") + managed_assets_path = "/workspace/primary/runs" + payload["invocation"]["authored"]["input_bindings"]["managed_assets_path"] = managed_assets_path + payload["invocation"]["effective_input_bindings"]["managed_assets_path"] = managed_assets_path + + with pytest.raises(ValidationError, match="managed_assets_path"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) def test_plan_rejects_otel_port_collision(multi_node_plan: ResolvedSlurmRunPlan) -> None: diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py index 1e4086f5e..dcf7addda 100644 --- a/packages/data-designer-slurm/tests/contracts/test_profiles.py +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -124,3 +124,12 @@ def test_profile_requires_explicit_version(profile_catalog: SlurmProfileCatalog) with pytest.raises(ValidationError, match="schema_version"): SlurmProfile.model_validate(payload) + + +def test_profile_rejects_mem_per_gpu_without_gres(profile_catalog: SlurmProfileCatalog) -> None: + payload = profile_catalog.clusters["primary"].model_dump(mode="json") + payload["gpu_request_mode"] = "visible" + payload["scheduler"]["mem_per_gpu"] = "80G" + + with pytest.raises(ValidationError, match="requires GRES"): + SlurmProfile.model_validate(payload) diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index 4a68d2cc6..a2e98e44a 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -9,7 +9,7 @@ "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "e28849e65152d52596c8d97737f218c616cf6355bd8cf7350839241c01ef3f78" + "sha256": "0b6455ac486e99a3b7a012fcdd0bdd406f02851ea4255834eec791482b63f816" }, "run_id": "run-single", "scheduler": { @@ -85,12 +85,12 @@ "run": { "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" + "sha256": "3ca8b4c53642796435ae782cb7710904681df6eb6750d414ad6701861eb85b01" }, "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "e28849e65152d52596c8d97737f218c616cf6355bd8cf7350839241c01ef3f78" + "sha256": "0b6455ac486e99a3b7a012fcdd0bdd406f02851ea4255834eec791482b63f816" }, "run_id": "run-single", "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index c80c95f5f..08d93d747 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -36,6 +36,16 @@ def test_renderer_matches_contract_bound_goldens( assert render_generation_attempt_script(plan, attempt_ordinal=1) == (GOLDEN_DIRECTORY / fixture_name).read_text() +def test_renderer_omits_unspecified_array_throttle(multi_node_plan: ResolvedSlurmRunPlan) -> None: + array_tasks = multi_node_plan.array_tasks.model_copy(update={"max_concurrent": None}) + plan = multi_node_plan.model_copy(update={"array_tasks": array_tasks}) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert "#SBATCH --array=0-1\n" in script + assert "#SBATCH --array=0-1%" not in script + + def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fields( single_node_plan: ResolvedSlurmRunPlan, ) -> None: @@ -98,7 +108,8 @@ def test_renderer_rejects_mem_per_gpu_without_a_slurm_gpu_request( "scheduler": SchedulerProfile(account="research", partition="batch", mem_per_gpu="80G"), } ) - plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + selected_profile = single_node_plan.selected_profile.model_copy(update={"profile": profile}) + plan = single_node_plan.model_copy(update={"selected_profile": selected_profile}) with pytest.raises(SlurmBatchRenderError, match="requires GRES"): render_generation_attempt_script(plan, attempt_ordinal=1) diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index 307ecf162..b611f847f 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -117,6 +117,14 @@ def test_compiler_resolves_explicit_hostname_and_default_profile_selection( } assert {plan.output.root for plan in plans} == {"/workspace/primary/runs/run-single/output"} assert all(plan.deployments[0].node_indices == (0,) for plan in plans) + assert all( + plan.invocation.effective_input_bindings.managed_assets_path == "/workspace/primary/managed-assets" + for plan in plans + ) + assert all( + plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == plan.client.authored.cpus + for plan in plans + ) def test_auto_gpu_resolution_is_explicit_and_scheduler_free( @@ -127,7 +135,7 @@ def test_auto_gpu_resolution_is_explicit_and_scheduler_free( ) -> None: selected = select_profile(profile_catalog, cluster="lab") runtime_bundle = single_node_plan.runtime_bundle.model_copy( - update={"path": "/workspace/lab/runtime/runtime.tar.gz"} + update={"path": f"/workspace/lab/runtime/{single_node_plan.runtime_bundle.sha256}.tar.gz"} ) with pytest.raises(SlurmConfigResolutionError, match="auto gpus_per_node"): @@ -168,7 +176,9 @@ def test_compiler_preserves_explicit_compatibility_run_values( "max_conversation_restarts": 3, "otel_metrics_port": 24000, "shutdown_error_rate": 0.25, + "non_inference_max_parallel_workers": 7, } + payload["invocation"]["input_bindings"]["managed_assets_path"] = "/workspace/explicit-assets" authored = DataDesignerSlurmConfig.model_validate(payload) effective = _resolve_fixture(authored, dependency_lock_single, single_node_plan) @@ -178,6 +188,8 @@ def test_compiler_preserves_explicit_compatibility_run_values( assert effective.invocation.effective_run_config["max_conversation_restarts"] == 3 assert effective.invocation.effective_run_config["otel_metrics_port"] == 24000 assert effective.invocation.effective_run_config["shutdown_error_rate"] == 0.25 + assert effective.invocation.effective_run_config["non_inference_max_parallel_workers"] == 7 + assert effective.invocation.effective_input_bindings.managed_assets_path == "/workspace/explicit-assets" @pytest.mark.parametrize( @@ -227,6 +239,7 @@ def test_compiler_rejects_tensor_parallelism_that_does_not_divide_gpu_shape( {"root": "/outside/output"}, {"root": "/workspace/primary/images/output"}, {"root": "/workspace/primary/runtime/output"}, + {"root": "/workspace/primary/managed-assets/output"}, {"root": "/workspace/primary/runs/other-run/output"}, {"root": "/workspace/primary/runs/run-single/shards"}, {"partitions": 9}, @@ -246,6 +259,42 @@ def test_resolution_rejects_invalid_output_destinations( _resolve_fixture(authored, dependency_lock_single, single_node_plan) +def test_resolution_rejects_output_overlapping_explicit_managed_assets( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["input_bindings"]["managed_assets_path"] = "/workspace/primary/custom-assets" + payload["output"]["root"] = "/workspace/primary/custom-assets/output" + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmConfigResolutionError, match="managed assets"): + _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + +@pytest.mark.parametrize( + "managed_assets_path", + [ + "/workspace/primary", + "/workspace/primary/images/assets", + "/workspace/primary/runs", + ], +) +def test_resolution_rejects_managed_assets_overlapping_workspace_state( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + managed_assets_path: str, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["input_bindings"]["managed_assets_path"] = managed_assets_path + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmConfigResolutionError, match="managed_assets_path"): + _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + def test_compiler_rejects_direct_effective_input_resolution_bypass( authored_run: DataDesignerSlurmConfig, dependency_lock: ResolvedDependencyLock, @@ -301,6 +350,19 @@ def test_compiler_rejects_direct_effective_input_resolution_bypass( ) ) + profile = effective.selected_profile.profile.model_copy( + update={ + "gpu_request_mode": "visible", + "scheduler": effective.selected_profile.profile.scheduler.model_copy(update={"mem_per_gpu": "80G"}), + } + ) + with pytest.raises(SlurmPlanCompilationError, match="requires GRES"): + SlurmRunCompiler.compile( + effective.model_copy( + update={"selected_profile": effective.selected_profile.model_copy(update={"profile": profile})} + ) + ) + def test_compiler_rejects_direct_effective_record_drift( authored_run_single: DataDesignerSlurmConfig, @@ -321,6 +383,15 @@ def test_compiler_rejects_direct_effective_record_drift( SlurmRunCompiler.compile( effective.model_copy(update={"output": effective.output.model_copy(update={"format": "jsonl"})}) ) + bindings = effective.invocation.effective_input_bindings.model_copy( + update={"managed_assets_path": "/workspace/other-assets"} + ) + with pytest.raises(SlurmPlanCompilationError, match="resolved invocation"): + SlurmRunCompiler.compile( + effective.model_copy( + update={"invocation": effective.invocation.model_copy(update={"effective_input_bindings": bindings})} + ) + ) def test_compiler_rejects_model_alias_missing_from_builder( @@ -768,16 +839,33 @@ def test_resolution_rejects_secret_values_in_sourced_builder_payload( assert error.value.__cause__ is None -def test_resolution_rejects_artifact_identity_mismatches( +def test_resolution_rejects_dependency_identity_mismatch( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, single_node_plan: ResolvedSlurmRunPlan, ) -> None: wrong_lock = dependency_lock_single.model_copy(update={"client_image_sha256": "a" * 64}) - wrong_runtime = ArtifactReference(path="/tmp/runtime.tar.gz", sha256="e" * 64) with pytest.raises(SlurmConfigResolutionError, match="client image"): _resolve_fixture(authored_run_single, wrong_lock, single_node_plan) + + +@pytest.mark.parametrize( + "runtime_path", + [ + "/tmp/" + "e" * 64 + ".tar.gz", + "/workspace/primary/runtime/runtime.tar.gz", + "/workspace/primary/runtime/" + "e" * 64 + ".tgz", + ], +) +def test_resolution_rejects_invalid_runtime_bundle_paths( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + runtime_path: str, +) -> None: + wrong_runtime = ArtifactReference(path=runtime_path, sha256="e" * 64) + with pytest.raises(SlurmConfigResolutionError, match="runtime bundle"): _resolve_fixture( authored_run_single, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 12a8dc129..73f1b36ba 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -10,10 +10,10 @@ set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" -readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" +readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="cd05642b6674bdef2f5f3389acb47a82a64c061f9f8501ba06ebb3022223c4da" +readonly DD_PLAN_SHA256="31edd3b3517be10e4b7841f542a65ccd024aaa08748d2770d71d6477600b2e67" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index a863e93e8..24fdb5ea8 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -10,10 +10,10 @@ set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" -readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" +readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="e28849e65152d52596c8d97737f218c616cf6355bd8cf7350839241c01ef3f78" +readonly DD_PLAN_SHA256="0b6455ac486e99a3b7a012fcdd0bdd406f02851ea4255834eec791482b63f816" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 939a0d5d1..137baaf63 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="c91039d39e3a7a2cbe727531495f430133f3b240a48fe5127f5b22c03cd3664f", + expected_fixture_sha256="874b3545cd1c9e601d5ab66eb2fd5c153012ce0467030507e46c80f853809bcc", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="6407ab2e57c8577d5ac1d1a14940bb8b5441c0f8f1c25b108acc2be41fdc4635", + expected_fixture_sha256="630308643ee26e55de5ad1720c046494f7debe8293096400d2095c114aef918a", ) @@ -55,7 +55,7 @@ def _assert_script_matches_plan( ) node_count = max(node_indices) + 1 array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}" - if plan.array_tasks.count > 1: + if plan.array_tasks.count > 1 and plan.array_tasks.max_concurrent is not None: array = f"{array}%{plan.array_tasks.max_concurrent}" plan_path = posixpath.join(posixpath.dirname(plan.authored_config.path), "resolved-plan.json") run_root = posixpath.dirname(plan.authored_config.path) From 25f7e24380f64896a0b133be15f7b2113a7fe909 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 28 Aug 2026 12:55:57 -0300 Subject: [PATCH 9/9] fix: address plan compiler review feedback --- .../src/data_designer/slurm/_errors.py | 22 ++++++++++++++++- .../data_designer/slurm/planning/compiler.py | 3 ++- .../data_designer/slurm/planning/models.py | 4 ++-- .../slurm/planning/resolution.py | 2 +- .../tests/config/test_loading_builder.py | 14 +++++++++++ .../contracts/golden/multi_node_plan.json | 2 +- .../contracts/golden/single_node_plan.json | 2 +- .../tests/contracts/test_planning_records.py | 6 ++--- .../golden/finalization_chain.json | 4 ++-- .../tests/planning/test_compiler.py | 24 +++++++++++++++---- .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 ++-- 13 files changed, 71 insertions(+), 20 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py index 5d2f2b7df..ed045dde9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py @@ -15,6 +15,19 @@ _LOCATION_SEGMENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") _SENSITIVE_LOCATION = re.compile(r"(?:api_?key|credential|password|secret|token)", re.IGNORECASE) +_DYNAMIC_LOCATION_OWNERS = frozenset( + { + "builder_payload", + "clusters", + "deployments", + "effective_run_config", + "environment", + "index_credentials", + "inline", + "model_concurrency", + "run_config", + } +) _ERROR_DESCRIPTIONS = { "extra_forbidden": "field is not permitted", "greater_than": "must be greater than the allowed minimum", @@ -32,6 +45,10 @@ "Value error, builder content digest does not match the resolved input": ( "resolved builder digest does not match its input" ), + "Value error, image inspection digest does not match the resolved SQSH": ( + "resolved image digest does not match its inspection record" + ), + "Value error, mem_per_gpu requires GRES GPU request mode": "mem_per_gpu requires GRES GPU request mode", "Value error, resolved model aliases do not match the inline builder": ( "resolved model aliases do not match the inline builder" ), @@ -61,18 +78,21 @@ def _format_error_detail(detail: dict[str, Any]) -> str: def _format_location(location: Iterable[object]) -> str: parts: list[str] = [] + owner: str | None = None for segment in location: if isinstance(segment, int): if parts: parts[-1] = f"{parts[-1]}[{segment}]" continue if ( - not isinstance(segment, str) + owner in _DYNAMIC_LOCATION_OWNERS + or not isinstance(segment, str) or _LOCATION_SEGMENT.fullmatch(segment) is None or _SENSITIVE_LOCATION.search(segment) is not None ): break parts.append(segment) + owner = segment return ".".join(parts) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index c93ec99a1..62fa8ce21 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -46,7 +46,8 @@ class SlurmRunCompiler: def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: """Return one immutable deterministic execution plan.""" try: - validate_effective_slurm_config(effective) + effective = EffectiveDataDesignerSlurmConfig.model_validate(effective.model_dump(mode="python")) + effective = validate_effective_slurm_config(effective) deployments = _compile_deployments(effective) client = _compile_client(effective, deployments) _validate_port_claims(effective, client, deployments) 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 98378fe19..8243a6aa3 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 @@ -407,8 +407,8 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: raise ValueError("managed_assets_path must not overlap package-managed workspace state") if "non_inference_max_parallel_workers" not in self.invocation.authored.run_config: workers = self.invocation.effective_run_config["non_inference_max_parallel_workers"] - if workers != self.client.authored.cpus: - raise ValueError("default non-inference worker count must match the client CPU count") + if workers != 4: + raise ValueError("default non-inference worker count must match the Data Designer default") expected_logical_names = tuple( f"{deployment.deployment_id}-logical-endpoint" for deployment in self.deployments diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py index 9cadf4a3e..586f8d5d3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -59,6 +59,7 @@ "display_tui": False, "max_conversation_correction_steps": 0, "max_conversation_restarts": 0, + "non_inference_max_parallel_workers": 4, "otel_metrics_port": None, "shutdown_error_rate": 1.0, } @@ -276,7 +277,6 @@ def _materialize_output(authored: DataDesignerSlurmConfig, run_root: str) -> Res def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: values = dict(authored.invocation.run_config) - values.setdefault("non_inference_max_parallel_workers", authored.client.cpus) authored_early_shutdown = {"disable_early_shutdown", "shutdown_error_rate", "shutdown_error_window"}.intersection( values ) diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 964cfd835..f76c8c8e5 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -115,6 +115,20 @@ def test_builder_validation_errors_hide_secret_inputs() -> None: assert error.value.__cause__ is None +def test_builder_validation_errors_hide_dynamic_mapping_keys() -> None: + secret = "sk_live_ABC123XYZ" + + with pytest.raises(SlurmConfigBuilderError) as error: + _config_builder().with_invocation( + num_records=1, + dataset_name="generated", + model_concurrency={secret: 0}, + ) + + assert "model_concurrency" in str(error.value) + assert secret not in str(error.value) + + def test_builder_custom_validation_errors_hide_secret_values() -> None: secret = "super-secret-token" 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 1d1d0fd96..c5e65841a 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 @@ -364,7 +364,7 @@ "max_conversation_correction_steps": 0, "max_conversation_restarts": 0, "max_in_flight_tasks": 1024, - "non_inference_max_parallel_workers": 32, + "non_inference_max_parallel_workers": 4, "otel_metrics_port": null, "preserve_dropped_columns": true, "progress_interval": 5.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 1ce8e8a5e..c2a978802 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 @@ -212,7 +212,7 @@ "max_conversation_correction_steps": 0, "max_conversation_restarts": 0, "max_in_flight_tasks": 1024, - "non_inference_max_parallel_workers": 32, + "non_inference_max_parallel_workers": 4, "otel_metrics_port": null, "preserve_dropped_columns": true, "progress_interval": 5.0, 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 63975bed8..5ae19bac4 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -111,13 +111,13 @@ def test_plan_rejects_unmaterialized_run_config(multi_node_plan: ResolvedSlurmRu ResolvedSlurmRunPlan.model_validate(payload) -def test_plan_derives_default_non_inference_worker_count_from_client_cpus( +def test_plan_requires_default_non_inference_worker_count( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: payload = multi_node_plan.model_dump(mode="json") - payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 4 + payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 32 - with pytest.raises(ValidationError, match="client CPU count"): + with pytest.raises(ValidationError, match="Data Designer default"): ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index a2e98e44a..3037df673 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -9,7 +9,7 @@ "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "0b6455ac486e99a3b7a012fcdd0bdd406f02851ea4255834eec791482b63f816" + "sha256": "4a64741df1a3eaba63d1cd19f08ce5dfe9017509f839a78678d17324e46d98b4" }, "run_id": "run-single", "scheduler": { @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "0b6455ac486e99a3b7a012fcdd0bdd406f02851ea4255834eec791482b63f816" + "sha256": "4a64741df1a3eaba63d1cd19f08ce5dfe9017509f839a78678d17324e46d98b4" }, "run_id": "run-single", "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index b611f847f..f9a2389fa 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -121,10 +121,7 @@ def test_compiler_resolves_explicit_hostname_and_default_profile_selection( plan.invocation.effective_input_bindings.managed_assets_path == "/workspace/primary/managed-assets" for plan in plans ) - assert all( - plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == plan.client.authored.cpus - for plan in plans - ) + assert all(plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == 4 for plan in plans) def test_auto_gpu_resolution_is_explicit_and_scheduler_free( @@ -481,6 +478,25 @@ def test_compiler_rejects_resolved_image_identity_drift( SlurmRunCompiler.compile(effective.model_copy(update={"deployment_images": deployment_images})) +@pytest.mark.parametrize("invalid_field", ["client_image", "runtime_bundle"]) +def test_compiler_revalidates_nested_effective_contracts( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + invalid_field: str, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + if invalid_field == "client_image": + invalid_value = effective.client_image.model_copy(update={"path": "/workspace/images/client.txt"}) + else: + invalid_value = effective.runtime_bundle.model_copy( + update={"path": f"/workspace/primary/runtime/../{effective.runtime_bundle.sha256}.tar.gz"} + ) + + with pytest.raises(SlurmPlanCompilationError, match="failed validation"): + SlurmRunCompiler.compile(effective.model_copy(update={invalid_field: invalid_value})) + + def test_compiler_rejects_otel_collision_before_runtime( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 73f1b36ba..aba4ed3d8 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="31edd3b3517be10e4b7841f542a65ccd024aaa08748d2770d71d6477600b2e67" +readonly DD_PLAN_SHA256="5b7622e41f710c3dca96201444753ffd5a36adc2803e8b81562979d1dfb0b670" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 24fdb5ea8..cdbe82d8a 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="0b6455ac486e99a3b7a012fcdd0bdd406f02851ea4255834eec791482b63f816" +readonly DD_PLAN_SHA256="4a64741df1a3eaba63d1cd19f08ce5dfe9017509f839a78678d17324e46d98b4" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 137baaf63..bcbeb3efa 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="874b3545cd1c9e601d5ab66eb2fd5c153012ce0467030507e46c80f853809bcc", + expected_fixture_sha256="c4ec9a60ffbc2feed6a34e1ca3b5363b47bba2add15f6cbb419e18d2c39c3ab7", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="630308643ee26e55de5ad1720c046494f7debe8293096400d2095c114aef918a", + expected_fixture_sha256="f967331f7c08f1e1adf24f42a13ddc7c196e9d26a34f9e3098f1533cc3348466", )