feat(slurm): wire retry and merge commands

packages/data-designer-slurm/src/data_designer/slurm/cli.py

@@ -5,6 +5,8 @@
 
 import re
 from collections.abc import Callable
+from enum import Enum
+from functools import partial
 from pathlib import Path
 from typing import NoReturn, TypeVar
 
@@ -33,6 +35,13 @@
     SlurmServiceErrorCode.INTERNAL: 1,
 }
 
+
+class _RetryResumeMode(str, Enum):
+    NEVER = "never"
+    ALWAYS = "always"
+    IF_POSSIBLE = "if_possible"
+
+
 app = typer.Typer(
     name="slurm",
     help="Run Data Designer workloads on Slurm",
@@ -98,6 +107,89 @@ def cancel_command(
     _emit_result(result)
 
 
+@app.command("retry")
+def retry_command(
+    run_or_job_id: str = typer.Argument(..., help="Managed run ID or Slurm array job ID"),
+    task_ids: list[int] | None = typer.Option(None, "--task-id", min=0, help="Array task ID to retry; repeatable"),
+    resume: _RetryResumeMode = typer.Option(_RetryResumeMode.IF_POSSIBLE, "--resume"),
+    dry_run: bool = typer.Option(False, "--dry-run"),
+    force: bool = typer.Option(False, "--force", help="Submit without confirmation"),
+    profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False),
+    cluster: str | None = typer.Option(None, "--cluster"),
+) -> None:
+    """Retry failed shards from immutable persisted run state."""
+    operation = SlurmServiceOperation.RETRY_RUN
+    shard_ids = None if task_ids is None else tuple(f"shard-{task_id:05d}" for task_id in task_ids)
+    service = _invoke(
+        operation,
+        lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster),
+    )
+
+    if not dry_run and not force:
+        planned = _invoke(
+            operation,
+            partial(
+                service.retry,
+                run_or_job_id,
+                shard_ids=shard_ids,
+                resume=resume.value,
+                dry_run=True,
+            ),
+        )
+        typer.echo(
+            f"Retry {', '.join(planned.shard_ids)} with resume={planned.effective_resume_mode}",
+            err=True,
+        )
+        try:
+            confirmed = click.confirm("Submit this retry?", default=False, err=True)
+        except click.Abort:
+            typer.echo(err=True)
+            _fail(
+                SlurmServiceError(
+                    SlurmServiceErrorCode.INVALID_REQUEST,
+                    operation,
+                    "interactive confirmation is unavailable; pass --force or --dry-run",
+                )
+            )
+        if not confirmed:
+            _emit_json({"operation": operation.value, "state": "declined"})
+            return
+        shard_ids = planned.shard_ids
+        resume = _RetryResumeMode(planned.effective_resume_mode)
+    result = _invoke(
+        operation,
+        partial(
+            service.retry,
+            run_or_job_id,
+            shard_ids=shard_ids,
+            resume=resume.value,
+            dry_run=dry_run,
+        ),
+    )
+    _emit_result(result)
+
+
+@app.command("merge")
+def merge_command(
+    input_path: Path = typer.Option(..., "--input-path", file_okay=False),
+    output_path: Path = typer.Option(..., "--output-path", file_okay=False),
+    num_partitions: int | None = typer.Option(None, "--num-partitions", min=1),
+    profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False),
+    cluster: str | None = typer.Option(None, "--cluster"),
+) -> None:
+    """Submit winner-driven collection as a zero-GPU Slurm job."""
+    operation = SlurmServiceOperation.COLLECT_RUN
+    result = _invoke(
+        operation,
+        lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).collect(
+            input_path,
+            destination=output_path,
+            num_partitions=num_partitions,
+        ),
+    )
+    _emit_result(result)
+
+
 @profile_app.command("init")
 def profile_init_command(
     workspace_root: Path = typer.Option(..., "--workspace-root", file_okay=False),

packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py

@@ -23,9 +23,11 @@
     create_slurm_profile_service,
 )
 from data_designer.slurm.services.results import (
+    SlurmCollectionExecution,
     SlurmPersistedAttemptStatus,
     SlurmPersistedRunStatus,
     SlurmPersistedShardStatus,
+    SlurmRetryExecution,
     SlurmRunCancellation,
     SlurmRunExecution,
 )
@@ -48,6 +50,7 @@
     "SlurmBatchScriptRenderer",
     "SlurmBenchmarkBackend",
     "SlurmBenchmarkService",
+    "SlurmCollectionExecution",
     "SlurmImageManager",
     "SlurmImageResolver",
     "SlurmImageService",
@@ -58,6 +61,7 @@
     "SlurmProfileMatch",
     "SlurmProfileService",
     "SlurmProfileValidation",
+    "SlurmRetryExecution",
     "SlurmRunArtifactPublisher",
     "SlurmRunBackend",
     "SlurmRunCancellation",

packages/data-designer-slurm/src/data_designer/slurm/services/errors.py

@@ -32,6 +32,8 @@ class SlurmServiceOperation(str, Enum):
     EXECUTE_RUN = "execute_run"
     STATUS_RUN = "status_run"
     CANCEL_RUN = "cancel_run"
+    RETRY_RUN = "retry_run"
+    COLLECT_RUN = "collect_run"
     INIT_PROFILE = "init_profile"
     VALIDATE_PROFILE = "validate_profile"
     RESOLVE_IMAGE = "resolve_image"

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",
 ]