From d93ead6c5947150308e2704a68158078194e6c52 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 10 Sep 2026 22:28:04 -0300 Subject: [PATCH 1/3] feat(slurm): add profile setup and image lifecycle wiring --- .../src/data_designer/slurm/cli.py | 52 +++ .../data_designer/slurm/services/__init__.py | 12 + .../data_designer/slurm/services/errors.py | 2 + .../slurm/services/image_lifecycle.py | 199 +++++++++ .../data_designer/slurm/services/images.py | 9 +- .../data_designer/slurm/services/profiles.py | 392 ++++++++++++++++++ .../data_designer/slurm/services/wiring.py | 25 +- .../tests/services/test_image_wiring.py | 359 ++++++++++++++++ .../tests/services/test_profile_service.py | 301 ++++++++++++++ .../tests/services/test_wiring.py | 13 +- .../data-designer-slurm/tests/test_cli.py | 89 +++- scripts/test_slurm_package_install.py | 25 ++ 12 files changed, 1456 insertions(+), 22 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py create mode 100644 packages/data-designer-slurm/tests/services/test_image_wiring.py create mode 100644 packages/data-designer-slurm/tests/services/test_profile_service.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index f2c9e12fd..ebe7de55e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -14,11 +14,13 @@ from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config from data_designer.slurm.contracts import canonical_json +from data_designer.slurm.images.records import validate_oci_source_for_lifecycle from data_designer.slurm.services import ( SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, create_slurm_image_service, + create_slurm_profile_service, create_slurm_run_service, ) @@ -37,7 +39,9 @@ no_args_is_help=True, ) image_app = typer.Typer(help="Manage verified Slurm images", no_args_is_help=True) +profile_app = typer.Typer(help="Initialize and validate Slurm profiles", no_args_is_help=True) app.add_typer(image_app, name="image") +app.add_typer(profile_app, name="profile") @app.callback() @@ -94,6 +98,46 @@ def cancel_command( _emit_result(result) +@profile_app.command("init") +def profile_init_command( + workspace_root: Path = typer.Option(..., "--workspace-root", file_okay=False), + image_build_partition: str = typer.Option(..., "--image-build-partition"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str = typer.Option("default", "--cluster"), + account: str | None = typer.Option(None, "--account"), + partition: str | None = typer.Option(None, "--partition"), + host_pattern: list[str] | None = typer.Option(None, "--host-pattern"), +) -> None: + """Create a safe portable starter profile without overwriting.""" + operation = SlurmServiceOperation.INIT_PROFILE + result = _invoke( + operation, + lambda: create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=workspace_root, + image_build_partition=image_build_partition, + cluster=cluster, + account=account, + partition=partition, + host_patterns=tuple(host_pattern or ()), + ), + ) + _emit_result(result) + + +@profile_app.command("validate") +def profile_validate_command( + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Validate strict loading, cluster selection, workspace, and Slurm facts.""" + operation = SlurmServiceOperation.VALIDATE_PROFILE + result = _invoke( + operation, + lambda: create_slurm_profile_service(profile_file=profile_file, cluster=cluster).validate(), + ) + _emit_result(result) + + @image_app.command("add") def image_add_command( source: str = typer.Argument(...), @@ -113,6 +157,14 @@ def add() -> BaseModel: operation, "OCI image source must be digest-qualified as name@sha256:", ) + try: + validate_oci_source_for_lifecycle(source) + except ValueError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "OCI image source must be a credential-free registry reference without a scheme", + ) from None request = ImageBuildRequest(name=name or _derive_image_name(source), kind=kind, source=source) return create_slurm_image_service(profile_file=profile_file, cluster=cluster).add(request, replace=replace) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py index 22b603c7a..5c0ae9c7e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py @@ -15,6 +15,13 @@ SlurmServiceOperation, ) from data_designer.slurm.services.images import SlurmImageManager, SlurmImageResolver, SlurmImageService +from data_designer.slurm.services.profiles import ( + SlurmProfileInitialization, + SlurmProfileMatch, + SlurmProfileService, + SlurmProfileValidation, + create_slurm_profile_service, +) from data_designer.slurm.services.results import ( SlurmPersistedAttemptStatus, SlurmPersistedRunStatus, @@ -47,6 +54,10 @@ "SlurmPersistedAttemptStatus", "SlurmPersistedRunStatus", "SlurmPersistedShardStatus", + "SlurmProfileInitialization", + "SlurmProfileMatch", + "SlurmProfileService", + "SlurmProfileValidation", "SlurmRunArtifactPublisher", "SlurmRunBackend", "SlurmRunCancellation", @@ -57,6 +68,7 @@ "SlurmServiceErrorCode", "SlurmServiceOperation", "create_slurm_image_service", + "create_slurm_profile_service", "create_slurm_run_service", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py index aa08fffed..7e7e19814 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py @@ -32,6 +32,8 @@ class SlurmServiceOperation(str, Enum): EXECUTE_RUN = "execute_run" STATUS_RUN = "status_run" CANCEL_RUN = "cancel_run" + INIT_PROFILE = "init_profile" + VALIDATE_PROFILE = "validate_profile" RESOLVE_IMAGE = "resolve_image" ADD_IMAGE = "add_image" LIST_IMAGES = "list_images" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py new file mode 100644 index 000000000..e864bf0ec --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Production composition for synchronous Slurm image lifecycle jobs.""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Sequence +from datetime import datetime, timedelta, timezone +from uuid import uuid4 + +from data_designer.slurm.config import ImageBuildRequest, SelectedSlurmProfile +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.images.errors import ( + ImageConflictError, + ImageLifecycleError, + ImageRegistryError, + ImageVerificationError, +) +from data_designer.slurm.images.lifecycle import ( + PreparedImageLifecycleJob, + cleanup_prepared_image_lifecycle, + prepare_image_lifecycle_job, + publish_completed_image_lifecycle, + submit_prepared_image_lifecycle, +) +from data_designer.slurm.images.records import RegisteredImage +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError +from data_designer.slurm.launcher.models import SlurmAccountingEntry, SlurmQueueEntry +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation +from data_designer.slurm.state import ( + SchedulerJobIdentity, + SchedulerObservation, + SchedulerObservationCollector, + SchedulerState, + SlurmStateError, +) +from data_designer.slurm.state.scheduler import is_scheduler_failure_state + +LifecycleIdFactory = Callable[[], str] +Clock = Callable[[], datetime] +Sleeper = Callable[[float], None] + +_POLL_INTERVAL_SECONDS = 300.0 +_ACCOUNTING_EXIT_LAG = timedelta(minutes=5) + + +class _TerminalLifecycleError(RuntimeError): + pass + + +class _RecordingObservationClient: + def __init__(self, launcher: SlurmCommandClient) -> None: + self._launcher = launcher + self.accounting: tuple[SlurmAccountingEntry, ...] = () + + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + return self._launcher.query_queue(selectors) + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + self.accounting = self._launcher.query_accounting(selectors) + return self.accounting + + +class SlurmImageLifecycleManager: + """Prepare, submit, reconcile, publish, and register one image.""" + + def __init__( + self, + selected_profile: SelectedSlurmProfile, + launcher: SlurmCommandClient, + *, + lifecycle_id_factory: LifecycleIdFactory | None = None, + clock: Clock | None = None, + sleep: Sleeper | None = None, + ) -> None: + self._profile = selected_profile + self._launcher = launcher + self._lifecycle_id_factory = lifecycle_id_factory or _new_lifecycle_id + self._clock = clock or _utc_now + self._sleep = sleep or time.sleep + + def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: + """Run one lifecycle job and publish only a successful verified result.""" + operation = SlurmServiceOperation.ADD_IMAGE + try: + prepared = prepare_image_lifecycle_job( + request, + self._profile, + lifecycle_id=self._lifecycle_id_factory(), + ) + except (ImageLifecycleError, OSError, ValueError): + raise _unavailable("image lifecycle job cannot be prepared") from None + + try: + receipt = submit_prepared_image_lifecycle(prepared, self._launcher) + except SlurmSubmissionError as error: + if not error.may_have_succeeded: + _cleanup_failed_lifecycle(prepared) + raise _unavailable("image lifecycle job cannot be submitted") from None + raise _unavailable( + f"image lifecycle submission outcome is unknown; lifecycle {prepared.plan.lifecycle_id} was retained" + ) from None + except (ImageLifecycleError, SlurmLauncherError, OSError, ValueError): + _cleanup_failed_lifecycle(prepared) + raise _unavailable("image lifecycle job cannot be submitted") from None + + try: + self._wait_for_success(receipt.job_id) + return publish_completed_image_lifecycle(prepared, replace=replace) + except (ImageConflictError, ImageVerificationError): + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + operation, + "image lifecycle result cannot be verified or registered", + ) from None + except ImageRegistryError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + "image registry operation failed", + ) from None + except ImageLifecycleError: + raise _unavailable("image lifecycle result cannot be published") from None + except _TerminalLifecycleError: + _cleanup_failed_lifecycle(prepared) + raise _unavailable(f"image lifecycle job {receipt.job_id} did not complete successfully") from None + except (SlurmLauncherError, SlurmStateError, OSError, ValueError): + self._cancel_and_cleanup(receipt.job_id, prepared) + raise _unavailable(f"image lifecycle job {receipt.job_id} did not complete successfully") from None + except BaseException: + self._cancel_and_cleanup(receipt.job_id, prepared) + raise + + def _wait_for_success(self, job_id: int) -> None: + client = _RecordingObservationClient(self._launcher) + observations = SchedulerObservationCollector(client) + previous: SchedulerObservation | None = None + completed_without_accounting_deadline: datetime | None = None + while True: + observed_at = self._clock() + observation = observations.collect( + (job_id,), + observed_at=observed_at, + previous={job_id: previous}, + )[0] + if observation.state is SchedulerState.COMPLETED: + accounting = next((entry for entry in client.accounting if entry.job_identity == job_id), None) + if accounting is None or accounting.state is not SchedulerState.COMPLETED: + if completed_without_accounting_deadline is None: + completed_without_accounting_deadline = observed_at + _ACCOUNTING_EXIT_LAG + elif observed_at > completed_without_accounting_deadline: + raise SlurmStateError("completed image lifecycle job has no exit evidence") + previous = observation + self._sleep(_POLL_INTERVAL_SECONDS) + continue + if ( + accounting.process_exit_code.exit_status != 0 + or accounting.process_exit_code.termination_signal != 0 + ): + raise _TerminalLifecycleError("completed image lifecycle job has no successful exit evidence") + return + if is_scheduler_failure_state(observation.state): + raise _TerminalLifecycleError("image lifecycle job did not complete successfully") + if observation.state is SchedulerState.UNKNOWN: + raise SlurmStateError("image lifecycle job has unknown scheduler state") + previous = observation + self._sleep(_POLL_INTERVAL_SECONDS) + + def _cancel_and_cleanup(self, job_id: int, prepared: PreparedImageLifecycleJob) -> None: + try: + self._launcher.cancel(job_id) + except (SlurmLauncherError, OSError, ValueError): + return + _cleanup_failed_lifecycle(prepared) + + +def _cleanup_failed_lifecycle(prepared: PreparedImageLifecycleJob) -> None: + try: + cleanup_prepared_image_lifecycle(prepared) + except ImageLifecycleError: + pass + + +def _unavailable(message: str) -> SlurmServiceError: + return SlurmServiceError(SlurmServiceErrorCode.UNAVAILABLE, SlurmServiceOperation.ADD_IMAGE, message) + + +def _new_lifecycle_id() -> Identifier: + return f"image-{uuid4().hex}" + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +__all__ = ["Clock", "LifecycleIdFactory", "Sleeper", "SlurmImageLifecycleManager"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/images.py b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py index 9a2a12253..531c468c5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/images.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py @@ -12,7 +12,7 @@ from data_designer.slurm.config import ImageKind, ImageRef from data_designer.slurm.config.images import ImageBuildRequest from data_designer.slurm.contracts import Identifier -from data_designer.slurm.images.records import RegisteredImage +from data_designer.slurm.images.records import RegisteredImage, validate_oci_source_for_lifecycle from data_designer.slurm.planning import ResolvedImage from data_designer.slurm.services.errors import ( SlurmServiceError, @@ -98,6 +98,13 @@ def add(self, request: ImageBuildRequest, *, replace: bool = False) -> Registere raise _make_invalid_request_error(operation, "request must be an ImageBuildRequest") if type(replace) is not bool: raise _make_invalid_request_error(operation, "replace must be a boolean") + try: + validate_oci_source_for_lifecycle(request.source) + except ValueError: + raise _make_invalid_request_error( + operation, + "OCI image source must be a credential-free registry reference without a scheme", + ) from None manager = self._require_manager(operation) def add_image() -> RegisteredImage: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py new file mode 100644 index 000000000..402834bc5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Package-owned Slurm profile initialization and validation.""" + +from __future__ import annotations + +import json +import os +import shlex +import socket +from collections.abc import Callable, Mapping +from fnmatch import fnmatchcase +from pathlib import Path + +import yaml +from pydantic import PositiveInt, ValidationError + +from data_designer.slurm.config import ( + DEFAULT_PROFILE_FILE_NAME, + PROFILE_FILE_ENVIRONMENT, + ProfileSelectionSource, + SlurmConfigLoadError, + SlurmProfile, + SlurmProfileCatalog, + load_profile_catalog, + resolve_profile, +) +from data_designer.slurm.contracts import ContractValue, Identifier, compute_canonical_json_sha256 +from data_designer.slurm.filesystem import create_restrictive_temporary_file, open_verified_directory +from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation + +HostnameResolver = Callable[[], tuple[str, ...]] + +_IMAGE_BUILD_CPUS = 2 +_IMAGE_BUILD_MEMORY = "8G" +_IMAGE_BUILD_TIME_LIMIT = "04:00:00" + + +class SlurmProfileInitialization(ContractValue): + """One newly created profile catalog.""" + + profile_file: str + validation_command: str + + +class SlurmProfileMatch(ContractValue): + """One cluster and its matching hostname patterns.""" + + cluster: Identifier + patterns: tuple[str, ...] + + +class SlurmProfileValidation(ContractValue): + """Effective local profile selection and workspace facts.""" + + profile_file: str | None + hostnames: tuple[str, ...] + matched_clusters: tuple[SlurmProfileMatch, ...] + default_cluster: Identifier | None + selected_cluster: Identifier | None + selection_source: ProfileSelectionSource + matched_pattern: str | None + workspace_root: str + image_root: str + registry_file: str + gpus_per_node: PositiveInt + + +class SlurmProfileService: + """Initialize and validate profiles through production loaders and selectors.""" + + def __init__( + self, + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + launcher: SlurmCommandClient | None = None, + hostname_resolver: HostnameResolver | None = None, + environ: Mapping[str, str] | None = None, + home_directory: str | Path | None = None, + ) -> None: + self._profile = profile + self._catalog = catalog + self._profile_file = profile_file + self._cluster = cluster + self._launcher = launcher or SlurmCommandClient() + self._hostname_resolver = hostname_resolver or _local_hostnames + self._environ = dict(os.environ if environ is None else environ) + self._home_directory = home_directory + + def initialize( + self, + *, + workspace_root: str | Path, + image_build_partition: str, + cluster: str = "default", + account: str | None = None, + partition: str | None = None, + host_patterns: tuple[str, ...] = (), + ) -> SlurmProfileInitialization: + """Create one deterministic starter catalog without overwriting a file.""" + operation = SlurmServiceOperation.INIT_PROFILE + try: + path = _resolve_profile_path( + self._profile_file, + environ=self._environ, + home_directory=self._home_directory, + ) + selected_patterns = host_patterns or _normalize_hostnames(self._hostname_resolver()) + payload = _starter_catalog_payload( + workspace_root=Path(workspace_root).expanduser().resolve(), + image_build_partition=image_build_partition, + cluster=cluster, + account=account, + partition=partition, + host_patterns=selected_patterns, + ) + SlurmProfileCatalog.model_validate(payload, strict=True) + _create_profile_file(path, _serialize_catalog(payload, suffix=path.suffix)) + except FileExistsError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + operation, + "profile file already exists", + ) from None + except FileNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "profile destination parent does not exist", + ) from None + except (SlurmConfigLoadError, ValidationError, ValueError, TypeError): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "profile initialization input is invalid", + ) from None + except OSError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "profile file cannot be created", + ) from None + command = shlex.join(("data-designer", "slurm", "profile", "validate", "--profile-file", path.as_posix())) + return SlurmProfileInitialization(profile_file=path.as_posix(), validation_command=command) + + def validate(self) -> SlurmProfileValidation: + """Load, select, and verify one effective production profile.""" + operation = SlurmServiceOperation.VALIDATE_PROFILE + try: + hostnames = _normalize_hostnames(self._hostname_resolver()) + selected = resolve_profile( + profile=self._profile, + catalog=self._catalog, + profile_file=self._profile_file, + cluster=self._cluster, + hostnames=hostnames, + environ=self._environ, + home_directory=self._home_directory, + ) + catalog = self._catalog + if selected.catalog_path is not None: + catalog = load_profile_catalog(selected.catalog_path) + if compute_canonical_json_sha256(catalog.model_dump(mode="json")) != selected.catalog_sha256: + raise SlurmConfigLoadError("profile catalog changed while it was being validated") + _validate_workspace(Path(selected.profile.workspace_root)) + gpus_per_node = self._resolve_gpu_count(selected.profile) + except SlurmServiceError: + raise + except (SlurmConfigLoadError, ValidationError, ValueError, TypeError): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "profile configuration cannot be resolved", + ) from None + except SlurmLauncherError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm is unavailable", + ) from None + except OSError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "profile workspace is unavailable", + ) from None + + store = ImageRegistryStore(selected.profile.workspace_root) + return SlurmProfileValidation( + profile_file=selected.catalog_path, + hostnames=hostnames, + matched_clusters=_matching_clusters(catalog, hostnames), + default_cluster=None if catalog is None else catalog.default_cluster, + selected_cluster=selected.cluster_name, + selection_source=selected.selection_source, + matched_pattern=selected.matched_pattern, + workspace_root=selected.profile.workspace_root, + image_root=store.image_root.as_posix(), + registry_file=store.registry_path.as_posix(), + gpus_per_node=gpus_per_node, + ) + + def _resolve_gpu_count(self, profile: SlurmProfile) -> int: + if profile.gpus_per_node != "auto": + return profile.gpus_per_node + counts = tuple(sorted(set(self._launcher.query_gpu_counts(partition=profile.scheduler.partition)))) + if len(counts) != 1: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.VALIDATE_PROFILE, + "eligible Slurm nodes do not report one GPU count", + ) + return counts[0] + + +def create_slurm_profile_service( + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + launcher: SlurmCommandClient | None = None, + hostname_resolver: HostnameResolver | None = None, + environ: Mapping[str, str] | None = None, + home_directory: str | Path | None = None, +) -> SlurmProfileService: + """Create the package-owned profile service.""" + return SlurmProfileService( + profile=profile, + catalog=catalog, + profile_file=profile_file, + cluster=cluster, + launcher=launcher, + hostname_resolver=hostname_resolver, + environ=environ, + home_directory=home_directory, + ) + + +def _starter_catalog_payload( + *, + workspace_root: Path, + image_build_partition: str, + cluster: str, + account: str | None, + partition: str | None, + host_patterns: tuple[str, ...], +) -> dict[str, object]: + scheduler = {key: value for key, value in (("account", account), ("partition", partition)) if value is not None} + profile: dict[str, object] = { + "schema_version": 1, + "host_patterns": list(host_patterns), + "gpus_per_node": "auto", + "workspace_root": workspace_root.as_posix(), + "image_build": { + "partition": image_build_partition, + "cpus_per_task": _IMAGE_BUILD_CPUS, + "memory": _IMAGE_BUILD_MEMORY, + "time_limit": _IMAGE_BUILD_TIME_LIMIT, + }, + "gpu_request_mode": "gres", + } + if scheduler: + profile["scheduler"] = scheduler + return { + "schema_version": 1, + "default_cluster": cluster, + "clusters": {cluster: profile}, + } + + +def _serialize_catalog(payload: dict[str, object], *, suffix: str) -> bytes: + if suffix == ".json": + return (json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n").encode() + return yaml.safe_dump(payload, allow_unicode=True, sort_keys=False).encode() + + +def _create_profile_file(path: Path, content: bytes) -> None: + with open_verified_directory(path.parent, resource_name="profile") as parent_descriptor: + descriptor, temporary_name = create_restrictive_temporary_file( + parent_descriptor, + prefix=f".{path.name}.", + suffix=".tmp", + ) + try: + with os.fdopen(descriptor, "wb") as output: + descriptor = -1 + output.write(content) + output.flush() + os.fsync(output.fileno()) + os.link( + temporary_name, + path.name, + src_dir_fd=parent_descriptor, + dst_dir_fd=parent_descriptor, + follow_symlinks=False, + ) + os.fsync(parent_descriptor) + finally: + if descriptor >= 0: + os.close(descriptor) + try: + os.unlink(temporary_name, dir_fd=parent_descriptor) + os.fsync(parent_descriptor) + except FileNotFoundError: + pass + + +def _validate_workspace(path: Path) -> None: + with open_verified_directory(path, resource_name="profile workspace") as descriptor: + temporary_descriptor, temporary_name = create_restrictive_temporary_file( + descriptor, + prefix=".data-designer-profile-validation-", + suffix=".tmp", + ) + try: + os.close(temporary_descriptor) + temporary_descriptor = -1 + os.unlink(temporary_name, dir_fd=descriptor) + os.fsync(descriptor) + finally: + if temporary_descriptor >= 0: + os.close(temporary_descriptor) + try: + os.unlink(temporary_name, dir_fd=descriptor) + except FileNotFoundError: + pass + + +def _matching_clusters( + catalog: SlurmProfileCatalog | None, + hostnames: tuple[str, ...], +) -> tuple[SlurmProfileMatch, ...]: + if catalog is None: + return () + matches = [] + for cluster, profile in sorted(catalog.clusters.items()): + patterns = tuple( + sorted( + pattern + for pattern in profile.host_patterns + if any(fnmatchcase(hostname, pattern.casefold()) for hostname in hostnames) + ) + ) + if patterns: + matches.append(SlurmProfileMatch(cluster=cluster, patterns=patterns)) + return tuple(matches) + + +def _resolve_profile_path( + explicit_path: str | Path | None, + *, + environ: Mapping[str, str], + home_directory: str | Path | None, +) -> Path: + source = explicit_path + if source is None: + source = environ.get(PROFILE_FILE_ENVIRONMENT) + if source == "": + raise SlurmConfigLoadError(f"{PROFILE_FILE_ENVIRONMENT} must not be empty") + if source is None: + home = Path.home() if home_directory is None else Path(home_directory) + source = home / DEFAULT_PROFILE_FILE_NAME + path = Path(source).expanduser().resolve() + if path.suffix not in {".json", ".yaml", ".yml"}: + raise SlurmConfigLoadError("configuration path must end in .json, .yaml, or .yml") + return path + + +def _normalize_hostnames(hostnames: tuple[str, ...]) -> tuple[str, ...]: + return tuple(dict.fromkeys(hostname.strip().casefold() for hostname in hostnames if hostname.strip())) + + +def _local_hostnames() -> tuple[str, ...]: + return socket.gethostname(), socket.getfqdn() + + +__all__ = [ + "HostnameResolver", + "SlurmProfileInitialization", + "SlurmProfileMatch", + "SlurmProfileService", + "SlurmProfileValidation", + "create_slurm_profile_service", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index b13834f56..f48757534 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -52,6 +52,7 @@ from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.services.artifacts import StateRunArtifactPublisher from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation +from data_designer.slurm.services.image_lifecycle import SlurmImageLifecycleManager from data_designer.slurm.services.images import SlurmImageService from data_designer.slurm.services.results import ( SlurmPersistedAttemptStatus, @@ -624,9 +625,10 @@ def cancel(self, run_id: Identifier) -> SlurmRunCancellation: class _RegistryImageBackend: - def __init__(self, workspace_root: str) -> None: + def __init__(self, workspace_root: str, lifecycle: SlurmImageLifecycleManager) -> None: self._verified = VerifiedImageRegistry(workspace_root) self._store = ImageRegistryStore(workspace_root) + self._lifecycle = lifecycle def resolve(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedImage: try: @@ -645,12 +647,7 @@ def resolve(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedI ) from None def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: - del request, replace - raise SlurmServiceError( - SlurmServiceErrorCode.UNAVAILABLE, - SlurmServiceOperation.ADD_IMAGE, - "image registration is not available; use a pre-registered image", - ) + return self._lifecycle.add(request, replace=replace) def list(self) -> tuple[RegisteredImage, ...]: return self._invoke_registry(SlurmServiceOperation.LIST_IMAGES, self._store.list_images) @@ -718,10 +715,22 @@ def create_slurm_image_service( catalog: SlurmProfileCatalog | None = None, profile_file: str | Path | None = None, cluster: str | None = None, + launcher: SlurmCommandClient | None = None, + lifecycle_id_factory: Callable[[], str] | None = None, + clock: Clock | None = None, + sleep: Callable[[float], None] | None = None, ) -> SlurmImageService: """Create the production image service for one selected cluster profile.""" selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) - backend = _RegistryImageBackend(selected.profile.workspace_root) + command_client = launcher or SlurmCommandClient() + lifecycle = SlurmImageLifecycleManager( + selected, + command_client, + lifecycle_id_factory=lifecycle_id_factory, + clock=clock, + sleep=sleep, + ) + backend = _RegistryImageBackend(selected.profile.workspace_root, lifecycle) return SlurmImageService(backend, backend) diff --git a/packages/data-designer-slurm/tests/services/test_image_wiring.py b/packages/data-designer-slurm/tests/services/test_image_wiring.py new file mode 100644 index 000000000..23ea118a9 --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -0,0 +1,359 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from collections.abc import Callable +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from slurm_test_fakes import FakeClock + +from data_designer.slurm.config import ( + ClientImageInspection, + ImageBuildProfile, + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, + SchedulerProfile, + SlurmProfile, +) +from data_designer.slurm.images.service import VerifiedImageRegistry +from data_designer.slurm.launcher.errors import SlurmSubmissionError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmProcessExitCode, + SlurmQueueEntry, +) +from data_designer.slurm.services import SlurmServiceError, SlurmServiceErrorCode, create_slurm_image_service +from data_designer.slurm.state import SchedulerState + +_JOB_ID = 42 +_LIFECYCLE_ID = "image-test" + + +class _Launcher: + def __init__( + self, + *, + queue_state: SchedulerState | None = None, + accounting_state: SchedulerState | None = SchedulerState.COMPLETED, + exit_status: int = 0, + termination_signal: int = 0, + on_submit: Callable[[], None] | None = None, + submission_error: SlurmSubmissionError | None = None, + ) -> None: + self.queue_state = queue_state + self.accounting_state = accounting_state + self.exit_status = exit_status + self.termination_signal = termination_signal + self.on_submit = on_submit + self.submission_error = submission_error + self.submissions: list[str] = [] + self.cancellations: list[int] = [] + + def submit_script(self, script: str, **_: object) -> SlurmJobSubmissionReceipt: + self.submissions.append(script) + if self.submission_error is not None: + raise self.submission_error + if self.on_submit is not None: + self.on_submit() + return SlurmJobSubmissionReceipt(job_id=_JOB_ID) + + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + del selectors + if self.queue_state is None: + return () + return (SlurmQueueEntry(job_identity=_JOB_ID, state=self.queue_state),) + + def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ...]: + del selectors + if self.accounting_state is None: + return () + return ( + SlurmAccountingEntry( + job_identity=_JOB_ID, + state=self.accounting_state, + process_exit_code=SlurmProcessExitCode( + exit_status=self.exit_status, + termination_signal=self.termination_signal, + ), + ), + ) + + def cancel(self, job_id: int) -> None: + self.cancellations.append(job_id) + + +class _DelayedAccountingLauncher(_Launcher): + def __init__(self, *, on_submit: Callable[[], None]) -> None: + super().__init__(queue_state=SchedulerState.COMPLETED, on_submit=on_submit) + self.queue_queries = 0 + self.accounting_queries = 0 + + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + self.queue_queries += 1 + if self.queue_queries == 1: + return super().query_queue(selectors) + return () + + def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ...]: + self.accounting_queries += 1 + if self.accounting_queries < 3: + return ( + SlurmAccountingEntry( + job_identity=_JOB_ID, + state=SchedulerState.RUNNING, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=0), + ), + ) + return super().query_accounting(selectors) + + +def test_default_image_add_runs_lifecycle_and_registers_existing_sqsh(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + launcher = _Launcher(on_submit=lambda: _write_inspection(workspace, content)) + service = _service(workspace, launcher) + + registered = service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert registered.path == source.as_posix() + assert registered.sqsh_sha256 == hashlib.sha256(content).hexdigest() + assert service.get("client") == registered + assert ( + VerifiedImageRegistry(workspace) + .resolve_for_planning( + ImageRef(name="client"), + expected_kind=ImageKind.CLIENT, + ) + .path + == source.as_posix() + ) + assert len(launcher.submissions) == 1 + assert "#SBATCH --partition=cpu" in launcher.submissions[0] + assert not _job_directory(workspace).exists() + + +def test_default_image_add_waits_for_completed_job_accounting_evidence(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + launcher = _DelayedAccountingLauncher(on_submit=lambda: _write_inspection(workspace, content)) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + + registered = _service(workspace, launcher, clock=clock).add( + ImageBuildRequest(name="client", kind="client", source=source.as_posix()) + ) + + assert registered.path == source.as_posix() + assert launcher.queue_queries == 3 + assert launcher.accounting_queries == 3 + assert clock.sleep_calls == [300.0, 300.0] + assert not _job_directory(workspace).exists() + + +def test_default_image_add_preserves_collision_then_replaces_explicitly(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + first = tmp_path / "first.sqsh" + second = tmp_path / "second.sqsh" + first.write_bytes(b"first") + second.write_bytes(b"second") + current_content = b"first" + launcher = _Launcher(on_submit=lambda: _write_inspection(workspace, current_content)) + service = _service(workspace, launcher) + original = service.add(ImageBuildRequest(name="client", kind="client", source=first.as_posix())) + current_content = b"second" + + with pytest.raises(SlurmServiceError) as caught: + service.add(ImageBuildRequest(name="client", kind="client", source=second.as_posix())) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert service.get("client") == original + + replacement = service.add( + ImageBuildRequest(name="client", kind="client", source=second.as_posix()), + replace=True, + ) + + assert replacement.path == second.as_posix() + assert service.get("client") == replacement + assert not _job_directory(workspace).exists() + + +@pytest.mark.parametrize( + ("state", "exit_status"), + ((SchedulerState.FAILED, 1), (SchedulerState.COMPLETED, 1)), + ids=("terminal-failure", "nonzero-exit"), +) +def test_default_image_add_fails_closed_and_cleans_terminal_jobs( + tmp_path: Path, + state: SchedulerState, + exit_status: int, +) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher(accounting_state=state, exit_status=exit_status) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher).add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert not _job_directory(workspace).exists() + assert not (workspace / "images" / "registry.yaml").exists() + + +def test_default_image_add_cleans_definitive_submission_failure(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher( + submission_error=SlurmSubmissionError("rejected", may_have_succeeded=False), + ) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher).add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert not _job_directory(workspace).exists() + + +def test_default_image_add_retains_ambiguous_submission_state(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher( + submission_error=SlurmSubmissionError("unknown", may_have_succeeded=True), + ) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher).add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert str(caught.value) == ("image lifecycle submission outcome is unknown; lifecycle image-test was retained") + assert _job_directory(workspace).is_dir() + + +def test_default_image_add_cancels_unknown_job_before_cleanup(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher(accounting_state=None) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher, clock=clock).add( + ImageBuildRequest(name="client", kind="client", source=source.as_posix()) + ) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert launcher.cancellations == [_JOB_ID] + assert clock.sleep_calls == [300.0, 300.0] + assert not _job_directory(workspace).exists() + + +def test_default_image_add_cancels_completed_job_without_exit_evidence(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher(queue_state=SchedulerState.COMPLETED, accounting_state=None) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher, clock=clock).add( + ImageBuildRequest(name="client", kind="client", source=source.as_posix()) + ) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert launcher.cancellations == [_JOB_ID] + assert clock.sleep_calls == [300.0, 300.0] + assert not _job_directory(workspace).exists() + + +def test_default_image_add_cancels_and_cleans_on_interrupt(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher(queue_state=SchedulerState.PENDING, accounting_state=None) + + def interrupt(_: float) -> None: + raise KeyboardInterrupt + + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=launcher, # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + sleep=interrupt, + ) + + with pytest.raises(KeyboardInterrupt): + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert launcher.cancellations == [_JOB_ID] + assert not _job_directory(workspace).exists() + + +def _service(workspace: Path, launcher: _Launcher, *, clock: FakeClock | None = None): + return create_slurm_image_service( + profile=_profile(workspace), + launcher=launcher, # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + clock=None if clock is None else clock.now, + sleep=None if clock is None else clock.sleep, + ) + + +def _profile(workspace: Path) -> SlurmProfile: + return SlurmProfile( + schema_version=1, + scheduler=SchedulerProfile(account="research", partition="gpu"), + gpus_per_node=8, + workspace_root=workspace.as_posix(), + image_build=ImageBuildProfile( + partition="cpu", + cpus_per_task=2, + memory="8G", + time_limit="04:00:00", + ), + ) + + +def _write_inspection(workspace: Path, content: bytes) -> None: + inspection = ImageInspectionRecord( + schema_version=1, + inspector_version="inspector-1", + sqsh_sha256=hashlib.sha256(content).hexdigest(), + inspection=ClientImageInspection( + kind="client", + python_implementation="cpython", + python_version="3.13.3", + python_abi="cp313", + distributions=tuple( + InstalledDistribution(name=name, version="0.9.2") + for name in ( + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + ) + ) + + (InstalledDistribution(name="pip", version="26.1"),), + installer_path="/usr/bin/pip", + installer_version="26.1", + ), + ) + output = _job_directory(workspace) / "output" / "inspection.json" + output.write_text(inspection.model_dump_json()) + + +def _job_directory(workspace: Path) -> Path: + return workspace / "images" / ".tmp" / "jobs" / _LIFECYCLE_ID diff --git a/packages/data-designer-slurm/tests/services/test_profile_service.py b/packages/data-designer-slurm/tests/services/test_profile_service.py new file mode 100644 index 000000000..933d6e04c --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_profile_service.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import stat +from pathlib import Path + +import pytest +import yaml + +from data_designer.slurm.config import ( + ImageBuildProfile, + ProfileSelectionSource, + SchedulerProfile, + SlurmProfile, + SlurmProfileCatalog, + load_profile_catalog, +) +from data_designer.slurm.services import SlurmServiceError, SlurmServiceErrorCode, create_slurm_profile_service + + +class _Launcher: + def __init__(self, gpu_counts: tuple[int, ...] = ()) -> None: + self.gpu_counts = gpu_counts + self.partitions: list[str | None] = [] + + def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: + self.partitions.append(partition) + return self.gpu_counts + + +def test_profile_init_creates_deterministic_private_starter_without_side_effects(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + workspace = tmp_path / "workspace" + launcher = _Launcher() + service = create_slurm_profile_service( + profile_file=profile_file, + launcher=launcher, # type: ignore[arg-type] + hostname_resolver=lambda: ("Ignored.EXAMPLE.test",), + ) + + result = service.initialize( + workspace_root=workspace, + image_build_partition="cpu", + cluster="primary", + account="research", + partition="gpu", + host_patterns=("Login.EXAMPLE.test",), + ) + + assert result.profile_file == profile_file.as_posix() + assert result.validation_command == f"data-designer slurm profile validate --profile-file {profile_file.as_posix()}" + assert stat.S_IMODE(profile_file.stat().st_mode) == 0o600 + assert not workspace.exists() + assert launcher.partitions == [] + assert load_profile_catalog(profile_file).model_dump(mode="json") == { + "schema_version": 1, + "default_cluster": "primary", + "clusters": { + "primary": { + "schema_version": 1, + "host_patterns": ["Login.EXAMPLE.test"], + "scheduler": { + "account": "research", + "partition": "gpu", + "mem_per_gpu": None, + "bin_path": None, + }, + "gpus_per_node": "auto", + "workspace_root": workspace.as_posix(), + "image_build": { + "partition": "cpu", + "cpus_per_task": 2, + "memory": "8G", + "time_limit": "04:00:00", + }, + "gpu_request_mode": "gres", + "container_mounts": [], + } + }, + } + assert "container_mounts" not in profile_file.read_text() + + +def test_profile_init_uses_explicit_environment_then_home_path_precedence(tmp_path: Path) -> None: + environment_file = tmp_path / "environment.yml" + explicit_file = tmp_path / "explicit.json" + environment = {"DATA_DESIGNER_SLURM_PROFILE_FILE": environment_file.as_posix()} + + explicit = create_slurm_profile_service( + profile_file=explicit_file, + environ=environment, + home_directory=tmp_path, + ).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + from_environment = create_slurm_profile_service( + environ=environment, + home_directory=tmp_path, + ).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + from_home = create_slurm_profile_service( + environ={}, + home_directory=tmp_path, + hostname_resolver=lambda: ("Node.EXAMPLE.test", "node.example.test", "node-short"), + ).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + ) + + assert explicit.profile_file == explicit_file.as_posix() + assert from_environment.profile_file == environment_file.as_posix() + assert from_home.profile_file == (tmp_path / ".data-designer-slurm-profile.yml").as_posix() + assert yaml.safe_load(explicit_file.read_text()) == json.loads(explicit_file.read_text()) + assert load_profile_catalog(from_home.profile_file).clusters["default"].host_patterns == [ + "node.example.test", + "node-short", + ] + + +def test_profile_init_refuses_to_overwrite_or_leave_temporary_files(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + profile_file.write_text("owned by caller\n") + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert profile_file.read_text() == "owned by caller\n" + assert tuple(tmp_path.iterdir()) == (profile_file,) + + +@pytest.mark.parametrize( + "profile_file", + (Path("profile.txt"), Path("missing/profile.yml")), + ids=("unsupported-suffix", "missing-parent"), +) +def test_profile_init_rejects_invalid_destination(tmp_path: Path, profile_file: Path) -> None: + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=tmp_path / profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_profile_validate_uses_strict_loader_selection_and_effective_checks(tmp_path: Path) -> None: + selected_workspace = tmp_path / "primary" + selected_workspace.mkdir() + catalog = _catalog(selected_workspace, tmp_path / "unused") + profile_file = tmp_path / "profile.yml" + profile_file.write_text(yaml.safe_dump(catalog.model_dump(mode="json"), sort_keys=False)) + launcher = _Launcher((8, 8)) + + result = create_slurm_profile_service( + profile_file=profile_file, + launcher=launcher, # type: ignore[arg-type] + hostname_resolver=lambda: ("LOGIN-01.EXAMPLE.TEST", "login-01.example.test"), + ).validate() + + assert result.profile_file == profile_file.as_posix() + assert result.hostnames == ("login-01.example.test",) + assert result.default_cluster == "fallback" + assert result.selected_cluster == "primary" + assert result.selection_source is ProfileSelectionSource.HOSTNAME + assert result.matched_pattern == "login-*.example.test" + assert [(match.cluster, match.patterns) for match in result.matched_clusters] == [ + ("primary", ("login-*.example.test",)) + ] + assert result.workspace_root == selected_workspace.as_posix() + assert result.image_root == (selected_workspace / "images").as_posix() + assert result.registry_file == (selected_workspace / "images" / "registry.yaml").as_posix() + assert result.gpus_per_node == 8 + assert launcher.partitions == ["gpu"] + assert tuple(selected_workspace.iterdir()) == () + assert not (tmp_path / "unused").exists() + + +def test_profile_validate_honors_explicit_cluster_and_fixed_gpu_count(tmp_path: Path) -> None: + workspace = tmp_path / "fallback" + workspace.mkdir() + catalog = _catalog(tmp_path / "unused", workspace) + launcher = _Launcher() + + result = create_slurm_profile_service( + catalog=catalog, + cluster="fallback", + launcher=launcher, # type: ignore[arg-type] + hostname_resolver=lambda: ("login-01.example.test",), + ).validate() + + assert result.selected_cluster == "fallback" + assert result.selection_source is ProfileSelectionSource.EXPLICIT + assert result.gpus_per_node == 4 + assert launcher.partitions == [] + assert [(match.cluster, match.patterns) for match in result.matched_clusters] == [ + ("primary", ("login-*.example.test",)) + ] + + +@pytest.mark.parametrize("gpu_counts", ((), (4, 8)), ids=("missing", "heterogeneous")) +def test_profile_validate_rejects_ambiguous_automatic_gpu_count( + tmp_path: Path, + gpu_counts: tuple[int, ...], +) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service( + profile=_profile(workspace), + launcher=_Launcher(gpu_counts), # type: ignore[arg-type] + ).validate() + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert str(caught.value) == "eligible Slurm nodes do not report one GPU count" + + +def test_profile_validate_rejects_ambiguous_selection_and_unavailable_workspace(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + overlapping = SlurmProfileCatalog( + schema_version=1, + default_cluster="first", + clusters={ + "first": _profile(workspace, host_patterns=["node-*"]), + "second": _profile(workspace, host_patterns=["*.example.test"]), + }, + ) + + with pytest.raises(SlurmServiceError) as ambiguous: + create_slurm_profile_service( + catalog=overlapping, + hostname_resolver=lambda: ("node-01.example.test",), + ).validate() + + missing = create_slurm_profile_service( + profile=_profile(tmp_path / "missing", gpus_per_node=4), + hostname_resolver=lambda: ("login",), + ) + with pytest.raises(SlurmServiceError) as unavailable: + missing.validate() + + assert ambiguous.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert unavailable.value.code is SlurmServiceErrorCode.UNAVAILABLE + + +def test_profile_validate_rejects_duplicate_keys_with_stable_error(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + profile_file.write_text("schema_version: 1\nschema_version: 1\n") + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=profile_file).validate() + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert str(caught.value) == "profile configuration cannot be resolved" + + +def _catalog(primary_workspace: Path, fallback_workspace: Path) -> SlurmProfileCatalog: + return SlurmProfileCatalog( + schema_version=1, + default_cluster="fallback", + clusters={ + "primary": _profile(primary_workspace, host_patterns=["login-*.example.test"]), + "fallback": _profile(fallback_workspace, gpus_per_node=4), + }, + ) + + +def _profile( + workspace: Path, + *, + host_patterns: list[str] | None = None, + gpus_per_node: int | str = "auto", +) -> SlurmProfile: + return SlurmProfile( + schema_version=1, + host_patterns=[] if host_patterns is None else host_patterns, + scheduler=SchedulerProfile(partition="gpu"), + gpus_per_node=gpus_per_node, + workspace_root=workspace.as_posix(), + image_build=ImageBuildProfile( + partition="cpu", + cpus_per_task=2, + memory="8G", + time_limit="04:00:00", + ), + ) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index bbe2b39d1..475bfd105 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -14,7 +14,6 @@ from data_designer.slurm.config import ( BuilderInput, DataDesignerSlurmConfig, - ImageBuildRequest, SecretRef, SlurmProfile, SlurmProfileCatalog, @@ -765,7 +764,7 @@ def test_production_status_maps_unknown_run_to_not_found( assert caught.value.code is SlurmServiceErrorCode.NOT_FOUND -def test_production_image_registry_operations_and_lifecycle_gap( +def test_production_image_registry_operations( tmp_path: Path, profile_catalog: SlurmProfileCatalog, authored_run_single: DataDesignerSlurmConfig, @@ -780,13 +779,3 @@ def test_production_image_registry_operations_and_lifecycle_gap( assert selected == removed == images[0] assert len(service.list()) == 1 - with pytest.raises(SlurmServiceError) as caught: - service.add( - ImageBuildRequest( - name="new-image", - kind="client", - source=f"registry.example/client@sha256:{'1' * 64}", - ) - ) - assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE - assert str(caught.value) == "image registration is not available; use a pre-registered image" diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 417657843..392e3e976 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -139,9 +139,96 @@ def test_image_add_rejects_mutable_oci_source(source: str) -> None: } +def test_profile_init_creates_starter_and_emits_validation_command(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + workspace = tmp_path / "workspace" + + result = CliRunner().invoke( + cli_module.create_cli(), + [ + "profile", + "init", + "--workspace-root", + str(workspace), + "--image-build-partition", + "cpu", + "--profile-file", + str(profile_file), + "--host-pattern", + "login.example.test", + ], + ) + + assert result.exit_code == 0, result.output + assert json.loads(result.stdout) == { + "profile_file": profile_file.as_posix(), + "validation_command": f"data-designer slurm profile validate --profile-file {profile_file.as_posix()}", + } + assert profile_file.is_file() + assert not workspace.exists() + + +def test_profile_validate_emits_selected_effective_paths(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + profile_file = tmp_path / "profile.json" + profile_file.write_text( + json.dumps( + { + "schema_version": 1, + "default_cluster": "local", + "clusters": { + "local": { + "schema_version": 1, + "gpus_per_node": 4, + "workspace_root": workspace.as_posix(), + "image_build": { + "partition": "cpu", + "cpus_per_task": 2, + "memory": "8G", + "time_limit": "04:00:00", + }, + } + }, + } + ) + ) + + result = CliRunner().invoke( + cli_module.create_cli(), + ["profile", "validate", "--profile-file", str(profile_file), "--cluster", "local"], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload["profile_file"] == profile_file.as_posix() + assert payload["selected_cluster"] == "local" + assert payload["selection_source"] == "explicit" + assert payload["workspace_root"] == workspace.as_posix() + assert payload["image_root"] == (workspace / "images").as_posix() + assert payload["registry_file"] == (workspace / "images" / "registry.yaml").as_posix() + assert payload["gpus_per_node"] == 4 + assert tuple(workspace.iterdir()) == () + + +def test_image_add_rejects_credential_bearing_oci_source() -> None: + source = f"https://user:secret@registry.example/image@sha256:{'a' * 64}" + + result = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source, "--kind", "client"]) + + assert result.exit_code == 2 + error = json.loads(result.stderr)["error"] + assert error == { + "code": "invalid_request", + "message": "OCI image source must be a credential-free registry reference without a scheme", + "operation": "add_image", + } + assert "secret" not in result.stderr + + def test_cli_exposes_only_m2_run_commands() -> None: result = CliRunner().invoke(cli_module.create_cli(), ["--help"]) assert result.exit_code == 0 - assert all(command in result.stdout for command in ("execute", "status", "cancel", "image")) + assert all(command in result.stdout for command in ("execute", "status", "cancel", "image", "profile")) assert all(command not in result.stdout for command in ("retry", "merge", "benchmark")) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 412a1b00b..6be54bcdd 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -113,6 +113,8 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non import data_designer.config import data_designer.engine import data_designer.interface +from pathlib import Path +from tempfile import TemporaryDirectory from typer.testing import CliRunner from data_designer.cli.main import app @@ -171,6 +173,29 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non assert StateArtifactReference is ContractArtifactReference assert StateRecordRange is ContractRecordRange assert StateResumeWorkspace is ContractResumeWorkspace +profile_help_result = CliRunner().invoke(app, ["slurm", "profile", "--help"]) +assert profile_help_result.exit_code == 0, profile_help_result.output +with TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + profile_file = root / "profile.yml" + profile_init_result = CliRunner().invoke( + app, + [ + "slurm", + "profile", + "init", + "--workspace-root", + str(root / "workspace"), + "--image-build-partition", + "cpu", + "--profile-file", + str(profile_file), + "--host-pattern", + "login.example.test", + ], + ) + assert profile_init_result.exit_code == 0, profile_init_result.output + assert profile_file.is_file() """ run([str(python), "-c", statement], cwd=cwd) From 15b162f244b7797b5ae88b6f254bd18259782ce0 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 11 Sep 2026 09:59:57 -0300 Subject: [PATCH 2/3] fix(slurm): address profile and lifecycle review --- .../data_designer/slurm/launcher/client.py | 14 ++-- .../data_designer/slurm/launcher/parsing.py | 21 ++++++ .../slurm/services/image_lifecycle.py | 24 ++++++- .../data_designer/slurm/services/profiles.py | 15 +++-- .../data-designer-slurm/tests/conftest.py | 2 +- .../tests/launcher/test_client.py | 4 +- .../tests/launcher/test_parsing.py | 17 +++++ .../tests/services/test_image_wiring.py | 64 ++++++++++++++++--- .../tests/services/test_profile_service.py | 38 +++++++++++ .../golden/slurm/sinfo_gres.txt | 2 +- .../tests/slurm_test_fakes/test_slurm.py | 2 +- 11 files changed, 173 insertions(+), 30 deletions(-) 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 cc40fbeff..e9c5fbf06 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 @@ -24,6 +24,7 @@ ) from data_designer.slurm.launcher.parsing import ( parse_accounting, + parse_default_partition_gpu_counts, parse_gpu_counts, parse_named_jobs, parse_queue, @@ -224,12 +225,13 @@ def release(self, job_id: int) -> None: self._run((self._executables.scontrol, "release", _format_job_id(job_id))) def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, ...]: - """Return configured GPU counts reported for eligible node groups.""" - command = [self._executables.sinfo, "--noheader", "--format=%G"] - if partition is not None: - if type(partition) is not str or _IDENTIFIER_PATTERN.fullmatch(partition) is None: - raise ValueError("Slurm partition must be a valid identifier") - command.append(f"--partition={partition}") + """Return configured GPU counts for the requested or default partition.""" + if partition is None: + command = (self._executables.sinfo, "--noheader", "--format=%P|%G") + return parse_default_partition_gpu_counts(self._run(command)) + if type(partition) is not str or _IDENTIFIER_PATTERN.fullmatch(partition) is None: + raise ValueError("Slurm partition must be a valid identifier") + command = (self._executables.sinfo, "--noheader", "--format=%G", f"--partition={partition}") return parse_gpu_counts(self._run(command)) def _run( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py index 32766d782..56e4f97a6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -146,6 +146,27 @@ def parse_gpu_counts(output: str) -> tuple[int, ...]: return tuple(counts) +def parse_default_partition_gpu_counts(output: str) -> tuple[int, ...]: + """Parse GPU counts from ``sinfo --format=%P|%G`` default-partition rows.""" + default_partition: str | None = None + resources: list[str] = [] + for line_number, line in _collect_nonempty_lines(output): + fields = tuple(field.strip() for field in line.split("|")) + if len(fields) != 2 or not all(fields): + raise SlurmCommandOutputError(f"sinfo line {line_number} must contain a partition and resources") + partition, gres = fields + if not partition.endswith("*"): + continue + partition = partition.removesuffix("*") + if _CLUSTER_NAME_PATTERN.fullmatch(partition) is None: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid default partition") + if default_partition is not None and partition != default_partition: + raise SlurmCommandOutputError("sinfo returned multiple default partitions") + default_partition = partition + resources.append(gres) + return parse_gpu_counts("\n".join(resources)) + + def parse_state(value: str) -> SchedulerState: """Normalize one Slurm long state spelling without guessing unknown states.""" normalized = value.strip().upper().removesuffix("+") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py index e864bf0ec..bd80f29e5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py @@ -37,7 +37,7 @@ SchedulerState, SlurmStateError, ) -from data_designer.slurm.state.scheduler import is_scheduler_failure_state +from data_designer.slurm.state.scheduler import is_scheduler_failure_state, is_scheduler_terminal_state LifecycleIdFactory = Callable[[], str] Clock = Callable[[], datetime] @@ -172,10 +172,30 @@ def _wait_for_success(self, job_id: int) -> None: def _cancel_and_cleanup(self, job_id: int, prepared: PreparedImageLifecycleJob) -> None: try: self._launcher.cancel(job_id) - except (SlurmLauncherError, OSError, ValueError): + if not self._wait_for_termination(job_id): + return + except (SlurmLauncherError, SlurmStateError, OSError, ValueError): return _cleanup_failed_lifecycle(prepared) + def _wait_for_termination(self, job_id: int) -> bool: + observations = SchedulerObservationCollector(self._launcher) + previous: SchedulerObservation | None = None + deadline = self._clock() + _ACCOUNTING_EXIT_LAG + while True: + observed_at = self._clock() + observation = observations.collect( + (job_id,), + observed_at=observed_at, + previous={job_id: previous}, + )[0] + if is_scheduler_terminal_state(observation.state): + return True + if observation.state is SchedulerState.UNKNOWN or observed_at >= deadline: + return False + previous = observation + self._sleep(_POLL_INTERVAL_SECONDS) + def _cleanup_failed_lifecycle(prepared: PreparedImageLifecycleJob) -> None: try: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py index 402834bc5..550af5562 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py @@ -173,18 +173,18 @@ def validate(self) -> SlurmProfileValidation: gpus_per_node = self._resolve_gpu_count(selected.profile) except SlurmServiceError: raise - except (SlurmConfigLoadError, ValidationError, ValueError, TypeError): - raise SlurmServiceError( - SlurmServiceErrorCode.INVALID_REQUEST, - operation, - "profile configuration cannot be resolved", - ) from None except SlurmLauncherError: raise SlurmServiceError( SlurmServiceErrorCode.UNAVAILABLE, operation, "Slurm is unavailable", ) from None + except (SlurmConfigLoadError, ValidationError, ValueError, TypeError): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "profile configuration cannot be resolved", + ) from None except OSError: raise SlurmServiceError( SlurmServiceErrorCode.UNAVAILABLE, @@ -368,7 +368,8 @@ def _resolve_profile_path( if source is None: home = Path.home() if home_directory is None else Path(home_directory) source = home / DEFAULT_PROFILE_FILE_NAME - path = Path(source).expanduser().resolve() + expanded = Path(source).expanduser() + path = expanded.parent.resolve() / expanded.name if path.suffix not in {".json", ".yaml", ".yml"}: raise SlurmConfigLoadError("configuration path must end in .json, .yaml, or .yml") return path diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py index a46fcb160..ed14a4820 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -197,7 +197,7 @@ def fake_slurm_runner() -> FakeSlurmRunner: ), ), sinfo_responses={ - ("sinfo", "--noheader", "--format=%G"): FakeCommandResponse( + ("sinfo", "--noheader", "--format=%P|%G"): FakeCommandResponse( stdout=(SLURM_GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() ) }, diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index e4be4ea4a..a9c247eff 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -236,11 +236,11 @@ def test_client_rejects_unbounded_or_invalid_job_selectors(fake_slurm_runner: Fa assert fake_slurm_runner.calls == [] -def test_client_queries_bounded_gpu_inventory(fake_slurm_runner: FakeSlurmRunner) -> None: +def test_client_queries_default_partition_gpu_inventory(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) assert client.query_gpu_counts() == (2,) - assert fake_slurm_runner.calls == [("sinfo", "--noheader", "--format=%G")] + assert fake_slurm_runner.calls == [("sinfo", "--noheader", "--format=%P|%G")] def test_client_queries_partition_scoped_gpu_inventory() -> None: diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index 7509f93c2..c6d1f955e 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -12,6 +12,7 @@ from data_designer.slurm.launcher.models import SlurmQueueEntry from data_designer.slurm.launcher.parsing import ( parse_accounting, + parse_default_partition_gpu_counts, parse_gpu_counts, parse_queue, parse_state, @@ -174,6 +175,22 @@ def test_parse_gpu_counts_normalizes_configured_gres(output: str, expected: tupl assert parse_gpu_counts(output) == expected +def test_parse_default_partition_gpu_counts_ignores_other_partitions() -> None: + output = "cpu|(null)\ngpu*|gpu:h100:8\nother|gpu:a100:4\ngpu*|gpu:h100:8\n" + + assert parse_default_partition_gpu_counts(output) == (8, 8) + + +@pytest.mark.parametrize( + "output", + ("gpu*|gpu:8|extra\n", "*|gpu:8\n", "gpu*|gpu:8\nother*|gpu:4\n"), + ids=("field-count", "missing-name", "multiple-defaults"), +) +def test_parse_default_partition_gpu_counts_rejects_malformed_rows(output: str) -> None: + with pytest.raises(SlurmCommandOutputError): + parse_default_partition_gpu_counts(output) + + @pytest.mark.parametrize( "output", ( diff --git a/packages/data-designer-slurm/tests/services/test_image_wiring.py b/packages/data-designer-slurm/tests/services/test_image_wiring.py index 23ea118a9..c48fbd275 100644 --- a/packages/data-designer-slurm/tests/services/test_image_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -30,7 +30,12 @@ SlurmProcessExitCode, SlurmQueueEntry, ) -from data_designer.slurm.services import SlurmServiceError, SlurmServiceErrorCode, create_slurm_image_service +from data_designer.slurm.services import ( + SlurmImageService, + SlurmServiceError, + SlurmServiceErrorCode, + create_slurm_image_service, +) from data_designer.slurm.state import SchedulerState _JOB_ID = 42 @@ -115,6 +120,31 @@ def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ... return super().query_accounting(selectors) +class _DelayedCancellationLauncher(_Launcher): + def __init__(self) -> None: + super().__init__(accounting_state=None) + + def cancel(self, job_id: int) -> None: + super().cancel(job_id) + self.accounting_state = SchedulerState.RUNNING + + def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ...]: + entries = super().query_accounting(selectors) + if self.cancellations: + self.accounting_state = SchedulerState.CANCELLED + return entries + + +class _ImmediateCancellationLauncher(_Launcher): + def cancel(self, job_id: int) -> None: + super().cancel(job_id) + self.accounting_state = SchedulerState.CANCELLED + + +def _interrupt(_: float) -> None: + raise KeyboardInterrupt + + def test_default_image_add_runs_lifecycle_and_registers_existing_sqsh(tmp_path: Path) -> None: workspace = tmp_path / "workspace" source = tmp_path / "client.sqsh" @@ -243,11 +273,11 @@ def test_default_image_add_retains_ambiguous_submission_state(tmp_path: Path) -> assert _job_directory(workspace).is_dir() -def test_default_image_add_cancels_unknown_job_before_cleanup(tmp_path: Path) -> None: +def test_default_image_add_waits_for_cancelled_job_before_cleanup(tmp_path: Path) -> None: workspace = tmp_path / "workspace" source = tmp_path / "client.sqsh" source.write_bytes(b"client") - launcher = _Launcher(accounting_state=None) + launcher = _DelayedCancellationLauncher() clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) with pytest.raises(SlurmServiceError) as caught: @@ -257,10 +287,27 @@ def test_default_image_add_cancels_unknown_job_before_cleanup(tmp_path: Path) -> assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE assert launcher.cancellations == [_JOB_ID] - assert clock.sleep_calls == [300.0, 300.0] + assert clock.sleep_calls == [300.0, 300.0, 300.0] assert not _job_directory(workspace).exists() +def test_default_image_add_retains_state_without_terminal_cancellation_evidence(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher(accounting_state=None) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + + with pytest.raises(SlurmServiceError): + _service(workspace, launcher, clock=clock).add( + ImageBuildRequest(name="client", kind="client", source=source.as_posix()) + ) + + assert launcher.cancellations == [_JOB_ID] + assert clock.sleep_calls == [300.0, 300.0, 300.0] + assert _job_directory(workspace).is_dir() + + def test_default_image_add_cancels_completed_job_without_exit_evidence(tmp_path: Path) -> None: workspace = tmp_path / "workspace" source = tmp_path / "client.sqsh" @@ -283,16 +330,13 @@ def test_default_image_add_cancels_and_cleans_on_interrupt(tmp_path: Path) -> No workspace = tmp_path / "workspace" source = tmp_path / "client.sqsh" source.write_bytes(b"client") - launcher = _Launcher(queue_state=SchedulerState.PENDING, accounting_state=None) - - def interrupt(_: float) -> None: - raise KeyboardInterrupt + launcher = _ImmediateCancellationLauncher(queue_state=SchedulerState.PENDING, accounting_state=None) service = create_slurm_image_service( profile=_profile(workspace), launcher=launcher, # type: ignore[arg-type] lifecycle_id_factory=lambda: _LIFECYCLE_ID, - sleep=interrupt, + sleep=_interrupt, ) with pytest.raises(KeyboardInterrupt): @@ -302,7 +346,7 @@ def interrupt(_: float) -> None: assert not _job_directory(workspace).exists() -def _service(workspace: Path, launcher: _Launcher, *, clock: FakeClock | None = None): +def _service(workspace: Path, launcher: _Launcher, *, clock: FakeClock | None = None) -> SlurmImageService: return create_slurm_image_service( profile=_profile(workspace), launcher=launcher, # type: ignore[arg-type] diff --git a/packages/data-designer-slurm/tests/services/test_profile_service.py b/packages/data-designer-slurm/tests/services/test_profile_service.py index 933d6e04c..1440ca609 100644 --- a/packages/data-designer-slurm/tests/services/test_profile_service.py +++ b/packages/data-designer-slurm/tests/services/test_profile_service.py @@ -18,6 +18,7 @@ SlurmProfileCatalog, load_profile_catalog, ) +from data_designer.slurm.launcher.errors import SlurmCommandOutputError from data_designer.slurm.services import SlurmServiceError, SlurmServiceErrorCode, create_slurm_profile_service @@ -31,6 +32,12 @@ def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: return self.gpu_counts +class _MalformedLauncher(_Launcher): + def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: + del partition + raise SlurmCommandOutputError("malformed") + + def test_profile_init_creates_deterministic_private_starter_without_side_effects(tmp_path: Path) -> None: profile_file = tmp_path / "profile.yml" workspace = tmp_path / "workspace" @@ -141,6 +148,23 @@ def test_profile_init_refuses_to_overwrite_or_leave_temporary_files(tmp_path: Pa assert tuple(tmp_path.iterdir()) == (profile_file,) +def test_profile_init_refuses_to_follow_dangling_destination_symlink(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + target = tmp_path / "target.yml" + profile_file.symlink_to(target) + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert profile_file.is_symlink() + assert not target.exists() + + @pytest.mark.parametrize( "profile_file", (Path("profile.txt"), Path("missing/profile.yml")), @@ -269,6 +293,20 @@ def test_profile_validate_rejects_duplicate_keys_with_stable_error(tmp_path: Pat assert str(caught.value) == "profile configuration cannot be resolved" +def test_profile_validate_reports_malformed_slurm_output_as_unavailable(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service( + profile=_profile(workspace), + launcher=_MalformedLauncher(), # type: ignore[arg-type] + ).validate() + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert str(caught.value) == "Slurm is unavailable" + + def _catalog(primary_workspace: Path, fallback_workspace: Path) -> SlurmProfileCatalog: return SlurmProfileCatalog( schema_version=1, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt index 25175f090..3e5fe27b7 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt @@ -1 +1 @@ -gpu:2 +batch*|gpu:2 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py index c47b21432..3051ca5ee 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -245,7 +245,7 @@ def test_fake_slurm_runner_exposes_retry_terminal_spellings() -> None: def test_fake_slurm_runner_bounds_sinfo_queries(fake_slurm_runner: FakeSlurmRunner) -> None: - response = fake_slurm_runner.run(("/usr/bin/sinfo", "--noheader", "--format=%G"), check=True) + response = fake_slurm_runner.run(("/usr/bin/sinfo", "--noheader", "--format=%P|%G"), check=True) assert response.stdout == (GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() with pytest.raises(AssertionError, match="unexpected sinfo query"): From 18035e2390a140020d0f9edba0488be8697277b3 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 11 Sep 2026 10:15:27 -0300 Subject: [PATCH 3/3] fix(slurm): require terminal cancellation evidence --- .../src/data_designer/slurm/services/image_lifecycle.py | 6 ++++-- .../data-designer-slurm/tests/services/test_image_wiring.py | 6 +++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py index bd80f29e5..a335e248f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py @@ -179,7 +179,8 @@ def _cancel_and_cleanup(self, job_id: int, prepared: PreparedImageLifecycleJob) _cleanup_failed_lifecycle(prepared) def _wait_for_termination(self, job_id: int) -> bool: - observations = SchedulerObservationCollector(self._launcher) + client = _RecordingObservationClient(self._launcher) + observations = SchedulerObservationCollector(client) previous: SchedulerObservation | None = None deadline = self._clock() + _ACCOUNTING_EXIT_LAG while True: @@ -189,7 +190,8 @@ def _wait_for_termination(self, job_id: int) -> bool: observed_at=observed_at, previous={job_id: previous}, )[0] - if is_scheduler_terminal_state(observation.state): + accounting = next((entry for entry in client.accounting if entry.job_identity == job_id), None) + if accounting is not None and is_scheduler_terminal_state(accounting.state): return True if observation.state is SchedulerState.UNKNOWN or observed_at >= deadline: return False diff --git a/packages/data-designer-slurm/tests/services/test_image_wiring.py b/packages/data-designer-slurm/tests/services/test_image_wiring.py index c48fbd275..e7b858353 100644 --- a/packages/data-designer-slurm/tests/services/test_image_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -308,7 +308,7 @@ def test_default_image_add_retains_state_without_terminal_cancellation_evidence( assert _job_directory(workspace).is_dir() -def test_default_image_add_cancels_completed_job_without_exit_evidence(tmp_path: Path) -> None: +def test_default_image_add_retains_completed_job_without_accounting_exit_evidence(tmp_path: Path) -> None: workspace = tmp_path / "workspace" source = tmp_path / "client.sqsh" source.write_bytes(b"client") @@ -322,8 +322,8 @@ def test_default_image_add_cancels_completed_job_without_exit_evidence(tmp_path: assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE assert launcher.cancellations == [_JOB_ID] - assert clock.sleep_calls == [300.0, 300.0] - assert not _job_directory(workspace).exists() + assert clock.sleep_calls == [300.0, 300.0, 300.0] + assert _job_directory(workspace).is_dir() def test_default_image_add_cancels_and_cleans_on_interrupt(tmp_path: Path) -> None: