From d0170d7cdfd22b51883dd9457f23268981aae950 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 10 Sep 2026 17:57:48 -0300 Subject: [PATCH 1/5] feat(slurm): wire retry and merge commands --- .../src/data_designer/slurm/cli.py | 49 +++- .../data_designer/slurm/services/__init__.py | 4 + .../data_designer/slurm/services/errors.py | 2 + .../data_designer/slurm/services/results.py | 54 +++- .../slurm/services/retry_collection.py | 263 ++++++++++++++++++ .../src/data_designer/slurm/services/run.py | 113 +++++++- .../data_designer/slurm/services/wiring.py | 42 ++- .../src/data_designer/slurm/state/observer.py | 16 +- .../src/data_designer/slurm/state/retry.py | 22 ++ .../tests/services/test_services.py | 100 +++++++ .../tests/services/test_wiring.py | 177 +++++++++++- .../data-designer-slurm/tests/test_cli.py | 101 ++++++- scripts/test_slurm_package_install.py | 2 + 13 files changed, 932 insertions(+), 13 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.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..ede8fae9b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -6,7 +6,7 @@ import re from collections.abc import Callable from pathlib import Path -from typing import NoReturn, TypeVar +from typing import Literal, NoReturn, TypeVar import click import typer @@ -94,6 +94,53 @@ 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: Literal["never", "always", "if_possible"] = typer.Option("if_possible", "--resume"), + dry_run: bool = typer.Option(False, "--dry-run"), + force: bool = typer.Option(False, "--force"), + 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) + result = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).retry( + run_or_job_id, + shard_ids=shard_ids, + resume=resume, + dry_run=dry_run, + force=force, + ), + ) + _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 = typer.Option(1, "--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) + + @image_app.command("add") def image_add_command( source: str = typer.Argument(...), 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..6993ff282 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 @@ -16,9 +16,11 @@ ) from data_designer.slurm.services.images import SlurmImageManager, SlurmImageResolver, SlurmImageService from data_designer.slurm.services.results import ( + SlurmCollectionExecution, SlurmPersistedAttemptStatus, SlurmPersistedRunStatus, SlurmPersistedShardStatus, + SlurmRetryExecution, SlurmRunCancellation, SlurmRunExecution, ) @@ -41,12 +43,14 @@ "SlurmBatchScriptRenderer", "SlurmBenchmarkBackend", "SlurmBenchmarkService", + "SlurmCollectionExecution", "SlurmImageManager", "SlurmImageResolver", "SlurmImageService", "SlurmPersistedAttemptStatus", "SlurmPersistedRunStatus", "SlurmPersistedShardStatus", + "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 aa08fffed..38092514b 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" RESOLVE_IMAGE = "resolve_image" ADD_IMAGE = "add_image" LIST_IMAGES = "list_images" 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..da9483473 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py @@ -0,0 +1,263 @@ +# 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.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, +) + +_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, + force: bool, + ) -> SlurmRetryExecution: + """Render or submit one sparse retry from persisted run state.""" + del force + operation = SlurmServiceOperation.RETRY_RUN + run_id = self._resolve_run_reference(run_or_job_id, operation) + try: + effective_resume_mode = self._resolve_retry_resume_mode(run_id, resume) + coordinator = SlurmRetryCoordinator(self._workspace_root, run_id, self._launcher) + if dry_run: + plan, batch_script = coordinator.preview( + shard_ids=shard_ids, + effective_resume_mode=effective_resume_mode, + observed_at=self._clock(), + ) + 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=self._clock(), + ), + 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, + ) -> 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() + if num_partitions != plan.output.partitions: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "num_partitions must match the persisted run output partitions", + ) + status = SlurmCollectionCoordinator(self._workspace_root, run_id, self._launcher).submit( + destination=destination, + submitted_at=self._clock(), + ) + if status.scheduler is None: + raise SlurmStateError("collection did not return an accepted Slurm job") + 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=num_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: + code = ( + SlurmServiceErrorCode.INVALID_REQUEST + if str(error).startswith("collection destination") + else SlurmServiceErrorCode.CONFLICT + ) + raise SlurmServiceError(code, 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 + + def _resolve_retry_resume_mode( + self, + run_id: Identifier, + requested: Literal["never", "always", "if_possible"], + ) -> Literal["never", "always"]: + if requested != "if_possible": + return requested + pinned = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan().invocation.authored.resume + return "always" if pinned == "if_possible" else pinned + + 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 + if direct.is_dir(): + 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) + if not matches: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "managed Slurm job was not found") + if len(matches) != 1: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + operation, + "Slurm job ID matches multiple managed runs", + ) + return matches[0] + + 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 + + +__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..f56daf990 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,14 @@ from __future__ import annotations +from collections.abc import Sequence 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 +22,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 +72,26 @@ 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, + force: bool, + ) -> SlurmRetryExecution: + """Render or submit one sparse retry.""" + + def collect( + self, + input_path: Path, + *, + destination: Path, + num_partitions: int, + ) -> SlurmCollectionExecution: + """Submit or recover one winner-driven collection.""" + class SlurmRunService: """Coordinate public run operations through package-owned boundaries. @@ -206,6 +230,74 @@ 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, + force: 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 or type(force) is not bool: + raise _make_invalid_request_error(operation, "dry_run and force must be booleans") + backend = self._require_backend(operation) + + def retry_run() -> SlurmRetryExecution: + result = backend.retry( + normalized_reference, + shard_ids=normalized_shards, + resume=resume, + dry_run=dry_run, + force=force, + ) + if not isinstance(result, SlurmRetryExecution): + raise TypeError("run backend returned an invalid retry result") + if normalized_shards is not None and result.shard_ids != normalized_shards: + 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 + + return _invoke_service_backend(operation, retry_run) + + def collect( + self, + input_path: str | Path, + *, + destination: str | Path, + num_partitions: int = 1, + ) -> 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 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) + + def collect_run() -> SlurmCollectionExecution: + result = backend.collect( + normalized_input, + destination=normalized_destination, + num_partitions=num_partitions, + ) + if not isinstance(result, SlurmCollectionExecution): + raise TypeError("run backend returned an invalid collection result") + if result.output_path != normalized_destination.as_posix() or result.num_partitions != num_partitions: + raise TypeError("run backend returned collection intent that does not match the request") + return result + + return _invoke_service_backend(operation, collect_run) + def _require_backend(self, operation: SlurmServiceOperation) -> SlurmRunBackend: if self._backend is None: raise SlurmServiceError( @@ -221,3 +313,20 @@ def _validate_run_id(run_id: object, operation: SlurmServiceOperation) -> Identi 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 b13834f56..94c21ac93 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 @@ -54,12 +54,15 @@ from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation 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 ( @@ -273,6 +276,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, @@ -622,6 +630,36 @@ 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, + force: bool, + ) -> SlurmRetryExecution: + return self._retry_collection.retry( + run_or_job_id, + shard_ids=shard_ids, + resume=resume, + dry_run=dry_run, + force=force, + ) + + def collect( + self, + input_path: Path, + *, + destination: Path, + num_partitions: int, + ) -> SlurmCollectionExecution: + return self._retry_collection.collect( + input_path, + destination=destination, + num_partitions=num_partitions, + ) + class _RegistryImageBackend: def __init__(self, workspace_root: str) -> None: 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..5ab547040 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 @@ -92,6 +92,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 +121,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,6 +180,8 @@ def _refresh_shard( self, expected: _ShardSnapshot, batch: _ObservationBatch, + *, + persist: bool, ) -> ShardStatus: try: with self._storage.acquire_shard_lock(expected.shard.shard_id): @@ -193,6 +203,7 @@ def _refresh_shard( batch, attempt, winner, + persist=persist, ) for attempt in attempts ) @@ -214,13 +225,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..14a0a11c5 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,28 @@ 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'") + with self._retries.acquire_lock(): + 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 _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/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index 80a5532e3..77426691b 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,102 @@ 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", + force=True, + ) + + assert result is expected + backend.retry.assert_called_once_with( + "42", + shard_ids=("shard-00001", "shard-00003"), + resume="always", + dry_run=False, + force=True, + ) + + +@pytest.mark.parametrize( + ("shard_ids", "resume", "dry_run", "force"), + [ + ((), "never", False, False), + (("shard-00001", "shard-00001"), "never", False, False), + (None, "sometimes", False, False), + (None, "never", 1, False), + (None, "never", False, 1), + ], +) +def test_run_service_rejects_invalid_retry_actions( + shard_ids: object, + resume: object, + dry_run: object, + force: 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, + force=force, + ) + + 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, + ) + + +@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 bbe2b39d1..9a5251e50 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, @@ -28,11 +30,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, ) @@ -40,6 +44,7 @@ AttemptLifecycleState, AttemptManifest, AttemptTerminalClassification, + CollectionState, RunManifest, SchedulerIdentity, SchedulerState, @@ -55,6 +60,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: @@ -65,6 +71,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 @@ -79,7 +87,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) @@ -103,6 +112,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__( @@ -484,6 +502,163 @@ 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})} + ) + _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) + 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="never", dry_run=True) + assert not (tmp_path / "runs" / "run-wired" / "retries").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="never") + repeated = service.retry("run-wired", shard_ids=("shard-00000",), resume="never") + + assert preview.state == "dry_run" + assert preview.job_id is None + assert preview.shard_ids == ("shard-00000",) + 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_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_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") + + assert unmanaged.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert destination.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/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 417657843..83a382d86 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, bool]] = [] + self.collection_calls: list[tuple[Path, Path, int]] = [] 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, + force: bool, + ) -> SlurmRetryExecution: + self.retry_calls.append((run_or_job_id, shard_ids, resume, dry_run, force)) + assert shard_ids is not None + return SlurmRetryExecution( + run_id="run-0001", + state="dry_run", + shard_ids=shard_ids, + attempt_ids=tuple("attempt-0002" for _ in shard_ids), + effective_resume_mode="always", + batch_script="#!/bin/bash\n", + ) + + def collect( + self, + input_path: Path, + *, + destination: Path, + num_partitions: int, + ) -> 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=num_partitions, + ) + def test_execute_emits_deterministic_json_and_forwards_actions( tmp_path: Path, @@ -65,6 +107,59 @@ 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", "--force"], + ) + + 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, 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"), [ @@ -139,9 +234,9 @@ def test_image_add_rejects_mutable_oci_source(source: str) -> None: } -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")) - 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")) + assert "benchmark" not in result.stdout diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 412a1b00b..ee6d91b20 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -133,6 +133,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 From b1c3f5cdafaa2f3e047b9b464d9dc10c081bd3d3 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 10 Sep 2026 18:16:29 -0300 Subject: [PATCH 2/5] fix(slurm): harden retry and collection actions --- .../src/data_designer/slurm/cli.py | 37 ++++-- .../slurm/services/retry_collection.py | 99 +++++++++++---- .../src/data_designer/slurm/services/run.py | 17 ++- .../data_designer/slurm/services/wiring.py | 4 +- .../data_designer/slurm/state/collection.py | 2 +- .../src/data_designer/slurm/state/retry.py | 9 +- .../tests/services/test_services.py | 36 ++++-- .../tests/services/test_wiring.py | 120 +++++++++++++++++- .../tests/state/test_retry_collection.py | 2 + .../data-designer-slurm/tests/test_cli.py | 59 +++++++-- 10 files changed, 301 insertions(+), 84 deletions(-) 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 ede8fae9b..105d39283 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -5,8 +5,9 @@ import re from collections.abc import Callable +from enum import Enum from pathlib import Path -from typing import Literal, NoReturn, TypeVar +from typing import NoReturn, TypeVar import click import typer @@ -31,6 +32,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,24 +106,31 @@ def cancel_command( 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: Literal["never", "always", "if_possible"] = typer.Option("if_possible", "--resume"), + resume: _RetryResumeMode = typer.Option(_RetryResumeMode.IF_POSSIBLE, "--resume"), dry_run: bool = typer.Option(False, "--dry-run"), - force: bool = typer.Option(False, "--force"), + 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) - result = _invoke( - operation, - lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).retry( + + def retry(*, preview: bool) -> BaseModel: + return create_slurm_run_service(profile_file=profile_file, cluster=cluster).retry( run_or_job_id, shard_ids=shard_ids, - resume=resume, - dry_run=dry_run, - force=force, - ), + resume=resume.value, + dry_run=preview, + ) + + if not dry_run and not force: + _invoke(operation, lambda: retry(preview=True)) + if not click.confirm("Submit this retry?", default=False, err=True): + raise typer.Exit() + result = _invoke( + operation, + lambda: retry(preview=dry_run), ) _emit_result(result) @@ -124,7 +139,7 @@ def retry_command( 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 = typer.Option(1, "--num-partitions", min=1), + 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: 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 index da9483473..b8c969a34 100644 --- 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 @@ -15,6 +15,7 @@ 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 ( @@ -26,6 +27,8 @@ StateCorruptionError, StateNotFoundError, ) +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.outputs import RetryPlan _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) @@ -45,21 +48,33 @@ def retry( shard_ids: tuple[ShardId, ...] | None, resume: Literal["never", "always", "if_possible"], dry_run: bool, - force: bool, ) -> SlurmRetryExecution: """Render or submit one sparse retry from persisted run state.""" - del force operation = SlurmServiceOperation.RETRY_RUN run_id = self._resolve_run_reference(run_or_job_id, operation) try: - effective_resume_mode = self._resolve_retry_resume_mode(run_id, resume) + observed_at = self._clock() coordinator = SlurmRetryCoordinator(self._workspace_root, run_id, self._launcher) - if dry_run: - plan, batch_script = coordinator.preview( + resolved_plan = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan() + effective_resume_mode = self._resolve_retry_resume_mode(resolved_plan, resume) + preview: tuple[RetryPlan, str] | None = None + if effective_resume_mode is None: + preview = coordinator.preview( shard_ids=shard_ids, - effective_resume_mode=effective_resume_mode, - observed_at=self._clock(), + effective_resume_mode="never", + observed_at=observed_at, ) + shard_ids = tuple(shard.shard_id for shard in preview[0].planned_shards) + effective_resume_mode = self._resolve_if_possible_resume_mode(resolved_plan, preview[0], operation) + 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", @@ -73,7 +88,7 @@ def retry( coordinator.retry( shard_ids=shard_ids, effective_resume_mode=effective_resume_mode, - observed_at=self._clock(), + observed_at=observed_at, ), key=lambda attempt: attempt.shard_id, ) @@ -109,32 +124,41 @@ def collect( input_path: Path, *, destination: Path, - num_partitions: int, + 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() - if num_partitions != plan.output.partitions: + 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 SlurmStateError("collection did not return an accepted Slurm job") + 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=num_partitions, + num_partitions=effective_partitions, ) except SlurmServiceError: raise @@ -145,12 +169,7 @@ def collect( "run state or shard winner was not found", ) from None except StateConflictError as error: - code = ( - SlurmServiceErrorCode.INVALID_REQUEST - if str(error).startswith("collection destination") - else SlurmServiceErrorCode.CONFLICT - ) - raise SlurmServiceError(code, operation, str(error)) from None + raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, str(error)) from None except StateCorruptionError: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -170,15 +189,31 @@ def collect( "collection state cannot be read", ) from None + @staticmethod def _resolve_retry_resume_mode( - self, - run_id: Identifier, + plan: ResolvedSlurmRunPlan, requested: Literal["never", "always", "if_possible"], - ) -> Literal["never", "always"]: + ) -> Literal["never", "always"] | None: if requested != "if_possible": return requested - pinned = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan().invocation.authored.resume - return "always" if pinned == "if_possible" else pinned + 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, @@ -187,7 +222,8 @@ def _resolve_run_reference( ) -> Identifier: workspace_root = Path(self._workspace_root) direct = workspace_root / "runs" / run_or_job_id - if direct.is_dir(): + 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") @@ -216,15 +252,18 @@ def _resolve_run_reference( attempt.scheduler is not None and attempt.scheduler.array_job_id == job_id for attempt in attempts ): matches.append(candidate_id) - if not matches: + 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(matches) != 1: + if len(targets) != 1: raise SlurmServiceError( SlurmServiceErrorCode.CONFLICT, operation, - "Slurm job ID matches multiple managed runs", + "numeric reference matches multiple managed runs", ) - return matches[0] + return targets.pop() def _resolve_run_input_path( self, @@ -260,4 +299,8 @@ def _has_cause(error: BaseException, expected_type: type[BaseException]) -> bool 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 f56daf990..3fddb1803 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 @@ -79,7 +79,6 @@ def retry( shard_ids: tuple[ShardId, ...] | None, resume: Literal["never", "always", "if_possible"], dry_run: bool, - force: bool, ) -> SlurmRetryExecution: """Render or submit one sparse retry.""" @@ -88,7 +87,7 @@ def collect( input_path: Path, *, destination: Path, - num_partitions: int, + num_partitions: int | None, ) -> SlurmCollectionExecution: """Submit or recover one winner-driven collection.""" @@ -237,7 +236,6 @@ def retry( shard_ids: Sequence[ShardId] | None = None, resume: Literal["never", "always", "if_possible"] = "if_possible", dry_run: bool = False, - force: bool = False, ) -> SlurmRetryExecution: """Render or submit retry attempts for failed shards.""" operation = SlurmServiceOperation.RETRY_RUN @@ -245,8 +243,8 @@ def retry( 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 or type(force) is not bool: - raise _make_invalid_request_error(operation, "dry_run and force must be booleans") + if type(dry_run) is not bool: + raise _make_invalid_request_error(operation, "dry_run must be a boolean") backend = self._require_backend(operation) def retry_run() -> SlurmRetryExecution: @@ -255,7 +253,6 @@ def retry_run() -> SlurmRetryExecution: shard_ids=normalized_shards, resume=resume, dry_run=dry_run, - force=force, ) if not isinstance(result, SlurmRetryExecution): raise TypeError("run backend returned an invalid retry result") @@ -272,13 +269,13 @@ def collect( input_path: str | Path, *, destination: str | Path, - num_partitions: int = 1, + 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 type(num_partitions) is not int or num_partitions <= 0: + 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() @@ -292,7 +289,9 @@ def collect_run() -> SlurmCollectionExecution: ) if not isinstance(result, SlurmCollectionExecution): raise TypeError("run backend returned an invalid collection result") - if result.output_path != normalized_destination.as_posix() or result.num_partitions != num_partitions: + if result.output_path != normalized_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 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 94c21ac93..eaee7d60c 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 @@ -637,14 +637,12 @@ def retry( shard_ids: tuple[ShardId, ...] | None, resume: Literal["never", "always", "if_possible"], dry_run: bool, - force: bool, ) -> SlurmRetryExecution: return self._retry_collection.retry( run_or_job_id, shard_ids=shard_ids, resume=resume, dry_run=dry_run, - force=force, ) def collect( @@ -652,7 +650,7 @@ def collect( input_path: Path, *, destination: Path, - num_partitions: int, + num_partitions: int | None, ) -> SlurmCollectionExecution: return self._retry_collection.collect( input_path, 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..e194367b9 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 @@ -295,7 +295,7 @@ def _validate_existing_destination( resolved = self._destinations.validate_persisted(resolved_plan, plan) 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/retry.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py index 14a0a11c5..0e4bff964 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 @@ -126,11 +126,10 @@ def preview( try: if effective_resume_mode not in {"never", "always"}: raise StateConflictError("effective resume mode must be 'never' or 'always'") - with self._retries.acquire_lock(): - 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) + 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: diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index 77426691b..b3df00eb7 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -245,7 +245,6 @@ def test_run_service_delegates_retry_with_stable_shard_order() -> None: "42", shard_ids=("shard-00003", "shard-00001"), resume="always", - force=True, ) assert result is expected @@ -254,25 +253,22 @@ def test_run_service_delegates_retry_with_stable_shard_order() -> None: shard_ids=("shard-00001", "shard-00003"), resume="always", dry_run=False, - force=True, ) @pytest.mark.parametrize( - ("shard_ids", "resume", "dry_run", "force"), + ("shard_ids", "resume", "dry_run"), [ - ((), "never", False, False), - (("shard-00001", "shard-00001"), "never", False, False), - (None, "sometimes", False, False), - (None, "never", 1, False), - (None, "never", False, 1), + ((), "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, - force: object, ) -> None: service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), Mock(spec=SlurmRunBackend)) @@ -282,7 +278,6 @@ def test_run_service_rejects_invalid_retry_actions( shard_ids=shard_ids, resume=resume, dry_run=dry_run, - force=force, ) assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST @@ -313,6 +308,27 @@ def test_run_service_normalizes_collection_paths(tmp_path: Path, monkeypatch: py ) +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)) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 9a5251e50..f801df4c2 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -461,7 +461,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() @@ -531,18 +534,20 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( ), ) - preview = service.retry("42", shard_ids=("shard-00000",), resume="never", dry_run=True) + 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="never") - repeated = service.retry("run-wired", shard_ids=("shard-00000",), resume="never") + 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" @@ -552,6 +557,87 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( 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, @@ -597,6 +683,29 @@ def test_production_retry_rejects_ambiguous_scheduler_job_id( 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, @@ -616,9 +725,12 @@ def test_collection_rejects_unmanaged_input_and_destination( 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 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 partitions.value.code is SlurmServiceErrorCode.INVALID_REQUEST def test_collection_returns_active_and_completed_jobs( 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..836d7609d 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -710,6 +710,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 83a382d86..d984b23ce 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -25,8 +25,8 @@ 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, bool]] = [] - self.collection_calls: list[tuple[Path, Path, int]] = [] + self.retry_calls: list[tuple[str, tuple[str, ...] | None, str, bool]] = [] + self.collection_calls: list[tuple[Path, Path, int | None]] = [] def execute( self, @@ -52,17 +52,17 @@ def retry( shard_ids: tuple[str, ...] | None, resume: str, dry_run: bool, - force: bool, ) -> SlurmRetryExecution: - self.retry_calls.append((run_or_job_id, shard_ids, resume, dry_run, force)) - assert shard_ids is not None + 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", - shard_ids=shard_ids, - attempt_ids=tuple("attempt-0002" for _ in shard_ids), + state="dry_run" if dry_run else "submitted", + shard_ids=selected, + attempt_ids=tuple("attempt-0002" for _ in selected), effective_resume_mode="always", - batch_script="#!/bin/bash\n", + job_id=None if dry_run else 43, + batch_script="#!/bin/bash\n" if dry_run else None, ) def collect( @@ -70,7 +70,7 @@ def collect( input_path: Path, *, destination: Path, - num_partitions: int, + num_partitions: int | None, ) -> SlurmCollectionExecution: self.collection_calls.append((input_path, destination, num_partitions)) return SlurmCollectionExecution( @@ -79,7 +79,7 @@ def collect( state=CollectionState.SUBMITTED, job_id=43, output_path=destination.resolve().as_posix(), - num_partitions=num_partitions, + num_partitions=1 if num_partitions is None else num_partitions, ) @@ -113,7 +113,7 @@ def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: result = CliRunner().invoke( cli_module.create_cli(), - ["retry", "42", "--task-id", "1", "--task-id", "3", "--resume", "always", "--dry-run", "--force"], + ["retry", "42", "--task-id", "1", "--task-id", "3", "--resume", "always", "--dry-run"], ) assert result.exit_code == 0 @@ -126,7 +126,40 @@ def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: "shard_ids": ["shard-00001", "shard-00003"], "state": "dry_run", } - assert service.retry_calls == [("42", ("shard-00001", "shard-00003"), "always", True, True)] + 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 service.retry_calls == [ + ("42", None, "if_possible", True), + ("42", None, "if_possible", 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_merge_emits_collection_job_and_forwards_paths(tmp_path: Path, monkeypatch) -> None: From 74f7577a6b6f33da82c261ca18d435c983596442 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 10 Sep 2026 18:30:30 -0300 Subject: [PATCH 3/5] fix(slurm): preserve retry confirmation intent --- .../src/data_designer/slurm/cli.py | 33 ++++++++++++++++--- .../slurm/services/retry_collection.py | 16 +++++---- .../src/data_designer/slurm/state/retry.py | 23 +++++++++++++ .../tests/services/test_wiring.py | 6 +++- .../data-designer-slurm/tests/test_cli.py | 31 ++++++++++++++++- 5 files changed, 96 insertions(+), 13 deletions(-) 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 105d39283..29819e19e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -16,6 +16,7 @@ from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config from data_designer.slurm.contracts import canonical_json from data_designer.slurm.services import ( + SlurmRetryExecution, SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, @@ -115,9 +116,13 @@ def retry_command( """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), + ) - def retry(*, preview: bool) -> BaseModel: - return create_slurm_run_service(profile_file=profile_file, cluster=cluster).retry( + def retry(*, preview: bool) -> SlurmRetryExecution: + return service.retry( run_or_job_id, shard_ids=shard_ids, resume=resume.value, @@ -125,9 +130,27 @@ def retry(*, preview: bool) -> BaseModel: ) if not dry_run and not force: - _invoke(operation, lambda: retry(preview=True)) - if not click.confirm("Submit this retry?", default=False, err=True): - raise typer.Exit() + planned = _invoke(operation, lambda: retry(preview=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, lambda: retry(preview=dry_run), 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 index b8c969a34..4732a957e 100644 --- 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 @@ -59,13 +59,17 @@ def retry( effective_resume_mode = self._resolve_retry_resume_mode(resolved_plan, resume) preview: tuple[RetryPlan, str] | None = None if effective_resume_mode is None: - preview = coordinator.preview( - shard_ids=shard_ids, - effective_resume_mode="never", - observed_at=observed_at, - ) + preview = coordinator.preview_active(shard_ids=shard_ids, observed_at=observed_at) + if preview 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) + else: + effective_resume_mode = preview[0].effective_resume_mode shard_ids = tuple(shard.shard_id for shard in preview[0].planned_shards) - effective_resume_mode = self._resolve_if_possible_resume_mode(resolved_plan, preview[0], operation) if dry_run: if preview is None or preview[0].effective_resume_mode != effective_resume_mode: preview = coordinator.preview( 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 0e4bff964..bc58d26ef 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 @@ -135,6 +135,29 @@ def preview( 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_status(retry_ids[-1]) + 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/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index f801df4c2..ae5da03ce 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -512,7 +512,10 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( 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(submission_job_ids=(42, 43)) @@ -525,6 +528,7 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( 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( diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index d984b23ce..a9e2407f4 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -146,9 +146,10 @@ def test_retry_auto_selects_tasks_and_confirms_submission(monkeypatch) -> None: "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", None, "if_possible", False), + ("42", ("shard-00000",), "always", False), ] @@ -162,6 +163,34 @@ def test_retry_force_skips_confirmation(monkeypatch) -> None: 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) From 768cb9d4984ceb54ff357c28aec501b9849aa438 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 10 Sep 2026 18:40:10 -0300 Subject: [PATCH 4/5] fix(slurm): recover active retry previews --- .../slurm/services/retry_collection.py | 25 ++++++------- .../src/data_designer/slurm/state/retry.py | 4 ++- .../slurm/state/retry_storage.py | 12 ++++++- .../tests/services/test_wiring.py | 35 +++++++++++++++++++ .../tests/state/test_retry_collection.py | 28 +++++++++++++++ 5 files changed, 90 insertions(+), 14 deletions(-) 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 index 4732a957e..46de9f076 100644 --- 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 @@ -57,18 +57,19 @@ def retry( 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: tuple[RetryPlan, str] | None = None - if effective_resume_mode is None: - preview = coordinator.preview_active(shard_ids=shard_ids, observed_at=observed_at) - if preview 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) - else: - effective_resume_mode = preview[0].effective_resume_mode + 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: 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 bc58d26ef..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 @@ -148,7 +148,9 @@ def preview_active( retry_ids = self._retries.list_retry_ids() if not retry_ids: return None - retry_status = self._retries.read_status(retry_ids[-1]) + 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 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/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index ae5da03ce..6d0b5ae2d 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -561,6 +561,41 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( 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, 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 836d7609d..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, From 44158adf98148b53126966ada0285832c05fec10 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 11 Sep 2026 09:59:20 -0300 Subject: [PATCH 5/5] fix(slurm): address retry and merge review feedback --- .../src/data_designer/slurm/cli.py | 29 ++++--- .../src/data_designer/slurm/services/run.py | 83 +++++++++++++------ .../data_designer/slurm/state/collection.py | 2 + .../data_designer/slurm/state/destinations.py | 20 ++++- .../src/data_designer/slurm/state/observer.py | 4 +- .../tests/launcher/test_collection.py | 30 +++++++ .../tests/services/test_wiring.py | 3 + .../tests/state/test_observer.py | 23 +++++ 8 files changed, 157 insertions(+), 37 deletions(-) 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 29819e19e..390a07048 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -6,6 +6,7 @@ import re from collections.abc import Callable from enum import Enum +from functools import partial from pathlib import Path from typing import NoReturn, TypeVar @@ -16,7 +17,6 @@ from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config from data_designer.slurm.contracts import canonical_json from data_designer.slurm.services import ( - SlurmRetryExecution, SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, @@ -121,16 +121,17 @@ def retry_command( lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster), ) - def retry(*, preview: bool) -> SlurmRetryExecution: - return service.retry( - run_or_job_id, - shard_ids=shard_ids, - resume=resume.value, - dry_run=preview, - ) - if not dry_run and not force: - planned = _invoke(operation, lambda: retry(preview=True)) + 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, @@ -153,7 +154,13 @@ def retry(*, preview: bool) -> SlurmRetryExecution: resume = _RetryResumeMode(planned.effective_resume_mode) result = _invoke( operation, - lambda: retry(preview=dry_run), + partial( + service.retry, + run_or_job_id, + shard_ids=shard_ids, + resume=resume.value, + dry_run=dry_run, + ), ) _emit_result(result) 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 3fddb1803..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 @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Sequence +from functools import partial from pathlib import Path from typing import Literal, Protocol @@ -247,22 +248,17 @@ def retry( raise _make_invalid_request_error(operation, "dry_run must be a boolean") backend = self._require_backend(operation) - def retry_run() -> SlurmRetryExecution: - result = backend.retry( + return _invoke_service_backend( + operation, + partial( + _retry_run, + backend, normalized_reference, shard_ids=normalized_shards, resume=resume, dry_run=dry_run, - ) - if not isinstance(result, SlurmRetryExecution): - raise TypeError("run backend returned an invalid retry result") - if normalized_shards is not None and result.shard_ids != normalized_shards: - 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 - - return _invoke_service_backend(operation, retry_run) + ), + ) def collect( self, @@ -281,21 +277,16 @@ def collect( normalized_destination = Path(destination).expanduser().resolve() backend = self._require_backend(operation) - def collect_run() -> SlurmCollectionExecution: - result = backend.collect( + return _invoke_service_backend( + operation, + partial( + _collect_run, + backend, normalized_input, destination=normalized_destination, num_partitions=num_partitions, - ) - if not isinstance(result, SlurmCollectionExecution): - raise TypeError("run backend returned an invalid collection result") - if result.output_path != normalized_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 - - return _invoke_service_backend(operation, collect_run) + ), + ) def _require_backend(self, operation: SlurmServiceOperation) -> SlurmRunBackend: if self._backend is None: @@ -307,6 +298,50 @@ 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) 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 e194367b9..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,6 +293,8 @@ 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 StateConflictError("persisted collection destination does not match the requested destination") 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 5ab547040..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 @@ -184,7 +185,8 @@ def _refresh_shard( 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( 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_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 6d0b5ae2d..97280a7c2 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -764,11 +764,14 @@ def test_collection_rejects_unmanaged_input_and_destination( 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 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)