fix(slurm): close one-node M2 integration

[andreatnvidia]

Summary

• replace host Python bootstrap with a Bash 4.3 controller that invokes Python only inside sealed client or serving images
• persist immutable inputs and submitted attempts before releasing held Slurm jobs
• add typed container phase/bootstrap manifests, scoped secrets and GPUs, readiness, cleanup, winner finalization, and vLLM queue backpressure
• package the Bash/runtime sources into the digest-bound runtime archive

Validation

• .venv/bin/ruff check --fix . && .venv/bin/ruff format .
• make check-slurm
• .venv/bin/pytest -q packages/data-designer-slurm/tests - 1309 passed
• make build-slurm
• make test-slurm-wheel-install
• GNU Bash 4.3.0 syntax and empty/sparse-array helper execution
• CPU allocation host-tool checks and zero-GPU Pyxis container smokes
• End-to-end sealed-image model-serving smokes, including concurrent shared-node allocations

Landing order

Land this integration prerequisite first. Then:

• feat: reconcile persisted Slurm state #913 should rebase and drop the overlapping resolved-plan provenance, one-node dataset lease/finalization, and launcher submission hunks while retaining fresh-process observation and reconciliation.
• feat(slurm): harden distributed allocation runtime #914 should rebase and drop the overlapping one-node queue-backpressure and runtime/controller/step hunks while retaining distributed topology, multi-node launch, and failure hardening.

Part of #868
Part of #874

[greptile-apps]

Greptile Summary

The PR completes the one-node Slurm M2 runtime by moving orchestration into a packaged Bash controller and sealed container phases.

• Persists run inputs and attempts before releasing held Slurm jobs.
• Adds translated container paths, scoped secrets and GPUs, readiness, cleanup, candidate finalization, and queue backpressure.
• Remaps shared-node ports from assigned GRES GPU identities and uses exclusive scheduling for visible-GPU profiles.
• The changes since the previous review align client endpoint validation with allocation-local port remapping.

Confidence Score: 5/5

The PR appears safe to merge; no new actionable issue or outstanding previous finding remains.

The current code translates attempt and dependency paths into container locations, finalizes the winner before marking success, removes unsupported client arguments, accepts GPU supersets, resolves nested mounts independently, and isolates ports through GRES-derived blocks or exclusive visible-GPU scheduling.

Important Files Changed

Filename	Overview
packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py	Remaps planned ports into deterministic allocation-local blocks derived from assigned GRES GPUs.
packages/data-designer-slurm/src/data_designer/slurm/client/execution.py	Translates mounted paths and validates model endpoints against the allocation-remapped plan.
packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py	Builds typed sealed-container steps with scoped secrets and allocation GPU metadata.
packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py	Publishes client candidates before atomically validating and finalizing the winning attempt.
packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py	Validates container-visible paths while retaining canonical host paths in allocation state.
packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh	Coordinates preflight, serving readiness, generation, cleanup, and terminal phases under Bash 4.3.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    Submit[Persist run and held attempt] --> Release[Release Slurm job]
    Release --> Bash[Bash allocation controller]
    Bash --> Prepare[Sealed client prepare phase]
    Prepare --> Preflight[Client preflight]
    Preflight --> Servers[Start serving containers]
    Servers --> Ready[Readiness checks]
    Ready --> Client[Run generation client]
    Client --> Candidate[Publish candidate result]
    Candidate --> Finalize[Validate and publish winner]
    Finalize --> Cleanup[Stop steps and finalize allocation]

Loading

_{Reviews (14): Last reviewed commit: "fix(slurm): validate allocation-local en..." | Re-trigger Greptile}

[nabinchha]

Thanks for putting this together, @andreatnvidia — this is a substantial step toward closing the one-node M2 path.

Summary

The change follows the stated direction: it moves host orchestration to Bash, runs Python only inside the sealed images, stages a digest-bound runtime bundle, and persists attempts before releasing held Slurm work. I confirmed the three existing P1 inline findings and have not repeated them below; they are valid merge blockers.

Findings

Critical — Let's fix these before merge

packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py:166 — The exported controller repeats the irreversible success-before-winner transition

