feat: add Slurm state writer

[nabinchha]

Summary

Adds the durable state foundation for Slurm execution so run intent, shards, attempts, and readiness can be reconstructed safely from compute-visible storage. The implementation composes focused writer, reader, storage, plan-validation, and filesystem responsibilities instead of concentrating them in store.py.

Related Issue

Part of #869 (869#1).

Merge Order and Dependencies

• This PR targets feat/slurm-execution, which includes the merged image-publication hardening from fix: harden Slurm image publication #907.
• feat: add one-node Slurm runtime #909 and feat: finalize Slurm shard winners #910 are stacked on this PR and must be rebased/retargeted after it merges.

Changes

Added

• Durable initialization and persistence for run, config, plan, shard, attempt, and readiness state.
• Shared descriptor-bound filesystem primitives for private directories, restrictive regular files, verified locks, and atomic publication/replacement.
• Public package-internal state temporary-file creation and recognition helpers, keeping the complete naming convention under one owner.
• A public package-internal persisted-plan validator that keeps state validation independent from the integration facade.
• Explicit schema-v1 compatibility enforcement and frozen canonical-byte fixtures.
• Adversarial coverage for corruption, unsafe filesystem objects, interrupted writes, concurrent writers, and fresh-process recovery.

Changed

• Split state responsibilities across SlurmStateWriter, StateReader, StateStorage, and focused filesystem/validation collaborators.
• Validate shared run/plan bindings in constant time for both full-audit and shard-local mutation paths; shard-local operations still inspect only their target shard.
• Map only lock-acquisition failures at the image registry boundary, preserving operation-body OSError exceptions unchanged.
• Use package-internal, non-underscored collaborator names so modules do not import objects that claim to be private.
• Centralize boundary error normalization while preserving original exceptions as context.

Fixed

• Removed the state-to-integration import cycle.
• Prevented unrelated shard corruption from blocking healthy shard mutations or readiness updates.
• Restored run.shard_count versus persisted-plan validation on the shard-local mutation path.
• Made interrupted publication recovery rebind the committed path before returning.
• Made lock-free attempt reads validate the exact opened file, closing replacement races.
• Fsync parent directories when immutable records already exist identically, preserving crash-safe idempotency.
• Read persisted JSON with exact newline handling so CRLF content cannot bypass canonical-byte validation.

Attention Areas

Reviewers: Please pay special attention to:

• state/store.py — public state-transition API and short mutation workflows.
• state/reader.py — target-shard and full-audit snapshot composition.
• state/storage.py — physical state layout and descriptor-bound storage operations.
• state/plan_validation.py — persisted-plan compatibility and state binding.
• filesystem.py and state/filesystem.py — shared filesystem invariants and state-specific policy.

Testing

• make check-slurm
• make test-slurm — 1,119 passed
• make test-slurm-wheel-install
• Corruption, recovery, concurrency, schema compatibility, layering, and shard-local regressions

Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs: N/A — this is an internal persistence boundary

---

Description updated with AI

[greptile-apps]

Greptile Summary

The PR adds durable, plan-bound Slurm run state with descriptor-bound filesystem operations and separates persistence responsibilities into focused writer, reader, storage, and validation layers.

• Persists run, configuration, plan, shard, attempt, and readiness records using restrictive files and atomic publication.
• Supports shard-local mutations while retaining an explicit full-run audit path.
• Repairs interrupted immutable publication and retries reads when concurrent atomic replacement changes the opened record.
• Extracts shared filesystem primitives for Slurm state and image-registry storage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
packages/data-designer-slurm/src/data_designer/slurm/filesystem.py	Adds shared descriptor-bound directory, lock, and temporary-file primitives with restrictive ownership policies.
packages/data-designer-slurm/src/data_designer/slurm/state/filesystem.py	Implements atomic state publication, replacement, consistent reads, and interrupted-publication recovery; the previously reported recovery races are addressed.
packages/data-designer-slurm/src/data_designer/slurm/state/storage.py	Defines the physical state layout and maps typed records onto verified filesystem operations.
packages/data-designer-slurm/src/data_designer/slurm/state/reader.py	Composes persisted records into target-shard or full-audit snapshots bound to the resolved plan.
packages/data-designer-slurm/src/data_designer/slurm/state/store.py	Adds durable initialization and shard-scoped mutation workflows over the reader, storage, and validation collaborators.
packages/data-designer-slurm/src/data_designer/slurm/state/plan_validation.py	Extracts persisted-plan compatibility and identity validation from the integration facade.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Writer[SlurmStateWriter] --> Validator[PersistedPlanStateValidator]
  Writer --> Storage[StateStorage]
  Reader[StateReader] --> Validator
  Reader --> Storage
  Storage --> StateFS[State filesystem policy]
  StateFS --> SharedFS[Descriptor-bound filesystem primitives]
  SharedFS --> Disk[(Compute-visible storage)]

Loading

_{Reviews (10): Last reviewed commit: "fix: align state boundary invariants" | Re-trigger Greptile}

[andreatnvidia]

One thing worth deciding now is how persisted state should behave across package upgrades. This parses with the current model and then requires an exact byte-for-byte reserialization, so even adding a defaulted field to a v1 config, plan, or state record makes older state look corrupt. We reproduced that internally by omitting the defaulted array_tasks field from a DataDesignerSlurmConfig: it still validates, but no longer serializes to the original bytes, and the digest checks would fail too. It would be good to make the policy explicit here, either with versioned readers and migrations or a clear unsupported-version check plus compatibility fixtures that should not be casually regenerated.

