Add Slurm integration (step operator + orchestrator) by htahir1 · Pull Request #5064 · zenml-io/zenml

htahir1 · 2026-07-10T22:56:16Z

What this does

Adds a Slurm integration so ZenML pipelines can run on an HPC cluster managed by the Slurm scheduler. It ships two stack components:

Step operator (slurm) — runs a single step as a Slurm batch job (delegating the rest of the pipeline to another orchestrator).
Orchestrator (slurm) — submits a whole pipeline as a graph of Slurm jobs wired together with sbatch --dependency=afterok chains.

Both run the step's standard ZenML Docker image on the compute node via a rootless HPC container runtime (Apptainer/Singularity by default, or NVIDIA Pyxis/enroot, or the Docker daemon where available), so code delivery uses ZenML's normal image build — nothing is shipped by the components. Submission goes over SSH to a login/controller node (reusing the ssh integration's SSHClient + connection mixin) or via local sbatch when the client already runs on the cluster.

Design decisions

Docker image, not a code tarball. An earlier approach shipped code as an archive on PYTHONPATH; that fought ZenML's entrypoint (which chdirs into the image's /app) and reimplemented delivery the framework already does. Switched to running the step image with an HPC container runtime.
No service connector. Follows the ssh integration's SecretField/connection-mixin pattern rather than a bespoke connector.
Detached orchestrator, modeled on the current ssh orchestrator (not the deprecated HyperAI). Submits the whole DAG upfront in topological order (via topsorted_layers); each step self-reports status, and the run id is the placeholder run id (unique per run), not the snapshot id.
Shared machinery factored out: flavors/base.py (config/settings mixins), slurm_job.py (container-command + sbatch-script builders, stage_and_submit, stack validator), and build_slurm_client (single transport factory) are shared by both components.

Reliability & security hardening

Following an external review, the failure/status and credential paths were hardened:

fetch_status reconciles a detached run from its Slurm jobs (squeue --jobs with a --states=all lookup for jobs that already left the active queue), catching failures that occur before the ZenML entrypoint starts (bad image, runtime error) which a step can never self-report. The missing-sentinel window is treated as still-running, not terminal.
Exact job-id status/cancellation; the Slurm job id is persisted as step-run metadata.
Private registry auth per runtime (Apptainer/Singularity *_DOCKER_USERNAME/PASSWORD, enroot .credentials, Docker config.json) via owner-only files — no passwords on any command line.
Credential lifecycle: owner-only staging (mkdir under umask 077, 0600 env file), an EXIT trap that scrubs credentials + records the exit code, staging rollback on submission failure, and an afterany cleanup/reaper job (depending on all step jobs) that scrubs credential files even for jobs cancelled before their trap runs.
Runtime safety: env-var names validated and values shell-escaped (Pyxis exports quoted values instead of sourceing a Docker-format file); Apptainer/Singularity run with --no-eval; CPU requests rounded up and memory emitted in MB so fractional/sub-GB resource settings produce valid #SBATCH directives.

Testing

43 unit tests cover config/validator, the client (submit/state/cancel by id), per-runtime container commands, registry auth (no password on the command line), env escaping, resource conversion, DAG dependency wiring, the reaper, and fetch_status. The Slurm CLI and SSH are never invoked — everything runs against an in-memory fake command runner.
ruff, ruff format, and mypy are clean on the package.
⚠️ Not yet validated on a real cluster. Integrations that talk to external services are typically verified locally/in CI against live infra; that step is pending a deployed Slurm cluster. Please run full CI.

Notes for reviewers

Closest reference points are the ssh orchestrator/step operator (SSH + Docker, detached submit) and ContainerizedOrchestrator/kubernetes for the image/DAG plumbing.
The run-metadata round-trip for the job-id map and the exact behavior on a real scheduler are the two things that can only be confirmed live.

🤖 Generated with Claude Code

