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 ebe7de55e..730ce513c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -5,6 +5,8 @@ import re from collections.abc import Callable +from enum import Enum +from functools import partial from pathlib import Path from typing import NoReturn, TypeVar @@ -33,6 +35,13 @@ SlurmServiceErrorCode.INTERNAL: 1, } + +class _RetryResumeMode(str, Enum): + NEVER = "never" + ALWAYS = "always" + IF_POSSIBLE = "if_possible" + + app = typer.Typer( name="slurm", help="Run Data Designer workloads on Slurm", @@ -98,6 +107,89 @@ def cancel_command( _emit_result(result) +@app.command("retry") +def retry_command( + run_or_job_id: str = typer.Argument(..., help="Managed run ID or Slurm array job ID"), + task_ids: list[int] | None = typer.Option(None, "--task-id", min=0, help="Array task ID to retry; repeatable"), + resume: _RetryResumeMode = typer.Option(_RetryResumeMode.IF_POSSIBLE, "--resume"), + dry_run: bool = typer.Option(False, "--dry-run"), + force: bool = typer.Option(False, "--force", help="Submit without confirmation"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Retry failed shards from immutable persisted run state.""" + operation = SlurmServiceOperation.RETRY_RUN + shard_ids = None if task_ids is None else tuple(f"shard-{task_id:05d}" for task_id in task_ids) + service = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster), + ) + + if not dry_run and not force: + planned = _invoke( + operation, + partial( + service.retry, + run_or_job_id, + shard_ids=shard_ids, + resume=resume.value, + dry_run=True, + ), + ) + typer.echo( + f"Retry {', '.join(planned.shard_ids)} with resume={planned.effective_resume_mode}", + err=True, + ) + try: + confirmed = click.confirm("Submit this retry?", default=False, err=True) + except click.Abort: + typer.echo(err=True) + _fail( + SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "interactive confirmation is unavailable; pass --force or --dry-run", + ) + ) + if not confirmed: + _emit_json({"operation": operation.value, "state": "declined"}) + return + shard_ids = planned.shard_ids + resume = _RetryResumeMode(planned.effective_resume_mode) + result = _invoke( + operation, + partial( + service.retry, + run_or_job_id, + shard_ids=shard_ids, + resume=resume.value, + dry_run=dry_run, + ), + ) + _emit_result(result) + + +@app.command("merge") +def merge_command( + input_path: Path = typer.Option(..., "--input-path", file_okay=False), + output_path: Path = typer.Option(..., "--output-path", file_okay=False), + num_partitions: int | None = typer.Option(None, "--num-partitions", min=1), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Submit winner-driven collection as a zero-GPU Slurm job.""" + operation = SlurmServiceOperation.COLLECT_RUN + result = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).collect( + input_path, + destination=output_path, + num_partitions=num_partitions, + ), + ) + _emit_result(result) + + @profile_app.command("init") def profile_init_command( workspace_root: Path = typer.Option(..., "--workspace-root", file_okay=False), 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 5c0ae9c7e..79722df62 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 @@ -23,9 +23,11 @@ create_slurm_profile_service, ) from data_designer.slurm.services.results import ( + SlurmCollectionExecution, SlurmPersistedAttemptStatus, SlurmPersistedRunStatus, SlurmPersistedShardStatus, + SlurmRetryExecution, SlurmRunCancellation, SlurmRunExecution, ) @@ -48,6 +50,7 @@ "SlurmBatchScriptRenderer", "SlurmBenchmarkBackend", "SlurmBenchmarkService", + "SlurmCollectionExecution", "SlurmImageManager", "SlurmImageResolver", "SlurmImageService", @@ -58,6 +61,7 @@ "SlurmProfileMatch", "SlurmProfileService", "SlurmProfileValidation", + "SlurmRetryExecution", "SlurmRunArtifactPublisher", "SlurmRunBackend", "SlurmRunCancellation", 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 7e7e19814..ecf336f49 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" + RETRY_RUN = "retry_run" + COLLECT_RUN = "collect_run" INIT_PROFILE = "init_profile" VALIDATE_PROFILE = "validate_profile" RESOLVE_IMAGE = "resolve_image" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/results.py b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py index 2417b000d..5d25a51f9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/results.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py @@ -7,12 +7,20 @@ from typing import Literal -from pydantic import PositiveInt, model_validator - -from data_designer.slurm.contracts import ContractValue, Identifier, Sha256Digest +from pydantic import PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ( + AttemptId, + ContractValue, + Identifier, + Sha256Digest, + ShardId, + validate_absolute_path, +) from data_designer.slurm.state import ( AttemptManifest, AttemptReadiness, + CollectionState, RunManifest, ShardManifest, ShardWinner, @@ -104,10 +112,50 @@ def validate_jobs(self) -> SlurmRunCancellation: return self +class SlurmRetryExecution(ContractValue): + """One rendered retry dry run or accepted sparse retry submission.""" + + run_id: Identifier + state: Literal["dry_run", "submitted"] + shard_ids: tuple[ShardId, ...] + attempt_ids: tuple[AttemptId, ...] + effective_resume_mode: Literal["never", "always"] + job_id: PositiveInt | None = None + batch_script: str | None = None + + @model_validator(mode="after") + def validate_execution(self) -> SlurmRetryExecution: + if not self.shard_ids or self.shard_ids != tuple(sorted(set(self.shard_ids))): + raise ValueError("retry shard IDs must be non-empty, sorted, and unique") + if len(self.attempt_ids) != len(self.shard_ids): + raise ValueError("retry attempt IDs must correspond to the selected shards") + if self.state == "dry_run": + if self.job_id is not None or not self.batch_script: + raise ValueError("dry-run retry requires only a rendered batch script") + elif self.job_id is None or self.batch_script is not None: + raise ValueError("submitted retry requires only a Slurm job ID") + return self + + +class SlurmCollectionExecution(ContractValue): + """One accepted or previously active collection submission.""" + + run_id: Identifier + collection_id: Identifier + state: CollectionState + job_id: PositiveInt + output_path: str + num_partitions: PositiveInt + + _output_path_is_absolute = field_validator("output_path")(validate_absolute_path) + + __all__ = [ + "SlurmCollectionExecution", "SlurmPersistedAttemptStatus", "SlurmPersistedRunStatus", "SlurmPersistedShardStatus", + "SlurmRetryExecution", "SlurmRunCancellation", "SlurmRunExecution", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py new file mode 100644 index 000000000..46de9f076 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py @@ -0,0 +1,311 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Production retry and collection adapters for the public run service.""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime +from pathlib import Path +from typing import Literal + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier, ShardId +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation +from data_designer.slurm.services.results import SlurmCollectionExecution, SlurmRetryExecution +from data_designer.slurm.state import ( + SlurmCollectionCoordinator, + SlurmRetryCoordinator, + SlurmStateError, + SlurmStateWriter, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.outputs import RetryPlan + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class RunRetryCollectionBackend: + """Adapt public run operations to persisted retry and collection capabilities.""" + + def __init__(self, workspace_root: str, launcher: SlurmCommandClient, clock: Callable[[], datetime]) -> None: + self._workspace_root = workspace_root + self._launcher = launcher + self._clock = clock + + def retry( + self, + run_or_job_id: Identifier, + *, + shard_ids: tuple[ShardId, ...] | None, + resume: Literal["never", "always", "if_possible"], + dry_run: bool, + ) -> SlurmRetryExecution: + """Render or submit one sparse retry from persisted run state.""" + operation = SlurmServiceOperation.RETRY_RUN + run_id = self._resolve_run_reference(run_or_job_id, operation) + try: + observed_at = self._clock() + coordinator = SlurmRetryCoordinator(self._workspace_root, run_id, self._launcher) + resolved_plan = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan() + effective_resume_mode = self._resolve_retry_resume_mode(resolved_plan, resume) + preview = coordinator.preview_active(shard_ids=shard_ids, observed_at=observed_at) + if preview is not None and effective_resume_mode not in {None, preview[0].effective_resume_mode}: + preview = None + if preview is not None: + effective_resume_mode = preview[0].effective_resume_mode + shard_ids = tuple(shard.shard_id for shard in preview[0].planned_shards) + elif effective_resume_mode is None: + preview = coordinator.preview( + shard_ids=shard_ids, + effective_resume_mode="never", + observed_at=observed_at, + ) + effective_resume_mode = self._resolve_if_possible_resume_mode(resolved_plan, preview[0], operation) + shard_ids = tuple(shard.shard_id for shard in preview[0].planned_shards) + if dry_run: + if preview is None or preview[0].effective_resume_mode != effective_resume_mode: + preview = coordinator.preview( + shard_ids=shard_ids, + effective_resume_mode=effective_resume_mode, + observed_at=observed_at, + ) + assert preview is not None + plan, batch_script = preview + return SlurmRetryExecution( + run_id=run_id, + state="dry_run", + shard_ids=tuple(shard.shard_id for shard in plan.planned_shards), + attempt_ids=tuple(shard.attempt_id for shard in plan.planned_shards), + effective_resume_mode=effective_resume_mode, + batch_script=batch_script, + ) + attempts = tuple( + sorted( + coordinator.retry( + shard_ids=shard_ids, + effective_resume_mode=effective_resume_mode, + observed_at=observed_at, + ), + key=lambda attempt: attempt.shard_id, + ) + ) + scheduler_ids = {attempt.scheduler.array_job_id for attempt in attempts if attempt.scheduler is not None} + if not attempts or len(scheduler_ids) != 1 or any(attempt.scheduler is None for attempt in attempts): + raise SlurmStateError("retry did not return one accepted sparse array") + return SlurmRetryExecution( + run_id=run_id, + state="submitted", + shard_ids=tuple(attempt.shard_id for attempt in attempts), + attempt_ids=tuple(attempt.attempt_id for attempt in attempts), + effective_resume_mode=effective_resume_mode, + job_id=scheduler_ids.pop(), + ) + except StateNotFoundError: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "run state was not found") from None + except StateConflictError as error: + raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, str(error)) from None + except StateCorruptionError: + raise SlurmServiceError(SlurmServiceErrorCode.INTERNAL, operation, "retry state is inconsistent") from None + except SlurmStateError as error: + if _has_cause(error, SlurmLauncherError): + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm retry submission is unavailable", + ) from None + raise SlurmServiceError(SlurmServiceErrorCode.INTERNAL, operation, "retry state cannot be read") from None + + def collect( + self, + input_path: Path, + *, + destination: Path, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + """Submit or recover one winner-driven collection.""" + operation = SlurmServiceOperation.COLLECT_RUN + run_id = self._resolve_run_input_path(input_path, operation) + try: + plan = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan() + try: + CollectionDestinationResolver().resolve(plan, destination) + except StateConflictError as error: + raise SlurmServiceError(SlurmServiceErrorCode.INVALID_REQUEST, operation, str(error)) from None + if num_partitions is not None and num_partitions != plan.output.partitions: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "num_partitions must match the persisted run output partitions", + ) + effective_partitions = plan.output.partitions + status = SlurmCollectionCoordinator(self._workspace_root, run_id, self._launcher).submit( + destination=destination, + submitted_at=self._clock(), + ) + if status.scheduler is None: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm collection submission was not accepted", + ) + return SlurmCollectionExecution( + run_id=run_id, + collection_id=status.collection_id, + state=status.state, + job_id=status.scheduler, + output_path=destination.as_posix(), + num_partitions=effective_partitions, + ) + except SlurmServiceError: + raise + except StateNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.NOT_FOUND, + operation, + "run state or shard winner was not found", + ) from None + except StateConflictError as error: + raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, str(error)) from None + except StateCorruptionError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + "collection state is inconsistent", + ) from None + except SlurmStateError as error: + if _has_cause(error, SlurmLauncherError): + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm collection submission is unavailable", + ) from None + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + "collection state cannot be read", + ) from None + + @staticmethod + def _resolve_retry_resume_mode( + plan: ResolvedSlurmRunPlan, + requested: Literal["never", "always", "if_possible"], + ) -> Literal["never", "always"] | None: + if requested != "if_possible": + return requested + pinned = plan.invocation.authored.resume + return None if pinned == "if_possible" else pinned + + @staticmethod + def _resolve_if_possible_resume_mode( + plan: ResolvedSlurmRunPlan, + retry_plan: RetryPlan, + operation: SlurmServiceOperation, + ) -> Literal["never", "always"]: + workspaces = {shard.shard_id: Path(shard.resume_workspace.path) for shard in plan.shards} + availability = {_has_resume_data(workspaces[shard.shard_id]) for shard in retry_plan.planned_shards} + if len(availability) != 1: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "retry selection mixes resumable and fresh shards; choose --resume never or --resume always", + ) + return "always" if availability.pop() else "never" + + def _resolve_run_reference( + self, + run_or_job_id: Identifier, + operation: SlurmServiceOperation, + ) -> Identifier: + workspace_root = Path(self._workspace_root) + direct = workspace_root / "runs" / run_or_job_id + direct_match = run_or_job_id if direct.is_dir() else None + if direct_match is not None and not run_or_job_id.isdecimal(): + return run_or_job_id + if not run_or_job_id.isdecimal() or int(run_or_job_id) <= 0: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "run state was not found") + job_id = int(run_or_job_id) + matches: list[Identifier] = [] + runs_root = workspace_root / "runs" + if runs_root.is_dir(): + for candidate in sorted(runs_root.iterdir(), key=lambda path: path.name): + if not candidate.is_dir(): + continue + try: + candidate_id = _IDENTIFIER_ADAPTER.validate_python(candidate.name, strict=True) + writer = SlurmStateWriter(workspace_root, candidate_id) + attempts = tuple( + attempt for shard in writer.load_shards() for attempt in writer.load_attempts(shard.shard_id) + ) + except (ValidationError, StateNotFoundError): + continue + except SlurmStateError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + "managed run state cannot be searched safely", + ) from None + if any( + attempt.scheduler is not None and attempt.scheduler.array_job_id == job_id for attempt in attempts + ): + matches.append(candidate_id) + targets = set(matches) + if direct_match is not None: + targets.add(direct_match) + if not targets: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "managed Slurm job was not found") + if len(targets) != 1: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + operation, + "numeric reference matches multiple managed runs", + ) + return targets.pop() + + def _resolve_run_input_path( + self, + input_path: Path, + operation: SlurmServiceOperation, + ) -> Identifier: + runs_root = (Path(self._workspace_root) / "runs").resolve() + if input_path.parent != runs_root: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "input_path must identify a managed run directory", + ) + try: + run_id = _IDENTIFIER_ADAPTER.validate_python(input_path.name, strict=True) + except ValidationError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "input_path must contain a valid managed run ID", + ) from None + if not input_path.is_dir(): + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "run state was not found") + return run_id + + +def _has_cause(error: BaseException, expected_type: type[BaseException]) -> bool: + cause = error.__cause__ + while cause is not None: + if isinstance(cause, expected_type): + return True + cause = cause.__cause__ + return False + + +def _has_resume_data(path: Path) -> bool: + return not path.is_symlink() and path.is_dir() and next(path.iterdir(), None) is not None + + +__all__ = ["RunRetryCollectionBackend"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py index 68eeaa90a..f40ea2614 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -5,13 +5,15 @@ from __future__ import annotations +from collections.abc import Sequence +from functools import partial from pathlib import Path -from typing import Protocol +from typing import Literal, Protocol from pydantic import TypeAdapter, ValidationError from data_designer.slurm.config import DataDesignerSlurmConfig -from data_designer.slurm.contracts import Identifier +from data_designer.slurm.contracts import Identifier, ShardId from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.services.errors import ( SlurmServiceError, @@ -21,12 +23,15 @@ _make_invalid_request_error, ) from data_designer.slurm.services.results import ( + SlurmCollectionExecution, SlurmPersistedRunStatus, + SlurmRetryExecution, SlurmRunCancellation, SlurmRunExecution, ) _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_SHARD_IDS_ADAPTER = TypeAdapter(tuple[ShardId, ...]) class SlurmRunPlanner(Protocol): @@ -68,6 +73,25 @@ def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: def cancel(self, run_id: Identifier) -> SlurmRunCancellation: """Request cancellation of active jobs.""" + def retry( + self, + run_or_job_id: Identifier, + *, + shard_ids: tuple[ShardId, ...] | None, + resume: Literal["never", "always", "if_possible"], + dry_run: bool, + ) -> SlurmRetryExecution: + """Render or submit one sparse retry.""" + + def collect( + self, + input_path: Path, + *, + destination: Path, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + """Submit or recover one winner-driven collection.""" + class SlurmRunService: """Coordinate public run operations through package-owned boundaries. @@ -206,6 +230,64 @@ def cancel_run() -> SlurmRunCancellation: return _invoke_service_backend(operation, cancel_run) + def retry( + self, + run_or_job_id: Identifier, + *, + shard_ids: Sequence[ShardId] | None = None, + resume: Literal["never", "always", "if_possible"] = "if_possible", + dry_run: bool = False, + ) -> SlurmRetryExecution: + """Render or submit retry attempts for failed shards.""" + operation = SlurmServiceOperation.RETRY_RUN + normalized_reference = _validate_run_id(run_or_job_id, operation) + normalized_shards = _validate_shard_ids(shard_ids, operation) + if resume not in {"never", "always", "if_possible"}: + raise _make_invalid_request_error(operation, "resume must be 'never', 'always', or 'if_possible'") + if type(dry_run) is not bool: + raise _make_invalid_request_error(operation, "dry_run must be a boolean") + backend = self._require_backend(operation) + + return _invoke_service_backend( + operation, + partial( + _retry_run, + backend, + normalized_reference, + shard_ids=normalized_shards, + resume=resume, + dry_run=dry_run, + ), + ) + + def collect( + self, + input_path: str | Path, + *, + destination: str | Path, + num_partitions: int | None = None, + ) -> SlurmCollectionExecution: + """Submit or recover collection for one managed run directory.""" + operation = SlurmServiceOperation.COLLECT_RUN + if not isinstance(input_path, str | Path) or not isinstance(destination, str | Path): + raise _make_invalid_request_error(operation, "input_path and destination must be paths") + if num_partitions is not None and (type(num_partitions) is not int or num_partitions <= 0): + raise _make_invalid_request_error(operation, "num_partitions must be a positive integer") + normalized_input = Path(input_path).expanduser().resolve() + normalized_destination = Path(destination).expanduser().resolve() + backend = self._require_backend(operation) + + return _invoke_service_backend( + operation, + partial( + _collect_run, + backend, + normalized_input, + destination=normalized_destination, + num_partitions=num_partitions, + ), + ) + def _require_backend(self, operation: SlurmServiceOperation) -> SlurmRunBackend: if self._backend is None: raise SlurmServiceError( @@ -216,8 +298,69 @@ def _require_backend(self, operation: SlurmServiceOperation) -> SlurmRunBackend: return self._backend +def _retry_run( + backend: SlurmRunBackend, + run_or_job_id: Identifier, + *, + shard_ids: tuple[ShardId, ...] | None, + resume: Literal["never", "always", "if_possible"], + dry_run: bool, +) -> SlurmRetryExecution: + result = backend.retry( + run_or_job_id, + shard_ids=shard_ids, + resume=resume, + dry_run=dry_run, + ) + if not isinstance(result, SlurmRetryExecution): + raise TypeError("run backend returned an invalid retry result") + if shard_ids is not None and result.shard_ids != shard_ids: + raise TypeError("run backend returned retry shards that do not match the request") + if result.state != ("dry_run" if dry_run else "submitted"): + raise TypeError("run backend returned a retry state that does not match the request") + return result + + +def _collect_run( + backend: SlurmRunBackend, + input_path: Path, + *, + destination: Path, + num_partitions: int | None, +) -> SlurmCollectionExecution: + result = backend.collect( + input_path, + destination=destination, + num_partitions=num_partitions, + ) + if not isinstance(result, SlurmCollectionExecution): + raise TypeError("run backend returned an invalid collection result") + if result.output_path != destination.as_posix() or ( + num_partitions is not None and result.num_partitions != num_partitions + ): + raise TypeError("run backend returned collection intent that does not match the request") + return result + + def _validate_run_id(run_id: object, operation: SlurmServiceOperation) -> Identifier: try: return _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) except ValidationError: raise _make_invalid_request_error(operation, "run_id must be a valid identifier") from None + + +def _validate_shard_ids( + shard_ids: Sequence[ShardId] | None, + operation: SlurmServiceOperation, +) -> tuple[ShardId, ...] | None: + if shard_ids is None: + return None + if isinstance(shard_ids, str | bytes): + raise _make_invalid_request_error(operation, "shard_ids must be a sequence of shard identifiers") + try: + normalized = _SHARD_IDS_ADAPTER.validate_python(tuple(shard_ids), strict=True) + except (TypeError, ValidationError): + raise _make_invalid_request_error(operation, "shard_ids must contain valid shard identifiers") from None + if not normalized or len(normalized) != len(set(normalized)): + raise _make_invalid_request_error(operation, "shard_ids must be non-empty and unique when provided") + return tuple(sorted(normalized)) 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 f48757534..691849bc4 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 @@ -12,7 +12,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import Protocol, TypeVar +from typing import Literal, Protocol, TypeVar from uuid import uuid4 from pydantic import JsonValue @@ -37,7 +37,7 @@ load_builder_payload, resolve_profile, ) -from data_designer.slurm.contracts import Identifier +from data_designer.slurm.contracts import Identifier, ShardId from data_designer.slurm.images.errors import ImageConflictError, ImageNotFoundError, SlurmImageError from data_designer.slurm.images.records import RegisteredImage from data_designer.slurm.images.registry import ImageRegistryStore @@ -55,12 +55,15 @@ from data_designer.slurm.services.image_lifecycle import SlurmImageLifecycleManager from data_designer.slurm.services.images import SlurmImageService from data_designer.slurm.services.results import ( + SlurmCollectionExecution, SlurmPersistedAttemptStatus, SlurmPersistedRunStatus, SlurmPersistedShardStatus, + SlurmRetryExecution, SlurmRunCancellation, SlurmRunExecution, ) +from data_designer.slurm.services.retry_collection import RunRetryCollectionBackend from data_designer.slurm.services.run import SlurmRunService from data_designer.slurm.serving.resolver import resolve_vllm_server from data_designer.slurm.state import ( @@ -274,6 +277,11 @@ def __init__( self._publisher = publisher self._clock = clock self._source_environment = source_environment + self._retry_collection = RunRetryCollectionBackend( + selected_profile.profile.workspace_root, + launcher, + clock, + ) def execute( self, @@ -623,6 +631,34 @@ def cancel(self, run_id: Identifier) -> SlurmRunCancellation: ) return SlurmRunCancellation(run_id=run_id, job_ids=job_ids) + def retry( + self, + run_or_job_id: Identifier, + *, + shard_ids: tuple[ShardId, ...] | None, + resume: Literal["never", "always", "if_possible"], + dry_run: bool, + ) -> SlurmRetryExecution: + return self._retry_collection.retry( + run_or_job_id, + shard_ids=shard_ids, + resume=resume, + dry_run=dry_run, + ) + + def collect( + self, + input_path: Path, + *, + destination: Path, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + return self._retry_collection.collect( + input_path, + destination=destination, + num_partitions=num_partitions, + ) + class _RegistryImageBackend: def __init__(self, workspace_root: str, lifecycle: SlurmImageLifecycleManager) -> None: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py index 396ccdcf5..78b1b0fab 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -293,9 +293,11 @@ def _validate_existing_destination( plan = self._load_bound_plan(status) resolved_plan = self._reader.load_resolved_plan() resolved = self._destinations.validate_persisted(resolved_plan, plan) + if requested_destination is not None and Path(requested_destination).as_posix() != resolved.host_path: + raise StateConflictError("persisted collection destination does not match the requested destination") requested = self._destinations.resolve(resolved_plan, requested_destination) if requested != resolved: - raise StateCorruptionError("persisted collection destination does not match the requested destination") + raise StateConflictError("persisted collection destination does not match the requested destination") if plan.host_destination != resolved.host_path or plan.container_destination != resolved.container_path: raise StateCorruptionError("persisted collection destination does not match the resolved plan") if status.state is CollectionState.SUCCEEDED: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py index 78c8a2bc4..1c738a2d6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -10,7 +10,7 @@ from pathlib import Path from data_designer.slurm.config import ContainerMount -from data_designer.slurm.contracts import is_path_below, validate_absolute_path +from data_designer.slurm.contracts import is_path_below, paths_overlap, validate_absolute_path from data_designer.slurm.images.records import validate_enroot_mount_path from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.state.errors import StateConflictError @@ -41,6 +41,7 @@ def resolve( except ValueError as error: raise StateConflictError("collection destination must be a normalized absolute path") from error workspace_root = plan.selected_profile.profile.workspace_root + _validate_destination_overlap(plan, host_path) workspace_mount = ContainerMount(source=workspace_root, target=workspace_root, read_only=False) authorized = { (mount.source, mount.target): mount @@ -89,4 +90,21 @@ def validate_persisted( return destination +def _validate_destination_overlap(plan: ResolvedSlurmRunPlan, host_path: str) -> None: + workspace_root = plan.selected_profile.profile.workspace_root + reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) + if any(paths_overlap(host_path, path) for path in reserved): + raise StateConflictError("collection destination must not overlap package-managed workspace state") + managed_assets_path = plan.invocation.effective_input_bindings.managed_assets_path + assert managed_assets_path is not None + if paths_overlap(host_path, managed_assets_path): + raise StateConflictError("collection destination must not overlap managed assets") + runs_root = posixpath.join(workspace_root, "runs") + run_output_root = posixpath.join(runs_root, plan.run_id, "output") + if paths_overlap(host_path, runs_root) and not ( + host_path == run_output_root or is_path_below(host_path, run_output_root) + ): + raise StateConflictError("collection destination must not overlap package-managed run state") + + __all__ = ["CollectionDestination", "CollectionDestinationResolver"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py index 4b018d726..1c49d2952 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -5,6 +5,7 @@ from __future__ import annotations +from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -92,6 +93,13 @@ def refresh(self, *, observed_at: datetime | None = None) -> RunStatus: SlurmStateError: If scheduler evidence cannot be queried or state cannot be reconstructed safely. """ + return self._observe(observed_at=observed_at, persist=True) + + def observe(self, *, observed_at: datetime | None = None) -> RunStatus: + """Return current scheduler status without persisting observations.""" + return self._observe(observed_at=observed_at, persist=False) + + def _observe(self, *, observed_at: datetime | None, persist: bool) -> RunStatus: timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at _validate_observed_at(timestamp) run, plan, shards = self._reader.load_context() @@ -114,6 +122,7 @@ def refresh(self, *, observed_at: datetime | None = None) -> RunStatus: self._refresh_shard( _ShardSnapshot(run, plan, shard, attempts_by_shard[shard.shard_id]), batch, + persist=persist, ) for shard in shards ) @@ -172,9 +181,12 @@ def _refresh_shard( self, expected: _ShardSnapshot, batch: _ObservationBatch, + *, + persist: bool, ) -> ShardStatus: try: - with self._storage.acquire_shard_lock(expected.shard.shard_id): + shard_lock = self._storage.acquire_shard_lock(expected.shard.shard_id) if persist else nullcontext() + with shard_lock: current_run, current_plan, current_shard = self._reader.load_shard_context(expected.shard.shard_id) attempts = self._reader.load_validated_shard_attempts(current_run, current_plan, current_shard) self._require_unchanged_context( @@ -193,6 +205,7 @@ def _refresh_shard( batch, attempt, winner, + persist=persist, ) for attempt in attempts ) @@ -214,13 +227,16 @@ def _build_attempt_status( batch: _ObservationBatch, attempt: AttemptManifest, winner: ShardWinner | None, + *, + persist: bool, ) -> AttemptStatus: readiness = self._reader.load_optional_readiness(snapshot.plan, attempt) result = self._reader.load_optional_attempt_result(snapshot.plan, snapshot.shard, attempt) scheduler = None if attempt.scheduler is not None: scheduler = batch.current[attempt.scheduler] - self._persist_observation(attempt, batch.previous[attempt.scheduler], scheduler) + if persist: + self._persist_observation(attempt, batch.previous[attempt.scheduler], scheduler) effective_state = reconcile_attempt_observation( attempt, readiness, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py index 607ab3fc9..a8b9188cd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -114,6 +114,52 @@ def retry( except (OSError, ValidationError, ValueError) as error: raise SlurmStateError(f"cannot retry persisted run {self._run_id!r}") from error + def preview( + self, + *, + shard_ids: Sequence[ShardId] | None = None, + effective_resume_mode: Literal["never", "always"], + observed_at: datetime | None = None, + ) -> tuple[RetryPlan, str]: + """Render a retry plan without persisting or submitting it.""" + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + try: + if effective_resume_mode not in {"never", "always"}: + raise StateConflictError("effective resume mode must be 'never' or 'always'") + status = self._reconciler.observe(observed_at=timestamp) + selected = _select_retryable_shards(status, shard_ids) + plan = self._build_retry_plan(status, selected, effective_resume_mode, timestamp) + return plan, render_generation_retry_script(self._reader.load_resolved_plan(status.run), plan) + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot preview retry for persisted run {self._run_id!r}") from error + + def preview_active( + self, + *, + shard_ids: Sequence[ShardId] | None = None, + observed_at: datetime | None = None, + ) -> tuple[RetryPlan, str] | None: + """Render the latest active retry without persisting observations.""" + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + try: + status = self._reconciler.observe(observed_at=timestamp) + retry_ids = self._retries.list_retry_ids() + if not retry_ids: + return None + retry_status = self._retries.read_optional_status(retry_ids[-1]) + if retry_status is None: + return None + plan = self._load_bound_plan(retry_status) + if self._load_active_retry(status, shard_ids, plan.effective_resume_mode) is None: + return None + return plan, render_generation_retry_script(self._reader.load_resolved_plan(status.run), plan) + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot preview active retry for persisted run {self._run_id!r}") from error + def _settle_pending_retry(self, updated_at: datetime) -> None: retry_ids = self._retries.list_retry_ids() if not retry_ids: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py index 489af26b0..cbb10803d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py @@ -49,7 +49,10 @@ def acquire_lock(self) -> Iterator[None]: def get_next_retry_id(self) -> Identifier: """Return the next monotonic retry identity.""" - return f"retry-{len(self.list_retry_ids()) + 1:04d}" + retry_ids = self.list_retry_ids() + if retry_ids and self.read_optional_status(retry_ids[-1]) is None: + return retry_ids[-1] + return f"retry-{len(retry_ids) + 1:04d}" def discard_incomplete_tail(self) -> None: """Discard one trailing journal that cannot have reached submission.""" @@ -156,6 +159,13 @@ def read_status(self, retry_id: Identifier) -> RetryStatus: raise OSError("retry status identity does not match its persisted location") return status + def read_optional_status(self, retry_id: Identifier) -> RetryStatus | None: + """Return retry status when its journal reached status publication.""" + try: + return self.read_status(retry_id) + except FileNotFoundError: + return None + @contextmanager def _open_retries_directory(self) -> Iterator[int]: with self._state.open_run_directory() as run_descriptor: diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index a79a8187a..478e202d9 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -191,6 +191,36 @@ def test_destination_reauthorizes_explicit_path_through_workspace_mapping( assert resolver.validate_persisted(plan, collection) == destination +@pytest.mark.parametrize( + ("destination_suffix", "error"), + [ + ("images/collected", "package-managed workspace state"), + ("runtime/collected", "package-managed workspace state"), + ("benchmarks/collected", "package-managed workspace state"), + ("managed-assets/collected", "managed assets"), + ("runs/run-001/collections", "package-managed run state"), + ("runs/run-001/retries", "package-managed run state"), + ("runs/other-run/output", "package-managed run state"), + ], +) +def test_destination_rejects_package_managed_paths( + multi_node_plan: ResolvedSlurmRunPlan, + destination_suffix: str, + error: str, +) -> None: + workspace_root = multi_node_plan.selected_profile.profile.workspace_root + requested = (Path(workspace_root) / destination_suffix).as_posix() + + with pytest.raises(StateConflictError, match=error): + CollectionDestinationResolver().resolve(multi_node_plan, requested) + + +def test_destination_allows_current_run_output_subtree(multi_node_plan: ResolvedSlurmRunPlan) -> None: + requested = (Path(multi_node_plan.output.root) / "collected").as_posix() + + assert CollectionDestinationResolver().resolve(multi_node_plan, requested).host_path == requested + + def test_destination_requires_one_unique_most_specific_mapping( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index 80a5532e3..b3df00eb7 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -5,6 +5,7 @@ import pickle from pathlib import Path +from unittest.mock import Mock import pytest from slurm_test_fakes import ( @@ -29,9 +30,11 @@ SlurmBatchScriptRenderer, SlurmBenchmarkBackend, SlurmBenchmarkService, + SlurmCollectionExecution, SlurmImageManager, SlurmImageResolver, SlurmImageService, + SlurmRetryExecution, SlurmRunArtifactPublisher, SlurmRunBackend, SlurmRunExecution, @@ -41,6 +44,7 @@ SlurmServiceErrorCode, SlurmServiceOperation, ) +from data_designer.slurm.state import CollectionState GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "rendered" @@ -224,6 +228,118 @@ def cancel(self, run_id): assert backend.calls == [(authored_run_single, tmp_path.resolve(), True, True)] +def test_run_service_delegates_retry_with_stable_shard_order() -> None: + backend = Mock(spec=SlurmRunBackend) + expected = SlurmRetryExecution( + run_id="run-0001", + state="submitted", + shard_ids=("shard-00001", "shard-00003"), + attempt_ids=("attempt-0002", "attempt-0004"), + effective_resume_mode="always", + job_id=43, + ) + backend.retry.return_value = expected + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) + + result = service.retry( + "42", + shard_ids=("shard-00003", "shard-00001"), + resume="always", + ) + + assert result is expected + backend.retry.assert_called_once_with( + "42", + shard_ids=("shard-00001", "shard-00003"), + resume="always", + dry_run=False, + ) + + +@pytest.mark.parametrize( + ("shard_ids", "resume", "dry_run"), + [ + ((), "never", False), + (("shard-00001", "shard-00001"), "never", False), + (None, "sometimes", False), + (None, "never", 1), + ], +) +def test_run_service_rejects_invalid_retry_actions( + shard_ids: object, + resume: object, + dry_run: object, +) -> None: + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), Mock(spec=SlurmRunBackend)) + + with pytest.raises(SlurmServiceError) as caught: + service.retry( # type: ignore[arg-type] + "run-0001", + shard_ids=shard_ids, + resume=resume, + dry_run=dry_run, + ) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.RETRY_RUN + + +def test_run_service_normalizes_collection_paths(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.chdir(tmp_path) + backend = Mock(spec=SlurmRunBackend) + expected = SlurmCollectionExecution( + run_id="run-0001", + collection_id="collection-0001", + state=CollectionState.SUBMITTED, + job_id=44, + output_path=(tmp_path / "collected").as_posix(), + num_partitions=2, + ) + backend.collect.return_value = expected + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) + + result = service.collect("runs/run-0001", destination="collected", num_partitions=2) + + assert result is expected + backend.collect.assert_called_once_with( + tmp_path / "runs/run-0001", + destination=tmp_path / "collected", + num_partitions=2, + ) + + +def test_run_service_uses_persisted_collection_partitions_when_omitted(tmp_path: Path) -> None: + backend = Mock(spec=SlurmRunBackend) + expected = SlurmCollectionExecution( + run_id="run-0001", + collection_id="collection-0001", + state=CollectionState.SUBMITTED, + job_id=44, + output_path=(tmp_path / "collected").as_posix(), + num_partitions=2, + ) + backend.collect.return_value = expected + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) + + assert service.collect(tmp_path / "runs/run-0001", destination=tmp_path / "collected") is expected + backend.collect.assert_called_once_with( + tmp_path / "runs/run-0001", + destination=tmp_path / "collected", + num_partitions=None, + ) + + +@pytest.mark.parametrize("num_partitions", [0, -1, True, 1.5]) +def test_run_service_rejects_invalid_collection_partitions(num_partitions: object) -> None: + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), Mock(spec=SlurmRunBackend)) + + with pytest.raises(SlurmServiceError) as caught: + service.collect("runs/run-0001", destination="collected", num_partitions=num_partitions) # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.COLLECT_RUN + + def test_run_service_normalizes_and_redacts_unexpected_backend_errors( authored_run_single: DataDesignerSlurmConfig, ) -> None: diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 475bfd105..cd453f598 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -7,9 +7,11 @@ from collections.abc import Mapping from datetime import datetime, timedelta, timezone from pathlib import Path +from types import SimpleNamespace import pytest +import data_designer.slurm.services.retry_collection as retry_collection_module from data_designer.slurm.client.dependencies import ResolvedClientDependencies from data_designer.slurm.config import ( BuilderInput, @@ -27,11 +29,13 @@ SlurmJobSubmissionReceipt, SlurmProcessExitCode, SlurmQueueEntry, + SlurmSubmissionMatch, ) from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.services import ( SlurmServiceError, SlurmServiceErrorCode, + SlurmServiceOperation, create_slurm_image_service, create_slurm_run_service, ) @@ -39,6 +43,7 @@ AttemptLifecycleState, AttemptManifest, AttemptTerminalClassification, + CollectionState, RunManifest, SchedulerIdentity, SchedulerState, @@ -54,6 +59,7 @@ def __init__( self, gpu_counts: tuple[int, ...] = (), *, + submission_job_ids: tuple[int, ...] = (42,), cancel_error: Exception | None = None, release_error: Exception | None = None, ) -> None: @@ -64,6 +70,8 @@ def __init__( self.exported_environments: list[dict[str, str]] = [] self.queue_entries: tuple[SlurmQueueEntry, ...] = () self.accounting_entries: tuple[SlurmAccountingEntry, ...] = () + self.submission_matches: tuple[SlurmSubmissionMatch, ...] = () + self.submission_job_ids = submission_job_ids self.gpu_counts = gpu_counts self.cancel_error = cancel_error self.release_error = release_error @@ -78,7 +86,8 @@ def submit_script( self.submissions.append(script) self.held_submissions.append(hold) self.exported_environments.append(dict(export_environment or {})) - return SlurmJobSubmissionReceipt(job_id=42) + index = min(len(self.submissions) - 1, len(self.submission_job_ids) - 1) + return SlurmJobSubmissionReceipt(job_id=self.submission_job_ids[index]) def cancel(self, job_id: int) -> None: self.cancellations.append(job_id) @@ -102,6 +111,15 @@ def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ... del selectors return self.accounting_entries + def query_submissions_by_name( + self, + job_name: str, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + del job_name, submitted_after + return self.submission_matches + class _Publisher: def __init__( @@ -442,7 +460,10 @@ def test_status_does_not_expire_requeue_window_from_another_attempt_clock( single_node_plan: ResolvedSlurmRunPlan, ) -> None: authored = authored_run_single.model_copy( - update={"array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2})} + update={ + "array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2}), + "invocation": authored_run_single.invocation.model_copy(update={"resume": "if_possible"}), + } ) _register_images(tmp_path, authored, single_node_plan) launcher = _Launcher() @@ -483,6 +504,314 @@ def test_status_does_not_expire_requeue_window_from_another_attempt_clock( assert requeued.shards[0].attempts[0].attempt.state is AttemptLifecycleState.PENDING +def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={ + "array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2}), + "invocation": authored_run_single.invocation.model_copy(update={"resume": "if_possible"}), + } + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher(submission_job_ids=(42, 43)) + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: now, + package_version="0.9.2", + ) + service.execute(authored, source_root=tmp_path) + assert SlurmStateWriter(tmp_path, "run-wired").load_resolved_plan().invocation.authored.resume == "if_possible" + failed = SchedulerIdentity(array_job_id=42, array_task_id=0) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=failed, + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ), + ) + + preview = service.retry("42", shard_ids=("shard-00000",), resume="if_possible", dry_run=True) + assert not (tmp_path / "runs" / "run-wired" / "retries").exists() + assert not (tmp_path / "runs" / "run-wired" / "retry.lock").exists() + assert not ( + tmp_path / "runs" / "run-wired" / "shards" / "shard-00000" / "attempts" / "attempt-0001" / "scheduler.json" + ).exists() + assert len(launcher.submissions) == 1 + submitted = service.retry("42", shard_ids=("shard-00000",), resume="if_possible") + repeated = service.retry("run-wired", shard_ids=("shard-00000",), resume="if_possible") + + assert preview.state == "dry_run" + assert preview.job_id is None + assert preview.shard_ids == ("shard-00000",) + assert preview.effective_resume_mode == "never" + assert preview.batch_script is not None + assert "#SBATCH --array=0" in preview.batch_script + assert submitted.state == "submitted" + assert submitted.job_id == 43 + assert submitted.shard_ids == ("shard-00000",) + assert repeated == submitted + assert len(launcher.submissions) == 2 + + +def test_production_retry_preview_recovers_active_pinned_mode( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher(submission_job_ids=(42, 43)) + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: now, + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=SchedulerIdentity(array_job_id=42, array_task_id=0), + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ), + ) + submitted = service.retry("run-wired", resume="never") + + preview = service.retry("run-wired", resume="never", dry_run=True) + + assert preview.state == "dry_run" + assert preview.shard_ids == submitted.shard_ids + assert preview.attempt_ids == submitted.attempt_ids + assert preview.effective_resume_mode == "never" + assert len(launcher.submissions) == 2 + + +def test_production_retry_if_possible_reuses_populated_workspace( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={"invocation": authored_run_single.invocation.model_copy(update={"resume": "if_possible"})} + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: now, + package_version="0.9.2", + ) + service.execute(authored, source_root=tmp_path) + plan = SlurmStateWriter(tmp_path, "run-wired").load_resolved_plan() + resume_workspace = Path(plan.shards[0].resume_workspace.path) + resume_workspace.mkdir(parents=True) + (resume_workspace / "partial.parquet").touch() + failed = SchedulerIdentity(array_job_id=42, array_task_id=0) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=failed, + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ), + ) + + preview = service.retry("run-wired", resume="if_possible", dry_run=True) + + assert preview.effective_resume_mode == "always" + assert preview.shard_ids == ("shard-00000",) + + +def test_production_retry_if_possible_rejects_mixed_workspace_availability( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={ + "array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2}), + "invocation": authored_run_single.invocation.model_copy(update={"resume": "if_possible"}), + } + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: now, + package_version="0.9.2", + ) + service.execute(authored, source_root=tmp_path) + plan = SlurmStateWriter(tmp_path, "run-wired").load_resolved_plan() + resume_workspace = Path(plan.shards[0].resume_workspace.path) + resume_workspace.mkdir(parents=True) + (resume_workspace / "partial.parquet").touch() + launcher.accounting_entries = tuple( + SlurmAccountingEntry( + job_identity=SchedulerIdentity(array_job_id=42, array_task_id=task_id), + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ) + for task_id in range(2) + ) + + with pytest.raises(SlurmServiceError, match="mixes resumable and fresh shards") as caught: + service.retry("run-wired", resume="if_possible", dry_run=True) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_production_retry_rejects_no_retryable_shards( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + + with pytest.raises(SlurmServiceError, match="no retryable shards") as caught: + service.retry("run-wired", resume="never") + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert caught.value.operation is SlurmServiceOperation.RETRY_RUN + + +def test_production_retry_rejects_ambiguous_scheduler_job_id( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + run_ids = iter(("run-one", "run-two")) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: next(run_ids), + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + service.execute(authored_run_single, source_root=tmp_path) + + with pytest.raises(SlurmServiceError, match="multiple managed runs") as caught: + service.retry("42", resume="never") + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + + +def test_production_retry_rejects_numeric_run_and_job_ambiguity( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + run_ids = iter(("run-one", "42")) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(submission_job_ids=(42, 43)), # type: ignore[arg-type] + run_id_factory=lambda: next(run_ids), + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + service.execute(authored_run_single, source_root=tmp_path) + + with pytest.raises(SlurmServiceError, match="numeric reference") as caught: + service.retry("42", resume="never") + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + + +def test_collection_rejects_unmanaged_input_and_destination( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + + with pytest.raises(SlurmServiceError) as unmanaged: + service.collect(tmp_path / "run-wired", destination=tmp_path / "collected") + with pytest.raises(SlurmServiceError) as destination: + service.collect(tmp_path / "runs/run-wired", destination=tmp_path.parent / "collected") + with pytest.raises(SlurmServiceError) as managed_state: + service.collect(tmp_path / "runs/run-wired", destination=tmp_path / "runs/run-wired/collections") + with pytest.raises(SlurmServiceError) as partitions: + service.collect(tmp_path / "runs/run-wired", destination=tmp_path / "collected", num_partitions=2) + + assert unmanaged.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert destination.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert managed_state.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert partitions.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_collection_returns_active_and_completed_jobs( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + statuses = iter( + ( + SimpleNamespace(collection_id="collection-0001", state=CollectionState.RUNNING, scheduler=43), + SimpleNamespace(collection_id="collection-0001", state=CollectionState.SUCCEEDED, scheduler=43), + ) + ) + + class Coordinator: + def submit(self, *, destination: Path, submitted_at: datetime) -> SimpleNamespace: + del destination, submitted_at + return next(statuses) + + monkeypatch.setattr(retry_collection_module, "SlurmCollectionCoordinator", lambda *_: Coordinator()) + output_path = tmp_path / "collected" + + active = service.collect(tmp_path / "runs/run-wired", destination=output_path) + completed = service.collect(tmp_path / "runs/run-wired", destination=output_path) + + assert active.state is CollectionState.RUNNING + assert completed.state is CollectionState.SUCCEEDED + assert active.job_id == completed.job_id == 43 + + def test_auto_gpu_resolution_rejects_mixed_node_shapes( tmp_path: Path, profile_catalog: SlurmProfileCatalog, diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py index 0a49de6be..41bc6dc73 100644 --- a/packages/data-designer-slurm/tests/state/test_observer.py +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -442,6 +442,29 @@ def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[S SlurmStateReconciler(case.workspace, case.plan.run_id, MutatingClient()).refresh(observed_at=observed_at) +def test_observe_does_not_create_a_missing_shard_lock_or_persist_scheduler_evidence( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + shard_root = case.writer.run_root / "shards/shard-00000" + (shard_root / "shard.lock").unlink() + client = _StaticSchedulerClient( + (SlurmQueueEntry(job_identity=scheduler, state=SchedulerState.RUNNING),), + (), + ) + + status = SlurmStateReconciler(case.workspace, case.plan.run_id, client).observe( + observed_at=case.created_at + timedelta(minutes=3) + ) + + assert status.shards[0].attempts[0].effective_state is EffectiveAttemptState.RUNNING + assert not (shard_root / "shard.lock").exists() + assert not (shard_root / "attempts/attempt-0001/scheduler.json").exists() + + def test_terminal_observation_remains_authoritative_during_later_accounting_gap() -> None: task = SchedulerIdentity(array_job_id=4101, array_task_id=0) first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 441744846..87ee5f001 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -129,6 +129,34 @@ def test_retry_refreshes_failed_shard_and_publishes_exact_next_attempt( assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 +def test_active_retry_preview_ignores_statusless_journal_tail( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + storage.ensure_retry("retry-0001") + + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + assert coordinator.preview_active(observed_at=case.created_at + timedelta(minutes=5)) is None + preview, _ = coordinator.preview( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + assert preview.retry_id == "retry-0001" + assert storage.get_retry_root("retry-0001").is_dir() + + def test_retry_rejects_a_nonterminal_explicit_shard( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -710,6 +738,8 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( assert submitted.state is CollectionState.SUBMITTED assert submitted.scheduler == 5101 + with pytest.raises(StateConflictError, match="requested destination"): + coordinator.submit(destination=Path(case.plan.output.root).parent / "other") script = cast(str, runner.inputs[-1]) assert "data_designer.slurm.state.collection_worker" in script assert "--gpus" not in script diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 392e3e976..aad4367ee 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -12,16 +12,21 @@ import data_designer.slurm.cli as cli_module from data_designer.slurm.config import DataDesignerSlurmConfig from data_designer.slurm.services import ( + SlurmCollectionExecution, + SlurmRetryExecution, SlurmRunExecution, SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, ) +from data_designer.slurm.state import CollectionState class _RunService: def __init__(self) -> None: self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool, bool]] = [] + self.retry_calls: list[tuple[str, tuple[str, ...] | None, str, bool]] = [] + self.collection_calls: list[tuple[Path, Path, int | None]] = [] def execute( self, @@ -40,6 +45,43 @@ def execute( batch_script="#!/bin/bash\n", ) + def retry( + self, + run_or_job_id: str, + *, + shard_ids: tuple[str, ...] | None, + resume: str, + dry_run: bool, + ) -> SlurmRetryExecution: + self.retry_calls.append((run_or_job_id, shard_ids, resume, dry_run)) + selected = ("shard-00000",) if shard_ids is None else shard_ids + return SlurmRetryExecution( + run_id="run-0001", + state="dry_run" if dry_run else "submitted", + shard_ids=selected, + attempt_ids=tuple("attempt-0002" for _ in selected), + effective_resume_mode="always", + job_id=None if dry_run else 43, + batch_script="#!/bin/bash\n" if dry_run else None, + ) + + def collect( + self, + input_path: Path, + *, + destination: Path, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + self.collection_calls.append((input_path, destination, num_partitions)) + return SlurmCollectionExecution( + run_id="run-0001", + collection_id="collection-0001", + state=CollectionState.SUBMITTED, + job_id=43, + output_path=destination.resolve().as_posix(), + num_partitions=1 if num_partitions is None else num_partitions, + ) + def test_execute_emits_deterministic_json_and_forwards_actions( tmp_path: Path, @@ -65,6 +107,121 @@ def test_execute_emits_deterministic_json_and_forwards_actions( assert service.calls == [(authored_run_single, tmp_path, True, True)] +def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke( + cli_module.create_cli(), + ["retry", "42", "--task-id", "1", "--task-id", "3", "--resume", "always", "--dry-run"], + ) + + assert result.exit_code == 0 + assert json.loads(result.stdout) == { + "attempt_ids": ["attempt-0002", "attempt-0002"], + "batch_script": "#!/bin/bash\n", + "effective_resume_mode": "always", + "job_id": None, + "run_id": "run-0001", + "shard_ids": ["shard-00001", "shard-00003"], + "state": "dry_run", + } + assert service.retry_calls == [("42", ("shard-00001", "shard-00003"), "always", True)] + + +def test_retry_auto_selects_tasks_and_confirms_submission(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42"], input="y\n") + + assert result.exit_code == 0 + assert json.loads(result.stdout.splitlines()[-1]) == { + "attempt_ids": ["attempt-0002"], + "batch_script": None, + "effective_resume_mode": "always", + "job_id": 43, + "run_id": "run-0001", + "shard_ids": ["shard-00000"], + "state": "submitted", + } + assert "Submit this retry?" in result.stderr + assert "Retry shard-00000 with resume=always" in result.stderr + assert service.retry_calls == [ + ("42", None, "if_possible", True), + ("42", ("shard-00000",), "always", False), + ] + + +def test_retry_force_skips_confirmation(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42", "--force"]) + + assert result.exit_code == 0 + assert service.retry_calls == [("42", None, "if_possible", False)] + + +def test_retry_decline_emits_stable_result(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42"], input="n\n") + + assert result.exit_code == 0 + assert json.loads(result.stdout.splitlines()[-1]) == {"operation": "retry_run", "state": "declined"} + assert service.retry_calls == [("42", None, "if_possible", True)] + + +def test_retry_noninteractive_confirmation_is_invalid_request(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42"], input="") + + assert result.exit_code == 2 + assert json.loads(result.stderr.splitlines()[-1]) == { + "error": { + "code": "invalid_request", + "message": "interactive confirmation is unavailable; pass --force or --dry-run", + "operation": "retry_run", + } + } + assert service.retry_calls == [("42", None, "if_possible", True)] + + +def test_merge_emits_collection_job_and_forwards_paths(tmp_path: Path, monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + input_path = tmp_path / "runs/run-0001" + output_path = tmp_path / "collected" + + result = CliRunner().invoke( + cli_module.create_cli(), + [ + "merge", + "--input-path", + str(input_path), + "--output-path", + str(output_path), + "--num-partitions", + "2", + ], + ) + + assert result.exit_code == 0 + assert json.loads(result.stdout) == { + "collection_id": "collection-0001", + "job_id": 43, + "num_partitions": 2, + "output_path": output_path.as_posix(), + "run_id": "run-0001", + "state": "submitted", + } + assert service.collection_calls == [(input_path, output_path, 2)] + + @pytest.mark.parametrize( ("code", "exit_code"), [ @@ -226,9 +383,11 @@ def test_image_add_rejects_credential_bearing_oci_source() -> None: assert "secret" not in result.stderr -def test_cli_exposes_only_m2_run_commands() -> None: +def test_cli_exposes_m3c_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", "profile")) - assert all(command not in result.stdout for command in ("retry", "merge", "benchmark")) + assert all( + command in result.stdout for command in ("execute", "status", "cancel", "retry", "merge", "image", "profile") + ) + assert "benchmark" not in result.stdout diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 6be54bcdd..b5f54d241 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -135,6 +135,8 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non statement += f""" slurm_help_result = CliRunner().invoke(app, ["slurm", "--help"]) assert slurm_help_result.exit_code == 0, (slurm_help_result.output, repr(slurm_help_result.exception)) +assert all(command in slurm_help_result.output for command in ("execute", "status", "cancel", "retry", "merge", "image")) +assert "benchmark" not in slurm_help_result.output assert "data_designer.slurm.cli" in sys.modules assert version("data-designer-slurm") == {version!r} from data_designer.slurm.contracts import ArtifactReference as ContractArtifactReference