[nabinchha]

Fixed in e253271. Persisted records now inspect schema_version before model parsing and report an explicit unsupported-version error. For schema v1, non-canonical bytes that only parse because of newly added defaults are rejected with the documented no-migration policy instead of being treated as silently compatible. A frozen compatibility fingerprint covers the authored config, resolved plan, and all state golden records, and regressions cover both an unknown version and the omitted-array_tasks case.

[andreatnvidia]

The dependency direction feels a little fragile here: state imports integration, while integration imports several state.* modules. The leaf imports and lazy writer export keep it working today, but a future eager export or another planning/client import could turn it into an import cycle. Maybe the persisted plan/state checks could live in a small module under state, with integration.py building on that for finalization validation?

[nabinchha]

Fixed in e253271. Persisted plan/state checks now live in state/plan_validation.py as PersistedPlanStateValidator. StateReader and SlurmStateWriter depend only on that state-layer validator; integration.PlanStateValidator builds on it for client/finalization checks. A subprocess regression asserts importing state.store does not import data_designer.slurm.integration.

[andreatnvidia]

This currently turns a shard-local operation into a full-run scan. Every attempt or readiness change reads and validates attempts for all shards; create and update also hold the run lock while doing it. With a large array, that will serialize the main write path and make unrelated shard state part of each operation. Since the plan already guarantees one array-task index per shard, a shard-local validation path seems sufficient here, leaving the all-run check for status or auditing.

[nabinchha]

Fixed in e253271. Attempt creation/update and readiness publication now acquire only the target shard lock, load one persisted shard, and validate only that shard's attempt sequence. The all-run validator remains available for explicit status/audit snapshots. A multi-shard regression corrupts an unrelated attempt directory and confirms target-shard attempt and readiness writes still succeed while direct reads of the corrupt shard fail.

[andreatnvidia]

We now have two versions of the same descriptor-bound filesystem building blocks, here and in images/filesystem.py, and they have already started to differ in link, permission, and replacement checks. The .state temporary-name convention is also defined separately in storage.py and in the generator here. This is security-sensitive code, so sharing the low-level primitives and a single temporary-name predicate would make future hardening much easier to keep consistent.

[nabinchha]

Fixed in e253271. A shared data_designer.slurm.filesystem module now owns private modes, descriptor-bound directory opens, single-link lock acquisition and replacement checks, restrictive temporary-file creation, stable file facts, and the single managed-temporary-name predicate. Both image and state filesystem adapters compose those primitives, while retaining their domain-specific record/publication policies.

[andreatnvidia]

[P2] This may become pretty expensive once #909 starts running larger Slurm arrays. Every readiness publication calls load_validated_attempts(), which reads and validates every attempt for every shard. Since each allocation publishes readiness several times, N active shards can produce roughly quadratic shared-filesystem work, and corruption in one unrelated shard prevents the others from reporting progress. Could readiness updates validate only the current shard and attempt, leaving run-wide checks for operations that actually need them?

[nabinchha]

Fixed in e253271. Readiness publication now loads and validates only its current shard and attempt under that shard's lock; it no longer calls the all-run attempt scan. The same shard-local path is used by attempt create/update. The regression deliberately corrupts another shard and verifies readiness progress for the target shard remains publishable.

[andreatnvidia]

The shared predicate helps, but the .state temporary-file convention is still defined twice. These prefix and suffix values are duplicated in state/filesystem.py, where the files are created. If they drift, recoverable leftovers become unexpected records and the attempt directory is reported as corrupt. Could state/filesystem.py expose a state-specific predicate or the constants so creation and recovery share the complete convention?

[nabinchha]

Fixed in d741a9d. state/filesystem.py now owns the complete state temporary convention through the public package-internal create_state_temporary_file(...) and is_state_temporary_name(...) pair. Creation, interrupted-publication repair, and StateStorage recovery all use that single owner; storage.py no longer duplicates the prefix/suffix. Added test_incomplete_attempt_recognizes_created_state_temporary_file.

[andreatnvidia]

One small regression from the extraction: because the yield is inside this try, an OSError from the protected registry operation is now reported as cannot lock image registry target. Previously only acquisition failures were translated, while errors from the operation kept their original context. Could we limit this exception mapping to lock acquisition so future failures are not misdiagnosed as lock problems?

[nabinchha]

Fixed in d741a9d. The image adapter now maps only failures from entering the shared lock context to ImageRegistryError; the protected operation runs outside that exception-mapping block and retains its original exception and context. Added test_registry_lock_preserves_operation_oserror.

[andreatnvidia]

The shard-local path looks much better, but it drops one constant-time run/plan check along with the full scan. validate_plan_shards() checks that run.shard_count matches the plan, while validate_plan_shard() does not. I reproduced the latter accepting an inflated count even though load_shards() rejects the same run as corrupt, so mutations can continue against state that status or audit will not load. Could we share that binding check between both paths? It keeps the hot path shard-local.

[nabinchha]

Fixed in d741a9d. validate_plan_shards(...) and shard-local validate_plan_shard(...) now share one constant-time run binding check covering run identity, authored config, plan reference, and shard_count. The hot path still reads only its target shard. Added test_shard_mutation_rejects_run_count_drift_without_scanning_other_shards.

[andreatnvidia]

Looks good from my side now. The follow-up closes the remaining state-boundary and filesystem gaps cleanly, with focused coverage for the cases we discussed. Thanks for turning these around so quickly!