Adds a `slurm` integration with a step operator that runs individual pipeline steps as Slurm batch jobs: - SlurmClient wraps the Slurm CLI (sbatch --parsable, squeue, scancel) behind a transport abstraction: LocalSlurmCommandRunner (client already on a login node) and SSHSlurmCommandRunner (reusing the ssh integration's SSHClient; the config mirrors its connection fields). - The operator is container-free by design: Slurm compute nodes cannot run Docker, so jobs activate a user-provided Python environment (env_setup_command) and unpack a code archive shipped at submission time into a per-run directory on the cluster. - Submission is stateless: jobs are named after the step run id, so get_status/cancel work by job name. Queue states come from squeue; after a job leaves the queue, a sentinel exit-code file written by the job script's EXIT trap decides success vs failure, so Slurm accounting (sacct) is never required. - ResourceSettings map onto sbatch directives (--cpus-per-task, --mem, --gres=gpu:N) plus partition/time/account/qos settings and a raw directive escape hatch. All unit tests drive the client and operator through an in-memory fake runner, so they run without a cluster. Live validation on a real Slurm cluster is the next step and will be reported on the issue. Refs #5063 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Add `set -eo pipefail` after the EXIT trap so a failing env_setup_command or a failed code-archive extraction aborts the job (with the failure code captured by the sentinel) instead of falling through to run the step in a broken environment. Verified by executing the generated script directly in bash (the #SBATCH directives are comments, so it is runnable without Slurm). Refs #5063 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ields Follow the ssh step operator's pattern: SlurmStepOperatorConfig now inherits SSHConnectionConfigMixin rather than re-declaring the SSH connection fields. This fixes a latent bug: SSHClient reads config.connection_timeout and config.keepalive_interval, which the hand-copied config never declared, so the ssh transport would have crashed with AttributeError on connect. The mixin provides them (and keeps every SSH field in sync with the client we already reuse). hostname/username are relaxed to optional (the mixin makes them required) because the `local` transport uses no SSH; the ssh transport still enforces them via the existing validator. The now-unnecessary cast in the SSH runner is removed. Decision recorded: no Slurm service connector. SSH-key auth is served by SecretFields + ZenML secret references exactly as the ssh operators do; a connector (HyperAI-style) would not fit the SSH-key model well. Refs #5063 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Surveyed the 11 other step operators; two gaps found and fixed: - Add a `validator` (every modern step operator has one). It rejects a local artifact store, since the cluster reads inputs / writes outputs there. Unlike the Docker-based operators it requires no container registry or image builder - the operator is container-free - matching the databricks operator, the only other container-free one. - Pass `info` (not `info.step_run`) to `get_settings` in the script builder, matching every other operator's `submit`. Verified consistent with convention (no change needed): not overriding `get_docker_builds` is correct - the StackComponent default returns [], and container-free operators (databricks) rely on it; the submit/get_status/cancel lifecycle plus the inherited `wait` matches the modern interface; and the default entrypoint config is right for the container-free path. Refs #5063 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ball The step operator previously shipped the step code as a tarball extracted on the cluster and put on PYTHONPATH. That fought the ZenML entrypoint (which chdirs into the image's /app) and reimplemented code delivery the framework already does well. Instead, declare the step's Docker build via get_docker_builds and run that image on the compute node with a rootless HPC container runtime (Apptainer/Singularity by default, or NVIDIA Pyxis, or the Docker daemon where available), so code delivery uses ZenML's standard, tested image path and no code is shipped by the operator. Security is handled up front: the environment file holding the step's credentials is written owner-only (0600) into a per-run directory created atomically owner-only (mkdir -m 700), passed to the container via --env-file (or, for pyxis, --container-env variable names sourced from that 0600 file) so values never appear on a command line visible to other cluster users, and scrubbed by the job's EXIT trap the moment the job ends. All interpolated paths and values go through shlex.quote. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Introduce a Slurm orchestrator that submits a whole pipeline as a graph of Slurm jobs. Each step becomes its own `sbatch` job running the step's ZenML Docker image on a compute node, and the pipeline DAG is reproduced with job dependencies (`--dependency=afterok`, plus `--kill-on-invalid-dep` so a failed upstream propagates instead of leaving jobs pending forever). Submission is fully detached: once the jobs are queued the scheduler runs them and each step reports its own status to the ZenML server, so the orchestrator never polls. Jobs are submitted upstream-first (topological order via the existing `topsorted_layers`) so each step's dependencies reference already-submitted job ids, and a mid-submission failure best-effort cancels the jobs already queued. The orchestrator and the existing step operator run the same kind of Slurm job, so the shared machinery is extracted to avoid divergence: - `flavors/base.py` holds the transport/runtime/`sbatch` settings and the SSH connection mixin both components inherit. - `slurm_job.py` owns the container-command and sbatch-script builders, the stack validator, and `stage_and_submit`, which centralizes the security-sensitive staging sequence (atomic `mkdir -m 700`, 0600 env file, 0700 script, submit) so the file-mode invariants cannot drift between the two components. - `build_slurm_client` in `slurm_client.py` is the single transport factory both components call. `SlurmClient.submit` gains optional dependency support for the DAG edges. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…shot The Slurm orchestrator was modelled on the HyperAI orchestrator, which is deprecated. Re-checked against the current orchestrators (the generic `ssh` orchestrator that HyperAI defers to, plus Kubernetes/Vertex/SageMaker/AzureML) and corrected the run-id handling to match them: - Use `placeholder_run.id` as the orchestrator run id, not `snapshot.id`. A snapshot can be executed many times, so `snapshot.id` is not unique per run: two runs of the same pipeline would collide on `get_orchestrator_run_id`, the staging directory, and the job names. Every current orchestrator derives the run id from the placeholder run; only the deprecated HyperAI used the snapshot id. - Drop the local `os.environ[...] = run_id` write. No current orchestrator does this; the run id only needs to reach `get_orchestrator_run_id` inside each step's job, which happens via the per-step environment injection that was already in place. (Another HyperAI-ism.) - Reject scheduled pipelines with a clear error, matching the `ssh` orchestrator, instead of silently ignoring the schedule. Also drop now-unused `Tuple` imports left over from the earlier client-factory refactor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Addresses an external review of the detached failure/status and credential lifecycle, which were not production-safe. Status and cancellation: - Persist the Slurm job id as step-run metadata and query/cancel by exact id instead of reconstructing the job name. - Add orchestrator `fetch_status` so a detached run is reconciled from its Slurm jobs, catching failures that happen before the ZenML entrypoint starts (bad image, container runtime error) which the steps can never self-report. - Reconcile job state by id via `squeue --jobs` with a `scontrol show job` fallback (controller memory), closing the window where a finished job has left the queue before its shared-filesystem sentinel is visible. Treat that gap as still-running rather than terminally cancelled. Credentials: - Add per-runtime private registry authentication (Apptainer/Singularity `*_DOCKER_USERNAME/PASSWORD`, enroot `.credentials`, Docker `config.json`) written as owner-only files, so private ECR/GCR/ACR/Docker Hub images pull without exposing passwords on any command line. - Roll back staged files (`rm -rf` the run dir) on any staging/submission failure, and submit an `afterany` cleanup job depending on all step jobs to scrub credential files even for jobs cancelled before their EXIT trap runs. Runtime safety: - Shell-escape and validate environment variable names; the pyxis path now `export`s quoted values instead of sourcing a Docker-format file, and Apptainer/Singularity run with `--no-eval`. - Round CPU requests up and emit memory in MB so fractional/sub-GB resource settings produce valid sbatch directives. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The review-hardening changes added `ValueError` input guards and `except Exception: ... raise` cleanup re-raises without updating the docstring `Raises:` sections, and moved a `raise` out of `_submit_step` into `stage_and_submit`, which the CI `pydoclint` check flagged (DOC501/DOC502/ DOC503). Align every `Raises:` section with its function body: - Document the numeric-id `ValueError` on `SlurmClient.submit`/`get_job_state`/ `cancel` and the enroot `ValueError` on `build_registry_auth`. - Document the re-raised `Exception` on `stage_and_submit`, `SlurmOrchestrator.submit_pipeline`, and `SlurmStepOperator.submit`. - Drop the stale `Raises:` from `_submit_step`, whose raise now lives in `stage_and_submit`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The CI lint runs ruff's F401/F841 check over the tests too; UUID was imported but only uuid4 is used. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

coderabbitai · 2026-07-13T09:34:28Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: a25dacd6-d9cb-4dad-afc3-4b879edd608c

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch feature/slurm-integration

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

github-actions · 2026-07-13T10:06:56Z

Documentation Link Check Results

✅ Absolute links check passed
✅ Relative links check passed
_{Last checked: 2026-07-17 14:28:29 UTC}

schustmi · 2026-07-16T07:54:51Z

+        container_registry = stack.container_registry
+        assert container_registry is not None
+        try:
+            for step_name in self._sorted_step_names(steps):


steps are already sorted

schustmi · 2026-07-16T07:55:07Z

+    @staticmethod
+    def _sorted_step_names(
+        steps: Dict[str, "Step"],
+    ) -> List[str]:
+        """Return the step invocation ids in a topological order.
+
+        Jobs must be submitted upstream-first so that a step's Slurm
+        dependencies reference the already-submitted job ids of its parents.
+
+        Args:
+            steps: The steps keyed by invocation id.
+
+        Returns:
+            The invocation ids in a valid topological order.
+        """
+
+        def parents(name: str) -> List[str]:
+            return [u for u in steps[name].spec.upstream_steps if u in steps]
+
+        def children(name: str) -> List[str]:
+            return [
+                other
+                for other in steps
+                if name in steps[other].spec.upstream_steps
+            ]
+
+        layers = topsorted_layers(
+            nodes=list(steps),
+            get_node_id_fn=lambda name: name,
+            get_parent_nodes=parents,
+            get_child_nodes=children,
+        )
+        return [name for layer in layers for name in layer]


Suggested change

@staticmethod

def _sorted_step_names(

steps: Dict[str, "Step"],

) -> List[str]:

"""Return the step invocation ids in a topological order.

Jobs must be submitted upstream-first so that a step's Slurm

dependencies reference the already-submitted job ids of its parents.

Args:

steps: The steps keyed by invocation id.

Returns:

The invocation ids in a valid topological order.

"""

def parents(name: str) -> List[str]:

return [u for u in steps[name].spec.upstream_steps if u in steps]

def children(name: str) -> List[str]:

return [

other

for other in steps

if name in steps[other].spec.upstream_steps

]

layers = topsorted_layers(

nodes=list(steps),

get_node_id_fn=lambda name: name,

get_parent_nodes=parents,

get_child_nodes=children,

)

return [name for layer in layers for name in layer]

schustmi · 2026-07-16T07:58:43Z

+        job_id: Optional[str] = None
+        sensitive_paths: List[str] = []
+        try:
+            job_id, sensitive_paths = self._submit_container_job(


Submitting a slurm job from an active slurm job seems like a horrible idea which will lead to starvation.

The new job will be queued, and the active job will be waiting and blocking resources while all other jobs in the queue are executing.

schustmi · 2026-07-16T08:00:39Z

+        run_dir = self._isolated_run_dir(
+            run_id=run_id, step_run_id=str(step_run_info.step_run_id)
+        )
+        client = build_slurm_client(self.config)


Shouldn't this always use a local runner, because by definition we're not running in the slurm cluster already?

Address review: submitting per-step Slurm jobs from inside an active orchestration job starves under per-user job limits and blocks the parent allocation while children queue. Dynamic pipelines now run as a single orchestration job sized by the pipeline-level resource settings, with every step executing inline - the isolated-step machinery is removed, and no sbatch ever happens from within the cluster. Steps that need their own allocation go through the Slurm step operator instead. Also per review, drop the redundant topological sort: snapshot step configurations are already topologically ordered by contract. RL-training readiness, kept additive: - Record the raw Slurm terminal state (TIMEOUT, NODE_FAIL, OUT_OF_MEMORY, PREEMPTED, ...) as step-run metadata in the step operator, so an infrastructure kill is distinguishable from a step that ran and exited non-zero (no exit-code sentinel exists in either case). - Document the srun-wrappable contract of build_container_command for a future multi-node command-step mode. - Document container_mounts as the supported shared-filesystem bridge for tools that manage their own environment (e.g. a prime-rl checkout with its uv venv). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Between a job leaving squeue --states=all and its exit-code sentinel becoming visible on the shared filesystem, the run's outcome is not knowable. Dynamic reconciliation previously passed cleanup_complete=True and declared FAILED on first sight of that window, terminally mis-marking runs whose sentinel merely lagged the queue purge; the static DAG path already treats this case as unknown. Mirror it, and warn on each poll so a job killed before its EXIT trap ran (NODE_FAIL) is visible to operators instead of silently pending. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…gration

The isolated-step removal orphaned the ORCHESTRATOR_DOCKER_IMAGE_KEY import, and "mis-mark" trips the typos checker. The develop merge also resolves the GitBook redirect check, which compared this branch's docs tree against a develop that had since gained the DigitalOcean pages. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The orchestration job's EXIT trap scrubs its staged credentials on clean exits, but a hard kill (NODE_FAIL, OOM, hard timeout) never runs it, stranding the env file and registry auth on the shared filesystem. Submit the same afterany cleanup job static runs get; if the reaper cannot be queued, cancel the orchestration job and scrub immediately rather than run without one. The cleanup job's completion marker also restores conclusive terminal resolution for dynamic reconciliation: a job gone from the queue with no exit sentinel stays unknown only until the cleanup job has run, then becomes FAILED - the same semantics the static DAG path has, replacing the previous never-resolving warning-only behavior. Also drop the typing.Any import orphaned by the isolated-step removal (CI lints with --isolated, which local per-file ignores had masked). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

pydoclint (DOC503) requires the bare re-raise in the credential-reaper error path to appear in the Raises section. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

ruff format --check in CI flags the get_job_states signature; the local ruff version had accepted it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

htahir1 and others added 8 commits July 10, 2026 11:40

github-actions Bot added internal To filter out internal PRs and issues enhancement New feature or request labels Jul 10, 2026

htahir1 added the release-notes Release notes will be attached and used publicly for this PR. label Jul 10, 2026

htahir1 mentioned this pull request Jul 10, 2026

Add a Slurm integration (step operator over SSH/local sbatch) #5063

Open

htahir1 and others added 5 commits July 10, 2026 17:39

Remove unused UUID import in the Slurm orchestrator tests

27cfa64

The CI lint runs ruff's F401/F841 check over the tests too; UUID was imported but only uuid4 is used. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Improve Slurm lifecycle reconciliation

d503952

Add dynamic pipeline support for Slurm

404072f

Fix Slurm failure modes and cancel cleanup

7f2de6a

Fix Slurm isolated metadata transaction

6da70f4

htahir1 added 2 commits July 13, 2026 11:39

Add Slurm integration docs

04483ce

Fix Slurm dynamic status polling

a929e65

schustmi requested changes Jul 16, 2026

View reviewed changes

htahir1 and others added 6 commits July 17, 2026 10:25

Merge remote-tracking branch 'origin/develop' into feature/slurm-inte…

4775de9

…gration

Document the cleanup-submission re-raise in the dynamic docstring

6f172da

pydoclint (DOC503) requires the bare re-raise in the credential-reaper error path to appear in the Raises section. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Format slurm_client to the CI ruff line length

1590627

ruff format --check in CI flags the get_job_states signature; the local ruff version had accepted it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add Slurm integration (step operator + orchestrator)#5064

Add Slurm integration (step operator + orchestrator)#5064
htahir1 wants to merge 23 commits into
developfrom
feature/slurm-integration

htahir1 commented Jul 10, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 13, 2026 •

edited

Loading

Review skipped

Uh oh!

github-actions Bot commented Jul 13, 2026 •

edited

Loading

Uh oh!

schustmi Jul 16, 2026

Uh oh!

schustmi Jul 16, 2026

Uh oh!

schustmi Jul 16, 2026

Uh oh!

schustmi Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

htahir1 commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this does

Design decisions

Reliability & security hardening

Testing

Notes for reviewers

Uh oh!

coderabbitai Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation Link Check Results

Uh oh!

schustmi Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

schustmi Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

schustmi Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

schustmi Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

htahir1 commented Jul 10, 2026 •

edited

Loading

coderabbitai Bot commented Jul 13, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading