From 48644f22e81385c9bbb608671f9d65bf5183166a Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 27 Aug 2026 16:17:22 -0600 Subject: [PATCH 1/5] feat: add Slurm image lifecycle jobs Stage digest-bound OCI import and existing-SQSH inspection plans beneath the selected workspace. Reuse safe batch rendering and isolated Slurm submission while shipping a standalone factual inspector for target images. Part of #867 Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/images/errors.py | 4 + .../data_designer/slurm/images/lifecycle.py | 282 ++++++++++++ .../src/data_designer/slurm/images/records.py | 91 +++- .../slurm/images/resources/__init__.py | 6 + .../slurm/images/resources/enroot.rc | 6 + .../slurm/images/resources/inspect_image.py | 137 ++++++ .../src/data_designer/slurm/launcher/batch.py | 57 +++ .../data_designer/slurm/launcher/renderer.py | 60 +-- .../tests/images/test_lifecycle.py | 435 ++++++++++++++++++ 9 files changed, 1026 insertions(+), 52 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/resources/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/resources/enroot.rc create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/batch.py create mode 100644 packages/data-designer-slurm/tests/images/test_lifecycle.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py index 7a6d426bf..db5807e21 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py @@ -28,3 +28,7 @@ class ImageConflictError(SlurmImageError): class ImageVerificationError(SlurmImageError): """Raised when an image file or inspection record fails verification.""" + + +class ImageLifecycleError(SlurmImageError): + """Raised when a CPU Slurm image lifecycle job cannot be prepared or submitted.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py new file mode 100644 index 000000000..efd3df653 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Structured CPU Slurm jobs for OCI import and existing-SQSH inspection.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from dataclasses import dataclass +from importlib import resources +from pathlib import Path + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.config import ImageBuildRequest, SelectedSlurmProfile +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.images.errors import ImageLifecycleError +from data_designer.slurm.images.filesystem import ensure_private_directory +from data_designer.slurm.images.records import ( + ImageLifecycleOperation, + ImageLifecyclePlan, + validate_oci_source_for_lifecycle, +) +from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt + +_IMAGE_JOB_TIME_LIMIT = "03:55:00" +_INSPECTOR_FILENAME = "inspect_image.py" +_ENROOT_RC_FILENAME = "enroot.rc" +_PLAN_FILENAME = "image-lifecycle-plan.json" +_SCRIPT_FILENAME = "image-lifecycle.sbatch" +_RESOURCE_PACKAGE = "data_designer.slurm.images.resources" +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +@dataclass(frozen=True, slots=True) +class PreparedImageLifecycleJob: + """Persisted, checksum-bound files ready for one Slurm submission.""" + + plan: ImageLifecyclePlan + plan_file: ArtifactReference + script_file: ArtifactReference + + +def prepare_image_lifecycle_job( + request: ImageBuildRequest, + selected_profile: SelectedSlurmProfile, + *, + lifecycle_id: Identifier, +) -> PreparedImageLifecycleJob: + """Stage one deterministic image plan, runtime, and batch script beneath the selected workspace.""" + try: + lifecycle_id = _IDENTIFIER_ADAPTER.validate_python(lifecycle_id, strict=True) + except ValidationError as error: + raise ImageLifecycleError("image lifecycle ID is invalid") from error + try: + validate_oci_source_for_lifecycle(request.source) + except ValueError as error: + raise ImageLifecycleError( + "OCI image source must be a credential-free registry reference without a scheme" + ) from error + workspace_root = Path(selected_profile.profile.workspace_root) + if any(delimiter in workspace_root.as_posix() for delimiter in (":", ",")): + raise ImageLifecycleError("selected workspace cannot be represented as an Enroot mount") + image_root = workspace_root / "images" + temporary_root = image_root / ".tmp" + job_root = temporary_root / "jobs" + job_directory = job_root / lifecycle_id + try: + ensure_private_directory(image_root, parents=True) + ensure_private_directory(temporary_root, parents=False) + ensure_private_directory(job_root, parents=False) + job_directory.mkdir(mode=0o700) + inspector_script = _stage_resource(job_directory, _INSPECTOR_FILENAME) + enroot_rc = _stage_resource(job_directory, _ENROOT_RC_FILENAME) + is_existing_sqsh = request.source.endswith(".sqsh") + plan = ImageLifecyclePlan( + schema_version=1, + lifecycle_id=lifecycle_id, + request=request, + selected_profile=selected_profile, + operation=( + ImageLifecycleOperation.INSPECT_SQSH if is_existing_sqsh else ImageLifecycleOperation.IMPORT_OCI + ), + job_directory=job_directory.as_posix(), + sqsh_path=(request.source if is_existing_sqsh else (job_directory / "candidate.sqsh").as_posix()), + inspection_output_path=(job_directory / "inspection.json").as_posix(), + inspector_script=inspector_script, + enroot_rc=enroot_rc, + source_oci_digest=None if is_existing_sqsh else request.source.rpartition("@sha256:")[2], + ) + plan_file = _write_file(job_directory / _PLAN_FILENAME, plan.serialize_json().encode(), mode=0o600) + script_file = _write_file( + job_directory / _SCRIPT_FILENAME, + render_image_lifecycle_script(plan).encode(), + mode=0o500, + ) + except (OSError, ValueError) as error: + raise ImageLifecycleError(f"cannot prepare image lifecycle job {lifecycle_id!r}") from error + return PreparedImageLifecycleJob(plan=plan, plan_file=plan_file, script_file=script_file) + + +def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: + """Render one thin CPU batch entrypoint from structured image lifecycle intent.""" + try: + plan = ImageLifecyclePlan.model_validate(plan.model_dump(mode="python"), strict=True) + except ValueError as error: + raise ImageLifecycleError("cannot render an invalid image lifecycle plan") from error + + profile = plan.selected_profile.profile + directives = render_batch_directives( + ( + ("job-name", f"dd-image-{plan.request.kind}"), + ("account", profile.scheduler.account), + ("partition", profile.image_build.partition), + ("nodes", "1"), + ("ntasks", "1"), + ("cpus-per-task", "1"), + ("time", _IMAGE_JOB_TIME_LIMIT), + ("chdir", plan.job_directory), + ("output", f"{plan.job_directory}/slurm-%j.out"), + ("error", f"{plan.job_directory}/slurm-%j.err"), + ) + ) + source_block = "" + if plan.operation is ImageLifecycleOperation.IMPORT_OCI: + source_block = f"""readonly DD_OCI_SOURCE={quote_shell_value(_format_enroot_oci_uri(plan.request.source))} +if [[ -e "${{DD_IMAGE_SQSH}}" ]]; then + printf '%s\\n' 'candidate SQSH path already exists' >&2 + exit 73 +fi +enroot import -o "${{DD_IMAGE_SQSH}}" "${{DD_OCI_SOURCE}}" +""" + + return f"""#!/usr/bin/env bash +{directives} +set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" + +readonly DD_JOB_DIR={quote_shell_value(plan.job_directory)} +readonly DD_IMAGE_KIND={quote_shell_value(plan.request.kind)} +readonly DD_IMAGE_SQSH={quote_shell_value(plan.sqsh_path)} +readonly DD_INSPECTION_OUTPUT={quote_shell_value(plan.inspection_output_path)} +readonly DD_INSPECTOR={quote_shell_value(plan.inspector_script.path)} +readonly DD_INSPECTOR_SHA256={quote_shell_value(plan.inspector_script.sha256)} +readonly DD_ENROOT_RC={quote_shell_value(plan.enroot_rc.path)} +readonly DD_ENROOT_RC_SHA256={quote_shell_value(plan.enroot_rc.sha256)} + +compute_file_sha256() {{ + local actual_sha256 + actual_sha256="$(sha256sum < "$1")" + printf '%s\\n' "${{actual_sha256%% *}}" +}} + +verify_sha256() {{ + [[ "$(compute_file_sha256 "$2")" == "$1" ]] +}} + +verify_sha256 "${{DD_INSPECTOR_SHA256}}" "${{DD_INSPECTOR}}" +verify_sha256 "${{DD_ENROOT_RC_SHA256}}" "${{DD_ENROOT_RC}}" +install -d -m 0700 "${{DD_JOB_DIR}}/enroot/cache" "${{DD_JOB_DIR}}/enroot/data" "${{DD_JOB_DIR}}/enroot/tmp" +export ENROOT_CACHE_PATH="${{DD_JOB_DIR}}/enroot/cache" +export ENROOT_DATA_PATH="${{DD_JOB_DIR}}/enroot/data" +export ENROOT_TEMP_PATH="${{DD_JOB_DIR}}/enroot/tmp" + +{source_block}if [[ ! -f "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then + printf '%s\\n' 'SQSH path must be a regular non-symlink file' >&2 + exit 66 +fi +readonly DD_SQSH_SHA256="$(compute_file_sha256 "${{DD_IMAGE_SQSH}}")" +if [[ ! ${{DD_SQSH_SHA256}} =~ ^[0-9a-f]{{64}}$ ]]; then + printf '%s\\n' 'SQSH checksum output is invalid' >&2 + exit 65 +fi +if [[ ! ${{SLURM_JOB_ID:-}} =~ ^[1-9][0-9]*$ ]]; then + printf '%s\\n' 'SLURM_JOB_ID must be a positive integer' >&2 + exit 64 +fi +readonly DD_CONTAINER_NAME="dd-image-${{SLURM_JOB_ID}}" + +cleanup() {{ + enroot remove -f "${{DD_CONTAINER_NAME}}" >/dev/null 2>&1 || true +}} +trap cleanup EXIT + +enroot create -f --name "${{DD_CONTAINER_NAME}}" "${{DD_IMAGE_SQSH}}" +ENROOT_LOGIN_SHELL=no ENROOT_MOUNT_HOME=no enroot start --root \ + --rc "${{DD_ENROOT_RC}}" \ + --mount "${{DD_INSPECTOR}}:/opt/data-designer-slurm/inspect_image.py:x-create=file,bind,ro" \ + --mount "${{DD_JOB_DIR}}:/opt/data-designer-slurm/job:x-create=dir,bind" \ + "${{DD_CONTAINER_NAME}}" -- /bin/sh -c ' +python_path="$(command -v python3 || command -v python || true)" +if [ -z "${{python_path}}" ]; then + printf "%s\\n" "target image does not contain Python" >&2 + exit 69 +fi +exec "${{python_path}}" "$@" +' dd-image-inspector \ + /opt/data-designer-slurm/inspect_image.py \ + "${{DD_IMAGE_KIND}}" \ + "${{DD_SQSH_SHA256}}" \ + /opt/data-designer-slurm/job/inspection.json + +[[ -s "${{DD_INSPECTION_OUTPUT}}" ]] +""" + + +def submit_prepared_image_lifecycle( + prepared: PreparedImageLifecycleJob, + client: SlurmCommandClient, +) -> SlurmJobSubmissionReceipt: + """Verify and submit one prepared image lifecycle script.""" + expected_artifacts = ( + ( + "plan", + prepared.plan_file, + Path(prepared.plan.job_directory) / _PLAN_FILENAME, + prepared.plan.serialize_json().encode(), + ), + ( + "script", + prepared.script_file, + Path(prepared.plan.job_directory) / _SCRIPT_FILENAME, + render_image_lifecycle_script(prepared.plan).encode(), + ), + ) + for label, artifact, expected_path, expected_content in expected_artifacts: + expected_sha256 = hashlib.sha256(expected_content).hexdigest() + if ( + artifact.path != expected_path.as_posix() + or artifact.sha256 != expected_sha256 + or _compute_regular_file_sha256(expected_path) != expected_sha256 + ): + raise ImageLifecycleError(f"prepared image lifecycle {label} no longer matches its digest") + return client.submit(prepared.script_file.path) + + +def _stage_resource(job_directory: Path, filename: str) -> ArtifactReference: + content = resources.files(_RESOURCE_PACKAGE).joinpath(filename).read_bytes() + return _write_file(job_directory / filename, content, mode=0o500) + + +def _format_enroot_oci_uri(source: str) -> str: + registry_or_namespace, separator, remainder = source.partition("/") + if not separator: + return f"docker://docker.io#library/{source}" + if "." in registry_or_namespace or ":" in registry_or_namespace or registry_or_namespace == "localhost": + return f"docker://{registry_or_namespace}#{remainder}" + return f"docker://docker.io#{source}" + + +def _write_file(path: Path, content: bytes, *, mode: int) -> ArtifactReference: + descriptor: int | None = None + try: + descriptor = os.open( + path, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + mode, + ) + os.fchmod(descriptor, mode) + with os.fdopen(descriptor, "wb") as output: + descriptor = None + output.write(content) + output.flush() + os.fsync(output.fileno()) + finally: + if descriptor is not None: + os.close(descriptor) + return ArtifactReference(path=path.as_posix(), sha256=hashlib.sha256(content).hexdigest()) + + +def _compute_regular_file_sha256(path: Path) -> str: + try: + path_status = path.lstat() + if not stat.S_ISREG(path_status.st_mode): + raise ImageLifecycleError("prepared image lifecycle file is not regular") + return hashlib.sha256(path.read_bytes()).hexdigest() + except OSError as error: + raise ImageLifecycleError("cannot read prepared image lifecycle file") from error diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/records.py b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py index 4e6589ec1..34304f7e4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py @@ -1,14 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Immutable records persisted by the Slurm image registry.""" +"""Immutable records for Slurm image lifecycle and registry operations.""" from __future__ import annotations +import posixpath +from enum import Enum + from pydantic import field_validator, model_validator -from data_designer.slurm.config import ImageInspectionRecord, ImageKind +from data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, ImageKind, SelectedSlurmProfile from data_designer.slurm.contracts import ( + ArtifactReference, ContractRecord, ContractValue, Identifier, @@ -18,6 +22,89 @@ ) +class ImageLifecycleOperation(str, Enum): + """Compute-node operation required by one image request.""" + + IMPORT_OCI = "import_oci" + INSPECT_SQSH = "inspect_sqsh" + + +class ImageLifecyclePlan(ContractRecord): + """Immutable inputs for one CPU Slurm image lifecycle job.""" + + lifecycle_id: Identifier + request: ImageBuildRequest + selected_profile: SelectedSlurmProfile + operation: ImageLifecycleOperation + job_directory: str + sqsh_path: str + inspection_output_path: str + inspector_script: ArtifactReference + enroot_rc: ArtifactReference + source_oci_digest: Sha256Digest | None = None + + _paths_are_absolute = field_validator("job_directory", "sqsh_path", "inspection_output_path")( + validate_absolute_path + ) + + @model_validator(mode="after") + def validate_lifecycle(self) -> ImageLifecyclePlan: + validate_oci_source_for_lifecycle(self.request.source) + workspace_root = self.selected_profile.profile.workspace_root + expected_job_directory = posixpath.join( + workspace_root, + "images", + ".tmp", + "jobs", + self.lifecycle_id, + ) + if self.job_directory != expected_job_directory: + raise ValueError("image lifecycle job directory must derive from the selected workspace") + if any(delimiter in self.job_directory for delimiter in (":", ",")): + raise ValueError("image lifecycle workspace path cannot be represented as an Enroot mount") + if self.inspection_output_path != posixpath.join(self.job_directory, "inspection.json"): + raise ValueError("image lifecycle inspection output must belong to its job directory") + expected_runtime_artifacts = ( + (self.inspector_script, "inspect_image.py"), + (self.enroot_rc, "enroot.rc"), + ) + for artifact, filename in expected_runtime_artifacts: + if artifact.path != posixpath.join(self.job_directory, filename): + raise ValueError("image lifecycle runtime artifacts must use their package-owned job paths") + + if self.request.source.endswith(".sqsh"): + if self.operation is not ImageLifecycleOperation.INSPECT_SQSH: + raise ValueError("existing SQSH requests require inspection without import") + if self.sqsh_path != self.request.source: + raise ValueError("existing SQSH plans must inspect the authored source path in place") + if self.source_oci_digest is not None: + raise ValueError("existing SQSH plans must not contain an OCI source digest") + else: + if self.operation is not ImageLifecycleOperation.IMPORT_OCI: + raise ValueError("OCI requests require an import operation") + if self.sqsh_path != posixpath.join(self.job_directory, "candidate.sqsh"): + raise ValueError("OCI import output must remain attempt-local until publication") + expected_source_digest = self.request.source.rpartition("@sha256:")[2] + if self.source_oci_digest != expected_source_digest: + raise ValueError("OCI source digest does not match the authored source") + return self + + +def validate_oci_source_for_lifecycle(source: str) -> str: + """Reject OCI source forms that could persist credentials or ambiguous schemes.""" + if source.endswith(".sqsh"): + return source + repository, separator, _digest = source.rpartition("@sha256:") + if ( + separator != "@sha256:" + or "@" in repository + or "://" in repository + or any(delimiter in repository for delimiter in ("?", "#")) + ): + raise ValueError("OCI image source must be a credential-free registry reference without a scheme") + return source + + class RegisteredImage(ContractRecord): """One immutable alias binding for a verified SQSH artifact.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/__init__.py new file mode 100644 index 000000000..d3be375ef --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Standalone resources staged for image inspection jobs.""" + +from __future__ import annotations diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/enroot.rc b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/enroot.rc new file mode 100644 index 000000000..9e31e54c5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/enroot.rc @@ -0,0 +1,6 @@ +#!/bin/sh +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[ "$1" = "--" ] && shift +exec "$@" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py new file mode 100644 index 000000000..b543f3e89 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Emit factual image metadata using only the target image's Python standard library.""" + +from __future__ import annotations + +import importlib.metadata +import json +import os +import platform +import re +import shutil +import sys +from collections.abc import Iterable +from pathlib import Path + +INSPECTOR_VERSION = "inspector-1" +_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") +_REQUIRED_CLIENT_DISTRIBUTIONS = ( + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", +) + + +def inspect_image(kind: str, sqsh_sha256: str) -> dict[str, object]: + """Return one JSON-compatible digest-bound inspection record.""" + if _SHA256_PATTERN.fullmatch(sqsh_sha256) is None: + raise ValueError("SQSH digest must be lowercase SHA-256 text") + if kind == "client": + inspection = _inspect_client() + elif kind == "serving": + inspection = _inspect_serving() + else: + raise ValueError("image kind must be 'client' or 'serving'") + return { + "schema_version": 1, + "inspector_version": INSPECTOR_VERSION, + "sqsh_sha256": sqsh_sha256, + "inspection": inspection, + } + + +def write_inspection(path: Path, record: dict[str, object]) -> None: + """Atomically write one restrictive inspection record.""" + path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + temporary_path = path.with_name(f".{path.name}.{os.getpid()}.tmp") + descriptor = os.open(temporary_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as output: + json.dump(record, output, allow_nan=False, ensure_ascii=False, indent=2, sort_keys=True) + output.write("\n") + output.flush() + os.fsync(output.fileno()) + os.replace(temporary_path, path) + directory_descriptor = os.open(path.parent, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(directory_descriptor) + finally: + os.close(directory_descriptor) + finally: + try: + temporary_path.unlink() + except FileNotFoundError: + pass + + +def main(arguments: list[str] | None = None) -> int: + """Run the standalone inspector entrypoint.""" + values = sys.argv[1:] if arguments is None else arguments + if len(values) != 3: + raise ValueError("expected KIND SQSH_SHA256 OUTPUT_PATH") + kind, sqsh_sha256, output_path = values + output = Path(output_path) + if not output.is_absolute(): + raise ValueError("inspection output path must be absolute") + write_inspection(output, inspect_image(kind, sqsh_sha256)) + return 0 + + +def _inspect_client() -> dict[str, object]: + versions = _list_distribution_versions(importlib.metadata.distributions()) + missing = tuple(name for name in _REQUIRED_CLIENT_DISTRIBUTIONS if name not in versions) + if missing: + raise RuntimeError(f"required client distributions are not installed: {', '.join(missing)}") + installer_path = shutil.which("pip") + if installer_path is None or not Path(installer_path).is_absolute(): + raise RuntimeError("required executable 'pip' is not installed at an absolute path") + cache_tag = sys.implementation.cache_tag + if cache_tag is None: + raise RuntimeError("active Python does not expose an ABI cache tag") + python_abi = f"cp{cache_tag.removeprefix('cpython-')}" if cache_tag.startswith("cpython-") else cache_tag + return { + "kind": "client", + "python_implementation": platform.python_implementation().casefold(), + "python_version": platform.python_version(), + "python_abi": python_abi, + "distributions": [{"name": name, "version": version} for name, version in sorted(versions.items())], + "installer_path": Path(installer_path).as_posix(), + "installer_version": versions["pip"], + } + + +def _inspect_serving() -> dict[str, object]: + executable_path = shutil.which("vllm") + if executable_path is None or not Path(executable_path).is_absolute(): + raise RuntimeError("required executable 'vllm' is not installed at an absolute path") + try: + runtime_version = importlib.metadata.version("vllm") + except importlib.metadata.PackageNotFoundError as error: + raise RuntimeError("required distribution 'vllm' is not installed") from error + return { + "kind": "serving", + "server_type": "vllm", + "runtime_version": runtime_version, + "executable_path": Path(executable_path).as_posix(), + } + + +def _list_distribution_versions(distributions: Iterable[importlib.metadata.Distribution]) -> dict[str, str]: + versions: dict[str, str] = {} + for distribution in distributions: + raw_name = distribution.metadata.get("Name") + if not raw_name: + raise RuntimeError("installed distribution is missing its canonical name") + name = re.sub(r"[-_.]+", "-", raw_name).casefold() + existing = versions.setdefault(name, distribution.version) + if existing != distribution.version: + raise RuntimeError(f"installed distribution {name!r} has conflicting versions") + return versions + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/batch.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/batch.py new file mode 100644 index 000000000..7342d3922 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/batch.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared safe rendering primitives for package-owned Slurm batch scripts.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +from data_designer.slurm.launcher.errors import SlurmBatchRenderError + +_DIRECTIVE_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9-]*$") +_DIRECTIVE_TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/,%+-]*$") + + +@dataclass(frozen=True) +class BatchDirective: + """One validated ``#SBATCH`` option.""" + + name: str + value: str + + def render(self) -> str: + """Render the directive as one non-executable scheduler line.""" + if type(self.name) is not str or _DIRECTIVE_NAME_PATTERN.fullmatch(self.name) is None: + raise SlurmBatchRenderError("batch directive name is invalid") + if type(self.value) is not str: + raise SlurmBatchRenderError("batch directive value must be text") + reject_control_characters(self.value, field_name=f"--{self.name} value") + value = self.value if _DIRECTIVE_TOKEN_PATTERN.fullmatch(self.value) else _quote_sbatch_option_value(self.value) + return f"#SBATCH --{self.name}={value}" + + +def render_batch_directives(values: tuple[tuple[str, str | None], ...]) -> str: + """Render ordered optional directive values as safe scheduler lines.""" + return "\n".join(BatchDirective(name=name, value=value).render() for name, value in values if value is not None) + + +def quote_shell_value(value: str) -> str: + """Quote one validated value for literal use in package-owned Bash.""" + if type(value) is not str: + raise SlurmBatchRenderError("shell value must be text") + reject_control_characters(value, field_name="shell value") + escaped = value.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`") + return f'"{escaped}"' + + +def reject_control_characters(value: str, *, field_name: str) -> None: + """Reject characters that can split a directive or shell assignment.""" + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise SlurmBatchRenderError(f"{field_name} must not contain control characters") + + +def _quote_sbatch_option_value(value: str) -> str: + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped}"' 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..0afb1f424 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 @@ -6,33 +6,11 @@ from __future__ import annotations import posixpath -import re -from dataclasses import dataclass +from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives from data_designer.slurm.launcher.errors import SlurmBatchRenderError from data_designer.slurm.planning import ResolvedSlurmRunPlan -_DIRECTIVE_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9-]*$") -_DIRECTIVE_TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/,%+-]*$") - - -@dataclass(frozen=True) -class _BatchDirective: - """One validated ``#SBATCH`` option.""" - - name: str - value: str - - def render(self) -> str: - """Render the directive as one non-executable scheduler line.""" - if type(self.name) is not str or _DIRECTIVE_NAME_PATTERN.fullmatch(self.name) is None: - raise SlurmBatchRenderError("batch directive name is invalid") - if type(self.value) is not str: - raise SlurmBatchRenderError("batch directive value must be text") - _reject_control_characters(self.value, field_name=f"--{self.name} value") - value = self.value if _DIRECTIVE_TOKEN_PATTERN.fullmatch(self.value) else _quote_sbatch_option_value(self.value) - return f"#SBATCH --{self.name}={value}" - def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: """Render a resolved generation plan as one thin deterministic entrypoint.""" @@ -41,8 +19,7 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi run_root = posixpath.dirname(plan.authored_config.path) plan_path = posixpath.join(run_root, "resolved-plan.json") - directives = _build_generation_directives(plan) - directive_text = "\n".join(directive.render() for directive in directives) + directive_text = render_batch_directives(_build_generation_directives(plan)) attempt = f"{attempt_ordinal:04d}" return f"""#!/usr/bin/env bash @@ -50,12 +27,12 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" -readonly DD_RUNTIME_ARCHIVE={_quote_shell_value(plan.runtime_bundle.path)} -readonly DD_RUNTIME_SHA256={_quote_shell_value(plan.runtime_bundle.sha256)} -readonly DD_PLAN={_quote_shell_value(plan_path)} -readonly DD_PLAN_SHA256={_quote_shell_value(plan.compute_sha256())} -readonly DD_RUN_ROOT={_quote_shell_value(run_root)} -readonly DD_ATTEMPT_ORDINAL={_quote_shell_value(attempt)} +readonly DD_RUNTIME_ARCHIVE={quote_shell_value(plan.runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={quote_shell_value(plan.runtime_bundle.sha256)} +readonly DD_PLAN={quote_shell_value(plan_path)} +readonly DD_PLAN_SHA256={quote_shell_value(plan.compute_sha256())} +readonly DD_RUN_ROOT={quote_shell_value(run_root)} +readonly DD_ATTEMPT_ORDINAL={quote_shell_value(attempt)} verify_sha256() {{ local actual_sha256 @@ -83,7 +60,7 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi """ -def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDirective, ...]: +def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[tuple[str, str | None], ...]: node_indices = ( plan.client.host_node_index, *(index for deployment in plan.deployments for index in deployment.node_indices), @@ -112,21 +89,4 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) if plan.submission.comment is not None: values.append(("comment", plan.submission.comment)) - return tuple(_BatchDirective(name=name, value=value) for name, value in values if value is not None) - - -def _quote_sbatch_option_value(value: str) -> str: - _reject_control_characters(value, field_name="batch directive value") - escaped = value.replace("\\", "\\\\").replace('"', '\\"') - return f'"{escaped}"' - - -def _quote_shell_value(value: str) -> str: - _reject_control_characters(value, field_name="shell value") - escaped = value.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`") - return f'"{escaped}"' - - -def _reject_control_characters(value: str, *, field_name: str) -> None: - if any(ord(character) < 32 or ord(character) == 127 for character in value): - raise SlurmBatchRenderError(f"{field_name} must not contain control characters") + return tuple(values) diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py new file mode 100644 index 000000000..19ce0af64 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -0,0 +1,435 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import stat +import subprocess +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from pydantic import ValidationError +from slurm_test_fakes import FakeSlurmJob, FakeSlurmRunner + +from data_designer.slurm.config import ( + ImageBuildProfile, + ImageBuildRequest, + ImageInspectionRecord, + SchedulerProfile, + SelectedSlurmProfile, + SlurmProfile, + injected_profile, +) +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.images.errors import ImageLifecycleError +from data_designer.slurm.images.lifecycle import ( + prepare_image_lifecycle_job, + render_image_lifecycle_script, + submit_prepared_image_lifecycle, +) +from data_designer.slurm.images.records import ImageLifecycleOperation, ImageLifecyclePlan +from data_designer.slurm.images.resources import inspect_image as resource_inspector +from data_designer.slurm.launcher.client import SlurmCommandClient + + +def test_prepare_oci_import_stages_checksum_bound_job_beneath_selected_workspace(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source_digest = "a" * 64 + request = ImageBuildRequest( + name="serving", + kind="serving", + source=f"registry.example.test/team/vllm@sha256:{source_digest}", + ) + + prepared = prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-0001") + + job_directory = workspace / "images" / ".tmp" / "jobs" / "image-job-0001" + assert prepared.plan.operation is ImageLifecycleOperation.IMPORT_OCI + assert prepared.plan.source_oci_digest == source_digest + assert prepared.plan.sqsh_path == (job_directory / "candidate.sqsh").as_posix() + assert prepared.plan.inspection_output_path == (job_directory / "inspection.json").as_posix() + assert prepared.plan_file.path == (job_directory / "image-lifecycle-plan.json").as_posix() + assert prepared.script_file.path == (job_directory / "image-lifecycle.sbatch").as_posix() + assert ImageLifecyclePlan.model_validate_json(Path(prepared.plan_file.path).read_text()) == prepared.plan + for artifact in ( + prepared.plan.inspector_script, + prepared.plan.enroot_rc, + prepared.plan_file, + prepared.script_file, + ): + content = Path(artifact.path).read_bytes() + assert hashlib.sha256(content).hexdigest() == artifact.sha256 + assert stat.S_IMODE(Path(prepared.plan.inspector_script.path).stat().st_mode) == 0o500 + assert stat.S_IMODE(Path(prepared.plan.enroot_rc.path).stat().st_mode) == 0o500 + assert stat.S_IMODE(job_directory.stat().st_mode) == 0o700 + assert stat.S_IMODE(Path(prepared.plan_file.path).stat().st_mode) == 0o600 + assert stat.S_IMODE(Path(prepared.script_file.path).stat().st_mode) == 0o500 + + +def test_prepare_existing_sqsh_inspects_in_place_without_oci_identity(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_path = tmp_path / "external" / "client.sqsh" + request = ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()) + + prepared = prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-0002") + + assert prepared.plan.operation is ImageLifecycleOperation.INSPECT_SQSH + assert prepared.plan.sqsh_path == image_path.as_posix() + assert prepared.plan.source_oci_digest is None + script = Path(prepared.script_file.path).read_text() + assert "enroot import" not in script + assert f'readonly DD_IMAGE_SQSH="{image_path.as_posix()}"' in script + + +def test_image_lifecycle_renderer_uses_cpu_profile_and_quotes_workspace_paths(tmp_path: Path) -> None: + workspace = tmp_path / "workspace with spaces" + source = f"registry.example.test/team/image:release@sha256:{'b' * 64}" + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="serving", kind="serving", source=source), + _get_selected_profile(workspace), + lifecycle_id="image-job-0003", + ) + + script = render_image_lifecycle_script(prepared.plan) + + assert "#SBATCH --account=research" in script + assert "#SBATCH --partition=image-build" in script + assert "#SBATCH --nodes=1" in script + assert "#SBATCH --gres=" not in script + expected_uri = f"docker://registry.example.test#team/image:release@sha256:{'b' * 64}" + assert f'readonly DD_OCI_SOURCE="{expected_uri}"' in script + assert 'enroot import -o "${DD_IMAGE_SQSH}" "${DD_OCI_SOURCE}"' in script + assert "inspect_image.py:x-create=file,bind,ro" in script + assert "/opt/data-designer-slurm/job:x-create=dir,bind" in script + completed = subprocess.run(("bash", "-n"), input=script, capture_output=True, text=True, check=False) + assert completed.returncode == 0, completed.stderr + + +@pytest.mark.parametrize("source_kind", ("oci", "existing"), ids=("oci-import", "existing-sqsh")) +def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( + tmp_path: Path, + source_kind: str, +) -> None: + if source_kind == "oci": + source = f"registry.example.test/client@sha256:{'a' * 64}" + else: + image_path = tmp_path / "client.sqsh" + image_path.write_bytes(b"client image") + source = image_path.as_posix() + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=source), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id=f"image-job-{source_kind}-execution", + ) + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_enroot = fake_bin / "enroot" + fake_enroot.write_text( + "#!/usr/bin/env bash\n" + "set -Eeuo pipefail\n" + 'if [[ "$1" == "import" ]]; then\n' + " printf '%s\\n' 'imported image' > \"$3\"\n" + 'elif [[ "$1" == "start" ]]; then\n' + " printf '%s\\n' '{}' > \"${DD_TEST_INSPECTION_OUTPUT}\"\n" + "fi\n" + ) + fake_enroot.chmod(0o700) + script = render_image_lifecycle_script(prepared.plan).replace( + 'export PATH="', + f'export PATH="{fake_bin.as_posix()}:', + 1, + ) + + completed = subprocess.run( + ("bash",), + input=script, + capture_output=True, + text=True, + check=False, + env={ + "DD_TEST_INSPECTION_OUTPUT": prepared.plan.inspection_output_path, + "SLURM_JOB_ID": "5101", + }, + ) + + assert completed.returncode == 0, completed.stderr + assert Path(prepared.plan.inspection_output_path).read_text() == "{}\n" + if source_kind == "oci": + assert Path(prepared.plan.sqsh_path).read_text() == "imported image\n" + + +@pytest.mark.parametrize( + "source", + ( + f"user:token@registry.example.test/image@sha256:{'b' * 64}", + f"docker://registry.example.test/image@sha256:{'b' * 64}", + f"registry.example.test/image?token=value@sha256:{'b' * 64}", + ), + ids=("credentials", "scheme", "query"), +) +def test_prepare_rejects_oci_sources_that_could_persist_credentials(tmp_path: Path, source: str) -> None: + request = ImageBuildRequest(name="serving", kind="serving", source=source) + + with pytest.raises(ImageLifecycleError, match="credential-free"): + prepare_image_lifecycle_job( + request, + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-secret", + ) + + assert not (tmp_path / "workspace" / "images").exists() + + +@pytest.mark.parametrize( + ("source", "expected_uri"), + ( + ( + f"vllm/vllm-openai@sha256:{'a' * 64}", + f"docker://docker.io#vllm/vllm-openai@sha256:{'a' * 64}", + ), + ( + f"example@sha256:{'a' * 64}", + f"docker://docker.io#library/example@sha256:{'a' * 64}", + ), + ( + f"registry.example.test:5000/team/image@sha256:{'a' * 64}", + f"docker://registry.example.test:5000#team/image@sha256:{'a' * 64}", + ), + ), + ids=("docker-hub-namespace", "docker-hub-library", "explicit-registry"), +) +def test_renderer_normalizes_digest_qualified_sources_for_enroot( + tmp_path: Path, + source: str, + expected_uri: str, +) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="serving", kind="serving", source=source), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-uri", + ) + + assert f'readonly DD_OCI_SOURCE="{expected_uri}"' in render_image_lifecycle_script(prepared.plan) + + +def test_prepare_refuses_duplicate_and_invalid_lifecycle_ids(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + profile = _get_selected_profile(workspace) + request = ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()) + prepare_image_lifecycle_job(request, profile, lifecycle_id="image-job-0004") + + with pytest.raises(ImageLifecycleError, match="cannot prepare"): + prepare_image_lifecycle_job(request, profile, lifecycle_id="image-job-0004") + with pytest.raises(ImageLifecycleError, match="ID is invalid"): + prepare_image_lifecycle_job(request, profile, lifecycle_id="../escape") # type: ignore[arg-type] + assert not (workspace / "images" / ".tmp" / "escape").exists() + + +def test_prepare_rejects_workspace_paths_that_enroot_cannot_mount(tmp_path: Path) -> None: + workspace = tmp_path / "workspace:unsafe" + request = ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()) + + with pytest.raises(ImageLifecycleError, match="Enroot mount"): + prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-unsafe") + + assert not (workspace / "images").exists() + + +def test_prepare_rejects_symlinked_package_owned_image_root(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (workspace / "images").symlink_to(outside, target_is_directory=True) + request = ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()) + + with pytest.raises(ImageLifecycleError, match="cannot prepare"): + prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-symlink") + + assert tuple(outside.iterdir()) == () + + +@pytest.mark.parametrize( + ("field", "value", "match"), + ( + ("operation", "inspect_sqsh", "require an import"), + ("sqsh_path", "/workspace/other.sqsh", "attempt-local"), + ("source_oci_digest", "f" * 64, "does not match"), + ("inspection_output_path", "/workspace/inspection.json", "job directory"), + ), +) +def test_image_lifecycle_plan_rejects_mismatched_oci_facts( + tmp_path: Path, + field: str, + value: str, + match: str, +) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name="serving", + kind="serving", + source=f"registry.example.test/vllm@sha256:{'e' * 64}", + ), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-contract", + ) + payload = prepared.plan.model_dump(mode="json") + payload[field] = value + + with pytest.raises(ValidationError, match=match): + ImageLifecyclePlan.model_validate_json(json.dumps(payload)) + + +def test_image_lifecycle_plan_rejects_non_package_inspector_path(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-runtime-path", + ) + payload = prepared.plan.model_dump(mode="json") + payload["inspector_script"]["path"] = f"{prepared.plan.job_directory}/other.py" + + with pytest.raises(ValidationError, match="package-owned job paths"): + ImageLifecyclePlan.model_validate_json(json.dumps(payload)) + + +def test_image_lifecycle_plan_rejects_credential_bearing_oci_source(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name="serving", + kind="serving", + source=f"registry.example.test/vllm@sha256:{'e' * 64}", + ), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-source-contract", + ) + payload = prepared.plan.model_dump(mode="json") + payload["request"]["source"] = f"user:token@registry.example.test/vllm@sha256:{'e' * 64}" + + with pytest.raises(ValidationError, match="credential-free"): + ImageLifecyclePlan.model_validate_json(json.dumps(payload)) + + +def test_submit_prepared_image_lifecycle_uses_isolated_sbatch_client(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-0005", + ) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + + receipt = submit_prepared_image_lifecycle(prepared, SlurmCommandClient(runner)) + + assert receipt.job_id == 5101 + assert runner.calls == [ + ( + "sbatch", + "--parsable", + "--export=NIL", + prepared.script_file.path, + ) + ] + + +@pytest.mark.parametrize("artifact_name", ("plan_file", "script_file"), ids=("plan", "script")) +def test_submit_rejects_modified_prepared_files(tmp_path: Path, artifact_name: str) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-0006", + ) + artifact = getattr(prepared, artifact_name) + artifact_path = Path(artifact.path) + artifact_path.chmod(0o700) + artifact_path.write_text("modified\n") + + with pytest.raises(ImageLifecycleError, match=f"{artifact_name.removesuffix('_file')} no longer matches"): + submit_prepared_image_lifecycle(prepared, SlurmCommandClient(FakeSlurmRunner())) + + +def test_submit_rejects_modified_script_even_when_artifact_digest_is_rebound(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-0007", + ) + script_path = Path(prepared.script_file.path) + script_path.chmod(0o700) + modified_content = b"#!/usr/bin/env bash\nexit 0\n" + script_path.write_bytes(modified_content) + rebound = replace( + prepared, + script_file=ArtifactReference( + path=script_path.as_posix(), + sha256=hashlib.sha256(modified_content).hexdigest(), + ), + ) + + with pytest.raises(ImageLifecycleError, match="script no longer matches"): + submit_prepared_image_lifecycle(rebound, SlurmCommandClient(FakeSlurmRunner())) + + +def test_standalone_client_inspector_emits_contract_compatible_json(tmp_path: Path) -> None: + distributions = tuple( + SimpleNamespace(metadata={"Name": name}, version="1.0.0") + for name in ( + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", + ) + ) + with ( + patch.object(resource_inspector.importlib.metadata, "distributions", return_value=distributions), + patch.object(resource_inspector.shutil, "which", return_value="/usr/bin/pip"), + ): + record = resource_inspector.inspect_image("client", "c" * 64) + output_path = tmp_path / "inspection.json" + + resource_inspector.write_inspection(output_path, record) + + inspection = ImageInspectionRecord.model_validate_json(output_path.read_text()) + assert inspection.inspection.kind == "client" + assert inspection.sqsh_sha256 == "c" * 64 + assert stat.S_IMODE(output_path.stat().st_mode) == 0o600 + + +def test_standalone_serving_inspector_emits_contract_compatible_json() -> None: + def find_executable(name: str) -> str | None: + return "/usr/local/bin/vllm" if name == "vllm" else None + + with ( + patch.object(resource_inspector.importlib.metadata, "version", return_value="0.21.0"), + patch.object(resource_inspector.shutil, "which", side_effect=find_executable), + ): + payload = resource_inspector.inspect_image("serving", "d" * 64) + + inspection = ImageInspectionRecord.model_validate_json(json.dumps(payload)) + assert inspection.inspection.kind == "serving" + assert inspection.inspection.runtime_version == "0.21.0" # type: ignore[union-attr] + + +@pytest.mark.parametrize( + ("kind", "sqsh_sha256", "match"), + (("unknown", "f" * 64, "image kind"), ("client", "not-a-digest", "SHA-256")), +) +def test_standalone_inspector_rejects_invalid_invocation(kind: str, sqsh_sha256: str, match: str) -> None: + with pytest.raises(ValueError, match=match): + resource_inspector.inspect_image(kind, sqsh_sha256) + + +def _get_selected_profile(workspace: Path) -> SelectedSlurmProfile: + profile = SlurmProfile( + schema_version=1, + scheduler=SchedulerProfile(account="research", partition="gpu"), + gpus_per_node=8, + workspace_root=workspace.as_posix(), + image_build=ImageBuildProfile(partition="image-build"), + ) + return injected_profile(profile) From 0c8a6cb95bca98086867627b623b11adf7d9750b Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Fri, 28 Aug 2026 13:22:36 -0600 Subject: [PATCH 2/5] fix: harden Slurm image inspection Restrict the container to a dedicated output mount, isolate Enroot user configuration, and bind vLLM metadata to its owning distribution.\n\nPart of #867 Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/images/lifecycle.py | 17 ++++-- .../src/data_designer/slurm/images/records.py | 4 +- .../slurm/images/resources/inspect_image.py | 39 +++++++++++-- .../tests/images/test_lifecycle.py | 57 ++++++++++++++++--- 4 files changed, 97 insertions(+), 20 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py index efd3df653..d707c16f0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -74,6 +74,8 @@ def prepare_image_lifecycle_job( ensure_private_directory(temporary_root, parents=False) ensure_private_directory(job_root, parents=False) job_directory.mkdir(mode=0o700) + inspection_directory = job_directory / "output" + ensure_private_directory(inspection_directory, parents=False) inspector_script = _stage_resource(job_directory, _INSPECTOR_FILENAME) enroot_rc = _stage_resource(job_directory, _ENROOT_RC_FILENAME) is_existing_sqsh = request.source.endswith(".sqsh") @@ -87,7 +89,7 @@ def prepare_image_lifecycle_job( ), job_directory=job_directory.as_posix(), sqsh_path=(request.source if is_existing_sqsh else (job_directory / "candidate.sqsh").as_posix()), - inspection_output_path=(job_directory / "inspection.json").as_posix(), + inspection_output_path=(inspection_directory / "inspection.json").as_posix(), inspector_script=inspector_script, enroot_rc=enroot_rc, source_oci_digest=None if is_existing_sqsh else request.source.rpartition("@sha256:")[2], @@ -143,6 +145,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: readonly DD_JOB_DIR={quote_shell_value(plan.job_directory)} readonly DD_IMAGE_KIND={quote_shell_value(plan.request.kind)} readonly DD_IMAGE_SQSH={quote_shell_value(plan.sqsh_path)} +readonly DD_INSPECTION_DIRECTORY={quote_shell_value(Path(plan.inspection_output_path).parent.as_posix())} readonly DD_INSPECTION_OUTPUT={quote_shell_value(plan.inspection_output_path)} readonly DD_INSPECTOR={quote_shell_value(plan.inspector_script.path)} readonly DD_INSPECTOR_SHA256={quote_shell_value(plan.inspector_script.sha256)} @@ -161,8 +164,14 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: verify_sha256 "${{DD_INSPECTOR_SHA256}}" "${{DD_INSPECTOR}}" verify_sha256 "${{DD_ENROOT_RC_SHA256}}" "${{DD_ENROOT_RC}}" -install -d -m 0700 "${{DD_JOB_DIR}}/enroot/cache" "${{DD_JOB_DIR}}/enroot/data" "${{DD_JOB_DIR}}/enroot/tmp" +install -d -m 0700 \ + "${{DD_JOB_DIR}}/enroot/cache" \ + "${{DD_JOB_DIR}}/enroot/config" \ + "${{DD_JOB_DIR}}/enroot/data" \ + "${{DD_JOB_DIR}}/enroot/tmp" \ + "${{DD_INSPECTION_DIRECTORY}}" export ENROOT_CACHE_PATH="${{DD_JOB_DIR}}/enroot/cache" +export ENROOT_CONFIG_PATH="${{DD_JOB_DIR}}/enroot/config" export ENROOT_DATA_PATH="${{DD_JOB_DIR}}/enroot/data" export ENROOT_TEMP_PATH="${{DD_JOB_DIR}}/enroot/tmp" @@ -190,7 +199,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: ENROOT_LOGIN_SHELL=no ENROOT_MOUNT_HOME=no enroot start --root \ --rc "${{DD_ENROOT_RC}}" \ --mount "${{DD_INSPECTOR}}:/opt/data-designer-slurm/inspect_image.py:x-create=file,bind,ro" \ - --mount "${{DD_JOB_DIR}}:/opt/data-designer-slurm/job:x-create=dir,bind" \ + --mount "${{DD_INSPECTION_DIRECTORY}}:/opt/data-designer-slurm/output:x-create=dir,bind" \ "${{DD_CONTAINER_NAME}}" -- /bin/sh -c ' python_path="$(command -v python3 || command -v python || true)" if [ -z "${{python_path}}" ]; then @@ -202,7 +211,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: /opt/data-designer-slurm/inspect_image.py \ "${{DD_IMAGE_KIND}}" \ "${{DD_SQSH_SHA256}}" \ - /opt/data-designer-slurm/job/inspection.json + /opt/data-designer-slurm/output/inspection.json [[ -s "${{DD_INSPECTION_OUTPUT}}" ]] """ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/records.py b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py index 34304f7e4..252223920 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py @@ -62,8 +62,8 @@ def validate_lifecycle(self) -> ImageLifecyclePlan: raise ValueError("image lifecycle job directory must derive from the selected workspace") if any(delimiter in self.job_directory for delimiter in (":", ",")): raise ValueError("image lifecycle workspace path cannot be represented as an Enroot mount") - if self.inspection_output_path != posixpath.join(self.job_directory, "inspection.json"): - raise ValueError("image lifecycle inspection output must belong to its job directory") + if self.inspection_output_path != posixpath.join(self.job_directory, "output", "inspection.json"): + raise ValueError("image lifecycle inspection output must belong to its dedicated output directory") expected_runtime_artifacts = ( (self.inspector_script, "inspect_image.py"), (self.enroot_rc, "enroot.rc"), diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py index b543f3e89..0665138eb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py @@ -105,18 +105,16 @@ def _inspect_client() -> dict[str, object]: def _inspect_serving() -> dict[str, object]: - executable_path = shutil.which("vllm") - if executable_path is None or not Path(executable_path).is_absolute(): - raise RuntimeError("required executable 'vllm' is not installed at an absolute path") try: - runtime_version = importlib.metadata.version("vllm") + distribution = importlib.metadata.distribution("vllm") except importlib.metadata.PackageNotFoundError as error: raise RuntimeError("required distribution 'vllm' is not installed") from error + executable_path = _find_distribution_console_script(distribution, "vllm") return { "kind": "serving", "server_type": "vllm", - "runtime_version": runtime_version, - "executable_path": Path(executable_path).as_posix(), + "runtime_version": distribution.version, + "executable_path": executable_path, } @@ -133,5 +131,34 @@ def _list_distribution_versions(distributions: Iterable[importlib.metadata.Distr return versions +def _find_distribution_console_script(distribution: importlib.metadata.Distribution, name: str) -> str: + entry_points = tuple( + entry_point + for entry_point in distribution.entry_points + if entry_point.group == "console_scripts" and entry_point.name == name + ) + if len(entry_points) != 1: + raise RuntimeError(f"required distribution {name!r} does not expose one console script") + files = distribution.files + if files is None: + raise RuntimeError(f"required distribution {name!r} does not expose an installed-file inventory") + executable_paths: set[Path] = set() + for installed_file in files: + if Path(str(installed_file)).name != name: + continue + candidate = Path(distribution.locate_file(installed_file)) + if not candidate.is_absolute(): + continue + try: + resolved = candidate.resolve(strict=True) + except OSError: + continue + if resolved.is_file() and os.access(resolved, os.X_OK): + executable_paths.add(resolved) + if len(executable_paths) != 1: + raise RuntimeError(f"required distribution {name!r} does not own one executable console script") + return next(iter(executable_paths)).as_posix() + + if __name__ == "__main__": raise SystemExit(main()) diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index 19ce0af64..1a1f29c9e 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -52,7 +52,7 @@ def test_prepare_oci_import_stages_checksum_bound_job_beneath_selected_workspace assert prepared.plan.operation is ImageLifecycleOperation.IMPORT_OCI assert prepared.plan.source_oci_digest == source_digest assert prepared.plan.sqsh_path == (job_directory / "candidate.sqsh").as_posix() - assert prepared.plan.inspection_output_path == (job_directory / "inspection.json").as_posix() + assert prepared.plan.inspection_output_path == (job_directory / "output" / "inspection.json").as_posix() assert prepared.plan_file.path == (job_directory / "image-lifecycle-plan.json").as_posix() assert prepared.script_file.path == (job_directory / "image-lifecycle.sbatch").as_posix() assert ImageLifecyclePlan.model_validate_json(Path(prepared.plan_file.path).read_text()) == prepared.plan @@ -67,6 +67,7 @@ def test_prepare_oci_import_stages_checksum_bound_job_beneath_selected_workspace assert stat.S_IMODE(Path(prepared.plan.inspector_script.path).stat().st_mode) == 0o500 assert stat.S_IMODE(Path(prepared.plan.enroot_rc.path).stat().st_mode) == 0o500 assert stat.S_IMODE(job_directory.stat().st_mode) == 0o700 + assert stat.S_IMODE((job_directory / "output").stat().st_mode) == 0o700 assert stat.S_IMODE(Path(prepared.plan_file.path).stat().st_mode) == 0o600 assert stat.S_IMODE(Path(prepared.script_file.path).stat().st_mode) == 0o500 @@ -105,7 +106,9 @@ def test_image_lifecycle_renderer_uses_cpu_profile_and_quotes_workspace_paths(tm assert f'readonly DD_OCI_SOURCE="{expected_uri}"' in script assert 'enroot import -o "${DD_IMAGE_SQSH}" "${DD_OCI_SOURCE}"' in script assert "inspect_image.py:x-create=file,bind,ro" in script - assert "/opt/data-designer-slurm/job:x-create=dir,bind" in script + assert "/opt/data-designer-slurm/output:x-create=dir,bind" in script + assert 'export ENROOT_CONFIG_PATH="${DD_JOB_DIR}/enroot/config"' in script + assert f'--mount "{prepared.plan.job_directory}:' not in script completed = subprocess.run(("bash", "-n"), input=script, capture_output=True, text=True, check=False) assert completed.returncode == 0, completed.stderr @@ -260,7 +263,7 @@ def test_prepare_rejects_symlinked_package_owned_image_root(tmp_path: Path) -> N ("operation", "inspect_sqsh", "require an import"), ("sqsh_path", "/workspace/other.sqsh", "attempt-local"), ("source_oci_digest", "f" * 64, "does not match"), - ("inspection_output_path", "/workspace/inspection.json", "job directory"), + ("inspection_output_path", "/workspace/inspection.json", "output directory"), ), ) def test_image_lifecycle_plan_rejects_mismatched_oci_facts( @@ -400,19 +403,57 @@ def test_standalone_client_inspector_emits_contract_compatible_json(tmp_path: Pa assert stat.S_IMODE(output_path.stat().st_mode) == 0o600 -def test_standalone_serving_inspector_emits_contract_compatible_json() -> None: - def find_executable(name: str) -> str | None: - return "/usr/local/bin/vllm" if name == "vllm" else None +def test_standalone_serving_inspector_binds_version_and_executable_to_one_distribution(tmp_path: Path) -> None: + executable_path = tmp_path / "active-environment" / "bin" / "vllm" + executable_path.parent.mkdir(parents=True) + executable_path.write_text("#!/bin/sh\n") + executable_path.chmod(0o700) + installed_file = Path("../../../bin/vllm") + entry_point = SimpleNamespace(group="console_scripts", name="vllm") + distribution = SimpleNamespace( + version="0.21.0", + entry_points=(entry_point,), + files=(installed_file,), + locate_file=lambda _installed_file: executable_path, + ) with ( - patch.object(resource_inspector.importlib.metadata, "version", return_value="0.21.0"), - patch.object(resource_inspector.shutil, "which", side_effect=find_executable), + patch.object(resource_inspector.importlib.metadata, "distribution", return_value=distribution), + patch.object(resource_inspector.shutil, "which", return_value="/shadow/bin/vllm"), ): payload = resource_inspector.inspect_image("serving", "d" * 64) inspection = ImageInspectionRecord.model_validate_json(json.dumps(payload)) assert inspection.inspection.kind == "serving" assert inspection.inspection.runtime_version == "0.21.0" # type: ignore[union-attr] + assert inspection.inspection.executable_path == executable_path.as_posix() # type: ignore[union-attr] + + +@pytest.mark.parametrize( + ("entry_points", "files", "match"), + ( + ((), (), "does not expose one console script"), + ((SimpleNamespace(group="console_scripts", name="vllm"),), None, "installed-file inventory"), + ((SimpleNamespace(group="console_scripts", name="vllm"),), (Path("vllm"),), "does not own one executable"), + ), +) +def test_standalone_serving_inspector_rejects_unverifiable_distribution_console_script( + entry_points: tuple[SimpleNamespace, ...], + files: tuple[Path, ...] | None, + match: str, +) -> None: + distribution = SimpleNamespace( + version="0.21.0", + entry_points=entry_points, + files=files, + locate_file=lambda installed_file: installed_file, + ) + + with ( + patch.object(resource_inspector.importlib.metadata, "distribution", return_value=distribution), + pytest.raises(RuntimeError, match=match), + ): + resource_inspector.inspect_image("serving", "d" * 64) @pytest.mark.parametrize( From 2b2f9ad457fa2a8947925a39de5572da83b74641 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 31 Aug 2026 10:36:24 -0600 Subject: [PATCH 3/5] fix: secure Slurm lifecycle boundaries Reject Enroot-unsafe mount paths and dangling import targets. Submit descriptor-verified scripts through standard input and bind lifecycle resource limits to the selected profile. Resolve pip and vLLM console scripts from their owning distributions with shared inspection semantics. Part of #867 Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/config/profiles.py | 11 ++ .../data_designer/slurm/images/inspection.py | 43 +++--- .../data_designer/slurm/images/lifecycle.py | 48 ++++-- .../src/data_designer/slurm/images/records.py | 11 +- .../slurm/images/resources/inspect_image.py | 35 +++-- .../data_designer/slurm/launcher/client.py | 16 +- .../data_designer/slurm/launcher/runner.py | 16 +- .../contracts/golden/multi_node_plan.json | 9 +- .../contracts/golden/profile_catalog.json | 10 +- .../contracts/golden/single_node_plan.json | 9 +- .../tests/contracts/test_profiles.py | 3 + .../tests/images/test_inspection.py | 59 ++++++-- .../tests/images/test_lifecycle.py | 137 +++++++++++++++--- .../golden/finalization_chain.json | 4 +- .../tests/launcher/test_client.py | 22 +++ .../tests/launcher/test_runner.py | 38 +++++ .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../tests/slurm_test_fakes/images.py | 15 +- .../tests/slurm_test_fakes/slurm.py | 10 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- 21 files changed, 390 insertions(+), 114 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py index 87412d84d..79e8f39b5 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 @@ -34,6 +34,17 @@ class SchedulerProfile(AuthoredConfig): class ImageBuildProfile(AuthoredConfig): partition: Identifier + cpus_per_task: PositiveInt + memory: Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:K|M|G|T)$")] + time_limit: Annotated[str, StringConstraints(pattern=r"^[0-9]+:[0-5][0-9]:[0-5][0-9]$")] + + @field_validator("time_limit") + @classmethod + def validate_time_limit(cls, value: str) -> str: + hours, minutes, seconds = (int(component) for component in value.split(":")) + if hours == minutes == seconds == 0: + raise ValueError("image-build time limit must be positive") + return value class ContainerMount(AuthoredConfig): diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py index d8e1fbcfc..adc2eb3e9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py @@ -7,9 +7,7 @@ import importlib.metadata import platform -import shutil import sys -from pathlib import Path from typing import Protocol from packaging.utils import canonicalize_name @@ -23,6 +21,7 @@ ) from data_designer.slurm.contracts import Identifier, Sha256Digest from data_designer.slurm.images.errors import ImageInspectionError +from data_designer.slurm.images.resources.inspect_image import find_distribution_console_script INSPECTOR_VERSION: Identifier = "inspector-1" _REQUIRED_CLIENT_DISTRIBUTIONS = ( @@ -49,11 +48,8 @@ def get_python_abi(self) -> str: def list_distributions(self) -> tuple[InstalledDistribution, ...]: """Return the installed Python distribution inventory.""" - def get_distribution_version(self, name: str) -> str: - """Return one installed distribution version.""" - - def find_executable(self, name: str) -> str: - """Return one absolute executable path.""" + def get_distribution_console_script(self, name: str) -> tuple[str, str]: + """Return one distribution version and its owned console-script path.""" class SystemInspectionEnvironment: @@ -94,22 +90,17 @@ def list_distributions(self) -> tuple[InstalledDistribution, ...]: except ValidationError as error: raise ImageInspectionError("installed distribution inventory is invalid") from error - def get_distribution_version(self, name: str) -> str: - """Return one installed distribution version with a canonical missing-package error.""" + def get_distribution_console_script(self, name: str) -> tuple[str, str]: + """Return one distribution version and its owned console-script path.""" try: - return importlib.metadata.version(name) + distribution = importlib.metadata.distribution(name) except importlib.metadata.PackageNotFoundError as error: raise ImageInspectionError(f"required distribution {name!r} is not installed") from error - - def find_executable(self, name: str) -> str: - """Return one resolved executable path with a canonical missing-tool error.""" - path = shutil.which(name) - if path is None: - raise ImageInspectionError(f"required executable {name!r} is not installed") - executable = Path(path) - if not executable.is_absolute(): - raise ImageInspectionError(f"required executable {name!r} did not resolve to an absolute path") - return executable.as_posix() + try: + executable = find_distribution_console_script(distribution, name) + except (OSError, RuntimeError) as error: + raise ImageInspectionError(str(error)) from error + return (distribution.version, executable) class ClientImageInspector: @@ -130,14 +121,17 @@ def inspect(self, sqsh_sha256: Sha256Digest) -> ImageInspectionRecord: if missing_distributions: missing = ", ".join(repr(name) for name in missing_distributions) raise ImageInspectionError(f"required client distributions are not installed: {missing}") + installer_version, installer_path = environment.get_distribution_console_script("pip") + if installer_version != versions_by_name["pip"]: + raise ImageInspectionError("pip console script does not match the distribution inventory") inspection = ClientImageInspection( kind="client", python_implementation=environment.get_python_implementation(), python_version=environment.get_python_version(), python_abi=environment.get_python_abi(), distributions=distributions, - installer_path=environment.find_executable("pip"), - installer_version=versions_by_name["pip"], + installer_path=installer_path, + installer_version=installer_version, ) return ImageInspectionRecord( schema_version=1, @@ -161,11 +155,12 @@ def inspect(self, sqsh_sha256: Sha256Digest) -> ImageInspectionRecord: """Return a digest-bound serving inspection for the active image environment.""" environment = self._environment try: + runtime_version, executable_path = environment.get_distribution_console_script("vllm") inspection = ServingImageInspection( kind="serving", server_type="vllm", - runtime_version=environment.get_distribution_version("vllm"), - executable_path=environment.find_executable("vllm"), + runtime_version=runtime_version, + executable_path=executable_path, ) return ImageInspectionRecord( schema_version=1, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py index d707c16f0..89d3dd53f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -21,13 +21,13 @@ from data_designer.slurm.images.records import ( ImageLifecycleOperation, ImageLifecyclePlan, + validate_enroot_mount_path, validate_oci_source_for_lifecycle, ) from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives from data_designer.slurm.launcher.client import SlurmCommandClient from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt -_IMAGE_JOB_TIME_LIMIT = "03:55:00" _INSPECTOR_FILENAME = "inspect_image.py" _ENROOT_RC_FILENAME = "enroot.rc" _PLAN_FILENAME = "image-lifecycle-plan.json" @@ -63,8 +63,10 @@ def prepare_image_lifecycle_job( "OCI image source must be a credential-free registry reference without a scheme" ) from error workspace_root = Path(selected_profile.profile.workspace_root) - if any(delimiter in workspace_root.as_posix() for delimiter in (":", ",")): - raise ImageLifecycleError("selected workspace cannot be represented as an Enroot mount") + try: + validate_enroot_mount_path(workspace_root.as_posix()) + except ValueError as error: + raise ImageLifecycleError("selected workspace cannot be represented as an Enroot mount") from error image_root = workspace_root / "images" temporary_root = image_root / ".tmp" job_root = temporary_root / "jobs" @@ -113,6 +115,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: raise ImageLifecycleError("cannot render an invalid image lifecycle plan") from error profile = plan.selected_profile.profile + image_build = profile.image_build directives = render_batch_directives( ( ("job-name", f"dd-image-{plan.request.kind}"), @@ -120,8 +123,9 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: ("partition", profile.image_build.partition), ("nodes", "1"), ("ntasks", "1"), - ("cpus-per-task", "1"), - ("time", _IMAGE_JOB_TIME_LIMIT), + ("cpus-per-task", str(image_build.cpus_per_task)), + ("mem", image_build.memory), + ("time", image_build.time_limit), ("chdir", plan.job_directory), ("output", f"{plan.job_directory}/slurm-%j.out"), ("error", f"{plan.job_directory}/slurm-%j.err"), @@ -130,7 +134,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: source_block = "" if plan.operation is ImageLifecycleOperation.IMPORT_OCI: source_block = f"""readonly DD_OCI_SOURCE={quote_shell_value(_format_enroot_oci_uri(plan.request.source))} -if [[ -e "${{DD_IMAGE_SQSH}}" ]]; then +if [[ -e "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then printf '%s\\n' 'candidate SQSH path already exists' >&2 exit 73 fi @@ -174,6 +178,11 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: export ENROOT_CONFIG_PATH="${{DD_JOB_DIR}}/enroot/config" export ENROOT_DATA_PATH="${{DD_JOB_DIR}}/enroot/data" export ENROOT_TEMP_PATH="${{DD_JOB_DIR}}/enroot/tmp" +if [[ ! ${{SLURM_CPUS_PER_TASK:-}} =~ ^[1-9][0-9]*$ ]]; then + printf '%s\\n' 'SLURM_CPUS_PER_TASK must be a positive integer' >&2 + exit 64 +fi +export ENROOT_MAX_PROCESSORS="${{SLURM_CPUS_PER_TASK}}" {source_block}if [[ ! -f "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then printf '%s\\n' 'SQSH path must be a regular non-symlink file' >&2 @@ -236,15 +245,25 @@ def submit_prepared_image_lifecycle( render_image_lifecycle_script(prepared.plan).encode(), ), ) + verified_script: bytes | None = None for label, artifact, expected_path, expected_content in expected_artifacts: expected_sha256 = hashlib.sha256(expected_content).hexdigest() + actual_content = _read_regular_file(expected_path) if ( artifact.path != expected_path.as_posix() or artifact.sha256 != expected_sha256 - or _compute_regular_file_sha256(expected_path) != expected_sha256 + or actual_content != expected_content ): raise ImageLifecycleError(f"prepared image lifecycle {label} no longer matches its digest") - return client.submit(prepared.script_file.path) + if label == "script": + verified_script = actual_content + if verified_script is None: + raise ImageLifecycleError("prepared image lifecycle script was not verified") + try: + script_text = verified_script.decode("utf-8", errors="strict") + except UnicodeError as error: + raise ImageLifecycleError("prepared image lifecycle script is not valid UTF-8") from error + return client.submit_script(script_text) def _stage_resource(job_directory: Path, filename: str) -> ArtifactReference: @@ -281,11 +300,18 @@ def _write_file(path: Path, content: bytes, *, mode: int) -> ArtifactReference: return ArtifactReference(path=path.as_posix(), sha256=hashlib.sha256(content).hexdigest()) -def _compute_regular_file_sha256(path: Path) -> str: +def _read_regular_file(path: Path) -> bytes: + descriptor: int | None = None try: - path_status = path.lstat() + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + path_status = os.fstat(descriptor) if not stat.S_ISREG(path_status.st_mode): raise ImageLifecycleError("prepared image lifecycle file is not regular") - return hashlib.sha256(path.read_bytes()).hexdigest() + with os.fdopen(descriptor, "rb") as source: + descriptor = None + return source.read() except OSError as error: raise ImageLifecycleError("cannot read prepared image lifecycle file") from error + finally: + if descriptor is not None: + os.close(descriptor) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/records.py b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py index 252223920..a44d24bbf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py @@ -60,8 +60,7 @@ def validate_lifecycle(self) -> ImageLifecyclePlan: ) if self.job_directory != expected_job_directory: raise ValueError("image lifecycle job directory must derive from the selected workspace") - if any(delimiter in self.job_directory for delimiter in (":", ",")): - raise ValueError("image lifecycle workspace path cannot be represented as an Enroot mount") + validate_enroot_mount_path(self.job_directory) if self.inspection_output_path != posixpath.join(self.job_directory, "output", "inspection.json"): raise ValueError("image lifecycle inspection output must belong to its dedicated output directory") expected_runtime_artifacts = ( @@ -105,6 +104,14 @@ def validate_oci_source_for_lifecycle(source: str) -> str: return source +def validate_enroot_mount_path(path: str) -> str: + """Reject host paths that Enroot's fstab-style mount parser cannot represent safely.""" + validate_absolute_path(path) + if any(character.isspace() or character in {":", ",", "\\"} for character in path): + raise ValueError("image lifecycle workspace path cannot be represented as an Enroot mount") + return path + + class RegisteredImage(ContractRecord): """One immutable alias binding for a verified SQSH artifact.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py index 0665138eb..605494ef2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py @@ -10,7 +10,6 @@ import os import platform import re -import shutil import sys from collections.abc import Iterable from pathlib import Path @@ -82,13 +81,18 @@ def main(arguments: list[str] | None = None) -> int: def _inspect_client() -> dict[str, object]: - versions = _list_distribution_versions(importlib.metadata.distributions()) + distributions = tuple(importlib.metadata.distributions()) + versions = _list_distribution_versions(distributions) missing = tuple(name for name in _REQUIRED_CLIENT_DISTRIBUTIONS if name not in versions) if missing: raise RuntimeError(f"required client distributions are not installed: {', '.join(missing)}") - installer_path = shutil.which("pip") - if installer_path is None or not Path(installer_path).is_absolute(): - raise RuntimeError("required executable 'pip' is not installed at an absolute path") + installer_distributions = tuple( + distribution for distribution in distributions if _distribution_name(distribution) == "pip" + ) + if len(installer_distributions) != 1: + raise RuntimeError("required distribution 'pip' is not installed exactly once") + installer_distribution = installer_distributions[0] + installer_path = find_distribution_console_script(installer_distribution, "pip") cache_tag = sys.implementation.cache_tag if cache_tag is None: raise RuntimeError("active Python does not expose an ABI cache tag") @@ -99,8 +103,8 @@ def _inspect_client() -> dict[str, object]: "python_version": platform.python_version(), "python_abi": python_abi, "distributions": [{"name": name, "version": version} for name, version in sorted(versions.items())], - "installer_path": Path(installer_path).as_posix(), - "installer_version": versions["pip"], + "installer_path": installer_path, + "installer_version": installer_distribution.version, } @@ -109,7 +113,7 @@ def _inspect_serving() -> dict[str, object]: distribution = importlib.metadata.distribution("vllm") except importlib.metadata.PackageNotFoundError as error: raise RuntimeError("required distribution 'vllm' is not installed") from error - executable_path = _find_distribution_console_script(distribution, "vllm") + executable_path = find_distribution_console_script(distribution, "vllm") return { "kind": "serving", "server_type": "vllm", @@ -121,17 +125,22 @@ def _inspect_serving() -> dict[str, object]: def _list_distribution_versions(distributions: Iterable[importlib.metadata.Distribution]) -> dict[str, str]: versions: dict[str, str] = {} for distribution in distributions: - raw_name = distribution.metadata.get("Name") - if not raw_name: - raise RuntimeError("installed distribution is missing its canonical name") - name = re.sub(r"[-_.]+", "-", raw_name).casefold() + name = _distribution_name(distribution) existing = versions.setdefault(name, distribution.version) if existing != distribution.version: raise RuntimeError(f"installed distribution {name!r} has conflicting versions") return versions -def _find_distribution_console_script(distribution: importlib.metadata.Distribution, name: str) -> str: +def _distribution_name(distribution: importlib.metadata.Distribution) -> str: + raw_name = distribution.metadata.get("Name") + if not raw_name: + raise RuntimeError("installed distribution is missing its canonical name") + return re.sub(r"[-_.]+", "-", raw_name).casefold() + + +def find_distribution_console_script(distribution: importlib.metadata.Distribution, name: str) -> str: + """Resolve exactly one executable console script owned by one distribution.""" entry_points = tuple( entry_point for entry_point in distribution.entry_points diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py index 3c0d3e858..47d061ad1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -76,6 +76,16 @@ def submit(self, script_path: str | Path) -> SlurmJobSubmissionReceipt: output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) return parse_submission(output) + def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + """Submit verified batch-script text through standard input.""" + if type(script) is not str or not script or "\0" in script: + raise ValueError("batch script text must be non-empty UTF-8 text without NUL") + output = self._run( + (self._executables.sbatch, "--parsable", "--export=NIL"), + input_text=script, + ) + return parse_submission(output) + def query_queue(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmQueueEntry, ...]: """Return normalized active-queue rows for explicit managed jobs.""" requested = tuple(selectors) @@ -133,10 +143,12 @@ def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, command.append(f"--partition={partition}") return parse_gpu_counts(self._run(command)) - def _run(self, command: Sequence[str]) -> str: + def _run(self, command: Sequence[str], *, input_text: str | None = None) -> str: command_name = Path(command[0]).name try: - completed = self._runner.run(command) + completed = ( + self._runner.run(command) if input_text is None else self._runner.run(command, input_text=input_text) + ) except (OSError, subprocess.SubprocessError) as error: raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") from error returncode = getattr(completed, "returncode", None) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py index 3679d3252..9e2e7839f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py @@ -16,7 +16,7 @@ class CommandRunner(Protocol): """Minimal command boundary implemented by production and fake runners.""" - def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + def run(self, command: Sequence[str], *, input_text: str | None = None) -> subprocess.CompletedProcess[str]: """Execute one argument-vector command.""" ... @@ -51,8 +51,20 @@ def environment(self) -> Mapping[str, str]: """Return the allowlisted environment forwarded to child processes.""" return self._environment - def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + def run(self, command: Sequence[str], *, input_text: str | None = None) -> subprocess.CompletedProcess[str]: """Execute an argument vector with captured text output.""" + if input_text is not None: + return subprocess.run( + tuple(command), + check=False, + input=input_text, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + env=dict(self._environment), + timeout=self._timeout_seconds, + ) return subprocess.run( tuple(command), check=False, 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 baf366891..fd89a82d7 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 @@ -387,7 +387,7 @@ "schema_version": 1, "selected_profile": { "catalog_path": "/workspace/profile.json", - "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "catalog_sha256": "a707e8e3b6f52170554b743d252b90273fa0deeea8141d4a87abe96e2e93215d", "cluster_name": "primary", "matched_pattern": null, "profile": { @@ -405,7 +405,10 @@ "primary-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" }, "scheduler": { "account": "research", @@ -415,7 +418,7 @@ "schema_version": 1, "workspace_root": "/workspace/primary" }, - "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "profile_sha256": "80e50266bc5748469c0429533731e1973856f6f58259d89ab6cde6d4dc0e5bc0", "schema_version": 1, "selection_source": "explicit" }, diff --git a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json index 8d500c664..7fc47e93d 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json +++ b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json @@ -8,7 +8,10 @@ "lab-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 2, + "memory": "8G", + "partition": "cpu", + "time_limit": "02:00:00" }, "scheduler": { "account": "lab", @@ -33,7 +36,10 @@ "primary-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" }, "scheduler": { "account": "research", 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 cfc0b5824..351e6ea1d 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 @@ -235,7 +235,7 @@ "schema_version": 1, "selected_profile": { "catalog_path": "/workspace/profile.json", - "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "catalog_sha256": "a707e8e3b6f52170554b743d252b90273fa0deeea8141d4a87abe96e2e93215d", "cluster_name": "primary", "matched_pattern": null, "profile": { @@ -253,7 +253,10 @@ "primary-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" }, "scheduler": { "account": "research", @@ -263,7 +266,7 @@ "schema_version": 1, "workspace_root": "/workspace/primary" }, - "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "profile_sha256": "80e50266bc5748469c0429533731e1973856f6f58259d89ab6cde6d4dc0e5bc0", "schema_version": 1, "selection_source": "explicit" }, diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py index 1e4086f5e..276a680c1 100644 --- a/packages/data-designer-slurm/tests/contracts/test_profiles.py +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -99,6 +99,9 @@ def test_explicit_selection_rejects_unknown_cluster(profile_catalog: SlurmProfil lambda payload: payload["clusters"]["primary"].update(extra="unknown"), lambda payload: payload["clusters"]["primary"].update(workspace_root="relative"), lambda payload: payload["clusters"]["primary"].update(host_patterns=["login[broken"]), + lambda payload: payload["clusters"]["primary"]["image_build"].update(cpus_per_task=0), + lambda payload: payload["clusters"]["primary"]["image_build"].update(memory="0G"), + lambda payload: payload["clusters"]["primary"]["image_build"].update(time_limit="00:00:00"), lambda payload: payload["clusters"]["primary"].update( container_mounts=[ {"source": "/one", "target": "/same"}, diff --git a/packages/data-designer-slurm/tests/images/test_inspection.py b/packages/data-designer-slurm/tests/images/test_inspection.py index 5d12cad12..c7415a4e5 100644 --- a/packages/data-designer-slurm/tests/images/test_inspection.py +++ b/packages/data-designer-slurm/tests/images/test_inspection.py @@ -5,7 +5,7 @@ import importlib.metadata import re -import shutil +from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock @@ -19,6 +19,7 @@ ServingImageInspector, SystemInspectionEnvironment, ) +from data_designer.slurm.images.resources import inspect_image as resource_inspector def _get_client_distributions() -> tuple[InstalledDistribution, ...]: @@ -151,16 +152,15 @@ def test_client_inspector_sorts_distribution_inventory() -> None: ) -def test_client_inspector_uses_installer_version_from_distribution_inventory() -> None: +def test_client_inspector_rejects_installer_version_outside_distribution_inventory() -> None: environment = FakeInspectionEnvironment( distributions=_get_client_distributions(), distribution_versions={"pip": "unexpected"}, executables={"pip": "/usr/bin/pip"}, ) - inspection = ClientImageInspector(environment).inspect("a" * 64).inspection - - assert inspection.installer_version == "26.1" + with pytest.raises(ImageInspectionError, match="does not match"): + ClientImageInspector(environment).inspect("a" * 64) def test_system_inspection_environment_normalizes_distribution_inventory( @@ -207,23 +207,56 @@ def test_system_inspection_environment_normalizes_missing_distribution( ) -> None: monkeypatch.setattr( importlib.metadata, - "version", + "distribution", Mock(side_effect=importlib.metadata.PackageNotFoundError("missing")), ) with pytest.raises(ImageInspectionError, match="not installed"): - SystemInspectionEnvironment().get_distribution_version("missing") + SystemInspectionEnvironment().get_distribution_console_script("missing") -@pytest.mark.parametrize("path", (None, "relative/pip"), ids=("missing", "relative")) -def test_system_inspection_environment_requires_absolute_executable( +def test_package_and_standalone_serving_inspectors_share_distribution_owned_console_script_semantics( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, - path: str | None, ) -> None: - monkeypatch.setattr(shutil, "which", Mock(return_value=path)) + executable_path = tmp_path / "active-environment" / "bin" / "vllm" + executable_path.parent.mkdir(parents=True) + executable_path.write_text("#!/bin/sh\n") + executable_path.chmod(0o700) + shadow_path = tmp_path / "shadow" / "bin" / "vllm" + shadow_path.parent.mkdir(parents=True) + shadow_path.write_text("#!/bin/sh\n") + shadow_path.chmod(0o700) + distribution = SimpleNamespace( + metadata={"Name": "vllm"}, + version="0.21.0", + entry_points=(SimpleNamespace(group="console_scripts", name="vllm"),), + files=(Path("../../../bin/vllm"),), + locate_file=lambda _installed_file: executable_path, + ) + monkeypatch.setenv("PATH", shadow_path.parent.as_posix()) + monkeypatch.setattr(importlib.metadata, "distribution", Mock(return_value=distribution)) - with pytest.raises(ImageInspectionError): - SystemInspectionEnvironment().find_executable("pip") + package_record = ServingImageInspector().inspect("d" * 64) + standalone_record = ImageInspectionRecord.model_validate(resource_inspector.inspect_image("serving", "d" * 64)) + + assert package_record == standalone_record + assert package_record.inspection.executable_path == executable_path.as_posix() # type: ignore[union-attr] + + +def test_system_inspection_environment_rejects_unowned_console_script( + monkeypatch: pytest.MonkeyPatch, +) -> None: + distribution = SimpleNamespace( + version="26.1", + entry_points=(SimpleNamespace(group="console_scripts", name="pip"),), + files=(Path("pip"),), + locate_file=lambda installed_file: installed_file, + ) + monkeypatch.setattr(importlib.metadata, "distribution", Mock(return_value=distribution)) + + with pytest.raises(ImageInspectionError, match="does not own"): + SystemInspectionEnvironment().get_distribution_console_script("pip") def test_system_inspection_environment_reports_current_python_facts() -> None: diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index 1a1f29c9e..c519ceca7 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -7,6 +7,7 @@ import json import stat import subprocess +from collections.abc import Sequence from dataclasses import replace from pathlib import Path from types import SimpleNamespace @@ -87,8 +88,8 @@ def test_prepare_existing_sqsh_inspects_in_place_without_oci_identity(tmp_path: assert f'readonly DD_IMAGE_SQSH="{image_path.as_posix()}"' in script -def test_image_lifecycle_renderer_uses_cpu_profile_and_quotes_workspace_paths(tmp_path: Path) -> None: - workspace = tmp_path / "workspace with spaces" +def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_path: Path) -> None: + workspace = tmp_path / "workspace-safe" source = f"registry.example.test/team/image:release@sha256:{'b' * 64}" prepared = prepare_image_lifecycle_job( ImageBuildRequest(name="serving", kind="serving", source=source), @@ -101,6 +102,9 @@ def test_image_lifecycle_renderer_uses_cpu_profile_and_quotes_workspace_paths(tm assert "#SBATCH --account=research" in script assert "#SBATCH --partition=image-build" in script assert "#SBATCH --nodes=1" in script + assert "#SBATCH --cpus-per-task=2" in script + assert "#SBATCH --mem=8G" in script + assert "#SBATCH --time=03:55:00" in script assert "#SBATCH --gres=" not in script expected_uri = f"docker://registry.example.test#team/image:release@sha256:{'b' * 64}" assert f'readonly DD_OCI_SOURCE="{expected_uri}"' in script @@ -108,6 +112,7 @@ def test_image_lifecycle_renderer_uses_cpu_profile_and_quotes_workspace_paths(tm assert "inspect_image.py:x-create=file,bind,ro" in script assert "/opt/data-designer-slurm/output:x-create=dir,bind" in script assert 'export ENROOT_CONFIG_PATH="${DD_JOB_DIR}/enroot/config"' in script + assert 'export ENROOT_MAX_PROCESSORS="${SLURM_CPUS_PER_TASK}"' in script assert f'--mount "{prepared.plan.job_directory}:' not in script completed = subprocess.run(("bash", "-n"), input=script, capture_output=True, text=True, check=False) assert completed.returncode == 0, completed.stderr @@ -156,6 +161,7 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( check=False, env={ "DD_TEST_INSPECTION_OUTPUT": prepared.plan.inspection_output_path, + "SLURM_CPUS_PER_TASK": "2", "SLURM_JOB_ID": "5101", }, ) @@ -233,8 +239,14 @@ def test_prepare_refuses_duplicate_and_invalid_lifecycle_ids(tmp_path: Path) -> assert not (workspace / "images" / ".tmp" / "escape").exists() -def test_prepare_rejects_workspace_paths_that_enroot_cannot_mount(tmp_path: Path) -> None: - workspace = tmp_path / "workspace:unsafe" +@pytest.mark.parametrize( + "workspace_name", ("workspace unsafe", "workspace\\unsafe", "workspace:unsafe", "workspace,unsafe") +) +def test_prepare_rejects_workspace_paths_that_enroot_cannot_mount( + tmp_path: Path, + workspace_name: str, +) -> None: + workspace = tmp_path / workspace_name request = ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()) with pytest.raises(ImageLifecycleError, match="Enroot mount"): @@ -243,6 +255,34 @@ def test_prepare_rejects_workspace_paths_that_enroot_cannot_mount(tmp_path: Path assert not (workspace / "images").exists() +def test_rendered_oci_import_rejects_dangling_candidate_symlink(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name="client", + kind="client", + source=f"registry.example.test/client@sha256:{'a' * 64}", + ), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-dangling-candidate", + ) + candidate = Path(prepared.plan.sqsh_path) + outside = tmp_path / "outside.sqsh" + candidate.symlink_to(outside) + + completed = subprocess.run( + ("bash",), + input=render_image_lifecycle_script(prepared.plan), + capture_output=True, + text=True, + check=False, + env={"SLURM_CPUS_PER_TASK": "2", "SLURM_JOB_ID": "5101"}, + ) + + assert completed.returncode == 73 + assert candidate.is_symlink() + assert not outside.exists() + + def test_prepare_rejects_symlinked_package_owned_image_root(tmp_path: Path) -> None: workspace = tmp_path / "workspace" workspace.mkdir() @@ -334,9 +374,9 @@ def test_submit_prepared_image_lifecycle_uses_isolated_sbatch_client(tmp_path: P "sbatch", "--parsable", "--export=NIL", - prepared.script_file.path, ) ] + assert runner.inputs == [render_image_lifecycle_script(prepared.plan)] @pytest.mark.parametrize("artifact_name", ("plan_file", "script_file"), ids=("plan", "script")) @@ -377,21 +417,51 @@ def test_submit_rejects_modified_script_even_when_artifact_digest_is_rebound(tmp submit_prepared_image_lifecycle(rebound, SlurmCommandClient(FakeSlurmRunner())) -def test_standalone_client_inspector_emits_contract_compatible_json(tmp_path: Path) -> None: - distributions = tuple( - SimpleNamespace(metadata={"Name": name}, version="1.0.0") - for name in ( - "data-designer", - "data-designer-config", - "data-designer-engine", - "data-designer-slurm", - "pip", - ) +def test_submit_uses_verified_script_bytes_when_path_changes_during_submission(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-submission-race", ) - with ( - patch.object(resource_inspector.importlib.metadata, "distributions", return_value=distributions), - patch.object(resource_inspector.shutil, "which", return_value="/usr/bin/pip"), - ): + expected_script = render_image_lifecycle_script(prepared.plan) + runner = _MutatingSubmissionRunner(Path(prepared.script_file.path)) + + receipt = submit_prepared_image_lifecycle(prepared, SlurmCommandClient(runner)) + + assert receipt.job_id == 5101 + assert runner.command == ("sbatch", "--parsable", "--export=NIL") + assert runner.input_text == expected_script + assert Path(prepared.script_file.path).read_text() == "replaced after verification\n" + + +def test_standalone_client_inspector_binds_pip_to_its_distribution( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + executable_path = tmp_path / "active-environment" / "bin" / "pip" + executable_path.parent.mkdir(parents=True) + executable_path.write_text("#!/bin/sh\n") + executable_path.chmod(0o700) + shadow_path = tmp_path / "shadow" / "bin" / "pip" + shadow_path.parent.mkdir(parents=True) + shadow_path.write_text("#!/bin/sh\n") + shadow_path.chmod(0o700) + monkeypatch.setenv("PATH", shadow_path.parent.as_posix()) + installer_distribution = SimpleNamespace( + metadata={"Name": "pip"}, + version="26.1", + entry_points=(SimpleNamespace(group="console_scripts", name="pip"),), + files=(Path("../../../bin/pip"),), + locate_file=lambda _installed_file: executable_path, + ) + distributions = ( + *( + SimpleNamespace(metadata={"Name": name}, version="1.0.0") + for name in ("data-designer", "data-designer-config", "data-designer-engine", "data-designer-slurm") + ), + installer_distribution, + ) + with patch.object(resource_inspector.importlib.metadata, "distributions", return_value=distributions): record = resource_inspector.inspect_image("client", "c" * 64) output_path = tmp_path / "inspection.json" @@ -400,6 +470,8 @@ def test_standalone_client_inspector_emits_contract_compatible_json(tmp_path: Pa inspection = ImageInspectionRecord.model_validate_json(output_path.read_text()) assert inspection.inspection.kind == "client" assert inspection.sqsh_sha256 == "c" * 64 + assert inspection.inspection.installer_path == executable_path.as_posix() # type: ignore[union-attr] + assert inspection.inspection.installer_version == "26.1" # type: ignore[union-attr] assert stat.S_IMODE(output_path.stat().st_mode) == 0o600 @@ -419,7 +491,6 @@ def test_standalone_serving_inspector_binds_version_and_executable_to_one_distri with ( patch.object(resource_inspector.importlib.metadata, "distribution", return_value=distribution), - patch.object(resource_inspector.shutil, "which", return_value="/shadow/bin/vllm"), ): payload = resource_inspector.inspect_image("serving", "d" * 64) @@ -471,6 +542,30 @@ def _get_selected_profile(workspace: Path) -> SelectedSlurmProfile: scheduler=SchedulerProfile(account="research", partition="gpu"), gpus_per_node=8, workspace_root=workspace.as_posix(), - image_build=ImageBuildProfile(partition="image-build"), + image_build=ImageBuildProfile( + partition="image-build", + cpus_per_task=2, + memory="8G", + time_limit="03:55:00", + ), ) return injected_profile(profile) + + +class _MutatingSubmissionRunner: + def __init__(self, script_path: Path) -> None: + self._script_path = script_path + self.command: tuple[str, ...] | None = None + self.input_text: str | None = None + + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + ) -> subprocess.CompletedProcess[str]: + self.command = tuple(command) + self.input_text = input_text + self._script_path.chmod(0o700) + self._script_path.write_text("replaced after verification\n") + return subprocess.CompletedProcess(command, 0, stdout="5101\n", stderr="") 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 e3b564f14..a818ecd9c 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": "8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" + "sha256": "90672b2e9b3adab921261876a436b4c5cd8caa242f2662bdcb78cd48f08da3fb" }, "run_id": "run-single", "scheduler": { @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" + "sha256": "90672b2e9b3adab921261876a436b4c5cd8caa242f2662bdcb78cd48f08da3fb" }, "run_id": "run-single", "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 5155ddb02..183680d34 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -28,6 +28,18 @@ def test_client_submits_and_observes_one_managed_array(fake_slurm_runner: FakeSl ] +def test_client_submits_verified_script_text_through_standard_input() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + client = SlurmCommandClient(runner) + script = "#!/usr/bin/env bash\nexit 0\n" + + submission = client.submit_script(script) + + assert submission.job_id == 5101 + assert runner.calls == [("sbatch", "--parsable", "--export=NIL")] + assert runner.inputs == [script] + + def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) client.submit("run.sbatch") @@ -247,6 +259,16 @@ def test_client_rejects_invalid_script_path(fake_slurm_runner: FakeSlurmRunner, assert fake_slurm_runner.calls == [] +@pytest.mark.parametrize("script", ("", "bad\0script", b"not text")) +def test_client_rejects_invalid_script_text(fake_slurm_runner: FakeSlurmRunner, script: object) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="batch script text"): + client.submit_script(script) # type: ignore[arg-type] + + assert fake_slurm_runner.calls == [] + + @pytest.mark.parametrize("executable", ("", "sbatch --wait", "sbatch\n")) def test_executables_reject_invalid_tokens(executable: str) -> None: with pytest.raises(ValueError, match="Slurm executable"): diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index 6bab64a83..b41811178 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -70,6 +70,44 @@ def test_subprocess_runner_default_environment_forwards_only_search_path(monkeyp assert runner.environment == {"LC_ALL": "C", "PATH": "/workspace/slurm/bin:/usr/bin"} +def test_subprocess_runner_forwards_explicit_standard_input(monkeypatch: pytest.MonkeyPatch) -> None: + observed: dict[str, object] = {} + + def fake_run( + command: Sequence[str], + *, + check: bool, + input: str, + capture_output: bool, + text: bool, + encoding: str, + errors: str, + env: Mapping[str, str], + timeout: float, + ) -> subprocess.CompletedProcess[str]: + observed.update(command=command, check=check, input=input, env=env, timeout=timeout) + assert capture_output is text is True + assert encoding == "utf-8" + assert errors == "replace" + return subprocess.CompletedProcess(command, 0, stdout="5101\n", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + + completed = SubprocessRunner(environment={"PATH": "/usr/bin"}).run( + ("sbatch", "--parsable"), + input_text="#!/bin/sh\n", + ) + + assert completed.stdout == "5101\n" + assert observed == { + "command": ("sbatch", "--parsable"), + "check": False, + "input": "#!/bin/sh\n", + "env": {"LC_ALL": "C", "PATH": "/usr/bin"}, + "timeout": 30.0, + } + + def test_subprocess_runner_default_environment_replaces_empty_search_path(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("PATH", "") 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 d934fd637..3e8e266a8 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="fff7573a8b6eb2f71ae12f27900a85e1e121debfe58b0e94fefc3fb9362970ce" +readonly DD_PLAN_SHA256="a910fe84276a2196168ad1f548148c2086cd716f45a0ff8a020a473ff30cd69f" 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 93156136d..f0d87350b 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" +readonly DD_PLAN_SHA256="90672b2e9b3adab921261876a436b4c5cd8caa242f2662bdcb78cd48f08da3fb" 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/images.py b/packages/data-designer-slurm/tests/slurm_test_fakes/images.py index df0d8d278..7ab4c084b 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/images.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/images.py @@ -39,16 +39,9 @@ def list_distributions(self) -> tuple[InstalledDistribution, ...]: """Return the configured distribution inventory.""" return self.distributions - def get_distribution_version(self, name: str) -> str: - """Return one configured distribution version.""" + def get_distribution_console_script(self, name: str) -> tuple[str, str]: + """Return one configured distribution version and console-script path.""" try: - return self.distribution_versions[name] + return (self.distribution_versions[name], self.executables[name]) except KeyError: - raise ImageInspectionError(f"required distribution {name!r} is not installed") from None - - def find_executable(self, name: str) -> str: - """Return one configured executable path.""" - try: - return self.executables[name] - except KeyError: - raise ImageInspectionError(f"required executable {name!r} is not installed") from None + raise ImageInspectionError(f"required distribution console script {name!r} is not installed") from None diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py index 22da3c51f..5a8494cfc 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -102,13 +102,21 @@ def __init__( self._scripted_responses: dict[str, deque[FakeCommandResponse]] = {} self._sinfo_responses = dict(sinfo_responses or {}) self.calls: list[tuple[str, ...]] = [] + self.inputs: list[str | None] = [] - def run(self, command: Sequence[str], *, check: bool = False) -> subprocess.CompletedProcess[str]: + def run( + self, + command: Sequence[str], + *, + check: bool = False, + input_text: str | None = None, + ) -> subprocess.CompletedProcess[str]: """Run one fake Slurm command and optionally raise on failure.""" if not command: raise ValueError("command must not be empty") argv = tuple(command) self.calls.append(argv) + self.inputs.append(input_text) command_name = Path(argv[0]).name scripted = self._scripted_responses.get(command_name) if scripted: 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 a8c29e5ef..01d2b7a11 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="8ddf07c38a825a7487c487fddfe051f0a1940f63063725b54b32bfe4c03fd9ca", + expected_fixture_sha256="91dde4ea46a6290c406685f65151f503fe42a7bc2d69618e9840954f45c45042", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="17a4c2e16189d22dfdb6885bf76264844ad3168dea0cf94aef70948d5ab2e6b7", + expected_fixture_sha256="dd4f96ade817cf11ab4e3e03f794c97d39dcf817f1efcb39948dd2bc8359b26f", ) From 76620f0cedb0e5b6171797f951cb7c090f264998 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 31 Aug 2026 16:35:55 -0600 Subject: [PATCH 4/5] fix: validate image lifecycle compatibility Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/images/inspection.py | 10 +-- .../data_designer/slurm/images/lifecycle.py | 21 +++++- .../slurm/images/resources/inspect_image.py | 28 +++++--- .../tests/images/test_inspection.py | 30 +++++++-- .../tests/images/test_lifecycle.py | 65 ++++++++++++++++++- 5 files changed, 128 insertions(+), 26 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py index adc2eb3e9..d0bc57abf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py @@ -21,7 +21,10 @@ ) from data_designer.slurm.contracts import Identifier, Sha256Digest from data_designer.slurm.images.errors import ImageInspectionError -from data_designer.slurm.images.resources.inspect_image import find_distribution_console_script +from data_designer.slurm.images.resources.inspect_image import ( + find_distribution_console_script, + find_unique_distribution, +) INSPECTOR_VERSION: Identifier = "inspector-1" _REQUIRED_CLIENT_DISTRIBUTIONS = ( @@ -93,10 +96,7 @@ def list_distributions(self) -> tuple[InstalledDistribution, ...]: def get_distribution_console_script(self, name: str) -> tuple[str, str]: """Return one distribution version and its owned console-script path.""" try: - distribution = importlib.metadata.distribution(name) - except importlib.metadata.PackageNotFoundError as error: - raise ImageInspectionError(f"required distribution {name!r} is not installed") from error - try: + distribution = find_unique_distribution(importlib.metadata.distributions(), name) executable = find_distribution_console_script(distribution, name) except (OSError, RuntimeError) as error: raise ImageInspectionError(str(error)) from error diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py index 89d3dd53f..c650ea0df 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -34,6 +34,7 @@ _SCRIPT_FILENAME = "image-lifecycle.sbatch" _RESOURCE_PACKAGE = "data_designer.slurm.images.resources" _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_MINIMUM_ENROOT_MAJOR = 4 @dataclass(frozen=True, slots=True) @@ -132,14 +133,18 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: ) ) source_block = "" + existing_sqsh_preflight = "" if plan.operation is ImageLifecycleOperation.IMPORT_OCI: source_block = f"""readonly DD_OCI_SOURCE={quote_shell_value(_format_enroot_oci_uri(plan.request.source))} if [[ -e "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then printf '%s\\n' 'candidate SQSH path already exists' >&2 exit 73 fi +verify_enroot_compatibility enroot import -o "${{DD_IMAGE_SQSH}}" "${{DD_OCI_SOURCE}}" """ + else: + existing_sqsh_preflight = "verify_enroot_compatibility\n" return f"""#!/usr/bin/env bash {directives} @@ -166,6 +171,20 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: [[ "$(compute_file_sha256 "$2")" == "$1" ]] }} +verify_enroot_compatibility() {{ + local version major + version="$(enroot version)" + if [[ ! ${{version}} =~ ^([0-9]+)\\.([0-9]+)\\.([0-9]+)([-+][0-9A-Za-z.-]+)?$ ]]; then + printf '%s\\n' 'Enroot version output is invalid' >&2 + exit 78 + fi + major="${{BASH_REMATCH[1]}}" + if (( 10#${{major}} < {_MINIMUM_ENROOT_MAJOR} )); then + printf '%s\\n' 'Enroot 4 or newer is required for digest imports and processor-bounded extraction' >&2 + exit 78 + fi +}} + verify_sha256 "${{DD_INSPECTOR_SHA256}}" "${{DD_INSPECTOR}}" verify_sha256 "${{DD_ENROOT_RC_SHA256}}" "${{DD_ENROOT_RC}}" install -d -m 0700 \ @@ -188,7 +207,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: printf '%s\\n' 'SQSH path must be a regular non-symlink file' >&2 exit 66 fi -readonly DD_SQSH_SHA256="$(compute_file_sha256 "${{DD_IMAGE_SQSH}}")" +{existing_sqsh_preflight}readonly DD_SQSH_SHA256="$(compute_file_sha256 "${{DD_IMAGE_SQSH}}")" if [[ ! ${{DD_SQSH_SHA256}} =~ ^[0-9a-f]{{64}}$ ]]; then printf '%s\\n' 'SQSH checksum output is invalid' >&2 exit 65 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py index 605494ef2..01dc0f6b6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py @@ -86,12 +86,7 @@ def _inspect_client() -> dict[str, object]: missing = tuple(name for name in _REQUIRED_CLIENT_DISTRIBUTIONS if name not in versions) if missing: raise RuntimeError(f"required client distributions are not installed: {', '.join(missing)}") - installer_distributions = tuple( - distribution for distribution in distributions if _distribution_name(distribution) == "pip" - ) - if len(installer_distributions) != 1: - raise RuntimeError("required distribution 'pip' is not installed exactly once") - installer_distribution = installer_distributions[0] + installer_distribution = find_unique_distribution(distributions, "pip") installer_path = find_distribution_console_script(installer_distribution, "pip") cache_tag = sys.implementation.cache_tag if cache_tag is None: @@ -109,10 +104,7 @@ def _inspect_client() -> dict[str, object]: def _inspect_serving() -> dict[str, object]: - try: - distribution = importlib.metadata.distribution("vllm") - except importlib.metadata.PackageNotFoundError as error: - raise RuntimeError("required distribution 'vllm' is not installed") from error + distribution = find_unique_distribution(importlib.metadata.distributions(), "vllm") executable_path = find_distribution_console_script(distribution, "vllm") return { "kind": "serving", @@ -139,6 +131,22 @@ def _distribution_name(distribution: importlib.metadata.Distribution) -> str: return re.sub(r"[-_.]+", "-", raw_name).casefold() +def find_unique_distribution( + distributions: Iterable[importlib.metadata.Distribution], + name: str, +) -> importlib.metadata.Distribution: + """Select exactly one installed distribution by canonical name.""" + canonical_name = re.sub(r"[-_.]+", "-", name).casefold() + matches = tuple( + distribution for distribution in distributions if _distribution_name(distribution) == canonical_name + ) + if not matches: + raise RuntimeError(f"required distribution {canonical_name!r} is not installed") + if len(matches) != 1: + raise RuntimeError(f"required distribution {canonical_name!r} is not installed exactly once") + return matches[0] + + def find_distribution_console_script(distribution: importlib.metadata.Distribution, name: str) -> str: """Resolve exactly one executable console script owned by one distribution.""" entry_points = tuple( diff --git a/packages/data-designer-slurm/tests/images/test_inspection.py b/packages/data-designer-slurm/tests/images/test_inspection.py index c7415a4e5..c869d6c0c 100644 --- a/packages/data-designer-slurm/tests/images/test_inspection.py +++ b/packages/data-designer-slurm/tests/images/test_inspection.py @@ -205,11 +205,7 @@ def test_system_inspection_environment_rejects_invalid_distribution_inventory( def test_system_inspection_environment_normalizes_missing_distribution( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setattr( - importlib.metadata, - "distribution", - Mock(side_effect=importlib.metadata.PackageNotFoundError("missing")), - ) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=())) with pytest.raises(ImageInspectionError, match="not installed"): SystemInspectionEnvironment().get_distribution_console_script("missing") @@ -235,7 +231,7 @@ def test_package_and_standalone_serving_inspectors_share_distribution_owned_cons locate_file=lambda _installed_file: executable_path, ) monkeypatch.setenv("PATH", shadow_path.parent.as_posix()) - monkeypatch.setattr(importlib.metadata, "distribution", Mock(return_value=distribution)) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=(distribution,))) package_record = ServingImageInspector().inspect("d" * 64) standalone_record = ImageInspectionRecord.model_validate(resource_inspector.inspect_image("serving", "d" * 64)) @@ -248,17 +244,37 @@ def test_system_inspection_environment_rejects_unowned_console_script( monkeypatch: pytest.MonkeyPatch, ) -> None: distribution = SimpleNamespace( + metadata={"Name": "pip"}, version="26.1", entry_points=(SimpleNamespace(group="console_scripts", name="pip"),), files=(Path("pip"),), locate_file=lambda installed_file: installed_file, ) - monkeypatch.setattr(importlib.metadata, "distribution", Mock(return_value=distribution)) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=(distribution,))) with pytest.raises(ImageInspectionError, match="does not own"): SystemInspectionEnvironment().get_distribution_console_script("pip") +def test_package_and_standalone_client_inspectors_reject_duplicate_same_version_pip( + monkeypatch: pytest.MonkeyPatch, +) -> None: + distributions = ( + *( + SimpleNamespace(metadata={"Name": name}, version="1.0.0") + for name in ("data-designer", "data-designer-config", "data-designer-engine", "data-designer-slurm") + ), + SimpleNamespace(metadata={"Name": "pip"}, version="26.1"), + SimpleNamespace(metadata={"Name": "Pip"}, version="26.1"), + ) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=distributions)) + + with pytest.raises(ImageInspectionError, match="exactly once"): + ClientImageInspector().inspect("a" * 64) + with pytest.raises(RuntimeError, match="exactly once"): + resource_inspector.inspect_image("client", "a" * 64) + + def test_system_inspection_environment_reports_current_python_facts() -> None: environment = SystemInspectionEnvironment() diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index c519ceca7..bc03ef174 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -113,6 +113,8 @@ def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_ assert "/opt/data-designer-slurm/output:x-create=dir,bind" in script assert 'export ENROOT_CONFIG_PATH="${DD_JOB_DIR}/enroot/config"' in script assert 'export ENROOT_MAX_PROCESSORS="${SLURM_CPUS_PER_TASK}"' in script + assert 'version="$(enroot version)"' in script + assert "if (( 10#${major} < 4 )); then" in script assert f'--mount "{prepared.plan.job_directory}:' not in script completed = subprocess.run(("bash", "-n"), input=script, capture_output=True, text=True, check=False) assert completed.returncode == 0, completed.stderr @@ -140,7 +142,9 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( fake_enroot.write_text( "#!/usr/bin/env bash\n" "set -Eeuo pipefail\n" - 'if [[ "$1" == "import" ]]; then\n' + 'if [[ "$1" == "version" ]]; then\n' + " printf '%s\\n' '4.0.0'\n" + 'elif [[ "$1" == "import" ]]; then\n' " printf '%s\\n' 'imported image' > \"$3\"\n" 'elif [[ "$1" == "start" ]]; then\n' " printf '%s\\n' '{}' > \"${DD_TEST_INSPECTION_OUTPUT}\"\n" @@ -172,6 +176,59 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( assert Path(prepared.plan.sqsh_path).read_text() == "imported image\n" +@pytest.mark.parametrize("source_kind", ("oci", "existing"), ids=("digest-import", "existing-sqsh")) +def test_rendered_image_lifecycle_rejects_enroot_before_version_4_before_image_operations( + tmp_path: Path, + source_kind: str, +) -> None: + if source_kind == "oci": + source = f"registry.example.test/client@sha256:{'a' * 64}" + else: + image_path = tmp_path / "client.sqsh" + image_path.write_bytes(b"client image") + source = image_path.as_posix() + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=source), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id=f"image-job-{source_kind}-old-enroot", + ) + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_enroot = fake_bin / "enroot" + fake_enroot.write_text( + "#!/usr/bin/env bash\n" + "set -Eeuo pipefail\n" + 'printf \'%s\\n\' "$1" >> "${DD_TEST_ENROOT_LOG}"\n' + 'if [[ "$1" == "version" ]]; then\n' + " printf '%s\\n' '3.5.0'\n" + "fi\n" + ) + fake_enroot.chmod(0o700) + script = render_image_lifecycle_script(prepared.plan).replace( + 'export PATH="', + f'export PATH="{fake_bin.as_posix()}:', + 1, + ) + enroot_log = tmp_path / "enroot.log" + + completed = subprocess.run( + ("bash",), + input=script, + capture_output=True, + text=True, + check=False, + env={ + "DD_TEST_ENROOT_LOG": enroot_log.as_posix(), + "SLURM_CPUS_PER_TASK": "2", + "SLURM_JOB_ID": "5101", + }, + ) + + assert completed.returncode == 78 + assert "Enroot 4 or newer" in completed.stderr + assert enroot_log.read_text() == "version\n" + + @pytest.mark.parametrize( "source", ( @@ -483,6 +540,7 @@ def test_standalone_serving_inspector_binds_version_and_executable_to_one_distri installed_file = Path("../../../bin/vllm") entry_point = SimpleNamespace(group="console_scripts", name="vllm") distribution = SimpleNamespace( + metadata={"Name": "vllm"}, version="0.21.0", entry_points=(entry_point,), files=(installed_file,), @@ -490,7 +548,7 @@ def test_standalone_serving_inspector_binds_version_and_executable_to_one_distri ) with ( - patch.object(resource_inspector.importlib.metadata, "distribution", return_value=distribution), + patch.object(resource_inspector.importlib.metadata, "distributions", return_value=(distribution,)), ): payload = resource_inspector.inspect_image("serving", "d" * 64) @@ -514,6 +572,7 @@ def test_standalone_serving_inspector_rejects_unverifiable_distribution_console_ match: str, ) -> None: distribution = SimpleNamespace( + metadata={"Name": "vllm"}, version="0.21.0", entry_points=entry_points, files=files, @@ -521,7 +580,7 @@ def test_standalone_serving_inspector_rejects_unverifiable_distribution_console_ ) with ( - patch.object(resource_inspector.importlib.metadata, "distribution", return_value=distribution), + patch.object(resource_inspector.importlib.metadata, "distributions", return_value=(distribution,)), pytest.raises(RuntimeError, match=match), ): resource_inspector.inspect_image("serving", "d" * 64) From 69b7b86846d79d37a9af185a2ccdfa17b7415959 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 1 Sep 2026 15:00:44 -0600 Subject: [PATCH 5/5] fix Enroot lifecycle compatibility Initialize a private job-local home for environment-isolated lifecycle submissions. Apply operation-specific Enroot minimum versions and cover the supported invocation contracts. --- .../data_designer/slurm/images/lifecycle.py | 23 ++++++--- .../tests/images/test_lifecycle.py | 47 ++++++++++++++++--- 2 files changed, 57 insertions(+), 13 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py index c650ea0df..cfb317f7b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -34,7 +34,8 @@ _SCRIPT_FILENAME = "image-lifecycle.sbatch" _RESOURCE_PACKAGE = "data_designer.slurm.images.resources" _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) -_MINIMUM_ENROOT_MAJOR = 4 +_MINIMUM_ENROOT_OCI_VERSION = (4, 0) +_MINIMUM_ENROOT_SQSH_VERSION = (3, 5) @dataclass(frozen=True, slots=True) @@ -135,16 +136,20 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: source_block = "" existing_sqsh_preflight = "" if plan.operation is ImageLifecycleOperation.IMPORT_OCI: + minimum_major, minimum_minor = _MINIMUM_ENROOT_OCI_VERSION source_block = f"""readonly DD_OCI_SOURCE={quote_shell_value(_format_enroot_oci_uri(plan.request.source))} if [[ -e "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then printf '%s\\n' 'candidate SQSH path already exists' >&2 exit 73 fi -verify_enroot_compatibility +verify_enroot_compatibility {minimum_major} {minimum_minor} "digest-pinned OCI imports" enroot import -o "${{DD_IMAGE_SQSH}}" "${{DD_OCI_SOURCE}}" """ else: - existing_sqsh_preflight = "verify_enroot_compatibility\n" + minimum_major, minimum_minor = _MINIMUM_ENROOT_SQSH_VERSION + existing_sqsh_preflight = ( + f'verify_enroot_compatibility {minimum_major} {minimum_minor} "existing SQSH inspection"\n' + ) return f"""#!/usr/bin/env bash {directives} @@ -172,15 +177,19 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: }} verify_enroot_compatibility() {{ - local version major + local required_major="$1" required_minor="$2" purpose="$3" + local version major minor version="$(enroot version)" if [[ ! ${{version}} =~ ^([0-9]+)\\.([0-9]+)\\.([0-9]+)([-+][0-9A-Za-z.-]+)?$ ]]; then printf '%s\\n' 'Enroot version output is invalid' >&2 exit 78 fi major="${{BASH_REMATCH[1]}}" - if (( 10#${{major}} < {_MINIMUM_ENROOT_MAJOR} )); then - printf '%s\\n' 'Enroot 4 or newer is required for digest imports and processor-bounded extraction' >&2 + minor="${{BASH_REMATCH[2]}}" + if (( 10#${{major}} < 10#${{required_major}} )) || + (( 10#${{major}} == 10#${{required_major}} && 10#${{minor}} < 10#${{required_minor}} )); then + printf 'Enroot %s.%s or newer is required for %s\\n' \\ + "${{required_major}}" "${{required_minor}}" "${{purpose}}" >&2 exit 78 fi }} @@ -188,11 +197,13 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: verify_sha256 "${{DD_INSPECTOR_SHA256}}" "${{DD_INSPECTOR}}" verify_sha256 "${{DD_ENROOT_RC_SHA256}}" "${{DD_ENROOT_RC}}" install -d -m 0700 \ + "${{DD_JOB_DIR}}/home" \ "${{DD_JOB_DIR}}/enroot/cache" \ "${{DD_JOB_DIR}}/enroot/config" \ "${{DD_JOB_DIR}}/enroot/data" \ "${{DD_JOB_DIR}}/enroot/tmp" \ "${{DD_INSPECTION_DIRECTORY}}" +export HOME="${{DD_JOB_DIR}}/home" export ENROOT_CACHE_PATH="${{DD_JOB_DIR}}/enroot/cache" export ENROOT_CONFIG_PATH="${{DD_JOB_DIR}}/enroot/config" export ENROOT_DATA_PATH="${{DD_JOB_DIR}}/enroot/data" diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index bc03ef174..7edd21431 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -86,6 +86,10 @@ def test_prepare_existing_sqsh_inspects_in_place_without_oci_identity(tmp_path: script = Path(prepared.script_file.path).read_text() assert "enroot import" not in script assert f'readonly DD_IMAGE_SQSH="{image_path.as_posix()}"' in script + assert 'verify_enroot_compatibility 3 5 "existing SQSH inspection"' in script + assert "enroot start --root" in script + assert "inspect_image.py:x-create=file,bind,ro" in script + assert "/output:x-create=dir,bind" in script def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_path: Path) -> None: @@ -111,19 +115,26 @@ def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_ assert 'enroot import -o "${DD_IMAGE_SQSH}" "${DD_OCI_SOURCE}"' in script assert "inspect_image.py:x-create=file,bind,ro" in script assert "/opt/data-designer-slurm/output:x-create=dir,bind" in script + assert 'export HOME="${DD_JOB_DIR}/home"' in script assert 'export ENROOT_CONFIG_PATH="${DD_JOB_DIR}/enroot/config"' in script assert 'export ENROOT_MAX_PROCESSORS="${SLURM_CPUS_PER_TASK}"' in script + assert script.index('export HOME="${DD_JOB_DIR}/home"') < script.rindex("\nverify_enroot_compatibility 4 0") assert 'version="$(enroot version)"' in script - assert "if (( 10#${major} < 4 )); then" in script + assert 'verify_enroot_compatibility 4 0 "digest-pinned OCI imports"' in script assert f'--mount "{prepared.plan.job_directory}:' not in script completed = subprocess.run(("bash", "-n"), input=script, capture_output=True, text=True, check=False) assert completed.returncode == 0, completed.stderr -@pytest.mark.parametrize("source_kind", ("oci", "existing"), ids=("oci-import", "existing-sqsh")) +@pytest.mark.parametrize( + ("source_kind", "enroot_version"), + (("oci", "4.0.0"), ("existing", "3.5.0")), + ids=("oci-import-enroot-4", "existing-sqsh-enroot-3.5"), +) def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( tmp_path: Path, source_kind: str, + enroot_version: str, ) -> None: if source_kind == "oci": source = f"registry.example.test/client@sha256:{'a' * 64}" @@ -142,11 +153,17 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( fake_enroot.write_text( "#!/usr/bin/env bash\n" "set -Eeuo pipefail\n" + ': "${HOME:?HOME must be set}"\n' + 'printf \'%s\\n\' "${HOME}" >> "${DD_TEST_HOME_LOG}"\n' + 'printf \'%s\\n\' "$1" >> "${DD_TEST_ENROOT_LOG}"\n' 'if [[ "$1" == "version" ]]; then\n' - " printf '%s\\n' '4.0.0'\n" + f" printf '%s\\n' '{enroot_version}'\n" 'elif [[ "$1" == "import" ]]; then\n' " printf '%s\\n' 'imported image' > \"$3\"\n" 'elif [[ "$1" == "start" ]]; then\n' + ' [[ " $* " == *" --root "* ]]\n' + ' [[ "$*" == *":x-create=file,bind,ro"* ]]\n' + ' [[ "$*" == *":x-create=dir,bind"* ]]\n' " printf '%s\\n' '{}' > \"${DD_TEST_INSPECTION_OUTPUT}\"\n" "fi\n" ) @@ -156,6 +173,8 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( f'export PATH="{fake_bin.as_posix()}:', 1, ) + home_log = tmp_path / "home.log" + enroot_log = tmp_path / "enroot.log" completed = subprocess.run( ("bash",), @@ -164,6 +183,8 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( text=True, check=False, env={ + "DD_TEST_ENROOT_LOG": enroot_log.as_posix(), + "DD_TEST_HOME_LOG": home_log.as_posix(), "DD_TEST_INSPECTION_OUTPUT": prepared.plan.inspection_output_path, "SLURM_CPUS_PER_TASK": "2", "SLURM_JOB_ID": "5101", @@ -171,15 +192,27 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( ) assert completed.returncode == 0, completed.stderr + expected_home = Path(prepared.plan.job_directory) / "home" + assert set(home_log.read_text().splitlines()) == {expected_home.as_posix()} + assert stat.S_IMODE(expected_home.stat().st_mode) == 0o700 assert Path(prepared.plan.inspection_output_path).read_text() == "{}\n" if source_kind == "oci": + assert enroot_log.read_text().splitlines() == ["version", "import", "create", "start", "remove"] assert Path(prepared.plan.sqsh_path).read_text() == "imported image\n" + else: + assert enroot_log.read_text().splitlines() == ["version", "create", "start", "remove"] -@pytest.mark.parametrize("source_kind", ("oci", "existing"), ids=("digest-import", "existing-sqsh")) -def test_rendered_image_lifecycle_rejects_enroot_before_version_4_before_image_operations( +@pytest.mark.parametrize( + ("source_kind", "enroot_version", "required_version"), + (("oci", "3.5.0", "4.0"), ("existing", "3.4.1", "3.5")), + ids=("digest-import-before-4", "existing-sqsh-before-3.5"), +) +def test_rendered_image_lifecycle_rejects_unsupported_enroot_before_image_operations( tmp_path: Path, source_kind: str, + enroot_version: str, + required_version: str, ) -> None: if source_kind == "oci": source = f"registry.example.test/client@sha256:{'a' * 64}" @@ -200,7 +233,7 @@ def test_rendered_image_lifecycle_rejects_enroot_before_version_4_before_image_o "set -Eeuo pipefail\n" 'printf \'%s\\n\' "$1" >> "${DD_TEST_ENROOT_LOG}"\n' 'if [[ "$1" == "version" ]]; then\n' - " printf '%s\\n' '3.5.0'\n" + f" printf '%s\\n' '{enroot_version}'\n" "fi\n" ) fake_enroot.chmod(0o700) @@ -225,7 +258,7 @@ def test_rendered_image_lifecycle_rejects_enroot_before_version_4_before_image_o ) assert completed.returncode == 78 - assert "Enroot 4 or newer" in completed.stderr + assert f"Enroot {required_version} or newer" in completed.stderr assert enroot_log.read_text() == "version\n"