feat: add Slurm retry and deterministic collection

[nabinchha]

📋 Summary

Implements the #869 retry-and-collection slice for persisted Slurm runs. Failed shards can be retried from fresh reconciliation state, and exactly one validated winner per planned shard is collected by a zero-GPU CPU Slurm job into a deterministic, atomically published dataset.

This PR is stacked on #913. PR #913 must merge first. The branch was restacked after #910 merged into the shared Slurm base and remains exactly four slice commits above the synchronized #913 head (506de489). It consumes #913's public observation/reconciliation APIs and adds capability APIs only; CLI wiring remains owned by #874#3.

🔗 Related Issue

Closes #869

🔄 Changes

✨ Added

• Persisted retry plans/status with exact sparse shard ownership, resume-and-shard locking, scheduler-identity validation, ambiguous-submission recovery, and winner-aware idempotency.
• Persisted collection plans/status/results with normalized profile-authorized destinations and exact winner/result/attempt/candidate validation.
• A zero-GPU CPU collection launcher and worker that performs bounded-memory Parquet input reads and deterministic Parquet, CSV, or JSONL partitioning.
• Same-parent staging and atomic no-overwrite publication with descriptor/path rebinding, exact-stage cleanup, and interruption/collision recovery.
• Adversarial regression coverage for stale state, conflicting winners, journal interruption, concurrent submissions, path replacement, output mutation, and publication races.

🔧 Changed

• Normalized launcher submission failures into definite versus potentially-successful outcomes so persisted operations can recover safely.
• Exposed the new retry and collection capabilities through public Slurm state APIs without importing sibling-private names or changing CLI behavior.

🐛 Fixed

• Reconcile ambiguous retry and collection submissions from fresh processes by an immutable plan-derived job name, explicit current-UID ownership, and exact ordinary/sparse-array shape before binding the scheduler ID.
• Persist a bounded visibility deadline so accepted submissions recover without resubmission, while absent submissions can be retired and safely replaced.
• Fence delayed ambiguous allocations at the persisted retry-attempt or collection scheduler identity so only the recorded job can write, including when an earlier accepted allocation becomes visible after replacement.
• Bound partial sparse-array visibility: strict subsets wait through the reconciliation deadline, then retire as bounded absence so a retry cannot remain permanently PREPARED.
• Let a fast collection allocation wait outside the state lock for its exact scheduler binding instead of failing while the durable status is still PREPARED.

🧪 Testing

• Full Slurm suite after the feat: finalize Slurm shard winners #910/feat: reconcile persisted Slurm state #913 restack — 1,318 passed
• Focused retry, collection, submission-recovery, launcher, and store tests — 163 passed
• make check-slurm
• Isolated wheel build/install/import and lazy CLI extension smoke test passes
• Unit tests added/updated
• E2E tests added/updated (worker/launcher integration through Slurm fakes)

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated (not applicable; this implements the approved Stage 2 plan)

---

Description updated with AI

[greptile-apps]

Greptile Summary

The PR adds persisted retry and deterministic collection capabilities for Slurm runs, including recovery of ambiguous submissions.

• Reconciles prepared retry and collection submissions through exact job-name, owner, and scheduler-shape matching.
• Bounds incomplete scheduler visibility with a persisted reconciliation deadline.
• Adds sparse retry arrays, allocation identity fencing, deterministic collection, and atomic dataset publication.

Confidence Score: 5/5

