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 a5c8c91c0..f3fc57987 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/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/inspection.py b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py index d8e1fbcfc..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 @@ -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,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, + find_unique_distribution, +) INSPECTOR_VERSION: Identifier = "inspector-1" _REQUIRED_CLIENT_DISTRIBUTIONS = ( @@ -49,11 +51,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 +93,14 @@ 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) - 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() + 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 + 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 new file mode 100644 index 000000000..cfb317f7b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -0,0 +1,347 @@ +# 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_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 + +_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) +_MINIMUM_ENROOT_OCI_VERSION = (4, 0) +_MINIMUM_ENROOT_SQSH_VERSION = (3, 5) + + +@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) + 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" + 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) + 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") + 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=(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], + ) + 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 + image_build = profile.image_build + 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", 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"), + ) + ) + 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 {minimum_major} {minimum_minor} "digest-pinned OCI imports" +enroot import -o "${{DD_IMAGE_SQSH}}" "${{DD_OCI_SOURCE}}" +""" + else: + 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} +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_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)} +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_enroot_compatibility() {{ + 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]}}" + 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 +}} + +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" +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 + exit 66 +fi +{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 +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_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 + 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/output/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(), + ), + ) + 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 actual_content != expected_content + ): + raise ImageLifecycleError(f"prepared image lifecycle {label} no longer matches its digest") + 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: + 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 _read_regular_file(path: Path) -> bytes: + descriptor: int | None = None + try: + 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") + 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 4e6589ec1..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 @@ -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,96 @@ ) +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") + 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 = ( + (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 + + +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/__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..01dc0f6b6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py @@ -0,0 +1,181 @@ +# 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 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]: + 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_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: + 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": installer_path, + "installer_version": installer_distribution.version, + } + + +def _inspect_serving() -> dict[str, object]: + distribution = find_unique_distribution(importlib.metadata.distributions(), "vllm") + executable_path = find_distribution_console_script(distribution, "vllm") + return { + "kind": "serving", + "server_type": "vllm", + "runtime_version": distribution.version, + "executable_path": executable_path, + } + + +def _list_distribution_versions(distributions: Iterable[importlib.metadata.Distribution]) -> dict[str, str]: + versions: dict[str, str] = {} + for distribution in distributions: + 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 _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_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( + 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/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/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/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index e8941df39..e076f5cb1 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), @@ -113,21 +90,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/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 1468aac49..4d257a8fb 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 @@ -394,7 +394,7 @@ "schema_version": 1, "selected_profile": { "catalog_path": "/workspace/profile.json", - "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "catalog_sha256": "a707e8e3b6f52170554b743d252b90273fa0deeea8141d4a87abe96e2e93215d", "cluster_name": "primary", "matched_pattern": null, "profile": { @@ -412,7 +412,10 @@ "primary-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" }, "scheduler": { "account": "research", @@ -422,7 +425,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 b2747d4c0..19b781d06 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 @@ -240,7 +240,7 @@ "schema_version": 1, "selected_profile": { "catalog_path": "/workspace/profile.json", - "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "catalog_sha256": "a707e8e3b6f52170554b743d252b90273fa0deeea8141d4a87abe96e2e93215d", "cluster_name": "primary", "matched_pattern": null, "profile": { @@ -258,7 +258,10 @@ "primary-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" }, "scheduler": { "account": "research", @@ -268,7 +271,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 dcf7addda..c0c67c00a 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..c869d6c0c 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( @@ -205,25 +205,74 @@ 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, - "version", - 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_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, "distributions", 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( + 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, "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: 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..7edd21431 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -0,0 +1,663 @@ +# 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 collections.abc import Sequence +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 / "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 + 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((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 + + +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 + 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: + 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), + _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 --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 + 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 '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", "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}" + 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" + ': "${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' + 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" + ) + fake_enroot.chmod(0o700) + script = render_image_lifecycle_script(prepared.plan).replace( + 'export PATH="', + f'export PATH="{fake_bin.as_posix()}:', + 1, + ) + home_log = tmp_path / "home.log" + 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(), + "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", + }, + ) + + 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", "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}" + 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' + f" printf '%s\\n' '{enroot_version}'\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 f"Enroot {required_version} or newer" in completed.stderr + assert enroot_log.read_text() == "version\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() + + +@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"): + prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-unsafe") + + 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() + 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", "output 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", + ) + ] + assert runner.inputs == [render_image_lifecycle_script(prepared.plan)] + + +@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_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", + ) + 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" + + 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 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 + + +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( + metadata={"Name": "vllm"}, + 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, "distributions", return_value=(distribution,)), + ): + 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( + metadata={"Name": "vllm"}, + version="0.21.0", + entry_points=entry_points, + files=files, + locate_file=lambda installed_file: installed_file, + ) + + with ( + patch.object(resource_inspector.importlib.metadata, "distributions", return_value=(distribution,)), + pytest.raises(RuntimeError, match=match), + ): + resource_inspector.inspect_image("serving", "d" * 64) + + +@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", + 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 4568b5566..c0e9887b2 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": "14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" + "sha256": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" }, "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": "14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" + "sha256": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" }, "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 fc36170e8..3f75a41b6 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="cd95eb995dd9705d12336b9760ca1f354cdd7614dce992cc4f92a6faa47dbfdc" +readonly DD_PLAN_SHA256="902919292da35aca426a191bee775acf140b0f8b7440b55887fb2ba60ed98a52" 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 e3f8529d5..553fadd86 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" +readonly DD_PLAN_SHA256="cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" 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 9880d08e8..621d4a33e 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="65032718bfc8fbd6c60700add08ebc098ba2e64858c23808decbefe5a8153d91", + expected_fixture_sha256="8238dd22393a46538169a83046e91230166aaa0f96e0a91255adf6a6074bfa20", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="58c3a2eddc9a0d377eb0a90655e70a6b2484cade0c52b2a01b9b388dcd4667e2", + expected_fixture_sha256="6153c1769d9b90315c4e7acc767bd5c469d0a5e1f5fc7843baddc1a733cb8e9b", )