• What: run() persists the terminal SUCCEEDED attempt through _persist_terminal_outcome() and only then calls finalize_winner(). The existing inline thread identifies this ordering in the new Bash control path's entrypoint.py, but the same invariant breach remains in OneNodeAllocationController, which is still exported from data_designer.slurm.runtime.
• Why: If candidate verification or winner publication fails, the attempt remains immutable and successful without a winner. A caller cannot mark it failed or retry it, so this path can leave the same uncollectable shard state even after the entrypoint.py occurrence is fixed.
• Suggestion: Make success plus winner publication recoverable as one state protocol and use it from both controller paths (or remove the superseded exported controller if it is no longer supported). Please add a failure-injection test where finalize_winner() raises and assert that the attempt remains recoverable.

What Looks Good

• The digest-bound runtime archive and typed, secret-free bootstrap manifest create a clear trust boundary between submission and allocation execution.
• Holding the array until state publication is the right direction for preventing jobs from racing ahead of durable run state.
• The new runtime, launcher, service, state, and Bash-controller coverage is broad, and all changed Python files pass Ruff and formatting checks.

Residual Risk

The focused runtime/launcher/service/state suite passed locally (289 tests). The full Slurm suite reached 1,308 passing tests; the remaining package-entry-point test could not discover data-designer-slurm because this isolated review worktree was not installed into the shared virtual environment. The PR also notes that a full model-serving M2 smoke was not run, so the sealed-image/Pyxis execution path remains the main integration risk.

Verdict

Needs changes. Please resolve the three existing P1 threads and make the success/winner transition recoverable in the still-exported Python controller as well. After those state and path invariants are fixed, the overall implementation shape looks sound.

---

This review was generated by an AI assistant.

[nabinchha]

Thanks for the quick follow-up, @andreatnvidia — the new commit closes several of the concrete integration gaps from the first pass.

Summary

I re-reviewed the new head rather than carrying the earlier verdict forward. The container-path handling and partial-submission compensation are now materially better, and success/winner publication has been centralized in the state finalizer. Greptile has already captured the separate client-command blocker on this head, so I have not duplicated it below.

Findings

Critical — Let's fix this before merge

packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py:196 — The success/winner transition still has an unrecoverable process-death window

• What: _publish_verified_resolution() durably replaces the attempt with SUCCEEDED and then publishes winner.json as a second write. The rollback handles a Python exception from publish_winner(), but it cannot run if the control process is killed or the node disappears after replace_attempt() and before publish_winner(). On the Bash path, the EXIT cleanup also cannot repair this state: _fail() returns immediately when it reloads a SUCCEEDED attempt, and I could not find another production reconciliation path that resumes winner publication for a successful attempt without a winner.
• Why: This leaves a shard durably reported as successful but without the immutable record needed to collect its result. Because success is terminal and the cleanup path skips it, the run cannot converge after exactly the interruption this protocol is intended to survive.
• Suggestion: Make the two-record transition restartable—for example, persist a durable intermediate/finalization intent that reconciliation can resume, or use one authoritative commit record from which both views can be derived. Add an interruption-recovery test that stops between the successful-attempt write and winner publication; exception injection inside publish_winner() does not exercise this window.

What Looks Good

• Container phases now validate/read through the container-visible attempt directory while retaining host paths for the host-side state layout.
• Submission compensation now updates the attempts that were actually persisted and tolerates a missing later attempt after partial publication.
• The new exception-path tests demonstrate rollback for ordinary publication errors and convergence when the winner link was committed before an fsync error.
• All changed Python files pass Ruff and formatting checks, and the four affected test modules pass locally (124 tests).

Residual Risk

The full Slurm suite reached 1,314 passing tests; its only failure was the package-entry-point check because this isolated review worktree is not installed as data-designer-slurm in the shared environment. The exploratory complexity check also flags _SystemRunBackend.execute() at 13 versus the configured exploratory threshold of 10 after the added compensation branches; extracting the repeated cancel-and-record-failure flow would make those recovery paths easier to audit, but I am not treating that as a merge blocker.

Verdict

Needs changes. The new head addresses the earlier path and catchable-exception recovery issues, but the separate successful-attempt and winner writes remain non-atomic across process or node loss. The existing Greptile blocker on this head also remains outstanding.

---

This review was generated by an AI assistant.

[andreatnvidia]

