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..584bdf62b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py @@ -0,0 +1,163 @@ +# 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 +from collections.abc import Iterable, Mapping, Sequence +from types import UnionType +from typing import Annotated, Any, Literal, Union, get_args, get_origin + +import yaml +from pydantic import BaseModel, ValidationError + +_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, 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" + ), + "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, + models: type[BaseModel] | tuple[type[BaseModel], ...], +) -> str: + """Summarize validation without rendering user-controlled values.""" + candidates = models if isinstance(models, tuple) else (models,) + model_type = next((model for model in candidates if model.__name__ == error.title), None) + details = error.errors(include_url=False, include_context=False, include_input=False) + summaries = sorted({_format_error_detail(detail, model_type=model_type) for detail in details}) + count = error.error_count() + noun = "error" if count == 1 else "errors" + summary = f": {'; '.join(summaries)}" if summaries else "" + return f"{subject} failed validation ({count} {noun}{summary})" + + +def _format_error_detail(detail: dict[str, Any], *, model_type: type[BaseModel] | None) -> 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", ()), model_type=model_type) + return f"{location}: {description}" if location else description + + +def _format_location(location: Iterable[object], *, model_type: type[BaseModel] | None) -> str: + if model_type is None: + return "" + parts: list[str] = [] + schemas: tuple[object, ...] = (model_type,) + for segment in location: + schemas = tuple(candidate for schema in schemas for candidate in _expand_schema(schema)) + if any(_is_mapping_schema(schema) for schema in schemas): + break + if isinstance(segment, int): + item_schemas = tuple(item for schema in schemas for item in _sequence_item_schemas(schema, index=segment)) + if not parts or not item_schemas: + break + parts[-1] = f"{parts[-1]}[{segment}]" + schemas = item_schemas + continue + if not isinstance(segment, str): + break + field_schemas = tuple( + field.annotation + for schema in schemas + if isinstance(schema, type) + and issubclass(schema, BaseModel) + and (field := schema.model_fields.get(segment)) is not None + ) + if not field_schemas: + branch_schemas = _tagged_union_schemas(schemas, tag=segment) + if not branch_schemas: + break + schemas = branch_schemas + continue + parts.append(segment) + schemas = field_schemas + return ".".join(parts) + + +def _expand_schema(schema: object) -> tuple[object, ...]: + origin = get_origin(schema) + if origin is Annotated: + return _expand_schema(get_args(schema)[0]) + if origin in (Union, UnionType): + return tuple(candidate for item in get_args(schema) for candidate in _expand_schema(item)) + return (schema,) + + +def _is_mapping_schema(schema: object) -> bool: + origin = get_origin(schema) + candidate = origin or schema + return isinstance(candidate, type) and issubclass(candidate, Mapping) + + +def _tagged_union_schemas(schemas: tuple[object, ...], *, tag: str) -> tuple[object, ...]: + if len(schemas) < 2: + return () + return tuple( + schema + for schema in schemas + if isinstance(schema, type) + and issubclass(schema, BaseModel) + and ( + schema.__name__ == tag + or any(_annotation_contains_literal(field.annotation, value=tag) for field in schema.model_fields.values()) + ) + ) + + +def _annotation_contains_literal(annotation: object, *, value: str) -> bool: + return any( + get_origin(candidate) is Literal and value in get_args(candidate) for candidate in _expand_schema(annotation) + ) + + +def _sequence_item_schemas(schema: object, *, index: int) -> tuple[object, ...]: + origin = get_origin(schema) + candidate = origin or schema + if not isinstance(candidate, type) or not issubclass(candidate, Sequence): + return () + arguments = get_args(schema) + if not arguments: + return () + if origin is tuple and arguments[-1] is not Ellipsis: + return (arguments[index],) if index < len(arguments) else () + return (arguments[0],) + + +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 8107c1f8a..6d329c23d 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,9 @@ DataDesignerSlurmBenchmarkConfig, FixedRecordPolicy, ) +from data_designer.slurm.config.builder import DataDesignerSlurmConfigBuilder from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.config.errors import SlurmConfigBuilderError, SlurmConfigLoadError from data_designer.slurm.config.images import ( ClientImageInspection, ImageBuildRequest, @@ -24,6 +26,13 @@ InstalledDistribution, ServingImageInspection, ) +from data_designer.slurm.config.loading import ( + DEFAULT_PROFILE_FILE_NAME, + PROFILE_FILE_ENVIRONMENT, + load_profile_catalog, + load_run_config, + resolve_profile, +) from data_designer.slurm.config.profiles import ( ContainerMount, GpuRequestMode, @@ -70,6 +79,8 @@ "ContainerMount", "DataDesignerSlurmBenchmarkConfig", "DataDesignerSlurmConfig", + "DataDesignerSlurmConfigBuilder", + "DEFAULT_PROFILE_FILE_NAME", "DeploymentResources", "DeploymentTopology", "FixedRecordPolicy", @@ -87,6 +98,7 @@ "LocalStdioMCPProviderConfig", "OutputConfig", "ProfileSelectionSource", + "PROFILE_FILE_ENVIRONMENT", "QueueBackpressureConfig", "RemoteMCPProviderConfig", "SchedulerProfile", @@ -94,11 +106,16 @@ "SelectedSlurmProfile", "ServerDeploymentConfig", "ServingImageInspection", + "SlurmConfigBuilderError", + "SlurmConfigLoadError", "SlurmProfile", "SlurmProfileCatalog", "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..679d818b3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -0,0 +1,200 @@ +# 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 +from typing import TypeVar + +import yaml +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, + BuilderInput, + ClientConfig, + ClientDependencies, + DataDesignerSlurmConfig, + InputBindings, + InvocationConfig, + InvocationDiagnostics, + MCPProviderConfig, + OutputConfig, + ServerDeploymentConfig, + SubmissionConfig, +) + +_ConfigValueT = TypeVar("_ConfigValueT", bound=BaseModel) + + +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(_validate_model(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(_validate_model(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 = _validate_model( + InvocationConfig, + { + "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 = _validate_model( + ClientConfig, + { + "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(_validate_model(ServerDeploymentConfig, deployment)) + return self + + def with_array_tasks( + self, + *, + count: int, + max_concurrent: int | None = None, + ) -> DataDesignerSlurmConfigBuilder: + """Set deterministic horizontal sharding.""" + 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 = _validate_model(SubmissionConfig, values) + return self + + def with_output(self, **values: object) -> DataDesignerSlurmConfigBuilder: + """Set typed dataset output intent.""" + self._output = _validate_model(OutputConfig, 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 SlurmConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") + assert self._invocation is not None + assert self._client is not None + 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: + """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 SlurmConfigBuilderError("config path must end in .json, .yaml, or .yml") + try: + output_path.write_text(contents, encoding="utf-8") + except OSError: + raise SlurmConfigBuilderError(f"cannot write Slurm config {output_path}") from None + + +def _validate_model(config_type: type[_ConfigValueT], value: object) -> _ConfigValueT: + try: + return config_type.model_validate(value) + except ValidationError as error: + message = format_validation_error(error, subject=config_type.__name__, models=config_type) + 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 new file mode 100644 index 000000000..684465871 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -0,0 +1,227 @@ +# 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._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, + 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" + +_ConfigT = TypeVar("_ConfigT", DataDesignerSlurmConfig, SlurmProfileCatalog) +_HostnameResolver = Callable[[], tuple[str, ...]] + + +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 SlurmConfigLoadError("YAML merge keys are not supported") + key = loader.construct_object(key_node, deep=deep) + try: + duplicate = key in mapping + except TypeError: + raise SlurmConfigLoadError("configuration mapping keys must be scalar values") from None + if duplicate: + raise SlurmConfigLoadError("duplicate configuration key") + 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.""" + try: + sources = sum(source is not None for source in (profile, catalog, profile_file)) + if sources > 1: + raise SlurmConfigLoadError("profile, catalog, and profile_file are mutually exclusive") + if profile is not None: + if cluster is not None: + raise SlurmConfigLoadError("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) + ) + return select_profile( + catalog, + cluster=cluster, + hostnames=normalized_hostnames, + catalog_path=catalog_path, + ) + except SlurmConfigLoadError: + raise + except ValidationError as error: + message = format_validation_error(error, subject="profile selection", models=SelectedSlurmProfile) + raise SlurmConfigLoadError(message) from None + except ValueError as error: + raise SlurmConfigLoadError(str(error)) from None + + +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: + 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: + _reject_run_environment_interpolation(payload) + else: + _reject_environment_interpolation(payload) + return config_type.model_validate(payload) + except SlurmConfigLoadError: + raise + except ValidationError as error: + message = format_validation_error( + error, + subject=f"configuration file {resolved_path}", + models=config_type, + ) + 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]: + 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 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 SlurmConfigLoadError("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 SlurmConfigLoadError("duplicate configuration key") + result[key] = value + return result + + +def _reject_environment_interpolation(value: object) -> None: + if isinstance(value, str) and "${" in value: + raise SlurmConfigLoadError("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 _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, + *, + 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 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) + 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 SlurmConfigLoadError("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/config/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py index a6a0b98c3..a5c8c91c0 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 cd323b763..0d9380a92 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 @@ -306,11 +306,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 b8c0b200e..226b9899e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -36,11 +36,11 @@ # TODO: Remove these compatibility exports after Stage 2 branches import shared scalars from slurm.types. ModelAlias = str -_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: @@ -61,7 +61,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: @@ -94,8 +94,8 @@ class ContractValue(BaseModel): model_config = ConfigDict( extra="forbid", frozen=True, - allow_inf_nan=False, hide_input_in_errors=True, + allow_inf_nan=False, protected_namespaces=(), strict=True, validate_default=True, @@ -169,6 +169,26 @@ def compute_canonical_json_sha256(value: object) -> Sha256Digest: return hashlib.sha256(canonical_json(value)).hexdigest() +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() + + +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) + + # TODO: Remove after in-flight Stage 2 branches adopt the explicit canonical-JSON digest name. compute_sha256 = compute_canonical_json_sha256 @@ -284,6 +304,7 @@ class ResumeWorkspace(ContractValue): "Duration", "EnvironmentName", "Identifier", + "ModelAlias", "NetworkPort", "NonNegativeDuration", "RecordRange", @@ -293,6 +314,11 @@ class ResumeWorkspace(ContractValue): "ShardId", "canonical_json", "compute_canonical_json_sha256", + "compute_serialized_json_sha256", + "compute_sha256", + "derive_managed_assets_path", + "is_path_below", + "paths_overlap", "pretty_json", "validate_absolute_path", "validate_local_config_path", 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/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py index 160b7ef77..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 @@ -23,12 +23,10 @@ ResolvedTopology, ResumeWorkspace, ) -from data_designer.slurm.planning.validation import PlanContractError, validate_resolved_plan __all__ = [ "ArtifactReference", "LockedPackage", - "PlanContractError", "PlannedShard", "PortClaim", "RecordRange", @@ -43,5 +41,4 @@ "ResolvedSubmission", "ResolvedTopology", "ResumeWorkspace", - "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 new file mode 100644 index 000000000..964a282c3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal pure deterministic Slurm plan compilation.""" + +from __future__ import annotations + +import posixpath + +from pydantic import ValidationError + +from data_designer.slurm._errors import format_validation_error +from data_designer.slurm.contracts import ( + ArtifactReference, + RecordRange, + ResumeWorkspace, + compute_serialized_json_sha256, +) +from data_designer.slurm.planning.errors import SlurmPlanCompilationError, SlurmPlanContractError +from data_designer.slurm.planning.models import ( + PlannedShard, + PortClaim, + ResolvedClient, + ResolvedDeployment, + ResolvedSlurmRunPlan, + ResolvedTopology, +) +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 +_PORT_RANGE_SIZE = 1000 +_COMPILER_VALIDATION_MODELS = ( + ArtifactReference, + EffectiveDataDesignerSlurmConfig, + PlannedShard, + PortClaim, + RecordRange, + ResolvedClient, + ResolvedDeployment, + ResolvedSlurmRunPlan, + ResolvedTopology, + ResumeWorkspace, +) + + +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: + 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) + 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 (SlurmPlanCompilationError, SlurmPlanContractError): + raise + except ValidationError as error: + message = format_validation_error( + error, + subject="Slurm plan compilation", + models=_COMPILER_VALIDATION_MODELS, + ) + raise SlurmPlanCompilationError(message) from None + except ValueError as error: + raise SlurmPlanCompilationError(str(error)) from None + + +def _compile_deployments( + effective: EffectiveDataDesignerSlurmConfig, +) -> tuple[ResolvedDeployment, ...]: + if len(effective.authored.deployments) > _PORT_RANGE_SIZE: + 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( + zip(effective.authored.deployments, effective.deployment_images, strict=True) + ): + topology = ResolvedTopology.derive( + node_count=authored.resources.nodes, + gpus_per_node=effective.resolved_gpus_per_node, + tensor_parallel=authored.topology.tensor_parallel, + nodes_per_replica=authored.topology.nodes_per_replica, + ) + if topology.replicas_per_node_group > _PORT_RANGE_SIZE: + raise SlurmPlanCompilationError("replica lanes exceed the compiler-owned deployment port range") + 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 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 SlurmPlanCompilationError("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.invocation.effective_input_bindings.seed_path + if seed_path is not None: + partition = ArtifactReference( + path=posixpath.join(shard_root, "input-partition.json"), + sha256=compute_serialized_json_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 _run_root(effective: EffectiveDataDesignerSlurmConfig) -> str: + return posixpath.join(effective.selected_profile.profile.workspace_root, "runs", effective.run_id) 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 9b71cc31f..ca2a57c80 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 @@ -28,6 +28,7 @@ ArrayTasksConfig, ClientConfig, ClientDependencies, + InputBindings, InvocationConfig, ServerDeploymentConfig, SubmissionConfig, @@ -37,15 +38,20 @@ ContractRecord, ContractValue, Identifier, + ModelAlias, RecordRange, ResumeWorkspace, Sha256Digest, ShardId, - compute_canonical_json_sha256, + compute_serialized_json_sha256, + derive_managed_assets_path, + is_path_below, + paths_overlap, validate_absolute_path, validate_local_config_path, validate_plain_text, ) +from data_designer.slurm.planning.builder_identity import get_declared_model_aliases from data_designer.slurm.types import NetworkPort @@ -158,9 +164,9 @@ class ResolvedBuilderInput(ContractValue): authored_source: str | None = None source: ArtifactReference | None = None inline: dict[str, JsonValue] | None = None + # Both forms use deterministic persisted builder JSON bytes. content_sha256: Sha256Digest - model_aliases: tuple[str, ...] - referenced_model_aliases: tuple[str, ...] = () + model_aliases: tuple[ModelAlias, ...] @model_validator(mode="after") def validate_input(self) -> ResolvedBuilderInput: @@ -170,20 +176,15 @@ def validate_input(self) -> ResolvedBuilderInput: if self.authored_source is not None: raise ValueError("inline builder input cannot contain authored_source") validate_no_plaintext_secrets(self.inline, field_name="resolved inline builder input") - expected_digest = compute_canonical_json_sha256(self.inline) - model_aliases, referenced_aliases = _extract_builder_aliases(self.inline) - if self.model_aliases != model_aliases: + 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") - 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 @@ -191,6 +192,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") @@ -401,6 +403,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: @@ -413,20 +417,33 @@ 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 not set(self.builder.referenced_model_aliases).issubset(aliases): - raise ValueError("each referenced Data Designer model alias requires a deployment") + if set(aliases) != set(self.builder.model_aliases): + raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") node_indices = tuple(index for deployment in self.deployments for index in deployment.node_indices) if node_indices != tuple(range(len(node_indices))): raise ValueError("deployment nodes must be disjoint and contiguous in authored order") if self.client.host_node_index != self.deployments[0].node_indices[0]: raise ValueError("client must be colocated on the first node of the first deployment") + 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 != 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 @@ -453,15 +470,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 @@ -474,7 +496,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") @@ -510,47 +532,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 - - -def _extract_builder_aliases(builder: dict[str, JsonValue]) -> tuple[tuple[str, ...], tuple[str, ...]]: - 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[str] = [] - 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[str] = [] - - def collect(value: JsonValue, *, key: str | None = None) -> None: - if key == "model_configs": - return - if key == "model_alias" or (key is not None and key.endswith("_model_alias")): - if not isinstance(value, str): - raise ValueError(f"builder {key} must be a string") - referenced_aliases.append(value) - return - if key == "model_aliases": - if not isinstance(value, list) or any(not isinstance(alias, str) for alias in value): - raise ValueError("builder model_aliases must be a list of strings") - referenced_aliases.extend(value) - return - if isinstance(value, dict): - for child_key, child in value.items(): - collect(child, key=child_key) - elif isinstance(value, list): - for child in value: - collect(child) - - collect(data_designer) - return tuple(model_aliases), tuple(dict.fromkeys(referenced_aliases)) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py new file mode 100644 index 000000000..3fc47b05f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -0,0 +1,428 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal 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, 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, + 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, +) + +__all__: list[str] = [] + +_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, + "non_inference_max_parallel_workers": 4, + "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)]) +_RESOLUTION_VALIDATION_MODELS = ( + RunConfig, + ArtifactReference, + BuilderInput, + InputBindings, + ResolvedBuilderInput, + ResolvedInvocation, + ResolvedOutput, + ResolvedSubmission, +) + + +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, resolved_builder_payload = _resolve_builder( + authored, + run_root=run_root, + builder_payload=builder_payload, + ) + effective = 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=resolved_builder_payload, + invocation=_materialize_invocation(authored, workspace_root), + client_image=client_image, + deployment_images=deployment_images, + dependency_lock=dependency_lock, + 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: + message = format_validation_error( + error, + subject="Slurm configuration resolution", + models=(*_RESOLUTION_VALIDATION_MODELS, EffectiveDataDesignerSlurmConfig), + ) + raise SlurmConfigResolutionError(message) from None + except ValueError as error: + 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 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: + 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, 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) + 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) + 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") + 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 content-addressed 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": + 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, +) -> 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_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") + 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, + ), + validated_payload, + ) + + +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), + ) + + +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( + 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: + _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: + 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): + _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, + *, + 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): + 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_path_below(output_root, run_output_root) + ): + raise SlurmConfigResolutionError("output root must not overlap another package-managed run") + + +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 c47ebc512..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 @@ -3,22 +3,22 @@ from __future__ import annotations +import posixpath + from pydantic import JsonValue from data_designer.config import RunConfig from data_designer.slurm.config.images import ClientImageInspection from data_designer.slurm.config.run import DataDesignerSlurmConfig +from data_designer.slurm.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 ( ResolvedDependencyLock, ResolvedSlurmRunPlan, - _extract_builder_aliases, ) -class PlanContractError(ValueError): - """Raised when a resolved plan does not match its authored inputs.""" - - def validate_resolved_plan( authored: DataDesignerSlurmConfig, dependency_lock: ResolvedDependencyLock, @@ -32,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, @@ -45,22 +53,28 @@ 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, "resolved builder source does not match authored input", ) if builder_payload is None: - raise PlanContractError("sourced builder validation requires its resolved payload") - model_aliases, referenced_aliases = _extract_builder_aliases(builder_payload) + 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.referenced_model_aliases == referenced_aliases, - "resolved referenced 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 expected_partition = authored.submission.partition or plan.selected_profile.profile.scheduler.partition @@ -126,7 +140,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 new file mode 100644 index 000000000..e1cd641a2 --- /dev/null +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -0,0 +1,362 @@ +# 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, LLMTextColumnConfig, ModelConfig +from data_designer.slurm.config import ( + DEFAULT_PROFILE_FILE_NAME, + PROFILE_FILE_ENVIRONMENT, + DataDesignerSlurmConfig, + DataDesignerSlurmConfigBuilder, + ProfileSelectionSource, + SlurmConfigBuilderError, + SlurmConfigLoadError, + SlurmProfileCatalog, + load_profile_catalog, + load_run_config, + resolve_profile, +) + + +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( + 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(SlurmConfigBuilderError, 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(SlurmConfigBuilderError, match="must end"): + _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(SlurmConfigBuilderError): + 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" + + 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 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_validation_errors_preserve_discriminated_union_locations() -> None: + secret = "super-secret-token" + + with pytest.raises(SlurmConfigBuilderError) as error: + _config_builder().with_invocation( + num_records=1, + dataset_name="generated", + mcp_providers=[ + { + "provider_type": "streamable_http", + "name": "example", + "endpoint": f"https://example.com/?token={secret}", + } + ], + ) + + assert "mcp_providers[0].endpoint" 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" + + 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(SlurmConfigBuilderError, match="cannot write"): + _config_builder().write_config(path) + + +@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"), + (".yaml", "schema_version: 1\nbuilder:\n source: ${HOME}/builder.json\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(SlurmConfigLoadError, 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(SlurmConfigLoadError, match="root must be an object"): + load_run_config(json_path) + with pytest.raises(SlurmConfigLoadError, match="must end"): + 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(SlurmConfigLoadError) as error: + load_run_config(path) + + assert secret not in str(error.value) + 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_validation_errors_preserve_nested_schema_locations(tmp_path: Path) -> None: + secret = "super-secret-token" + payload = _config_builder().build().model_dump(mode="json") + payload["deployments"][0]["server"]["image"] = {"path": f"/images/{secret}.txt"} + path = tmp_path / "run.json" + path.write_text(json.dumps(payload)) + + with pytest.raises(SlurmConfigLoadError) as error: + load_run_config(path) + + assert "deployments[0].server.image.path" in str(error.value) + assert secret not in str(error.value) + + +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 None + + +@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, +) -> 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: pytest.fail("explicit cluster selection must not resolve hostnames"), + 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(SlurmConfigLoadError, match="mutually exclusive"): + resolve_profile(catalog=profile_catalog, profile_file="profile.json") + with pytest.raises(SlurmConfigLoadError, match="must not be empty"): + resolve_profile(environ={PROFILE_FILE_ENVIRONMENT: ""}) + with pytest.raises(SlurmConfigLoadError, match="cluster selection"): + resolve_profile(profile=profile_catalog.clusters["primary"], cluster="primary") + with pytest.raises(SlurmConfigLoadError, 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(SlurmConfigLoadError, match="multiple clusters"): + 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() + + 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/contracts/golden/authored_run_single.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json index 346b4500c..7a852dbcd 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 48427e849..1468aac49 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": [], @@ -31,7 +31,6 @@ "generator", "judge" ], - "referenced_model_aliases": [], "source": null }, "client": { @@ -355,6 +354,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, @@ -385,7 +388,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 4b86c8657..b2747d4c0 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,15 +1,15 @@ { "array_tasks": { "count": 1, - "max_concurrent": 1 + "max_concurrent": null }, "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "3d092f4a6d3731d048c519b5e5a41b5f0c2973a4b3aff6459710dfbe64f66e18" + "sha256": "5b7614af448e9c647a2244257da37923784c3897e221a15f0a050b3e5f0c0324" }, "builder": { "authored_source": null, - "content_sha256": "b3ef5fc1fe675a8e004633f84842ac60cf82d5ba3dc68b4d50ee4438448b0570", + "content_sha256": "0bac3a88498774a30b64f0a62617511c8e75601ca54e0b775d86484bb13788d8", "inline": { "data_designer": { "columns": [], @@ -25,7 +25,6 @@ "model_aliases": [ "generator" ], - "referenced_model_aliases": [], "source": null }, "client": { @@ -201,6 +200,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, @@ -231,7 +234,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 dcf9c876d..58651662a 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -695,6 +695,7 @@ def test_run_validates_public_run_config_and_shard_count(authored_run: DataDesig @pytest.mark.parametrize("readiness_path", ["health", "//other-host/health", "/health check"]) def test_small_config_values_validate_at_boundary(readiness_path: str) -> 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 cb907c010..fa45b8291 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -10,18 +10,18 @@ from pydantic import ValidationError from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig -from data_designer.slurm.contracts import compute_canonical_json_sha256 +from data_designer.slurm.contracts import compute_serialized_json_sha256 from data_designer.slurm.planning import ( ArtifactReference, - PlanContractError, ResolvedBuilderInput, ResolvedDependencyLock, ResolvedDeployment, ResolvedSlurmRunPlan, ResolvedSubmission, ResolvedTopology, - validate_resolved_plan, ) +from data_designer.slurm.planning.errors import SlurmPlanContractError +from data_designer.slurm.planning.validation import validate_resolved_plan def test_resolved_topology_derives_all_resource_fields() -> None: @@ -105,7 +105,7 @@ def test_resolved_builder_rejects_plaintext_secrets(multi_node_plan: ResolvedSlu inline = payload["inline"] assert isinstance(inline, dict) inline["plugin_config"] = {"api_key": "VERY_SECRET_VALUE"} - payload["content_sha256"] = compute_canonical_json_sha256(inline) + payload["content_sha256"] = compute_serialized_json_sha256(inline) with pytest.raises(ValidationError, match="plaintext values under secret-bearing keys") as error: ResolvedBuilderInput.model_validate_json(json.dumps(payload)) @@ -134,6 +134,12 @@ def test_resolved_builder_rejects_plaintext_secrets(multi_node_plan: ResolvedSlu 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: @@ -152,13 +158,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_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"] = 32 - with pytest.raises(ValidationError, match="RunConfig default"): + with pytest.raises(ValidationError, match="Data Designer default"): ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) @@ -166,12 +172,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: @@ -207,22 +232,12 @@ 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_canonical_json_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)) -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_canonical_json_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 @@ -273,19 +288,19 @@ 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_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"] = { "authored_source": "builder.json", - "source": {"path": "/workspace/primary/runs/run-001/builder.json", "sha256": "a" * 64}, + "source": {"path": "/workspace/primary/runs/run-001/builder-config.json", "sha256": builder_digest}, "inline": None, - "content_sha256": "a" * 64, + "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 ( @@ -299,6 +314,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", [ @@ -394,7 +429,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) @@ -413,7 +448,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) @@ -431,7 +466,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") @@ -454,7 +489,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( @@ -487,7 +522,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) @@ -506,7 +541,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) @@ -520,7 +555,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) @@ -533,5 +568,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/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 7e993b116..4568b5566 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": "897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" + "sha256": "14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" }, "run_id": "run-single", "scheduler": { @@ -85,12 +85,12 @@ "run": { "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "3d092f4a6d3731d048c519b5e5a41b5f0c2973a4b3aff6459710dfbe64f66e18" + "sha256": "5b7614af448e9c647a2244257da37923784c3897e221a15f0a050b3e5f0c0324" }, "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" + "sha256": "14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" }, "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 new file mode 100644 index 000000000..f9a2389fa --- /dev/null +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -0,0 +1,904 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from copy import deepcopy +from pathlib import Path + +import pytest + +from data_designer.config import ( + DataDesignerConfigBuilder, + DropColumnsProcessorConfig, + JudgeScoreProfilerConfig, + LLMTextColumnConfig, + ModelConfig, +) +from data_designer.slurm.config import ( + BuilderInput, + DataDesignerSlurmConfig, + InputBindings, + SlurmProfileCatalog, + select_profile, +) +from data_designer.slurm.contracts import compute_serialized_json_sha256 +from data_designer.slurm.planning import ( + ArtifactReference, + ResolvedDependencyLock, + ResolvedSlurmRunPlan, +) +from data_designer.slurm.planning.compiler import SlurmRunCompiler +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" + + +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 = 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() + 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( + SlurmRunCompiler.compile( + _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) + 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"] == 4 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": f"/workspace/lab/runtime/{single_node_plan.runtime_bundle.sha256}.tar.gz"} + ) + + with pytest.raises(SlurmConfigResolutionError, 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 = SlurmRunCompiler.compile( + _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, + "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) + + 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 + 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( + ("run_config", "expected_rate", "expected_window"), + [ + ({"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( + 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, + 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(SlurmPlanCompilationError, match="tensor_parallel"): + SlurmRunCompiler.compile(_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"}, + {"root": "/workspace/primary/managed-assets/output"}, + {"root": "/workspace/primary/runs/other-run/output"}, + {"root": "/workspace/primary/runs/run-single/shards"}, + {"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(SlurmConfigResolutionError, match="output"): + _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, + 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), + } + ) + ) + + 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, + 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"})}) + ) + 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( + 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(SlurmPlanCompilationError, match="exactly cover"): + SlurmRunCompiler.compile(_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(SlurmPlanCompilationError, match="exactly cover"): + SlurmRunCompiler.compile( + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + ) + + +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})) + + +@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, + 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(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( + 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 = 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), + (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_serialized_json_sha256( + { + "record_range": shard.record_range.model_dump(mode="json"), + "seed_path": "/datasets/seed.parquet", + } + ) + 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"), + [ + ( + { + "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(SlurmConfigResolutionError, 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(SlurmConfigResolutionError, match=field): + _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(SlurmConfigResolutionError, 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, + 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 = SlurmRunCompiler.compile(_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 = 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, + 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 = SlurmRunCompiler.compile( + _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 + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("bytes", "builder digest"), + ("aliases", "model aliases"), + ], +) +def test_compiler_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" + drifted = effective.model_copy(update={"builder_payload": drifted_payload}) + + with pytest.raises(SlurmPlanContractError, match=message): + SlurmRunCompiler.compile(drifted) + + +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"] + 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(SlurmConfigResolutionError, match="failed validation") as error: + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + +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}) + + 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, + 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 = SlurmRunCompiler.compile(_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/packages/data-designer-slurm/tests/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py index f19973cbd..fdd83ee99 100644 --- a/packages/data-designer-slurm/tests/serving/test_resolver.py +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -11,7 +11,7 @@ import data_designer.slurm.serving.resolver as resolver_module from data_designer.slurm.config import QueueBackpressureConfig -from data_designer.slurm.contracts import pretty_json +from data_designer.slurm.contracts import compute_serialized_json_sha256, pretty_json from data_designer.slurm.planning import ResolvedDeployment, ResolvedSlurmRunPlan from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment from data_designer.slurm.serving.resolver import ( @@ -261,6 +261,11 @@ def _multi_group_plan(plan: ResolvedSlurmRunPlan) -> ResolvedSlurmRunPlan: payload = plan.model_dump(mode="json") payload["deployments"] = payload["deployments"][:1] payload["client"]["ports"] = payload["client"]["ports"][:1] + payload["builder"]["inline"]["data_designer"]["model_configs"] = payload["builder"]["inline"]["data_designer"][ + "model_configs" + ][:1] + payload["builder"]["model_aliases"] = ["generator"] + payload["builder"]["content_sha256"] = compute_serialized_json_sha256(payload["builder"]["inline"]) deployment = payload["deployments"][0] deployment["authored"]["resources"]["nodes"] = 4 deployment["authored"]["topology"]["tensor_parallel"] = 4 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 b3c1e40c8..fc36170e8 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="00e206f4b759ca06336c02d3694f7289551d1a0b9ba50c05f37cdbf7a213756d" +readonly DD_PLAN_SHA256="cd95eb995dd9705d12336b9760ca1f354cdd7614dce992cc4f92a6faa47dbfdc" 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 115d998da..e3f8529d5 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="897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" +readonly DD_PLAN_SHA256="14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" 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 34d947543..9880d08e8 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="244fe9b160b361645e3f9be13d40d5d6157563a69fa37037fa974d2a3acead2f", + expected_fixture_sha256="65032718bfc8fbd6c60700add08ebc098ba2e64858c23808decbefe5a8153d91", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="37681e8805eb81ca9c4e1fc2da88b14031d7a150c9fd729104ed875ed2503224", + expected_fixture_sha256="58c3a2eddc9a0d377eb0a90655e70a6b2484cade0c52b2a01b9b388dcd4667e2", ) @@ -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)