The PR appears safe to merge because the previously reported permanent PREPARED-state failures are resolved and no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py	Implements deadline-bounded recovery with exact ordinary or sparse-array shape matching, resolving the previously reported permanent PREPARED states.
packages/data-designer-slurm/src/data_designer/slurm/state/retry.py	Reconciles prepared retry submissions before proceeding and binds recovered scheduler identities into persisted retry state.
packages/data-designer-slurm/src/data_designer/slurm/state/collection.py	Adds persisted collection submission, reconciliation, refresh, result validation, and failed-stage cleanup.
packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py	Adds exact-name current-user submission lookup across active and accounting scheduler views and classifies ambiguous submission failures.
packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py	Adds deterministic bounded-row Parquet merging and partitioned Parquet, CSV, or JSONL output.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  P[Persist PREPARED operation] --> S[Submit Slurm job]
  S -->|Receipt received| B[Bind scheduler identity]
  S -->|Ambiguous outcome| Q[Query squeue and sacct by exact name and UID]
  Q -->|Exact shape match| B
  Q -->|Absent or partial before deadline| W[Remain PREPARED and retry reconciliation]
  W --> Q
  Q -->|Bounded absence after deadline| F[Mark operation FAILED]
  F --> R[Permit replacement submission]
  B --> E[Execute identity-fenced retry or collection]

Loading

_{Reviews (8): Last reviewed commit: "fix(slurm): bound partial array recovery" | Re-trigger Greptile}

[greptile-apps]

Ambiguous retries remain permanently blocked

When sbatch raises SlurmSubmissionError with may_have_succeeded=True, _submit leaves the retry in PREPARED, and every subsequent retry call rejects that state before reconciliation. This permanently prevents the run from retrying any failed shard, even after the ambiguous job has terminated or was never submitted.

Prompt To Fix With AI

This is a comment left during a code review.
Path: packages/data-designer-slurm/src/data_designer/slurm/state/retry.py
Line: 110-111

Comment:
**Ambiguous retries remain permanently blocked**

When `sbatch` raises `SlurmSubmissionError` with `may_have_succeeded=True`, `_submit` leaves the retry in `PREPARED`, and every subsequent retry call rejects that state before reconciliation. This permanently prevents the run from retrying any failed shard, even after the ambiguous job has terminated or was never submitted.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[nabinchha]

Fixed in b994ac3. PREPARED retry and collection submissions now reconcile from a fresh process by a plan-derived exact job name across current-UID squeue and sacct views, validate ordinary-versus-array shape and the complete sparse task set, and persist a bounded visibility deadline. Exactly one match is rebound to its real scheduler ID; multiple or mismatched matches fail safely; authoritative absence after the deadline retires the prepared operation so a later call can proceed. Retry attempt identity and collection worker identity fence any delayed old allocation from writing, and the collection worker waits outside the state lock when it starts before its ID is durably bound. Regression coverage includes accepted receipt loss and fresh-process recovery, delayed visibility, definite absence, multiple/wrong-shape matches, concurrent collection startup, and late-allocation fencing. Validation: 47 focused recovery tests passed; full Slurm suite 1,214 passed; Ruff/format/strict complexity and isolated wheel install/import smoke passed.

[greptile-apps]

Partial array views block recovery

If a completed sparse-array task has left squeue before appearing in sacct, the merged task IDs are a strict subset of the planned shape. This exact comparison raises StateConflictError before the visibility-deadline branch and leaves the retry in PREPARED, so every subsequent retry remains blocked.

Prompt To Fix With AI

This is a comment left during a code review.
Path: packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py
Line: 61-62

Comment:
**Partial array views block recovery**

If a completed sparse-array task has left `squeue` before appearing in `sacct`, the merged task IDs are a strict subset of the planned shape. This exact comparison raises `StateConflictError` before the visibility-deadline branch and leaves the retry in `PREPARED`, so every subsequent retry remains blocked.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[nabinchha]

Fixed in 359b53fc. Strict-subset sparse-array observations are now treated as incomplete scheduler visibility: they remain pending through the reconciliation deadline, then resolve to bounded absence so the PREPARED retry can fail and be safely replaced instead of blocking forever. Non-subset/ordinary shape mismatches still fail closed. Regression coverage exercises both the at-deadline wait and post-deadline retirement; focused recovery/retry tests pass (47), and the full Slurm suite passes (1,282).