Thanks, @nabinchha. 3419ee51 addresses the first-pass findings.

We also completed a full one-node model-serving smoke through readiness, record generation, and winner publication. I’ll address the remaining review points and run a few more end-to-end smokes before merge.

[andreatnvidia]

Thanks, @nabinchha. I missed your second review earlier. 458dc28b adds automatic recovery for the success-to-winner interruption window, including process-death coverage, and addresses the remaining review and smoke-test findings. The full Slurm suite passes with 1,323 tests. I’ll run a few more end-to-end smokes on this head and report back.

[nabinchha]

Thanks for the follow-up, @andreatnvidia — the new head fixes both the unsupported client arguments and the process-death window from my previous review. I also verified that the full Slurm test suite now passes locally.

Summary

The recovery path is now restartable: a fresh writer can detect a successful attempt without a winner and complete publication during status reconciliation. The host/container path work also handles a non-identity workspace mount in the common single-mount case. One supported mount-layout edge case still loses the plan's most-specific mapping, though, so I don't think this head is ready to merge yet.

Findings

Warning — Please fix before merge

packages/data-designer-slurm/src/data_designer/slurm/state/storage.py:92 — Finalization ignores more-specific mounts below the workspace root

• What: StateStorage.get_local_path() maps every persisted logical path by stripping logical_workspace_root and appending it to one workspace_root. That is only equivalent to get_container_path() when every path below the workspace uses the same mount. Profiles currently allow nested mount sources, and the runtime deliberately selects the most-specific source. With a valid workspace mount /host/workspace -> /container/workspace plus /host/workspace/runs/run-1/shards/shard-00000/attempts -> /fast/attempts, the client writes the candidate to /fast/attempts/attempt-0001/dataset, while the finalizer resolves the same persisted path to /container/workspace/runs/run-1/shards/shard-00000/attempts/attempt-0001/dataset.
• Why: Winner finalization and resume-workspace handling then read or create the wrong directory. A run using a more-specific writable mount can complete generation but fail artifact verification/winner publication, even though the plan's mount mapping is valid and the client used it correctly.
• Suggestion: Make state-side path resolution use the resolved plan's full most-specific mount mapping (or reject overlapping/nested mount sources as an explicit profile invariant). Please add a regression test with a nested writable mount beneath the workspace and verify that client output, artifact verification, and resumed finalization all resolve to the same physical path.

What Looks Good

• resume_incomplete_finalization() plus the fresh-writer interruption test directly closes the durable SUCCEEDED-without-winner window from the previous head.
• The bootstrap command now removes only the unsupported shard/attempt arguments while retaining the plan, attempt directory, and endpoint arguments.
• Schema-digest normalization now matches the state verifier's metadata-independent contract.
• Scheduler reconciliation is conservative on unavailable scheduler queries and preserves terminal attempt immutability.

Verification

• Full Slurm suite: 1,323 passed.
• Ruff and formatting checks pass across the Slurm package.
• Current GitHub checks are green, and the PR is reported mergeable.

Verdict

Needs changes. The prior blockers are fixed, but state-side path resolution still diverges from the runtime's documented most-specific mount behavior for nested mounts.

---

This review was generated by an AI assistant.

[nabinchha]

Thanks for the careful follow-through, @andreatnvidia — the latest fixes close the remaining mount and runtime-path gaps cleanly.

Summary

This head now applies the resolved plan's most-specific mount mapping consistently across state finalization, client bootstrap artifacts, attempt-local overlays, and absolute serving-model paths. It also aligns the Bash and Python GPU preflight checks so exclusive allocations with extra visible GPUs are accepted as intended.

Findings

No new findings. The previous nested-mount warning and the existing GPU/dependency-path inline findings are resolved on this head.

What Looks Good

• State finalization now resolves logical dataset paths through the same plan-aware mapping used by the runtime and client, including interrupted-finalization recovery through a nested attempt mount.
• Client bootstrap resolves the plan, lock, input artifacts, wheels, and writable overlay independently, so a more-specific mount for one artifact no longer incorrectly relocates its siblings.
• The Bash and Python GPU checks now share the same minimum-count contract, with direct regression coverage for GPU supersets.
• Absolute vLLM model paths are translated into the serving container while Hugging Face model identifiers remain unchanged.

Residual Risk

Local validation cannot reproduce the real Slurm/Pyxis environment. The exact head should still receive the intended sealed-image cluster smoke, especially for the newly mapped absolute model path.

Verdict

Ship it. The reviewed blockers are addressed, the implementation matches the PR's one-node M2 integration goal, and the current Slurm test suite passes locally.

---

This review was generated by an AI assistant.

[nabinchha]

Thanks for the quick follow-up, @andreatnvidia — the GPU-owned port blocks close the shared-node collision identified on the previous head.

Summary

This revision keeps visible-GPU allocations exclusive while allowing GRES allocations to share a node, and maps each GRES allocation's runtime service claims into a deterministic 256-port block owned by its lowest allocated global GPU ID. The scheduler environment is carried consistently through preflight, bootstrap, entrypoint, and the exported controller path.

Findings

No merge-blocking findings on this head. The previous shared-node service-port P1 is resolved.

What Looks Good

• GPU 0 and GPU 1 receive disjoint service-port blocks, so independently submitted jobs sharing a node cannot select the same remapped runtime ports.
• Visible-GPU profiles retain #SBATCH --exclusive, preserving the contract for profiles where the scheduler cannot provide a stable GRES allocation identity.
• Invalid or missing SLURM_JOB_GPUS values fail preflight instead of silently falling back to a collision-prone identity.
• The new environment plumbing and remapped deployment selection are covered across the Python and Bash runtime paths.

Verification

• Full Slurm suite: 1,336 passed.
• Ruff and formatting checks pass for the changed Python files.
• GitHub reports the PR mergeable with all current required checks green; Greptile also completed successfully on this exact head.

Residual Risk

An explicitly configured OTEL metrics port is not remapped into the GPU-owned block. Two opt-in jobs choosing the same metrics port can cause one exporter to disable metrics, but that bind failure is deliberately non-fatal and Slurm defaults telemetry to disabled, so I do not consider it a merge blocker. A real shared-node Slurm/Pyxis smoke remains the best integration validation.

Verdict

Approved. The current head is good to merge.

---

This review was generated by an AI assistant.

[andreatnvidia]

@nabinchha, final smoke update: we completed end-to-end sealed-image serving for a single GPU, concurrent shards, and two simultaneous one-GPU tasks sharing a node. We also exercised interrupted winner publication, cancellation reconciliation, and nested mount layouts.

The resulting fixes cover restartable winner publication, consistent client and mount-path handling, partial-GPU scheduling without whole-node reservation, and deterministic GPU-owned service ports. The final runs published winners and produced the expected outputs. Could you take one final look at ae26a95f?

[nabinchha]

Thanks for the final pass, @andreatnvidia — carrying the allocation identity into the sealed client closes the last port-validation mismatch cleanly.

Summary

The current head makes client preflight reconstruct the same allocation-local plan used to build the runtime endpoints, with SLURM_JOB_GPUS forwarded through both the Bash manifest and the Python srun builder. The full PR now matches its one-node M2 integration intent across persistence, sealed-container execution, shared-node isolation, finalization, and recovery.

Findings

No findings.

What Looks Good

• resolve_allocation_plan() is now the single mapping operation used by both runtime deployment construction and client endpoint validation, while the persisted plan remains unchanged for provenance.
• The allocation identity is forwarded only where the zero-GPU client needs it, preserving the existing secret and GPU scoping of the sealed steps.
• The regression coverage rejects an endpoint from another GPU's port block and exercises both client-step construction paths.
• The earlier path, finalization, GPU-superset, nested-mount, and shared-node collision findings remain resolved on this head.

Verification

• Focused client/runtime validation: 37 passed.
• Full Slurm suite: 1,337 passed.
• Ruff and formatting checks pass for all seven Python files changed since the previous reviewed head.
• GitHub reports this exact head clean and mergeable with current checks green.

Residual Risk

I could not independently reproduce the real Slurm/Pyxis environment locally. Andre reports successful sealed-image smokes covering a single GPU, concurrent shards, two simultaneous one-GPU tasks sharing a node, interrupted winner publication, cancellation reconciliation, and nested mounts.

Verdict

Ship it. The current head is ready to merge.

---

This review was generated by an AI assistant.