diff --git a/docs/book/component-guide/orchestrators/README.md b/docs/book/component-guide/orchestrators/README.md index a204667fab9..15c04f2e4e8 100644 --- a/docs/book/component-guide/orchestrators/README.md +++ b/docs/book/component-guide/orchestrators/README.md @@ -35,6 +35,7 @@ Out of the box, ZenML comes with a `local` orchestrator already part of the defa | [SkypilotGCPOrchestrator](skypilot-vm.md) | `vm_gcp` | `skypilot[gcp]` | Runs your pipelines in GCP VMs using SkyPilot | | [SkypilotAzureOrchestrator](skypilot-vm.md) | `vm_azure` | `skypilot[azure]` | Runs your pipelines in Azure VMs using SkyPilot | | [SSHOrchestrator](ssh.md) | `ssh` | `ssh` | Runs your pipelines on a remote host via SSH and Docker Compose. | +| [SlurmOrchestrator](slurm.md) | `slurm` | `slurm` | Runs your pipelines as Slurm jobs on an HPC cluster. | | [HyperAIOrchestrator](hyperai.md) | `hyperai` | `hyperai` | _Deprecated_ — use the SSH orchestrator. Runs your pipeline in HyperAI.ai instances. | | [Custom Implementation](custom.md) | _custom_ | | Extend the orchestrator abstraction and provide your own implementation | diff --git a/docs/book/component-guide/orchestrators/slurm.md b/docs/book/component-guide/orchestrators/slurm.md new file mode 100644 index 00000000000..ce5a7d9b3ea --- /dev/null +++ b/docs/book/component-guide/orchestrators/slurm.md @@ -0,0 +1,229 @@ +--- +description: Orchestrating ZenML pipelines on Slurm clusters. +--- + +# Slurm Orchestrator + +The Slurm orchestrator is an [orchestrator](./) flavor that submits ZenML +pipeline runs to an existing [Slurm](https://slurm.schedmd.com/) cluster. Each +pipeline step runs as a Slurm batch job, using the standard ZenML container image +for that step and a container runtime available on the compute nodes. + +Use this orchestrator when you already have an HPC cluster managed by Slurm and +want ZenML to submit whole pipeline runs to it. If you only want to offload a +few individual steps while another orchestrator runs the rest of the pipeline, +use the [Slurm step operator](../step-operators/slurm.md) instead. + +{% hint style="warning" %} +The Slurm integration targets containerized, one-Slurm-job-per-step workloads +with a shared POSIX staging directory. It does not currently provide +first-class Slurm arrays, MPI or multi-node job helpers, `slurmrestd`, +accounting-based reconciliation, or scheduling support. +{% endhint %} + +### When to use it + +You should use the Slurm orchestrator if: + +* you have an existing Slurm cluster and want to run complete ZenML pipelines on + it. +* your cluster can pull container images through Apptainer, Singularity, Pyxis, + or Docker. +* the Slurm login node and compute nodes can access the same staging directory + configured as `workdir`. +* your pipeline artifacts live in a remote artifact store that compute nodes can + reach. + +Use another orchestrator if your cluster cannot run container images, if compute +nodes cannot reach your artifact store, or if you need native Slurm features +that are not exposed by this integration yet. + +### Prerequisites + +To use the Slurm orchestrator, you need: + +* The ZenML `slurm` integration installed. +* Access to a Slurm submission host with `sbatch`, `squeue`, `scontrol`, and + `scancel` available. +* A shared `workdir` path visible from both the submission host and compute + nodes. ZenML stores job scripts, owner-only env files, output logs, and + completion sentinels there. +* A remote artifact store, such as S3, GCS, Azure Blob Storage, or another + artifact store reachable from the compute nodes. +* A remote container registry and an image builder so ZenML can build and push + the step images that Slurm jobs pull. +* One supported compute-node container runtime: + * `apptainer` or `singularity`, using `docker://` image references. + * `pyxis`, using `srun --container-image`. + * `docker`, only on clusters where compute nodes expose a Docker daemon. + +### How it works + +The Slurm orchestrator connects to a Slurm submission host either over SSH or by +running Slurm commands locally: + +* `transport=ssh` connects to a login node with the same SSH configuration used + by the SSH integration. +* `transport=local` runs `sbatch`, `squeue`, `scontrol`, and `scancel` on the + current machine. Use this when the ZenML client or orchestration process + already runs on a Slurm login node. + +For static pipelines, ZenML submits one Slurm job per step and wires the jobs +together with `afterok` dependencies. Slurm handles the dependency graph, and +ZenML reconciles job state through batched `squeue` lookups plus sentinel files +written by the job scripts. + +For dynamic pipelines, ZenML submits a single orchestration job that runs +ZenML's dynamic runner inside the orchestrator image. Every step of the +pipeline executes inline within that job's allocation, so size the job with +pipeline-level resource settings. The orchestrator deliberately does not +submit per-step child jobs from the orchestration job: an active Slurm job +that queues and waits on new jobs deadlocks under per-user job limits and +blocks its own allocation while the children sit in the queue. If individual +steps need their own Slurm allocation, run them through the Slurm step +operator instead. + +The orchestrator supports `CONTINUE_ON_FAILURE`, `STOP_ON_FAILURE`, and +`FAIL_FAST` execution modes. In continue-on-failure mode, independent sibling +jobs are left running when another branch fails. In stop-on-failure and +fail-fast modes, unfinished jobs are cancelled. + +### How to use it + +Install the Slurm integration: + +```shell +zenml integration install slurm +``` + +Register the orchestrator for a remote login node: + +```shell +zenml orchestrator register \ + --flavor=slurm \ + --transport=ssh \ + --hostname= \ + --username= \ + --ssh_key_path= \ + --workdir=/shared/zenml-runs \ + --container_runtime=apptainer \ + --partition= \ + --account= +``` + +For a client that already runs on a Slurm login node, use local transport: + +```shell +zenml orchestrator register \ + --flavor=slurm \ + --transport=local \ + --workdir=/shared/zenml-runs \ + --container_runtime=apptainer +``` + +Create or update a stack with the Slurm orchestrator, a remote artifact store, a +container registry, and an image builder: + +```shell +zenml stack register \ + -o \ + -a \ + -c \ + -i \ + --set +``` + +You can now run a pipeline as usual: + +```shell +python file_that_runs_a_zenml_pipeline.py +``` + +### Configuration options + +Key registration-time options: + +| Option | Description | +|--------|-------------| +| `transport` | `ssh` to connect to a remote login node, or `local` to run Slurm commands on the current machine. | +| `hostname`, `username`, `port` | SSH connection details. Required for `transport=ssh`. | +| `ssh_key_path` | Path to a private key on the submitting machine. Use for normal local submissions. | +| `ssh_private_key`, `ssh_key_passphrase` | Private key content and optional passphrase. Use ZenML secrets for these values. | +| `verify_host_key`, `known_hosts_path` | Host-key verification settings inherited from the SSH integration. Verification is enabled by default. | +| `workdir` | Shared directory used to stage per-run Slurm files. Must be visible from the submission host and compute nodes. | +| `container_runtime` | One of `apptainer`, `singularity`, `pyxis`, or `docker`. | +| `partition`, `time_limit`, `account`, `qos` | Common Slurm scheduling directives mapped to `#SBATCH` lines. | +| `extra_sbatch_directives` | Additional raw `#SBATCH` options such as `--constraint=a100`, `--reservation=`, or `--exclusive`. | +| `container_mounts` | Host-path to container-path mounts, for example `{'/scratch/user': '/scratch'}`. | +| `container_run_args` | Additional arguments passed to the selected container runtime. | + +You can override Slurm job settings per pipeline or per step: + +```python +from zenml import step +from zenml.integrations.slurm.flavors.slurm_orchestrator_flavor import ( + SlurmOrchestratorSettings, +) + +gpu_settings = SlurmOrchestratorSettings( + partition="gpu", + time_limit="2:00:00", + account="ml-research", + extra_sbatch_directives=["--constraint=a100"], + container_mounts={"/scratch/$USER": "/scratch"}, +) + + +@step(settings={"orchestrator": gpu_settings}) +def train_model() -> None: + ... +``` + +ZenML maps standard step resource settings to Slurm directives: + +* `cpu_count` becomes `--cpus-per-task`. +* memory becomes `--mem`. +* `gpu_count` becomes `--gres=gpu:`. + +### Dynamic pipeline notes + +Dynamic pipelines run entirely inside one long-running orchestration job. Keep +these operational details in mind: + +* Size the orchestration job for the whole run: set resource settings at the + pipeline level (CPU, memory, GPUs) — every step executes inside this one + allocation. +* Configure `time_limit`, `partition`, `account`, `qos`, and any site-required + `extra_sbatch_directives` so the orchestration job is allowed to run for the + full duration of the pipeline, not just one step. +* The orchestration container must be able to reach the ZenML server and the + artifact store. It never talks to the Slurm scheduler itself. + +### Security and cleanup + +ZenML stages each Slurm job in an owner-only run directory under `workdir`. The +environment file and any registry authentication files are written with `0600` +permissions and are not passed as command-line arguments. Job scripts install an +`EXIT` trap that records the exit code and removes credential-bearing files. + +Static pipeline runs also submit an `afterany` cleanup job for staged sensitive +files. Cancellation paths explicitly scrub the known sensitive files when a +pending job is cancelled before its job script starts. + +### Limitations + +The current Slurm orchestrator is intentionally conservative: + +* It runs one containerized Slurm job per ZenML step. +* It does not manage ZenML schedules. Trigger scheduled runs from cron, CI, or + another scheduler. +* It does not provide first-class helpers for MPI, multi-node training, Slurm + arrays, pre-staged SIF images, or `sacct`/accounting-based reconciliation. +* Advanced cluster requirements should be passed through + `extra_sbatch_directives` and validated on a real cluster before relying on + them in production. + +Check out the [SDK docs](https://sdkdocs.zenml.io/latest/) for the full API +reference. + +
ZenML Scarf
diff --git a/docs/book/component-guide/step-operators/README.md b/docs/book/component-guide/step-operators/README.md index a6dda7d340e..6a95592b730 100644 --- a/docs/book/component-guide/step-operators/README.md +++ b/docs/book/component-guide/step-operators/README.md @@ -28,6 +28,7 @@ Step operators to execute steps on one of the big cloud providers are provided b | [Modal](modal.md) | `modal` | `modal` | ✅ | Uses Modal to execute steps | | [SageMaker](sagemaker.md) | `sagemaker` | `aws` | ✅ | Uses SageMaker to execute steps | | [Run:AI](runai.md) | `runai` | `runai` | ✅ | Uses Run:AI to execute steps | +| [Slurm](slurm.md) | `slurm` | `slurm` | ❌ | Runs selected steps as Slurm jobs on an HPC cluster | | [SSH](ssh.md) | `ssh` | `ssh` | ❌ | Runs steps on a remote host via SSH + Docker | | [Spark](spark-kubernetes.md) | `spark` | `spark` | ❌ | Uses Spark on Kubernetes to execute steps in a distributed manner | | [Vertex](vertex.md) | `vertex` | `gcp` | ✅ | Uses Vertex AI to execute steps | diff --git a/docs/book/component-guide/step-operators/slurm.md b/docs/book/component-guide/step-operators/slurm.md new file mode 100644 index 00000000000..7dd40b123fc --- /dev/null +++ b/docs/book/component-guide/step-operators/slurm.md @@ -0,0 +1,245 @@ +--- +description: Executing individual ZenML steps on Slurm clusters. +--- + +# Slurm Step Operator + +The Slurm step operator is a [step operator](./) flavor that runs selected ZenML +steps as Slurm batch jobs on an existing HPC cluster. It is useful when most of +your pipeline can run on another orchestrator, but one or more compute-heavy +steps need Slurm-managed CPU, memory, or GPU resources. + +{% hint style="info" %} +The Slurm step operator does not orchestrate dynamic pipelines. For full +pipeline execution, including dynamic pipelines, use the +[Slurm orchestrator](../orchestrators/slurm.md). +{% endhint %} + +{% hint style="warning" %} +The Slurm step operator runs a single containerized job per selected step. It +does not provide first-class Slurm arrays, MPI or multi-node helpers, +`slurmrestd`, accounting-based reconciliation, or scheduling support. +{% endhint %} + +### When to use it + +You should use the Slurm step operator if: + +* your organization already has a Slurm cluster. +* only selected steps need HPC resources. +* compute nodes can pull the step image from your container registry. +* compute nodes can read and write artifacts through your remote artifact store. +* a shared `workdir` path is visible from the Slurm submission host and compute + nodes. + +If you want every step in the pipeline to run as a Slurm job, use the +[Slurm orchestrator](../orchestrators/slurm.md) instead. + +### Prerequisites + +To use the Slurm step operator, you need: + +* The ZenML `slurm` integration installed. +* Access to a Slurm submission host with `sbatch`, `squeue`, `scontrol`, and + `scancel` available. +* A shared `workdir` path visible from both the submission host and compute + nodes. +* A remote artifact store reachable from compute nodes. +* A remote container registry and image builder. +* One supported compute-node container runtime: + * `apptainer` or `singularity`. + * `pyxis`. + * `docker`, only where compute nodes expose a Docker daemon. + +### How it works + +When a step uses the Slurm step operator, ZenML builds and pushes the step image +through the active stack's image builder and container registry. The step +operator then connects to the Slurm submission host, stages a job script and +owner-only environment file in `workdir`, and submits the job with `sbatch`. + +The job script runs the step container on the compute node and writes an exit +sentinel when it finishes. ZenML polls Slurm with `squeue`/`scontrol` and reads +the sentinel file after the job leaves the queue, so it does not require Slurm +accounting (`sacct`) to be enabled. + +### How to use it + +Install the Slurm integration: + +```shell +zenml integration install slurm +``` + +Register the step operator for a remote login node: + +```shell +zenml step-operator register \ + --flavor=slurm \ + --transport=ssh \ + --hostname= \ + --username= \ + --ssh_key_path= \ + --workdir=/shared/zenml-runs \ + --container_runtime=apptainer \ + --partition= \ + --account= +``` + +For a client or orchestrator process that already runs on a Slurm login node, +use local transport: + +```shell +zenml step-operator register \ + --flavor=slurm \ + --transport=local \ + --workdir=/shared/zenml-runs \ + --container_runtime=apptainer +``` + +Attach the step operator to a stack that also includes a remote artifact store, +a container registry, and an image builder: + +```shell +zenml stack update \ + -s \ + -a \ + -c \ + -i +``` + +Use the step operator on selected steps: + +```python +from zenml import step + + +@step(step_operator="") +def train_model() -> None: + ... +``` + +### Configuration options + +Key registration-time options: + +| Option | Description | +|--------|-------------| +| `transport` | `ssh` to connect to a remote login node, or `local` to run Slurm commands on the current machine. | +| `hostname`, `username`, `port` | SSH connection details. Required for `transport=ssh`. | +| `ssh_key_path` | Path to a private key on the machine that submits the step. | +| `ssh_private_key`, `ssh_key_passphrase` | Private key content and optional passphrase. Prefer ZenML secrets for these values when the submitting process runs remotely. | +| `verify_host_key`, `known_hosts_path` | Host-key verification settings inherited from the SSH integration. Verification is enabled by default. | +| `workdir` | Shared directory used to stage per-step Slurm files. Must be visible from the submission host and compute nodes. | +| `container_runtime` | One of `apptainer`, `singularity`, `pyxis`, or `docker`. | +| `partition`, `time_limit`, `account`, `qos` | Common Slurm scheduling directives mapped to `#SBATCH` lines. | +| `extra_sbatch_directives` | Additional raw `#SBATCH` options such as `--constraint=a100`, `--reservation=`, or `--exclusive`. | +| `container_mounts` | Host-path to container-path mounts, for example `{'/scratch/user': '/scratch'}`. | +| `container_run_args` | Additional arguments passed to the selected container runtime. | + +You can override Slurm job settings per step: + +```python +from zenml import step +from zenml.integrations.slurm.flavors.slurm_step_operator_flavor import ( + SlurmStepOperatorSettings, +) + +slurm_settings = SlurmStepOperatorSettings( + partition="gpu", + time_limit="4:00:00", + account="ml-research", + extra_sbatch_directives=["--constraint=a100"], + container_mounts={"/scratch/$USER": "/scratch"}, +) + + +@step( + step_operator="", + settings={"step_operator": slurm_settings}, +) +def train_model() -> None: + ... +``` + +ZenML maps standard step resource settings to Slurm directives: + +* `cpu_count` becomes `--cpus-per-task`. +* memory becomes `--mem`. +* `gpu_count` becomes `--gres=gpu:`. + +### Running tools from shared storage + +`container_mounts` is the supported bridge to software that lives on the +cluster's shared filesystem instead of inside the ZenML image. This is the +standard pattern for frameworks that manage their own environment (for +example, a [prime-rl](https://github.com/PrimeIntellect-ai/prime-rl) checkout +with its `uv`-managed virtual environment, which its docs require to be on a +shared filesystem visible from every node): bind-mount the checkout into the +container and have the step invoke it there. + +```python +slurm_settings = SlurmStepOperatorSettings( + partition="gpu", + time_limit="24:00:00", + # The checkout and its venv live on the shared filesystem; the step + # command runs them in place, e.g. `uv run --project /shared/prime-rl ...` + container_mounts={"/shared/prime-rl": "/shared/prime-rl"}, +) +``` + +The invoked binaries (`uv` in this example) must be on the `PATH` inside the +ZenML image, or referenced by absolute path under the mount. + +### Honest failure states + +When a job ends because the scheduler killed it (`TIMEOUT`, `NODE_FAIL`, +`OUT_OF_MEMORY`, `PREEMPTED`), the job's EXIT trap never runs, so no exit-code +sentinel exists. The step operator records the raw Slurm terminal state as +step-run metadata under the `slurm_state` key whenever it observes one, so +downstream consumers can distinguish an infrastructure kill from a step that +ran to completion and exited non-zero. + +### Container runtimes + +The selected `container_runtime` controls how the ZenML image is executed on the +compute node: + +* `apptainer` and `singularity` run `apptainer exec` or `singularity exec` with + `docker://`. +* `pyxis` uses `srun --container-image=` and passes selected environment + variables through `--container-env`. +* `docker` uses `docker run --rm`, which only works on clusters where Docker is + available and permitted on compute nodes. + +Private registry credentials, when configured on the active container registry, +are staged as owner-only files and consumed by the selected runtime. + +### Security and cleanup + +The step operator writes the environment file with `0600` permissions in an +owner-only staging directory. Secrets are passed through files or environment +variable names, not as visible command-line values. The generated Slurm script +installs an `EXIT` trap that writes the step exit code and removes credential +files when the job finishes. + +If a step is cancelled while still pending, the step operator explicitly removes +known sensitive files because the job script may never start and therefore may +never run its `EXIT` trap. + +### Limitations + +The Slurm step operator is intentionally narrow: + +* It is a per-step offload mechanism, not a dynamic pipeline engine. +* It requires a remote artifact store and remote container registry. +* It does not manage schedules. +* It does not provide first-class helpers for MPI, multi-node training, Slurm + arrays, pre-staged SIF images, or Slurm accounting. +* Site-specific requirements should be passed through + `extra_sbatch_directives` and validated on your cluster. + +Check out the [SDK docs](https://sdkdocs.zenml.io/latest/) for the full API +reference. + +
ZenML Scarf
diff --git a/docs/book/component-guide/toc.md b/docs/book/component-guide/toc.md index 2e7d8a2eb66..c282565e676 100644 --- a/docs/book/component-guide/toc.md +++ b/docs/book/component-guide/toc.md @@ -19,6 +19,7 @@ * [Airflow Orchestrator](orchestrators/airflow.md) * [Skypilot VM Orchestrator](orchestrators/skypilot-vm.md) * [SSH Orchestrator](orchestrators/ssh.md) + * [Slurm Orchestrator](orchestrators/slurm.md) * [HyperAI Orchestrator](orchestrators/hyperai.md) * [Lightning AI Orchestrator](orchestrators/lightning.md) * [Develop a custom orchestrator](orchestrators/custom.md) @@ -66,6 +67,7 @@ * [Kubernetes](step-operators/kubernetes.md) * [Run:AI](step-operators/runai.md) * [Modal](step-operators/modal.md) + * [Slurm](step-operators/slurm.md) * [SSH](step-operators/ssh.md) * [Spark](step-operators/spark-kubernetes.md) * [Develop a Custom Step Operator](step-operators/custom.md) diff --git a/src/zenml/integrations/constants.py b/src/zenml/integrations/constants.py index 2681a6ce244..ffe2fab110b 100644 --- a/src/zenml/integrations/constants.py +++ b/src/zenml/integrations/constants.py @@ -74,6 +74,7 @@ SKYPILOT_LAMBDA = "skypilot_lambda" SKYPILOT_KUBERNETES = "skypilot_kubernetes" SLACK = "slack" +SLURM = "slurm" SPARK = "spark" TEKTON = "tekton" TENSORBOARD = "tensorboard" diff --git a/src/zenml/integrations/slurm/__init__.py b/src/zenml/integrations/slurm/__init__.py new file mode 100644 index 00000000000..5b8e014e544 --- /dev/null +++ b/src/zenml/integrations/slurm/__init__.py @@ -0,0 +1,61 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Initialization of the Slurm integration. + +The Slurm integration runs ZenML pipelines on an HPC cluster managed by the +Slurm scheduler. It provides two stack components: + +- a step operator that runs a single step as a Slurm job (delegating the rest + of the pipeline to another orchestrator), and +- an orchestrator that submits the whole pipeline as a graph of Slurm jobs + wired together with ``sbatch`` dependencies. + +Both run the step's standard ZenML Docker 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 image build. Submission goes over SSH to a login/controller node, or +via local ``sbatch`` when the client already runs on the cluster. +""" + +from typing import List, Type + +from zenml.integrations.constants import SLURM +from zenml.integrations.integration import Integration +from zenml.stack import Flavor + +SLURM_STEP_OPERATOR_FLAVOR = "slurm" +SLURM_ORCHESTRATOR_FLAVOR = "slurm" + + +class SlurmIntegration(Integration): + """Definition of the Slurm integration for ZenML.""" + + NAME = SLURM + # paramiko is only needed for the ssh transport; kept in sync with the + # ssh integration so the two can share the same SSH client. + REQUIREMENTS = ["paramiko>=3.0.0,<6"] + + @classmethod + def flavors(cls) -> List[Type[Flavor]]: + """Declare the stack component flavors for the Slurm integration. + + Returns: + List of stack component flavors for this integration. + """ + from zenml.integrations.slurm.flavors import ( + SlurmOrchestratorFlavor, + SlurmStepOperatorFlavor, + ) + + return [SlurmStepOperatorFlavor, SlurmOrchestratorFlavor] diff --git a/src/zenml/integrations/slurm/flavors/__init__.py b/src/zenml/integrations/slurm/flavors/__init__.py new file mode 100644 index 00000000000..58ea0f88a2d --- /dev/null +++ b/src/zenml/integrations/slurm/flavors/__init__.py @@ -0,0 +1,34 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Slurm integration flavors.""" + +from zenml.integrations.slurm.flavors.slurm_orchestrator_flavor import ( # noqa + SlurmOrchestratorConfig, + SlurmOrchestratorFlavor, + SlurmOrchestratorSettings, +) +from zenml.integrations.slurm.flavors.slurm_step_operator_flavor import ( # noqa + SlurmStepOperatorConfig, + SlurmStepOperatorFlavor, + SlurmStepOperatorSettings, +) + +__all__ = [ + "SlurmOrchestratorConfig", + "SlurmOrchestratorFlavor", + "SlurmOrchestratorSettings", + "SlurmStepOperatorConfig", + "SlurmStepOperatorFlavor", + "SlurmStepOperatorSettings", +] diff --git a/src/zenml/integrations/slurm/flavors/base.py b/src/zenml/integrations/slurm/flavors/base.py new file mode 100644 index 00000000000..9a89e43e404 --- /dev/null +++ b/src/zenml/integrations/slurm/flavors/base.py @@ -0,0 +1,167 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Shared configuration for the Slurm step operator and orchestrator. + +Both components submit the same kind of Slurm job - the step's Docker image +run on a compute node with an HPC container runtime - so they share the +transport, container-runtime and ``sbatch`` settings defined here. +""" + +from enum import Enum +from typing import Dict, List, Optional + +from pydantic import Field, model_validator + +from zenml.config.base_settings import BaseSettings +from zenml.integrations.ssh.flavors.base import SSHConnectionConfigMixin + + +class SlurmTransport(str, Enum): + """How the client reaches the Slurm submission host.""" + + SSH = "ssh" + LOCAL = "local" + + +class SlurmContainerRuntime(str, Enum): + """Container runtime used to run the step image on the compute node. + + Slurm compute nodes rarely have a Docker daemon (no root), so the ZenML + step image is run with a rootless HPC container runtime by default. + """ + + APPTAINER = "apptainer" + SINGULARITY = "singularity" + PYXIS = "pyxis" + DOCKER = "docker" + + +class SlurmJobSettings(BaseSettings): + """Settings shared by every Slurm job ZenML submits. + + These map onto ``sbatch`` directives / container-run flags and can be + overridden per step. + """ + + partition: Optional[str] = Field( + default=None, + description="Slurm partition (queue) to submit the job to, passed as " + "`--partition`. Uses the cluster's default partition if unset. " + "Example: 'gpu'", + ) + time_limit: Optional[str] = Field( + default=None, + description="Wall-clock time limit for the job in Slurm time format, " + "passed as `--time`. Examples: '30:00' (30 minutes), '2:00:00' " + "(2 hours), '1-12:00:00' (1 day 12 hours). Uses the partition " + "default if unset", + ) + account: Optional[str] = Field( + default=None, + description="Slurm account to charge the job to, passed as " + "`--account`. Required on clusters that enforce accounting. " + "Example: 'ml-research'", + ) + qos: Optional[str] = Field( + default=None, + description="Quality-of-service level for the job, passed as " + "`--qos`. Example: 'normal'", + ) + extra_sbatch_directives: List[str] = Field( + default_factory=list, + description="Additional raw `#SBATCH` directive lines appended to " + "the generated job script, as an escape hatch for cluster-specific " + "options. Example: ['--constraint=a100', '--exclusive']", + ) + container_mounts: Dict[str, str] = Field( + default_factory=dict, + description="Host-path to container-path bind mounts applied to the " + "step container, e.g. to expose a shared scratch filesystem. " + "Example: {'/scratch/user': '/scratch'}", + ) + container_run_args: List[str] = Field( + default_factory=list, + description="Additional arguments passed to the container runtime " + "command (apptainer/singularity/docker/srun), as an escape hatch. " + "Example: ['--writable-tmpfs']", + ) + + +class SlurmConnectionConfig(SSHConnectionConfigMixin): + """Shared connection/runtime configuration for the Slurm components. + + Inherits the SSH connection fields (``ssh_key_path``, ``ssh_private_key``, + ``port``, ``connection_timeout``, ...) from the ssh integration's + connection mixin, so the same ``SSHClient`` consumes this config directly + and the fields never drift out of sync. + """ + + transport: SlurmTransport = Field( + default=SlurmTransport.SSH, + description="How to reach the Slurm submission host: 'ssh' connects " + "to a remote login/controller node, 'local' runs sbatch directly and " + "requires the client to already run on the cluster. Example: 'ssh'", + ) + container_runtime: SlurmContainerRuntime = Field( + default=SlurmContainerRuntime.APPTAINER, + description="Runtime used to run the step's Docker image on the " + "compute node: 'apptainer' or 'singularity' (rootless, pulls " + "`docker://` images; requires Apptainer 1.1+ or SingularityCE 3.10+ " + "for `--no-eval` support), 'pyxis' (NVIDIA " + "enroot/pyxis via `srun --container-image`), or 'docker' (only where " + "a Docker daemon is available). Example: 'apptainer'", + ) + workdir: str = Field( + description="Directory on the cluster used to stage the job script, " + "the environment file and the job output/exit-code sentinel, ideally " + "on a filesystem shared between the submission host and the compute " + "nodes. The step's code is not staged here; it lives in the container " + "image. Example: '/shared/zenml-runs'", + ) + # The connection mixin requires hostname/username, but the `local` + # transport does not use SSH, so these are relaxed to optional (the + # ignore below) and enforced for the `ssh` transport in the validator. + hostname: Optional[str] = Field( # type: ignore[assignment] + default=None, + description="Hostname or IP address of the Slurm login/controller " + "node, required for the 'ssh' transport. Example: " + "'login.cluster.example.com'", + ) + username: Optional[str] = Field( # type: ignore[assignment] + default=None, + description="SSH username on the login node, required for the 'ssh' " + "transport. Example: 'mlops'", + ) + + @model_validator(mode="after") + def _validate_transport_requirements(self) -> "SlurmConnectionConfig": + """Validates that the ssh transport has connection details. + + Returns: + The validated config. + + Raises: + ValueError: If the ssh transport is selected without a hostname + or username. + """ + if self.transport == SlurmTransport.SSH and not ( + self.hostname and self.username + ): + raise ValueError( + "The Slurm component with the 'ssh' transport needs " + "`hostname` and `username` to reach the login node. Set " + "them, or use transport='local' if the client runs on the " + "cluster itself." + ) + return self diff --git a/src/zenml/integrations/slurm/flavors/slurm_orchestrator_flavor.py b/src/zenml/integrations/slurm/flavors/slurm_orchestrator_flavor.py new file mode 100644 index 00000000000..91b9c958824 --- /dev/null +++ b/src/zenml/integrations/slurm/flavors/slurm_orchestrator_flavor.py @@ -0,0 +1,110 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Slurm orchestrator flavor.""" + +from typing import TYPE_CHECKING, Optional, Type + +from zenml.integrations.slurm import SLURM_ORCHESTRATOR_FLAVOR +from zenml.integrations.slurm.flavors.base import ( + SlurmConnectionConfig, + SlurmJobSettings, +) +from zenml.orchestrators import ( + BaseOrchestratorConfig, + BaseOrchestratorFlavor, +) + +if TYPE_CHECKING: + from zenml.integrations.slurm.orchestrators import SlurmOrchestrator + + +class SlurmOrchestratorSettings(SlurmJobSettings): + """Settings for the Slurm orchestrator.""" + + +class SlurmOrchestratorConfig( + BaseOrchestratorConfig, + SlurmConnectionConfig, + SlurmOrchestratorSettings, +): + """Configuration for the Slurm orchestrator.""" + + @property + def is_remote(self) -> bool: + """Whether this component runs remotely. + + Returns: + True, since jobs execute on the Slurm cluster. + """ + return True + + +class SlurmOrchestratorFlavor(BaseOrchestratorFlavor): + """Flavor of the Slurm orchestrator.""" + + @property + def name(self) -> str: + """Name of the flavor. + + Returns: + The name of the flavor. + """ + return SLURM_ORCHESTRATOR_FLAVOR + + @property + def docs_url(self) -> Optional[str]: + """A URL to point at docs explaining this flavor. + + Returns: + A flavor docs url. + """ + return self.generate_default_docs_url() + + @property + def sdk_docs_url(self) -> Optional[str]: + """A URL to point at SDK docs explaining this flavor. + + Returns: + A flavor SDK docs url. + """ + return self.generate_default_sdk_docs_url() + + @property + def logo_url(self) -> str: + """A URL to represent the flavor in the dashboard. + + Returns: + The flavor logo. + """ + return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/orchestrator/slurm.png" + + @property + def config_class(self) -> Type[SlurmOrchestratorConfig]: + """Config class for this flavor. + + Returns: + The config class. + """ + return SlurmOrchestratorConfig + + @property + def implementation_class(self) -> Type["SlurmOrchestrator"]: + """Implementation class for this flavor. + + Returns: + The implementation class. + """ + from zenml.integrations.slurm.orchestrators import SlurmOrchestrator + + return SlurmOrchestrator diff --git a/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py b/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py new file mode 100644 index 00000000000..7fbef9db2d4 --- /dev/null +++ b/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py @@ -0,0 +1,119 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Slurm step operator flavor.""" + +from typing import TYPE_CHECKING, Optional, Type + +from zenml.integrations.slurm import SLURM_STEP_OPERATOR_FLAVOR +from zenml.integrations.slurm.flavors.base import ( + SlurmConnectionConfig, + SlurmJobSettings, +) +from zenml.step_operators.base_step_operator import ( + BaseStepOperatorConfig, + BaseStepOperatorFlavor, +) + +if TYPE_CHECKING: + from zenml.integrations.slurm.step_operators import SlurmStepOperator + + +class SlurmStepOperatorSettings(SlurmJobSettings): + """Settings for the Slurm step operator.""" + + +class SlurmStepOperatorConfig( + BaseStepOperatorConfig, + SlurmConnectionConfig, + SlurmStepOperatorSettings, +): + """Configuration for the Slurm step operator.""" + + @property + def is_remote(self) -> bool: + """Whether this component runs remotely. + + Returns: + True, since steps execute on the Slurm cluster. + """ + return True + + +class SlurmStepOperatorFlavor(BaseStepOperatorFlavor): + """Flavor of the Slurm step operator.""" + + @property + def name(self) -> str: + """Name of the flavor. + + Returns: + The name of the flavor. + """ + return SLURM_STEP_OPERATOR_FLAVOR + + @property + def display_name(self) -> str: + """Display name of the flavor. + + Returns: + The display name of the flavor. + """ + return "Slurm" + + @property + def docs_url(self) -> Optional[str]: + """A URL to point at docs explaining this flavor. + + Returns: + A flavor docs url. + """ + return self.generate_default_docs_url() + + @property + def sdk_docs_url(self) -> Optional[str]: + """A URL to point at SDK docs explaining this flavor. + + Returns: + A flavor SDK docs url. + """ + return self.generate_default_sdk_docs_url() + + @property + def logo_url(self) -> str: + """A URL to represent the flavor in the dashboard. + + Returns: + The flavor logo. + """ + return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/step_operator/slurm.png" + + @property + def config_class(self) -> Type[SlurmStepOperatorConfig]: + """Config class for this flavor. + + Returns: + The config class. + """ + return SlurmStepOperatorConfig + + @property + def implementation_class(self) -> Type["SlurmStepOperator"]: + """Implementation class for this flavor. + + Returns: + The implementation class. + """ + from zenml.integrations.slurm.step_operators import SlurmStepOperator + + return SlurmStepOperator diff --git a/src/zenml/integrations/slurm/orchestrators/__init__.py b/src/zenml/integrations/slurm/orchestrators/__init__.py new file mode 100644 index 00000000000..5ffe6bd2ec8 --- /dev/null +++ b/src/zenml/integrations/slurm/orchestrators/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Initialization of the Slurm orchestrator.""" + +from zenml.integrations.slurm.orchestrators.slurm_orchestrator import ( # noqa + SlurmOrchestrator, +) + +__all__ = ["SlurmOrchestrator"] diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py new file mode 100644 index 00000000000..3698f8d375c --- /dev/null +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -0,0 +1,995 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Orchestrator that runs a whole pipeline as a graph of Slurm jobs. + +Each step is submitted as its own ``sbatch`` job running the step's ZenML +Docker image on a compute node, and the pipeline DAG is reproduced with Slurm +job dependencies (``--dependency=afterok``): a step's job only starts once all +of its upstream jobs have completed successfully. Submission is fully detached +- once the jobs are queued the scheduler runs them, and each step reports its +own status back to the ZenML server, so the orchestrator does not poll. + +Submission goes over SSH to a login/controller node, or via local ``sbatch`` +when the client already runs on the cluster. + +Dynamic pipelines run as a single Slurm orchestration job sized via the +pipeline-level resource settings; their steps execute inline inside that +job's allocation. Submitting per-step jobs from within an active Slurm job +would deadlock under per-user job limits and waste the parent allocation +while children wait in the queue, so the orchestrator deliberately does not +support isolated steps. +""" + +import os +import shlex +from typing import ( + TYPE_CHECKING, + Callable, + Dict, + List, + Optional, + Tuple, + cast, +) + +from zenml.config.base_settings import BaseSettings +from zenml.config.resource_settings import ResourceSettings +from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID +from zenml.entrypoints.step_entrypoint_configuration import ( + StepEntrypointConfiguration, +) +from zenml.enums import ExecutionMode, ExecutionStatus +from zenml.integrations.slurm.flavors.slurm_orchestrator_flavor import ( + SlurmOrchestratorConfig, + SlurmOrchestratorSettings, +) +from zenml.integrations.slurm.slurm_client import ( + PENDING_STATES, + RUNNING_STATES, + SlurmClient, + build_slurm_client, +) +from zenml.integrations.slurm.slurm_job import ( + CANCELLED_FILE, + DOCKER_CONFIG_DIR, + ENROOT_CONFIG_DIR, + ENV_FILE, + EXIT_CODE_FILE, + OUTPUT_FILE, + REGISTRY_AUTH_FILE, + REQUIRED_COMPONENTS, + SCRIPT_FILE, + build_container_command, + build_registry_auth, + build_sbatch_script, + serialize_environment, + stage_and_submit, + validate_remote_stack, +) +from zenml.logger import get_logger +from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult +from zenml.stack import StackValidator + +if TYPE_CHECKING: + from zenml.config.step_configurations import Step + from zenml.models import ( + PipelineRunResponse, + PipelineSnapshotResponse, + ) + from zenml.stack import Stack + +logger = get_logger(__name__) + +ENV_ZENML_SLURM_RUN_ID = "ZENML_SLURM_ORCHESTRATOR_RUN_ID" +SLURM_JOB_IDS_METADATA_KEY = "slurm_job_ids" +SLURM_CLEANUP_JOB_ID_METADATA_KEY = "slurm_cleanup_job_id" +SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY = "slurm_orchestration_job_id" +_CLEANUP_COMPLETE_FILE = "cleanup_complete" + + +class SlurmOrchestrator(ContainerizedOrchestrator): + """Orchestrator that submits a pipeline as a graph of Slurm jobs.""" + + @property + def config(self) -> SlurmOrchestratorConfig: + """Returns the config of this orchestrator. + + Returns: + The config. + """ + return cast(SlurmOrchestratorConfig, self._config) + + @property + def supported_execution_modes(self) -> List[ExecutionMode]: + """Execution modes supported by this orchestrator. + + Returns: + The supported execution modes. + """ + return [ + ExecutionMode.FAIL_FAST, + ExecutionMode.STOP_ON_FAILURE, + ExecutionMode.CONTINUE_ON_FAILURE, + ] + + @property + def settings_class(self) -> Optional[type[BaseSettings]]: + """Settings class for this orchestrator. + + Returns: + The settings class. + """ + return SlurmOrchestratorSettings + + @property + def validator(self) -> Optional[StackValidator]: + """Validate that the stack meets the orchestrator's requirements. + + Returns: + A stack validator. + """ + return StackValidator( + required_components=REQUIRED_COMPONENTS, + custom_validation_function=validate_remote_stack, + ) + + def get_orchestrator_run_id(self) -> str: + """Returns the active orchestrator run id. + + Returns: + The orchestrator run id. + + Raises: + RuntimeError: If the run id environment variable is not set, which + means this was not called from within a step's job. + """ + try: + return os.environ[ENV_ZENML_SLURM_RUN_ID] + except KeyError: + raise RuntimeError( + "Unable to read the orchestrator run id from environment " + f"variable {ENV_ZENML_SLURM_RUN_ID}." + ) + + def _job_name(self, run_id: str, step_name: str) -> str: + """Slurm job name for a step in a run. + + Args: + run_id: The orchestrator run id. + step_name: The step invocation id. + + Returns: + The Slurm job name. + """ + return f"zenml-{run_id}-{step_name}" + + def _run_dir(self, run_id: str, step_name: str) -> str: + """Per-step staging directory on the cluster. + + Args: + run_id: The orchestrator run id. + step_name: The step invocation id. + + Returns: + The absolute path of the step's run directory. + """ + workdir = self.config.workdir.rstrip("/") + return f"{workdir}/{run_id}/{step_name}" + + def _orchestration_run_dir(self, run_id: str) -> str: + """Staging directory for a dynamic orchestration job. + + Args: + run_id: The orchestrator run id. + + Returns: + The absolute path of the orchestration job's run directory. + """ + workdir = self.config.workdir.rstrip("/") + return f"{workdir}/{run_id}/orchestration" + + def submit_pipeline( + self, + snapshot: "PipelineSnapshotResponse", + stack: "Stack", + base_environment: Dict[str, str], + step_environments: Dict[str, Dict[str, str]], + placeholder_run: Optional["PipelineRunResponse"] = None, + ) -> Optional[SubmissionResult]: + """Submit the pipeline as a graph of dependent Slurm jobs. + + Args: + snapshot: The pipeline snapshot to submit. + stack: The stack the pipeline will run on. + base_environment: Base environment shared by all steps. Unused: the + per-step environments already include it. + step_environments: Environment variables per step. + placeholder_run: The placeholder run for the pipeline. + + Returns: + None, since the pipeline is submitted detached. + + Raises: + RuntimeError: If the pipeline has a schedule, which is not + supported. + Exception: Re-raised after cancelling already-queued jobs if a + step fails to submit. + """ + if snapshot.schedule: + raise RuntimeError( + "The Slurm orchestrator does not support scheduled pipelines. " + "Remove the schedule and trigger the pipeline directly (e.g. " + "from your own cron job or CI), or use an orchestrator that " + "supports scheduling." + ) + # The run id must be unique per run and identical for every step of the + # run; the placeholder run id satisfies both (the snapshot id does not - + # a snapshot can be executed many times). It is injected into each + # step's job environment below and read back by + # `get_orchestrator_run_id` inside the step's job. + assert placeholder_run is not None + run_id = str(placeholder_run.id) + + steps = snapshot.step_configurations + client = build_slurm_client(self.config) + submitted: Dict[str, str] = {} + sensitive_paths: List[str] = [] + cleanup_job_id: Optional[str] = None + cleanup_settings = cast( + SlurmOrchestratorSettings, self.get_settings(snapshot) + ) + container_registry = stack.container_registry + assert container_registry is not None + try: + # `step_configurations` is topologically sorted by contract, so + # every step's upstream job ids are already submitted when its + # `--dependency=afterok` list is built. + for step_name, step in steps.items(): + job_id, step_sensitive_paths = self._submit_step( + client=client, + snapshot=snapshot, + run_id=run_id, + step_name=step_name, + step=step, + step_environment=step_environments[step_name], + upstream_job_ids=[ + submitted[up] + for up in step.spec.upstream_steps + if up in submitted + ], + registry_uri=container_registry.config.uri, + registry_credentials=container_registry.credentials, + ) + submitted[step_name] = job_id + sensitive_paths.extend(step_sensitive_paths) + logger.info( + "Submitted step `%s` as Slurm job `%s`.", + step_name, + job_id, + ) + cleanup_job_id = self._submit_cleanup_job( + client=client, + run_id=run_id, + submitted=submitted, + sensitive_paths=sensitive_paths, + settings=cleanup_settings, + ) + except Exception: + # Best-effort cleanup: cancel the jobs already queued for this run + # so a partially-submitted pipeline does not leave work running. + if submitted: + for job_id in submitted.values(): + try: + client.cancel(job_id) + except Exception: + logger.warning( + "Failed to cancel Slurm job `%s` after a submission " + "failure.", + job_id, + ) + if sensitive_paths: + client.runner.run( + "rm -rf -- " + + " ".join( + shlex.quote(path) for path in sensitive_paths + ) + ) + raise + finally: + client.runner.close() + + logger.info( + "Submitted pipeline `%s` as %d Slurm jobs.", + snapshot.pipeline_configuration.name, + len(submitted), + ) + assert cleanup_job_id is not None + return SubmissionResult( + metadata={ + METADATA_ORCHESTRATOR_RUN_ID: run_id, + SLURM_JOB_IDS_METADATA_KEY: submitted, + SLURM_CLEANUP_JOB_ID_METADATA_KEY: cleanup_job_id, + } + ) + + def submit_dynamic_pipeline( + self, + snapshot: "PipelineSnapshotResponse", + stack: "Stack", + environment: Dict[str, str], + placeholder_run: Optional["PipelineRunResponse"] = None, + ) -> Optional[SubmissionResult]: + """Submit a dynamic pipeline as a Slurm orchestration job. + + Args: + snapshot: The pipeline snapshot to submit. + stack: The stack the pipeline will run on. + environment: Environment variables for the orchestration + environment. + placeholder_run: The placeholder run for the pipeline. + + Returns: + Submission metadata containing the orchestration job id. + + Raises: + RuntimeError: If the pipeline has a schedule, which is not + supported. + Exception: Re-raised after cancelling the orchestration job if + its credential-cleanup job cannot be submitted. + """ + from zenml.pipelines.dynamic.entrypoint_configuration import ( + DynamicPipelineEntrypointConfiguration, + ) + + if snapshot.schedule: + raise RuntimeError( + "The Slurm orchestrator does not support scheduled pipelines. " + "Remove the schedule and trigger the pipeline directly (e.g. " + "from your own cron job or CI), or use an orchestrator that " + "supports scheduling." + ) + + assert placeholder_run is not None + run_id = str(placeholder_run.id) + settings = cast(SlurmOrchestratorSettings, self.get_settings(snapshot)) + container_registry = stack.container_registry + assert container_registry is not None + command = ( + DynamicPipelineEntrypointConfiguration.get_entrypoint_command() + + DynamicPipelineEntrypointConfiguration.get_entrypoint_arguments( + snapshot_id=snapshot.id, + run_id=run_id, + ) + ) + + client = build_slurm_client(self.config) + try: + job_id, sensitive_paths = self._submit_container_job( + client=client, + run_id=run_id, + run_dir=self._orchestration_run_dir(run_id), + job_name=f"zenml-{run_id}-orchestration", + image=self.get_image(snapshot=snapshot), + entrypoint_command=command, + environment=environment, + # Every step of a dynamic pipeline runs inline inside this + # job, so its allocation must carry the pipeline's resources. + resources=snapshot.pipeline_configuration.resource_settings, + settings=settings, + registry_uri=container_registry.config.uri, + registry_credentials=container_registry.credentials, + ) + try: + # The job's own EXIT trap scrubs credentials on clean exits, + # but a hard kill (NODE_FAIL, OOM) never runs it - the + # `afterany` cleanup job is the reaper of last resort, same + # as for static runs. Its completion marker also makes a + # missing exit sentinel conclusive during reconciliation. + cleanup_job_id = self._submit_cleanup_job( + client=client, + run_id=run_id, + submitted={"orchestration": job_id}, + sensitive_paths=sensitive_paths, + settings=settings, + ) + except Exception: + # Without the reaper, a hard-killed job would strand + # credentials on the shared filesystem indefinitely - refuse + # to run in that state. + try: + client.cancel(job_id) + except Exception: + logger.warning( + "Failed to cancel Slurm orchestration job `%s` " + "after its cleanup job could not be submitted.", + job_id, + ) + if sensitive_paths: + client.runner.run( + "rm -rf -- " + + " ".join( + shlex.quote(path) for path in sensitive_paths + ) + ) + raise + finally: + client.runner.close() + + logger.info( + "Submitted dynamic pipeline `%s` as Slurm orchestration job `%s`.", + snapshot.pipeline_configuration.name, + job_id, + ) + return SubmissionResult( + metadata={ + METADATA_ORCHESTRATOR_RUN_ID: run_id, + SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: job_id, + SLURM_CLEANUP_JOB_ID_METADATA_KEY: cleanup_job_id, + } + ) + + def _submit_container_job( + self, + client: SlurmClient, + run_id: str, + run_dir: str, + job_name: str, + image: str, + entrypoint_command: List[str], + environment: Dict[str, str], + resources: ResourceSettings, + settings: SlurmOrchestratorSettings, + registry_uri: Optional[str], + registry_credentials: Optional[Tuple[str, str]], + dependencies: Optional[List[str]] = None, + ) -> Tuple[str, List[str]]: + """Stage and submit one containerized Slurm job. + + Args: + client: The Slurm client used to stage files and submit the job. + run_id: The orchestrator run id to inject into the job + environment. + run_dir: Per-job staging directory. + job_name: Slurm job name. + image: Container image to execute. + entrypoint_command: Full command executed inside the container. + environment: Runtime environment for the job. + resources: Resource requests for the Slurm job. + settings: Slurm settings for the job. + registry_uri: URI of the stack's container registry. + registry_credentials: Registry username and password. + dependencies: Optional Slurm job ids this job depends on. + + Returns: + The submitted Slurm job id and credential-bearing paths. + """ + env_file = f"{run_dir}/{ENV_FILE}" + runtime_environment = environment.copy() + runtime_environment[ENV_ZENML_SLURM_RUN_ID] = run_id + env_content = serialize_environment( + runtime_environment, runtime=self.config.container_runtime + ) + registry_auth = build_registry_auth( + runtime=self.config.container_runtime, + run_dir=run_dir, + registry_uri=registry_uri, + credentials=registry_credentials, + ) + container_command = build_container_command( + runtime=self.config.container_runtime, + image=image, + entrypoint_command=entrypoint_command, + env_file=env_file, + env_keys=sorted(runtime_environment), + use_gpu=bool(resources.gpu_count), + settings=settings, + registry_auth=registry_auth, + ) + script = build_sbatch_script( + job_name=job_name, + run_dir=run_dir, + container_command=container_command, + resources=resources, + settings=settings, + sensitive_paths=registry_auth.sensitive_paths, + ) + job_id = stage_and_submit( + client, + run_dir=run_dir, + env_content=env_content, + script=script, + dependencies=dependencies, + extra_files=registry_auth.files, + ) + return job_id, [env_file, *registry_auth.sensitive_paths] + + def _submit_step( + self, + client: SlurmClient, + snapshot: "PipelineSnapshotResponse", + run_id: str, + step_name: str, + step: "Step", + step_environment: Dict[str, str], + upstream_job_ids: List[str], + registry_uri: Optional[str], + registry_credentials: Optional[Tuple[str, str]], + ) -> Tuple[str, List[str]]: + """Stage and submit a single step as a Slurm job. + + Args: + client: The Slurm client used to stage files and submit the job. + snapshot: The pipeline snapshot. + run_id: The orchestrator run id. + step_name: The step invocation id. + step: The step configuration. + step_environment: Environment variables for the step. + upstream_job_ids: Slurm job ids this step depends on. + registry_uri: URI of the stack's container registry. + registry_credentials: Registry username and password. + + Returns: + The submitted Slurm job id and credential-bearing paths. + """ + settings = cast(SlurmOrchestratorSettings, self.get_settings(step)) + entrypoint_command = ( + StepEntrypointConfiguration.get_entrypoint_command() + + StepEntrypointConfiguration.get_entrypoint_arguments( + step_name=step_name, snapshot_id=snapshot.id + ) + ) + return self._submit_container_job( + client=client, + run_id=run_id, + run_dir=self._run_dir(run_id, step_name), + job_name=self._job_name(run_id, step_name), + image=self.get_image(snapshot=snapshot, step_name=step_name), + entrypoint_command=entrypoint_command, + environment=step_environment, + resources=step.config.resource_settings, + settings=settings, + registry_uri=registry_uri, + registry_credentials=registry_credentials, + dependencies=upstream_job_ids, + ) + + def _submit_cleanup_job( + self, + client: SlurmClient, + run_id: str, + submitted: Dict[str, str], + sensitive_paths: List[str], + settings: SlurmOrchestratorSettings, + ) -> str: + """Submit one credential cleanup job after all DAG jobs terminate. + + Args: + client: Slurm client used for submission. + run_id: Orchestrator run ID. + submitted: Mapping of step names to Slurm job IDs. + sensitive_paths: Credential-bearing files and directories. + settings: Pipeline-level Slurm settings for the cleanup job. + + Returns: + The cleanup Slurm job ID. + + Raises: + RuntimeError: If the cleanup job cannot be staged. + """ + job_ids = list(submitted.values()) + cleanup_dir = f"{self.config.workdir.rstrip('/')}/{run_id}/cleanup" + cleanup_marker = f"{cleanup_dir}/{_CLEANUP_COMPLETE_FILE}" + cleanup_command = "rm -rf -- " + " ".join( + shlex.quote(path) for path in sensitive_paths + ) + directives = [ + f"#SBATCH --job-name=zenml-{run_id}-cleanup", + f"#SBATCH --output={cleanup_dir}/{OUTPUT_FILE}", + ] + if settings.partition: + directives.append(f"#SBATCH --partition={settings.partition}") + if settings.account: + directives.append(f"#SBATCH --account={settings.account}") + if settings.qos: + directives.append(f"#SBATCH --qos={settings.qos}") + for directive in settings.extra_sbatch_directives: + directives.append(f"#SBATCH {directive}") + + script = f"""#!/bin/bash +{chr(10).join(directives)} + +set -eo pipefail +{cleanup_command} +touch {shlex.quote(cleanup_marker)} +""" + result = client.runner.run( + f"umask 077 && mkdir -p {shlex.quote(cleanup_dir)} && " + f"chmod 700 {shlex.quote(cleanup_dir)}" + ) + if result.exit_code != 0: + raise RuntimeError( + f"Failed to create cleanup directory `{cleanup_dir}`: " + f"{result.stderr.strip()}" + ) + script_path = f"{cleanup_dir}/{SCRIPT_FILE}" + client.runner.put_text(script_path, script, mode=0o700) + return client.submit( + script_path, + dependencies=job_ids, + dependency_type="afterany", + ) + + def fetch_status( + self, run: "PipelineRunResponse", include_steps: bool = False + ) -> Tuple[ + Optional[ExecutionStatus], Optional[Dict[str, ExecutionStatus]] + ]: + """Reconcile a detached pipeline run with its Slurm jobs. + + Args: + run: Pipeline run submitted by this orchestrator. + include_steps: Whether to return individual step statuses. + + Returns: + The pipeline status and optional step statuses. + """ + raw_job_ids = run.run_metadata.get(SLURM_JOB_IDS_METADATA_KEY) + if not isinstance(raw_job_ids, dict): + return self._fetch_dynamic_status(run) + + job_ids = { + str(name): str(job_id) for name, job_id in raw_job_ids.items() + } + cleanup_marker = ( + f"{self.config.workdir.rstrip('/')}/{run.id}/cleanup/" + f"{_CLEANUP_COMPLETE_FILE}" + ) + client = build_slurm_client(self.config) + try: + try: + client.runner.read_text(cleanup_marker) + except Exception: + cleanup_complete = False + else: + cleanup_complete = True + + job_states = client.get_job_states(list(job_ids.values())) + statuses = { + step_name: self._get_job_status( + client=client, + state=job_states[job_id], + run_dir=self._run_dir(str(run.id), step_name), + cleanup_complete=cleanup_complete, + ) + for step_name, job_id in job_ids.items() + } + + known_statuses = [status for status in statuses.values() if status] + pipeline_status: Optional[ExecutionStatus] = None + if not run.status.is_finished: + failed_or_cancelled = any( + status + in {ExecutionStatus.FAILED, ExecutionStatus.CANCELLED} + for status in known_statuses + ) + all_terminal = len(known_statuses) == len(statuses) and all( + status.is_finished for status in known_statuses + ) + if failed_or_cancelled: + if self._should_cancel_unfinished_jobs(run): + pipeline_status = ExecutionStatus.FAILED + self._cancel_unfinished_jobs( + client=client, + run_id=str(run.id), + job_ids=job_ids, + statuses=statuses, + ) + elif all_terminal: + pipeline_status = ExecutionStatus.FAILED + elif ExecutionStatus.RUNNING in known_statuses: + pipeline_status = ExecutionStatus.RUNNING + else: + pipeline_status = ExecutionStatus.PROVISIONING + elif len(known_statuses) == len(statuses) and all( + status == ExecutionStatus.COMPLETED + for status in known_statuses + ): + pipeline_status = ExecutionStatus.COMPLETED + elif ExecutionStatus.RUNNING in known_statuses: + pipeline_status = ExecutionStatus.RUNNING + elif known_statuses: + pipeline_status = ExecutionStatus.PROVISIONING + finally: + client.runner.close() + + step_statuses = None + if include_steps: + step_statuses = { + name: status for name, status in statuses.items() if status + } + return pipeline_status, step_statuses + + @staticmethod + def _should_cancel_unfinished_jobs(run: "PipelineRunResponse") -> bool: + """Check whether failure reconciliation should cancel sibling jobs. + + Args: + run: Pipeline run submitted by this orchestrator. + + Returns: + Whether unfinished jobs should be cancelled. + """ + return run.config.execution_mode in { + ExecutionMode.FAIL_FAST, + ExecutionMode.STOP_ON_FAILURE, + } + + def _fetch_dynamic_status( + self, run: "PipelineRunResponse" + ) -> Tuple[Optional[ExecutionStatus], None]: + """Reconcile a dynamic pipeline run with its orchestration job. + + Args: + run: Pipeline run submitted by this orchestrator. + + Returns: + The pipeline status and no step statuses. + """ + job_id = run.run_metadata.get(SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY) + if job_id is None: + logger.warning("No Slurm job metadata found for run `%s`.", run.id) + return None, None + + cleanup_marker = ( + f"{self.config.workdir.rstrip('/')}/{run.id}/cleanup/" + f"{_CLEANUP_COMPLETE_FILE}" + ) + client = build_slurm_client(self.config) + try: + try: + client.runner.read_text(cleanup_marker) + except Exception: + cleanup_complete = False + else: + cleanup_complete = True + + state = client.get_job_state(str(job_id)) + status = self._get_job_status( + client=client, + state=state, + run_dir=self._orchestration_run_dir(str(run.id)), + # The orchestration job writes its own exit-code sentinel + # via its EXIT trap; the `afterany` cleanup job's marker + # tells us when there is nothing left to wait for. Until + # then, a job that has vanished from the queue with no + # sentinel is unknown, not failed - mirroring the static + # DAG reconciliation. Declaring FAILED on first sight would + # terminally mislabel a run whose sentinel merely lags the + # queue purge (NFS/Lustre flush). + cleanup_complete=cleanup_complete, + ) + if status is None and state is None: + logger.warning( + "Slurm orchestration job `%s` for run `%s` is no longer " + "known to Slurm and has not written an exit-code " + "sentinel yet; keeping the run status unchanged until " + "the cleanup job resolves it.", + job_id, + run.id, + ) + finally: + client.runner.close() + + pipeline_status: Optional[ExecutionStatus] = None + if not run.status.is_finished and status: + if status in {ExecutionStatus.FAILED, ExecutionStatus.CANCELLED}: + pipeline_status = ExecutionStatus.FAILED + elif status == ExecutionStatus.COMPLETED: + pipeline_status = ExecutionStatus.COMPLETED + elif status == ExecutionStatus.RUNNING: + pipeline_status = ExecutionStatus.RUNNING + elif status == ExecutionStatus.QUEUED: + pipeline_status = ExecutionStatus.PROVISIONING + + return pipeline_status, None + + @staticmethod + def _get_job_status( + client: SlurmClient, + state: Optional[str], + run_dir: str, + cleanup_complete: bool, + ) -> Optional[ExecutionStatus]: + """Map one Slurm job and its sentinel to a ZenML status. + + Args: + client: Slurm client used to inspect the job. + state: Slurm job state fetched for the job. + run_dir: Per-step staging directory. + cleanup_complete: Whether the final cleanup job has run. + + Returns: + The reconciled status, or None while the outcome is unknown. + """ + if state in PENDING_STATES: + return ExecutionStatus.QUEUED + if state in RUNNING_STATES: + return ExecutionStatus.RUNNING + if state == "COMPLETED": + return ExecutionStatus.COMPLETED + if state is not None: + if state.startswith("CANCEL"): + return ExecutionStatus.CANCELLED + return ExecutionStatus.FAILED + + try: + client.runner.read_text(f"{run_dir}/{CANCELLED_FILE}") + except Exception: + pass + else: + return ExecutionStatus.CANCELLED + + try: + exit_code = client.runner.read_text( + f"{run_dir}/{EXIT_CODE_FILE}" + ).strip() + except Exception: + return ExecutionStatus.FAILED if cleanup_complete else None + return ( + ExecutionStatus.COMPLETED + if exit_code == "0" + else ExecutionStatus.FAILED + ) + + def _stop_run( + self, run: "PipelineRunResponse", graceful: bool = False + ) -> None: + """Cancel all Slurm jobs submitted for a pipeline run. + + Args: + run: Pipeline run submitted by this orchestrator. + graceful: Unused. Slurm does not expose a DAG-level graceful stop, + so both graceful and forceful stops cancel submitted jobs. + """ + _ = graceful + raw_job_ids = run.run_metadata.get(SLURM_JOB_IDS_METADATA_KEY) + static_job_ids = ( + {str(name): str(job_id) for name, job_id in raw_job_ids.items()} + if isinstance(raw_job_ids, dict) + else {} + ) + orchestration_job_id = run.run_metadata.get( + SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY + ) + orchestration_job_ids = ( + {"orchestration": str(orchestration_job_id)} + if orchestration_job_id + else {} + ) + if not (static_job_ids or orchestration_job_ids): + logger.warning("No Slurm job metadata found for run `%s`.", run.id) + return + + run_id = str(run.id) + client = build_slurm_client(self.config) + try: + self._cancel_jobs( + client=client, + job_ids=static_job_ids, + run_dir_for_key=lambda step_name: self._run_dir( + run_id, step_name + ), + ) + self._cancel_jobs( + client=client, + job_ids=orchestration_job_ids, + run_dir_for_key=lambda _: self._orchestration_run_dir(run_id), + ) + finally: + client.runner.close() + + def _cancel_unfinished_jobs( + self, + client: SlurmClient, + run_id: str, + job_ids: Dict[str, str], + statuses: Dict[str, Optional[ExecutionStatus]], + ) -> None: + """Cancel jobs that are not yet terminal. + + Args: + client: Slurm client used to cancel jobs. + run_id: Orchestrator run ID. + job_ids: Mapping of step names to Slurm job IDs. + statuses: Reconciled step statuses keyed by step name. + """ + jobs_to_cancel: Dict[str, str] = {} + for step_name, job_id in job_ids.items(): + status = statuses.get(step_name) + if not (status and status.is_finished): + jobs_to_cancel[step_name] = job_id + self._cancel_jobs( + client=client, + job_ids=jobs_to_cancel, + run_dir_for_key=lambda step_name: self._run_dir(run_id, step_name), + ) + + def _cancel_jobs( + self, + client: SlurmClient, + job_ids: Dict[str, str], + run_dir_for_key: Callable[[str], str], + ) -> None: + """Cancel Slurm jobs and persist cancellation sentinels. + + Args: + client: Slurm client used to cancel jobs. + job_ids: Mapping of stable job keys to Slurm job IDs. + run_dir_for_key: Function returning the staging directory for a + job key. + """ + if not job_ids: + return + + cancelled_steps = set(job_ids) + try: + client.cancel_jobs(list(job_ids.values())) + except Exception: + logger.warning( + "Failed to cancel Slurm jobs in one request; retrying " + "individually." + ) + cancelled_steps = set() + for step_name, job_id in job_ids.items(): + try: + client.cancel(job_id) + except Exception: + logger.warning( + "Failed to cancel Slurm job `%s` for step `%s`.", + job_id, + step_name, + ) + else: + cancelled_steps.add(step_name) + + for job_key in cancelled_steps: + try: + run_dir = run_dir_for_key(job_key) + client.runner.put_text( + f"{run_dir}/{CANCELLED_FILE}", "1\n", mode=0o600 + ) + self._cleanup_sensitive_files(client, run_dir) + except Exception: + logger.warning( + "Failed to write Slurm cancellation marker for job `%s`.", + job_key, + ) + + @staticmethod + def _cleanup_sensitive_files(client: SlurmClient, run_dir: str) -> None: + """Remove credential-bearing files for a submitted Slurm job. + + Args: + client: Slurm client used to remove the files. + run_dir: Per-job staging directory. + """ + paths = [ + f"{run_dir}/{ENV_FILE}", + f"{run_dir}/{REGISTRY_AUTH_FILE}", + f"{run_dir}/{DOCKER_CONFIG_DIR}", + f"{run_dir}/{ENROOT_CONFIG_DIR}", + ] + client.runner.run( + "rm -rf -- " + " ".join(shlex.quote(path) for path in paths) + ) diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py new file mode 100644 index 00000000000..9882fe30f1e --- /dev/null +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -0,0 +1,425 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Client for interacting with a Slurm cluster. + +The client wraps the Slurm CLI (``sbatch``, ``squeue``, ``scancel``) behind a +small transport abstraction so the same logic works whether the commands run +locally (the ZenML client already lives on a login node) or over SSH (the +common case: submitting to a remote cluster). A REST transport (``slurmrestd``) +can slot in later without touching the step operator. + +Only the CLI surface that is universally available is used: job state is read +from ``squeue`` while a job is queued or running, and from a sentinel exit-code +file written by the job script after it leaves the queue, so the client does +not depend on Slurm accounting (``sacct``) being configured. +""" + +import re +import shlex +import subprocess +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Sequence + +if TYPE_CHECKING: + from zenml.integrations.slurm.flavors.base import SlurmConnectionConfig + +# Slurm job states that mean the job is still on its way to running. +PENDING_STATES = { + "PENDING", + "CONFIGURING", + "REQUEUED", + "RESIZING", + "SUSPENDED", +} +# Slurm job states that mean the job is actively running or wrapping up. +RUNNING_STATES = {"RUNNING", "COMPLETING", "STAGE_OUT", "SIGNALING"} + + +@dataclass +class CommandResult: + """Result of a command executed on the submission host.""" + + exit_code: int + stdout: str + stderr: str + + +class SlurmCommandRunner(ABC): + """Transport abstraction: run commands and move files to the cluster.""" + + @abstractmethod + def run(self, command: str) -> CommandResult: + """Run a shell command on the submission host. + + Args: + command: The shell command to run. + + Returns: + The command result. + """ + + @abstractmethod + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: + """Write text content to a file on the submission host. + + Args: + remote_path: Destination path on the submission host. + content: The text content to write. + mode: File permissions. Defaults to owner read/write only, since + the job script and environment file may hold credentials on a + shared cluster filesystem. + """ + + @abstractmethod + def read_text(self, remote_path: str) -> str: + """Read a text file from the submission host. + + Args: + remote_path: Path of the file on the submission host. + + Returns: + The file content. + """ + + def close(self) -> None: + """Release any resources held by the runner.""" + + +class LocalSlurmCommandRunner(SlurmCommandRunner): + """Runs Slurm commands as local subprocesses. + + For setups where the ZenML client or orchestrator already runs on a + cluster login node with the Slurm CLI in ``PATH``. + """ + + def run(self, command: str) -> CommandResult: + """Run a shell command locally. + + Args: + command: The shell command to run. + + Returns: + The command result. + """ + completed = subprocess.run( # nosec B602 - commands are built by us + command, + shell=True, + capture_output=True, + text=True, + ) + return CommandResult( + exit_code=completed.returncode, + stdout=completed.stdout, + stderr=completed.stderr, + ) + + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: + """Write text content to the destination path. + + Args: + remote_path: Destination path. + content: The text content to write. + mode: File permissions to apply. + """ + import os + + flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC + flags |= getattr(os, "O_NOFOLLOW", 0) + file_descriptor = os.open(remote_path, flags, mode) + os.fchmod(file_descriptor, mode) + with os.fdopen(file_descriptor, "w") as f: + f.write(content) + + def read_text(self, remote_path: str) -> str: + """Read a text file. + + Args: + remote_path: Path of the file. + + Returns: + The file content. + """ + with open(remote_path) as f: + return f.read() + + +class SSHSlurmCommandRunner(SlurmCommandRunner): + """Runs Slurm commands on a remote login node over SSH. + + Reuses the SSH client from the ``ssh`` integration; the config object must + provide the same connection fields (``hostname``, ``username``, ...). + """ + + def __init__(self, config: "SlurmConnectionConfig") -> None: + """Initialize the runner with an open SSH connection. + + Args: + config: The Slurm component config holding SSH connection settings. + """ + from zenml.integrations.ssh.ssh_client import SSHClient + + # SlurmConnectionConfig inherits SSHConnectionConfigMixin, so the SSH + # client consumes it directly. + self._ssh = SSHClient(config) + self._ssh.__enter__() + + def run(self, command: str) -> CommandResult: + """Run a shell command on the remote host. + + Args: + command: The shell command to run. + + Returns: + The command result. + """ + result = self._ssh.exec(command) + return CommandResult( + exit_code=result.exit_code, + stdout=result.stdout, + stderr=result.stderr, + ) + + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: + """Write text content to a file on the remote host. + + Args: + remote_path: Destination path on the remote host. + content: The text content to write. + mode: File permissions to apply. + """ + self._ssh.put_text(remote_path, content, mode=mode) + + def read_text(self, remote_path: str) -> str: + """Read a text file from the remote host. + + Args: + remote_path: Path of the file on the remote host. + + Returns: + The file content. + """ + return self._ssh.read_text(remote_path) + + def close(self) -> None: + """Close the SSH connection.""" + self._ssh.__exit__(None, None, None) + + +class SlurmClient: + """Thin wrapper around the Slurm CLI on top of a command runner.""" + + def __init__(self, runner: SlurmCommandRunner) -> None: + """Initialize the client. + + Args: + runner: The transport used to reach the submission host. + """ + self.runner = runner + + def submit( + self, + script_path: str, + dependencies: Optional[List[str]] = None, + dependency_type: Literal["afterok", "afterany"] = "afterok", + ) -> str: + """Submit a batch script and return the Slurm job id. + + Args: + script_path: Path of the sbatch script on the submission host. + dependencies: Job ids this job must wait for. The job runs only + after all of them complete successfully (``afterok``); if any + cannot be satisfied, the job is cancelled rather than left + pending forever (``--kill-on-invalid-dep``), which propagates a + failure through the DAG. + dependency_type: Slurm dependency type to use. + + Returns: + The Slurm job id. + + Raises: + ValueError: If any dependency job id is not numeric. + RuntimeError: If the submission fails. + """ + command = "sbatch --parsable" + if dependencies: + if not all( + re.fullmatch(r"[0-9]+", job_id) for job_id in dependencies + ): + raise ValueError("Slurm dependency job IDs must be numeric.") + command += ( + f" --dependency={dependency_type}:{':'.join(dependencies)}" + ) + if dependency_type == "afterok": + command += " --kill-on-invalid-dep=yes" + result = self.runner.run(f"{command} {shlex.quote(script_path)}") + if result.exit_code != 0: + raise RuntimeError( + f"`sbatch` failed with exit code {result.exit_code}: " + f"{result.stderr.strip() or result.stdout.strip()}" + ) + # --parsable prints `jobid[;cluster]` + job_id = result.stdout.strip().split(";")[0] + if not re.fullmatch(r"[0-9]+", job_id): + raise RuntimeError( + f"`sbatch` returned an invalid job ID: {job_id!r}" + ) + return job_id + + def get_job_states( + self, job_ids: Sequence[str] + ) -> Dict[str, Optional[str]]: + """Get the queue states of multiple jobs by their IDs. + + Args: + job_ids: The Slurm job IDs to look up. + + Returns: + A mapping from job ID to the Slurm job state (e.g. ``PENDING``, + ``RUNNING``) or ``None`` if the job is no longer known to Slurm. + Jobs leave the ``squeue`` output once they reach a terminal state, + so ``None`` means the job finished, failed, or never existed - the + caller disambiguates via the exit-code sentinel file. + + Raises: + ValueError: If any job id is not numeric. + RuntimeError: If ``squeue`` fails. + """ + job_ids = list(dict.fromkeys(job_ids)) + if not all(re.fullmatch(r"[0-9]+", job_id) for job_id in job_ids): + raise ValueError("Slurm job IDs must be numeric.") + states: Dict[str, Optional[str]] = {job_id: None for job_id in job_ids} + if not job_ids: + return states + + jobs_arg = ",".join(job_ids) + result = self.runner.run( + f"squeue --noheader --format={shlex.quote('%i|%T')} " + f"--jobs={shlex.quote(jobs_arg)}" + ) + if result.exit_code != 0: + if "invalid job id" in result.stderr.lower(): + return states + raise RuntimeError( + f"`squeue` failed with exit code {result.exit_code}: " + f"{result.stderr.strip()}" + ) + + for line in result.stdout.strip().splitlines(): + job_id, separator, state = line.strip().partition("|") + if separator and job_id in states: + states[job_id] = state.strip() + + missing_job_ids = [ + job_id for job_id, state in states.items() if state is None + ] + if not missing_job_ids: + return states + + # Completed jobs can disappear from the default squeue result before + # a shared-filesystem sentinel becomes visible. Ask for all states in + # one supported batched lookup to close that race without sacct. + missing_jobs_arg = ",".join(missing_job_ids) + result = self.runner.run( + f"squeue --noheader --format={shlex.quote('%i|%T')} " + f"--states=all --jobs={shlex.quote(missing_jobs_arg)}" + ) + if result.exit_code != 0: + if "invalid job id" in result.stderr.lower(): + return states + raise RuntimeError( + f"`squeue --states=all` failed with exit code " + f"{result.exit_code}: {result.stderr.strip()}" + ) + for line in result.stdout.strip().splitlines(): + job_id, separator, state = line.strip().partition("|") + if separator and job_id in states: + states[job_id] = state.strip() + + return states + + def get_job_state(self, job_id: str) -> Optional[str]: + """Get the queue state of a job by its ID. + + Args: + job_id: The Slurm job ID to look up. + + Returns: + The Slurm job state (e.g. ``PENDING``, ``RUNNING``) or ``None`` + if the job is no longer known to Slurm. + """ + return self.get_job_states([job_id])[job_id] + + def cancel(self, job_id: str) -> None: + """Cancel a job by its ID. + + Args: + job_id: The Slurm job ID to cancel. + """ + self.cancel_jobs([job_id]) + + def cancel_jobs(self, job_ids: Sequence[str]) -> None: + """Cancel multiple jobs by their IDs. + + Args: + job_ids: Slurm job IDs to cancel. + + Raises: + ValueError: If any job id is not numeric. + RuntimeError: If ``scancel`` fails. + """ + job_ids = list(dict.fromkeys(job_ids)) + if not all(re.fullmatch(r"[0-9]+", job_id) for job_id in job_ids): + raise ValueError("Slurm job IDs must be numeric.") + if not job_ids: + return + + result = self.runner.run( + "scancel " + " ".join(shlex.quote(job_id) for job_id in job_ids) + ) + if result.exit_code != 0: + raise RuntimeError( + f"`scancel` failed with exit code {result.exit_code}: " + f"{result.stderr.strip()}" + ) + + +def build_slurm_client(config: "SlurmConnectionConfig") -> SlurmClient: + """Build a Slurm client for the transport configured on a component. + + Both the step operator and the orchestrator reach the cluster the same + way, so they share this factory. The caller owns the returned client and + must close its ``runner`` when done (``client.runner.close()``). + + Args: + config: The component config holding the transport and, for the ssh + transport, the SSH connection settings. + + Returns: + A Slurm client wrapping a runner for the configured transport. + """ + from zenml.integrations.slurm.flavors.base import SlurmTransport + + runner: SlurmCommandRunner + if config.transport == SlurmTransport.LOCAL: + runner = LocalSlurmCommandRunner() + else: + runner = SSHSlurmCommandRunner(config) + return SlurmClient(runner) diff --git a/src/zenml/integrations/slurm/slurm_job.py b/src/zenml/integrations/slurm/slurm_job.py new file mode 100644 index 00000000000..15ff51b93f1 --- /dev/null +++ b/src/zenml/integrations/slurm/slurm_job.py @@ -0,0 +1,444 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Shared Slurm job-script construction for the step operator and orchestrator. + +Both components run a ZenML step as a Slurm batch job: the step's Docker image +is executed on a compute node with a rootless HPC container runtime. The job +script and the per-runtime container command are identical between them, so +they are built here. + +Security: the environment file holds the step's credentials (ZenML store +token, etc.). It is passed to the container via ``--env-file`` / +``--container-env`` so the values never appear on a command line visible to +other cluster users, and the job's EXIT trap scrubs it the moment the job +ends - even if the submitting process dies. All interpolated paths and values +go through ``shlex.quote``. +""" + +import base64 +import json +import math +import posixpath +import re +import shlex +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Tuple + +from zenml.enums import StackComponentType +from zenml.integrations.slurm.flavors.base import ( + SlurmContainerRuntime, + SlurmJobSettings, +) + +if TYPE_CHECKING: + from zenml.config.resource_settings import ResourceSettings + from zenml.integrations.slurm.slurm_client import SlurmClient + from zenml.stack import Stack + +# Files written into a job's per-run directory on the cluster. +ENV_FILE = "env" +EXIT_CODE_FILE = "exit_code" +OUTPUT_FILE = "output.log" +SCRIPT_FILE = "job.sh" +CANCELLED_FILE = "cancelled" +DOCKER_CONFIG_DIR = "docker_config" +ENROOT_CONFIG_DIR = "enroot_config" +REGISTRY_AUTH_FILE = "registry_auth.sh" + +_ENV_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + +# The required stack components shared by both Slurm components. +REQUIRED_COMPONENTS = { + StackComponentType.CONTAINER_REGISTRY, + StackComponentType.IMAGE_BUILDER, +} + + +@dataclass(frozen=True) +class SlurmRegistryAuth: + """Files and shell setup required for a private registry pull.""" + + files: Dict[str, str] = field(default_factory=dict, repr=False) + shell_setup: str = "" + sensitive_paths: List[str] = field(default_factory=list) + + +def serialize_environment( + environment: Mapping[str, str], runtime: SlurmContainerRuntime +) -> str: + """Serialize a step environment safely for the selected runtime. + + Args: + environment: Environment variables passed to the step. + runtime: Container runtime that consumes the environment file. + + Returns: + Serialized environment file content. + + Raises: + ValueError: If an environment variable name is invalid. + """ + for key in environment: + if not _ENV_NAME_PATTERN.fullmatch(key): + raise ValueError(f"Invalid environment variable name: {key!r}.") + + if runtime == SlurmContainerRuntime.PYXIS: + return "".join( + f"export {key}={shlex.quote(value)}\n" + for key, value in sorted(environment.items()) + ) + + from zenml.integrations.ssh.utils import serialize_env_for_docker_env_file + + return serialize_env_for_docker_env_file(environment) + + +def build_registry_auth( + runtime: SlurmContainerRuntime, + run_dir: str, + registry_uri: Optional[str], + credentials: Optional[Tuple[str, str]], +) -> SlurmRegistryAuth: + """Build owner-only registry authentication files for a Slurm job. + + Args: + runtime: Container runtime that pulls the image. + run_dir: Per-job staging directory. + registry_uri: Registry URI from the active stack. + credentials: Registry username and password. + + Returns: + Authentication files and shell setup for the selected runtime. + + Raises: + ValueError: If an enroot registry hostname or username contains + whitespace. + """ + if not registry_uri or not credentials: + return SlurmRegistryAuth() + + username, password = credentials + if runtime == SlurmContainerRuntime.DOCKER: + config_dir = f"{run_dir}/{DOCKER_CONFIG_DIR}" + config_path = f"{config_dir}/config.json" + encoded = base64.b64encode( + f"{username}:{password}".encode("utf-8") + ).decode("ascii") + content = json.dumps( + {"auths": {registry_uri: {"auth": encoded}}}, sort_keys=True + ) + return SlurmRegistryAuth( + files={config_path: content}, + shell_setup=f"export DOCKER_CONFIG={shlex.quote(config_dir)}\n", + sensitive_paths=[config_dir], + ) + + if runtime == SlurmContainerRuntime.PYXIS: + config_dir = f"{run_dir}/{ENROOT_CONFIG_DIR}" + credentials_path = f"{config_dir}/.credentials" + registry = registry_uri.removeprefix("https://").removeprefix( + "http://" + ) + registry = registry.rstrip("/").split("/", 1)[0] + if any(character.isspace() for character in registry + username): + raise ValueError( + "Enroot registry hostnames and usernames cannot contain " + "whitespace." + ) + + def _netrc_quote(value: str) -> str: + return '"' + value.replace("\\", "\\\\").replace('"', '\\"') + '"' + + # Enroot matches these fields textually to infer the registry user, so + # only the password can use netrc quoting here. + content = ( + f"machine {registry} " + f"login {username} " + f"password {_netrc_quote(password)}\n" + ) + return SlurmRegistryAuth( + files={credentials_path: content}, + shell_setup=f"export ENROOT_CONFIG_PATH={shlex.quote(config_dir)}\n", + sensitive_paths=[config_dir], + ) + + auth_path = f"{run_dir}/{REGISTRY_AUTH_FILE}" + content = "".join( + [ + f"export APPTAINER_DOCKER_USERNAME={shlex.quote(username)}\n", + f"export APPTAINER_DOCKER_PASSWORD={shlex.quote(password)}\n", + f"export SINGULARITY_DOCKER_USERNAME={shlex.quote(username)}\n", + f"export SINGULARITY_DOCKER_PASSWORD={shlex.quote(password)}\n", + ] + ) + return SlurmRegistryAuth( + files={auth_path: content}, + shell_setup=f"source {shlex.quote(auth_path)}\n", + sensitive_paths=[auth_path], + ) + + +def validate_remote_stack(stack: "Stack") -> Tuple[bool, str]: + """Ensure the stack can feed a step running on a remote Slurm cluster. + + The cluster reads inputs and writes outputs to the artifact store and + pulls the step image from the container registry, so both must be remote. + + Args: + stack: The stack to validate. + + Returns: + Whether the stack is valid and an error message if it is not. + """ + if stack.artifact_store.config.is_local: + return False, ( + "The Slurm integration runs steps on a remote cluster that must " + "read inputs and write outputs to a shared artifact store, but " + f"the artifact store '{stack.artifact_store.name}' is local. " + "Please use a remote artifact store (S3, GCS, Azure Blob, etc.)." + ) + + container_registry = stack.container_registry + assert container_registry is not None + if container_registry.config.is_local: + return False, ( + "The Slurm integration runs the step's image on a remote cluster, " + "which must pull it from a registry, but the container registry " + f"'{container_registry.name}' is local. Please use a remote " + "container registry (ECR, GCR, ACR, DockerHub, etc.)." + ) + return True, "" + + +def build_container_command( + runtime: SlurmContainerRuntime, + image: str, + entrypoint_command: List[str], + env_file: str, + env_keys: List[str], + use_gpu: bool, + settings: SlurmJobSettings, + registry_auth: Optional[SlurmRegistryAuth] = None, +) -> str: + """Build the shell snippet that runs the step image on the node. + + Secrets are passed only via the env file (or, for pyxis, via + ``--container-env`` variable names whose values are sourced from the 0600 + env file), never on the command line. + + Contract: the returned snippet is one self-contained foreground command + that must stay wrappable with ``srun --nodes=N --ntasks-per-node=1`` (the + planned multi-node mode for command steps), so it must not assume it runs + exactly once on the batch node. The pyxis runtime already executes + through ``srun``. + + Args: + runtime: The container runtime to use. + image: Fully-qualified image reference to run. + entrypoint_command: The full step command (entrypoint + arguments). + env_file: Path to the owner-only environment file on the cluster. + env_keys: Names of the environment variables (for pyxis). + use_gpu: Whether the step requested GPUs. + settings: The resolved step settings (mounts, extra run args). + registry_auth: Optional private-registry authentication setup. + + Returns: + A shell snippet that runs the container in the foreground. + """ + entrypoint = " ".join(shlex.quote(p) for p in entrypoint_command) + extra = " ".join(shlex.quote(a) for a in settings.container_run_args) + shell_setup = registry_auth.shell_setup if registry_auth else "" + + if runtime in ( + SlurmContainerRuntime.APPTAINER, + SlurmContainerRuntime.SINGULARITY, + ): + binary = runtime.value # "apptainer" or "singularity" + parts = [binary, "exec", "--no-eval"] + if use_gpu: + parts.append("--nv") + parts += ["--env-file", shlex.quote(env_file)] + for host, container in settings.container_mounts.items(): + parts += ["--bind", shlex.quote(f"{host}:{container}")] + if extra: + parts.append(extra) + parts.append(shlex.quote(f"docker://{image}")) + return f"{shell_setup}{' '.join(parts)} {entrypoint}" + + if runtime == SlurmContainerRuntime.DOCKER: + parts = ["docker", "run", "--rm"] + if use_gpu: + parts += ["--gpus", "all"] + parts += ["--env-file", shlex.quote(env_file)] + for host, container in settings.container_mounts.items(): + parts += ["-v", shlex.quote(f"{host}:{container}")] + if extra: + parts.append(extra) + parts.append(shlex.quote(image)) + return f"{shell_setup}{' '.join(parts)} {entrypoint}" + + # Pyxis / enroot via srun. `--container-env` takes variable *names*; the + # values are sourced from the 0600 env file into the job shell, so they + # are not exposed on the command line. + srun = ["srun", f"--container-image={shlex.quote(image)}"] + if settings.container_mounts: + mounts = ",".join( + f"{host}:{container}" + for host, container in settings.container_mounts.items() + ) + srun.append(f"--container-mounts={shlex.quote(mounts)}") + if env_keys: + srun.append(f"--container-env={','.join(env_keys)}") + if extra: + srun.append(extra) + source_env = f"set -a\nsource {shlex.quote(env_file)}\nset +a\n" + return f"{shell_setup}{source_env}{' '.join(srun)} {entrypoint}" + + +def build_sbatch_script( + job_name: str, + run_dir: str, + container_command: str, + resources: "ResourceSettings", + settings: SlurmJobSettings, + sensitive_paths: Optional[List[str]] = None, +) -> str: + """Render the sbatch job script for a step. + + Args: + job_name: The Slurm job name. + run_dir: The per-run staging directory on the cluster. + container_command: The shell snippet that runs the step container. + resources: The step's resource settings (cpu/mem/gpu). + settings: The resolved Slurm job settings (partition, time, ...). + sensitive_paths: Credential-bearing paths removed when the job exits. + + Returns: + The job script content. + """ + directives = [ + f"#SBATCH --job-name={job_name}", + f"#SBATCH --output={run_dir}/{OUTPUT_FILE}", + ] + if settings.partition: + directives.append(f"#SBATCH --partition={settings.partition}") + if settings.time_limit: + directives.append(f"#SBATCH --time={settings.time_limit}") + if settings.account: + directives.append(f"#SBATCH --account={settings.account}") + if settings.qos: + directives.append(f"#SBATCH --qos={settings.qos}") + if resources.cpu_count: + directives.append( + f"#SBATCH --cpus-per-task={math.ceil(resources.cpu_count)}" + ) + if memory_mb := resources.get_memory(unit="MB"): + directives.append(f"#SBATCH --mem={math.ceil(memory_mb)}M") + if resources.gpu_count: + directives.append(f"#SBATCH --gres=gpu:{resources.gpu_count}") + for directive in settings.extra_sbatch_directives: + directives.append(f"#SBATCH {directive}") + + # The EXIT trap records the job outcome in a sentinel file that the caller + # reads after the job leaves the queue (so no dependency on Slurm + # accounting) and scrubs the credential-bearing env file, so secrets never + # outlive the job even if the submitting process dies. `set -e` aborts the + # job (with the failing code captured) if the container runtime fails. + cleanup_paths = [f"{run_dir}/{ENV_FILE}", *(sensitive_paths or [])] + cleanup_command = "rm -rf -- " + " ".join( + shlex.quote(path) for path in cleanup_paths + ) + exit_code_path = shlex.quote(f"{run_dir}/{EXIT_CODE_FILE}") + + return f"""#!/bin/bash +{chr(10).join(directives)} + +cleanup() {{ + ec=$? + trap - EXIT + echo "$ec" > {exit_code_path} + {cleanup_command} + exit "$ec" +}} +trap cleanup EXIT +set -eo pipefail + +{container_command} +""" + + +def stage_and_submit( + client: "SlurmClient", + run_dir: str, + env_content: str, + script: str, + dependencies: Optional[List[str]] = None, + extra_files: Optional[Dict[str, str]] = None, +) -> str: + """Stage a job's files on the cluster and submit it. + + This owns the security-sensitive part both components share: the run + directory is created owner-only and atomically (``mkdir -m 700``, so it is + never briefly readable by other cluster users), the credential-bearing + environment file is written owner-only (0600), and the job script 0700. + + Args: + client: The Slurm client for the cluster. + run_dir: The per-job staging directory on the cluster. + env_content: The serialized environment file content (holds secrets). + script: The sbatch job script content. + dependencies: Slurm job ids this job must wait for (for DAG edges). + extra_files: Additional owner-only files to stage for the job. + + Returns: + The submitted Slurm job id. + + Raises: + RuntimeError: If the run directory cannot be created on the cluster. + Exception: Re-raised after cleaning up staged files if staging or + submission fails. + """ + runner = client.runner + result = runner.run( + f"umask 077 && mkdir -p {shlex.quote(run_dir)} && " + f"chmod 700 {shlex.quote(run_dir)}" + ) + if result.exit_code != 0: + raise RuntimeError( + f"Failed to create run directory `{run_dir}` on the cluster: " + f"{result.stderr.strip()}" + ) + + try: + runner.put_text(f"{run_dir}/{ENV_FILE}", env_content, mode=0o600) + for path, content in (extra_files or {}).items(): + parent = posixpath.dirname(path) + if parent != run_dir: + result = runner.run( + f"umask 077 && mkdir -p {shlex.quote(parent)} && " + f"chmod 700 {shlex.quote(parent)}" + ) + if result.exit_code != 0: + raise RuntimeError( + f"Failed to create registry auth directory `{parent}`: " + f"{result.stderr.strip()}" + ) + runner.put_text(path, content, mode=0o600) + + script_path = f"{run_dir}/{SCRIPT_FILE}" + runner.put_text(script_path, script, mode=0o700) + return client.submit(script_path, dependencies=dependencies) + except Exception: + runner.run(f"rm -rf -- {shlex.quote(run_dir)}") + raise diff --git a/src/zenml/integrations/slurm/step_operators/__init__.py b/src/zenml/integrations/slurm/step_operators/__init__.py new file mode 100644 index 00000000000..7260ca2931c --- /dev/null +++ b/src/zenml/integrations/slurm/step_operators/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Initialization of the Slurm step operator.""" + +from zenml.integrations.slurm.step_operators.slurm_step_operator import ( # noqa + SlurmStepOperator, +) + +__all__ = ["SlurmStepOperator"] diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py new file mode 100644 index 00000000000..a5d2bb93f6e --- /dev/null +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -0,0 +1,428 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Step operator that runs individual steps as Slurm jobs. + +The step's ZenML Docker image is run 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 image +build - no code is shipped by the operator. Submission goes over SSH to a +login node, or via a local ``sbatch`` when the client already runs on the +cluster. + +The Slurm job ID is persisted as step-run metadata for exact status lookups and +cancellation. Once a job leaves the queue, a sentinel exit-code file written by +the job script disambiguates success from failure without requiring Slurm +accounting (``sacct``). +""" + +import shlex +from typing import TYPE_CHECKING, Dict, List, Optional, cast +from uuid import UUID + +from zenml.config.base_settings import BaseSettings +from zenml.config.build_configuration import BuildConfiguration +from zenml.enums import ExecutionStatus +from zenml.integrations.slurm.flavors.slurm_step_operator_flavor import ( + SlurmStepOperatorConfig, + SlurmStepOperatorSettings, +) +from zenml.integrations.slurm.slurm_client import ( + PENDING_STATES, + RUNNING_STATES, + SlurmCommandRunner, + build_slurm_client, +) +from zenml.integrations.slurm.slurm_job import ( + CANCELLED_FILE, + DOCKER_CONFIG_DIR, + ENROOT_CONFIG_DIR, + ENV_FILE, + EXIT_CODE_FILE, + OUTPUT_FILE, + REGISTRY_AUTH_FILE, + REQUIRED_COMPONENTS, + build_container_command, + build_registry_auth, + build_sbatch_script, + serialize_environment, + stage_and_submit, + validate_remote_stack, +) +from zenml.logger import get_logger +from zenml.orchestrators.publish_utils import publish_step_run_metadata +from zenml.stack import StackValidator +from zenml.step_operators import BaseStepOperator + +if TYPE_CHECKING: + from zenml.config.step_run_info import StepRunInfo + from zenml.metadata.metadata_types import MetadataType + from zenml.models import PipelineSnapshotBase, StepRunResponse + +logger = get_logger(__name__) + +SLURM_STEP_OPERATOR_DOCKER_IMAGE_KEY = "slurm_step_operator" +SLURM_JOB_ID_METADATA_KEY = "slurm_job_id" +SLURM_STATE_METADATA_KEY = "slurm_state" + +_JOB_NAME_PREFIX = "zenml" + + +class SlurmStepOperator(BaseStepOperator): + """Step operator that submits steps as Slurm batch jobs.""" + + @property + def config(self) -> SlurmStepOperatorConfig: + """Returns the config of this step operator. + + Returns: + The config. + """ + return cast(SlurmStepOperatorConfig, self._config) + + @property + def settings_class(self) -> Optional[type[BaseSettings]]: + """Settings class for this step operator. + + Returns: + The settings class. + """ + return SlurmStepOperatorSettings + + @property + def validator(self) -> Optional[StackValidator]: + """Validate that the stack meets the operator's requirements. + + Returns: + A stack validator. + """ + return StackValidator( + required_components=REQUIRED_COMPONENTS, + custom_validation_function=validate_remote_stack, + ) + + def get_docker_builds( + self, snapshot: "PipelineSnapshotBase" + ) -> List["BuildConfiguration"]: + """Declare the Docker builds for steps using this operator. + + Args: + snapshot: The pipeline snapshot. + + Returns: + One build configuration per step that uses this step operator. + """ + builds = [] + for step_name, step in snapshot.step_configurations.items(): + if step.config.uses_step_operator(self.name): + builds.append( + BuildConfiguration( + key=SLURM_STEP_OPERATOR_DOCKER_IMAGE_KEY, + settings=step.config.docker_settings, + step_name=step_name, + ) + ) + return builds + + @staticmethod + def _job_name(step_run_id: UUID) -> str: + """Deterministic Slurm job name for a step run. + + Args: + step_run_id: The step run id. + + Returns: + The Slurm job name. + """ + return f"{_JOB_NAME_PREFIX}-{step_run_id}" + + def _run_dir(self, step_run_id: UUID) -> str: + """Per-run staging directory on the cluster. + + Args: + step_run_id: The step run id. + + Returns: + The absolute path of the run directory. + """ + return ( + f"{self.config.workdir.rstrip('/')}/{self._job_name(step_run_id)}" + ) + + def submit( + self, + info: "StepRunInfo", + entrypoint_command: List[str], + environment: Dict[str, str], + ) -> None: + """Submit a step as a Slurm batch job and return immediately. + + Args: + info: Information about the step run. + entrypoint_command: Command that executes the step. + environment: Environment variables to set in the step operator + environment. + + Raises: + Exception: Re-raised after cancelling the job and cleaning up + staged credentials if submission or metadata publication fails. + """ + image = info.get_image(key=SLURM_STEP_OPERATOR_DOCKER_IMAGE_KEY) + run_dir = self._run_dir(info.step_run_id) + env_file = f"{run_dir}/{ENV_FILE}" + settings = cast(SlurmStepOperatorSettings, self.get_settings(info)) + use_gpu = bool(info.config.resource_settings.gpu_count) + + env_content = serialize_environment( + environment, runtime=self.config.container_runtime + ) + from zenml.client import Client + + container_registry = Client().active_stack.container_registry + assert container_registry is not None + registry_auth = build_registry_auth( + runtime=self.config.container_runtime, + run_dir=run_dir, + registry_uri=container_registry.config.uri, + credentials=container_registry.credentials, + ) + container_command = build_container_command( + runtime=self.config.container_runtime, + image=image, + entrypoint_command=entrypoint_command, + env_file=env_file, + env_keys=sorted(environment), + use_gpu=use_gpu, + settings=settings, + registry_auth=registry_auth, + ) + script = build_sbatch_script( + job_name=self._job_name(info.step_run_id), + run_dir=run_dir, + container_command=container_command, + resources=info.config.resource_settings, + settings=settings, + sensitive_paths=registry_auth.sensitive_paths, + ) + + client = build_slurm_client(self.config) + job_id: Optional[str] = None + try: + job_id = stage_and_submit( + client, + run_dir, + env_content, + script, + extra_files=registry_auth.files, + ) + metadata: Dict[str, "MetadataType"] = { + SLURM_JOB_ID_METADATA_KEY: job_id + } + publish_step_run_metadata( + step_run_id=info.step_run_id, + step_run_metadata={self.id: metadata}, + ) + info.step_run.run_metadata.update(metadata) + logger.info( + "Submitted step `%s` as Slurm job `%s` (job name `%s`).", + info.pipeline_step_name, + job_id, + self._job_name(info.step_run_id), + ) + except Exception: + if job_id is not None: + try: + client.cancel(job_id) + except Exception: + logger.warning( + "Failed to cancel Slurm job `%s` after step metadata " + "publication failed.", + job_id, + ) + self._cleanup_sensitive_files(client.runner, run_dir) + raise + finally: + client.runner.close() + + def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: + """Get the status of a step run from the Slurm queue. + + Args: + step_run: The step run to get the status of. + + Returns: + The execution status of the step run. + """ + job_id = step_run.run_metadata.get(SLURM_JOB_ID_METADATA_KEY) + if job_id is None: + logger.warning( + "No Slurm job ID recorded for step `%s`.", step_run.id + ) + return ExecutionStatus.FAILED + run_dir = self._run_dir(step_run.id) + client = build_slurm_client(self.config) + runner = client.runner + try: + state = client.get_job_state(str(job_id)) + if state in PENDING_STATES: + return ExecutionStatus.QUEUED + if state in RUNNING_STATES: + return ExecutionStatus.RUNNING + if state == "COMPLETED": + return ExecutionStatus.COMPLETED + if state is not None: + # Cancellation and failure states can linger in the queue + # output briefly before the job is purged from it. + self._publish_terminal_state(step_run, state) + if state.startswith("CANCEL"): + return ExecutionStatus.CANCELLED + return ExecutionStatus.FAILED + + # The job is no longer queued: the sentinel file written by the + # job script's EXIT trap holds the outcome. + try: + runner.read_text(f"{run_dir}/{CANCELLED_FILE}") + except Exception: + pass + else: + return ExecutionStatus.CANCELLED + + try: + exit_code = runner.read_text( + f"{run_dir}/{EXIT_CODE_FILE}" + ).strip() + except Exception: + logger.warning( + "Slurm job `%s` for step `%s` is no longer known to " + "Slurm and did not write an exit-code sentinel.", + job_id, + step_run.id, + ) + return ExecutionStatus.FAILED + + if exit_code == "0": + return ExecutionStatus.COMPLETED + + self._log_failure_output(runner, run_dir) + return ExecutionStatus.FAILED + finally: + runner.close() + + def _publish_terminal_state( + self, step_run: "StepRunResponse", state: str + ) -> None: + """Record the raw Slurm terminal state as step metadata. + + A scheduler-level kill (``TIMEOUT``, ``NODE_FAIL``, + ``OUT_OF_MEMORY``, ``PREEMPTED``) never runs the job's EXIT trap, + so the sentinel-based status collapses it into a plain failure. + The raw state is the only signal that distinguishes an + infrastructure kill from a step that ran and exited non-zero. + + Args: + step_run: The step run the state belongs to. + state: The Slurm terminal state as reported by the queue. + """ + if SLURM_STATE_METADATA_KEY in step_run.run_metadata: + return + try: + publish_step_run_metadata( + step_run_id=step_run.id, + step_run_metadata={self.id: {SLURM_STATE_METADATA_KEY: state}}, + ) + except Exception as e: + logger.warning( + "Failed to publish Slurm terminal state `%s` for step " + "`%s`: %s", + state, + step_run.id, + e, + ) + + def cancel(self, step_run: "StepRunResponse") -> None: + """Cancel the Slurm job of a step run. + + Args: + step_run: The step run to cancel. + """ + client = build_slurm_client(self.config) + try: + job_id = step_run.run_metadata.get(SLURM_JOB_ID_METADATA_KEY) + if job_id is None: + logger.warning( + "No Slurm job ID recorded for step `%s`.", step_run.id + ) + return + client.cancel(str(job_id)) + run_dir = self._run_dir(step_run.id) + client.runner.put_text( + f"{run_dir}/{CANCELLED_FILE}", "1\n", mode=0o600 + ) + self._cleanup_sensitive_files(client.runner, run_dir) + finally: + client.runner.close() + + def cleanup_step_submission(self, step_run: "StepRunResponse") -> None: + """Remove the per-run directory after the step has finished. + + The credential-bearing env file is already scrubbed by the job's EXIT + trap; this removes the remaining job script, output and sentinel. + + Args: + step_run: The finished step run. + """ + run_dir = self._run_dir(step_run.id) + client = build_slurm_client(self.config) + try: + client.runner.run(f"rm -rf -- {shlex.quote(run_dir)}") + except Exception as e: + logger.warning( + "Failed to clean up Slurm run directory `%s`: %s", run_dir, e + ) + finally: + client.runner.close() + + def _log_failure_output( + self, runner: SlurmCommandRunner, run_dir: str + ) -> None: + """Log the tail of a failed job's output file. + + Args: + runner: The command runner to read the output with. + run_dir: The per-run staging directory on the cluster. + """ + result = runner.run(f"tail -n 50 {shlex.quote(run_dir)}/{OUTPUT_FILE}") + if result.exit_code == 0 and result.stdout.strip(): + logger.error( + "Tail of the failed Slurm job output:\n%s", result.stdout + ) + + @staticmethod + def _cleanup_sensitive_files( + runner: SlurmCommandRunner, run_dir: str + ) -> None: + """Remove credential-bearing files for a submitted step. + + Args: + runner: Command runner used to remove the files. + run_dir: Per-step staging directory. + """ + paths = [ + f"{run_dir}/{ENV_FILE}", + f"{run_dir}/{REGISTRY_AUTH_FILE}", + f"{run_dir}/{DOCKER_CONFIG_DIR}", + f"{run_dir}/{ENROOT_CONFIG_DIR}", + ] + runner.run( + "rm -rf -- " + " ".join(shlex.quote(path) for path in paths) + ) diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py new file mode 100644 index 00000000000..f295f0a6cf0 --- /dev/null +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -0,0 +1,854 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Unit tests for the Slurm orchestrator. + +The Slurm CLI and SSH are never invoked: the orchestrator is driven through an +in-memory fake command runner, so the tests run without a cluster. +""" + +from datetime import datetime +from types import SimpleNamespace +from typing import Dict, List, Optional +from uuid import uuid4 + +import pytest + +from zenml.config.resource_settings import ResourceSettings +from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID +from zenml.enums import ExecutionMode, ExecutionStatus, StackComponentType +from zenml.integrations.slurm.flavors import ( + SlurmOrchestratorConfig, + SlurmOrchestratorFlavor, + SlurmOrchestratorSettings, +) +from zenml.integrations.slurm.orchestrators import SlurmOrchestrator +from zenml.integrations.slurm.orchestrators.slurm_orchestrator import ( + ENV_ZENML_SLURM_RUN_ID, + SLURM_CLEANUP_JOB_ID_METADATA_KEY, + SLURM_JOB_IDS_METADATA_KEY, + SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY, +) +from zenml.integrations.slurm.slurm_client import ( + CommandResult, + SlurmClient, + SlurmCommandRunner, +) + +SECRET_TOKEN = "super-secret-zenml-token" +IMAGE = "registry.example.com/zenml:abc123" + + +class FakeRunner(SlurmCommandRunner): + """In-memory fake host that hands out incrementing job ids.""" + + def __init__(self) -> None: + """Initialize the fake.""" + self.files: Dict[str, str] = {} + self.modes: Dict[str, int] = {} + self.commands: List[str] = [] + self._next_job_id = 1000 + + def run(self, command: str) -> CommandResult: + """Record the command; return a fresh job id for each sbatch. + + Args: + command: The command line. + + Returns: + The command result. + """ + self.commands.append(command) + if command.startswith("sbatch"): + job_id = str(self._next_job_id) + self._next_job_id += 1 + return CommandResult(exit_code=0, stdout=job_id, stderr="") + return CommandResult(exit_code=0, stdout="", stderr="") + + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: + """Record a written file and its mode. + + Args: + remote_path: Destination path. + content: The content written. + mode: The file mode. + """ + self.files[remote_path] = content + self.modes[remote_path] = mode + + def read_text(self, remote_path: str) -> str: + """Read a recorded file. + + Args: + remote_path: Path to read. + + Returns: + The recorded content. + + Raises: + FileNotFoundError: If nothing was recorded at that path. + """ + if remote_path not in self.files: + raise FileNotFoundError(remote_path) + return self.files[remote_path] + + +def _build_orchestrator( + config: Optional[SlurmOrchestratorConfig] = None, +) -> SlurmOrchestrator: + """Construct a Slurm orchestrator. + + Args: + config: The config for the orchestrator. + + Returns: + A SlurmOrchestrator instance. + """ + return SlurmOrchestrator( + name="", + id=uuid4(), + config=config + or SlurmOrchestratorConfig(transport="local", workdir="/runs"), + flavor="slurm", + type=StackComponentType.ORCHESTRATOR, + user=uuid4(), + created=datetime.now(), + updated=datetime.now(), + ) + + +def _step(upstream: List[str], gpu: int = 0) -> SimpleNamespace: + """Build a Step stand-in. + + Args: + upstream: Invocation ids this step depends on. + gpu: Number of GPUs the step requests. + + Returns: + A namespace usable where a Step is expected. + """ + return SimpleNamespace( + spec=SimpleNamespace(upstream_steps=upstream), + config=SimpleNamespace( + resource_settings=ResourceSettings(gpu_count=gpu or None) + ), + ) + + +def _snapshot( + steps: Dict[str, SimpleNamespace], + pipeline_resources: Optional[ResourceSettings] = None, +) -> SimpleNamespace: + """Build a PipelineSnapshotResponse stand-in. + + Args: + steps: The step configurations keyed by invocation id. + pipeline_resources: Pipeline-level resource settings. + + Returns: + A namespace usable where a snapshot is expected. + """ + return SimpleNamespace( + id=uuid4(), + build=SimpleNamespace( + get_image=lambda component_key, step: IMAGE, + ), + step_configurations=steps, + pipeline_configuration=SimpleNamespace( + name="p", + resource_settings=pipeline_resources or ResourceSettings(), + ), + schedule=None, + stack=SimpleNamespace(id=uuid4()), + ) + + +# --- flavor / config ---------------------------------------------------------- + + +def test_flavor_identity(): + """The flavor identifier and classes are Slurm's.""" + flavor = SlurmOrchestratorFlavor() + assert flavor.name == "slurm" + assert flavor.config_class is SlurmOrchestratorConfig + assert flavor.implementation_class is SlurmOrchestrator + + +def test_ssh_transport_requires_connection_details(): + """The ssh transport needs hostname and username; local does not.""" + with pytest.raises(ValueError): + SlurmOrchestratorConfig(transport="ssh", workdir="/s") + + SlurmOrchestratorConfig(transport="local", workdir="/s") + SlurmOrchestratorConfig( + transport="ssh", workdir="/s", hostname="h", username="u" + ) + + +def test_orchestrator_is_remote_and_detached(): + """The orchestrator runs remotely and does not wait for completion.""" + config = SlurmOrchestratorConfig(transport="local", workdir="/s") + assert config.is_remote is True + assert config.is_synchronous is False + + +def test_orchestrator_supports_failure_modes(): + """The orchestrator advertises the failure modes it implements.""" + op = _build_orchestrator() + assert op.supported_execution_modes == [ + ExecutionMode.FAIL_FAST, + ExecutionMode.STOP_ON_FAILURE, + ExecutionMode.CONTINUE_ON_FAILURE, + ] + + +# --- run id ------------------------------------------------------------------- + + +def test_get_orchestrator_run_id_reads_env(monkeypatch): + """The run id comes from the injected environment variable.""" + op = _build_orchestrator() + monkeypatch.setenv(ENV_ZENML_SLURM_RUN_ID, "the-run-id") + assert op.get_orchestrator_run_id() == "the-run-id" + + +def test_get_orchestrator_run_id_raises_when_unset(monkeypatch): + """Without the env var, the run id lookup fails loudly.""" + op = _build_orchestrator() + monkeypatch.delenv(ENV_ZENML_SLURM_RUN_ID, raising=False) + with pytest.raises(RuntimeError): + op.get_orchestrator_run_id() + + +# --- submit ------------------------------------------------------------------- + + +def _use_fake_client(monkeypatch, runner: FakeRunner) -> None: + """Route the orchestrator's client factory to the fake runner. + + Args: + monkeypatch: The pytest monkeypatch fixture. + runner: The fake runner to hand back. + """ + monkeypatch.setattr( + "zenml.integrations.slurm.orchestrators.slurm_orchestrator" + ".build_slurm_client", + lambda config: SlurmClient(runner), + ) + + +def _placeholder_run() -> SimpleNamespace: + """Build a PipelineRunResponse stand-in with a unique id. + + Returns: + A namespace usable where a placeholder run is expected. + """ + return SimpleNamespace(id=uuid4()) + + +def _stack( + credentials: Optional["tuple[str, str]"] = None, +) -> SimpleNamespace: + """Build a stack stand-in with a remote container registry. + + Args: + credentials: Optional registry username and password. + + Returns: + A namespace usable where a stack is expected. + """ + return SimpleNamespace( + container_registry=SimpleNamespace( + config=SimpleNamespace(uri="registry.example.com"), + credentials=credentials, + ) + ) + + +def _run( + run_id, + run_metadata: Dict[str, object], + execution_mode: ExecutionMode = ExecutionMode.CONTINUE_ON_FAILURE, +) -> SimpleNamespace: + """Build a PipelineRunResponse stand-in. + + Args: + run_id: The run ID. + run_metadata: Run metadata to attach. + execution_mode: Pipeline execution mode. + + Returns: + A namespace usable where a PipelineRunResponse is expected. + """ + return SimpleNamespace( + id=run_id, + status=ExecutionStatus.PROVISIONING, + config=SimpleNamespace(execution_mode=execution_mode), + run_metadata=run_metadata, + ) + + +def _submit_two_step_dag(monkeypatch) -> "tuple[FakeRunner, str]": + """Submit a load -> train DAG and return the fake host and run id. + + Args: + monkeypatch: The pytest monkeypatch fixture. + + Returns: + The fake runner recording everything the submission did, and the + orchestrator run id (the placeholder run id). + """ + op = _build_orchestrator() + op.get_settings = lambda _step: SlurmOrchestratorSettings() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + + snapshot = _snapshot({"load": _step([]), "train": _step(["load"], gpu=1)}) + placeholder = _placeholder_run() + op.submit_pipeline( + snapshot=snapshot, + stack=_stack(), + base_environment={}, + step_environments={ + "load": {"ZENML_STORE_API_KEY": SECRET_TOKEN}, + "train": {"ZENML_STORE_API_KEY": SECRET_TOKEN}, + }, + placeholder_run=placeholder, + ) + return runner, str(placeholder.id) + + +def test_submit_returns_slurm_metadata(monkeypatch): + """Submitted job ids and the run id are attached to the run metadata.""" + op = _build_orchestrator() + op.get_settings = lambda _step: SlurmOrchestratorSettings() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + snapshot = _snapshot({"load": _step([]), "train": _step(["load"])}) + placeholder = _placeholder_run() + + result = op.submit_pipeline( + snapshot=snapshot, + stack=_stack(), + base_environment={}, + step_environments={"load": {}, "train": {}}, + placeholder_run=placeholder, + ) + + assert result is not None + assert result.metadata == { + METADATA_ORCHESTRATOR_RUN_ID: str(placeholder.id), + SLURM_JOB_IDS_METADATA_KEY: {"load": "1000", "train": "1001"}, + SLURM_CLEANUP_JOB_ID_METADATA_KEY: "1002", + } + + +def test_submit_wires_dependencies_in_topological_order(monkeypatch): + """The dependent step's job waits on the parent's job id.""" + runner, _ = _submit_two_step_dag(monkeypatch) + sbatch = [c for c in runner.commands if c.startswith("sbatch")] + assert len(sbatch) == 3 + # load submitted first (job id 1000, no dependency), train second. + assert "--dependency" not in sbatch[0] + assert "--dependency=afterok:1000" in sbatch[1] + assert "--kill-on-invalid-dep=yes" in sbatch[1] + assert "--dependency=afterany:1000:1001" in sbatch[2] + cleanup_scripts = [ + content + for path, content in runner.files.items() + if path.endswith("/cleanup/job.sh") + ] + assert len(cleanup_scripts) == 1 + assert "/env" in cleanup_scripts[0] + assert "cleanup_complete" in cleanup_scripts[0] + + +def test_cleanup_job_inherits_extra_sbatch_directives(monkeypatch): + """Cleanup jobs preserve cluster-specific sbatch directives.""" + op = _build_orchestrator() + + def _settings(obj): + if hasattr(obj, "pipeline_configuration"): + return SlurmOrchestratorSettings( + extra_sbatch_directives=["--constraint=a100"] + ) + return SlurmOrchestratorSettings() + + op.get_settings = _settings + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + op.submit_pipeline( + snapshot=_snapshot({"load": _step([])}), + stack=_stack(), + base_environment={}, + step_environments={"load": {}}, + placeholder_run=_placeholder_run(), + ) + + cleanup_script = next( + content + for path, content in runner.files.items() + if path.endswith("/cleanup/job.sh") + ) + assert "#SBATCH --constraint=a100" in cleanup_script + + +def test_submit_injects_placeholder_run_id_into_step_env(monkeypatch): + """Each step's env file carries the placeholder run id as the run id.""" + runner, run_id = _submit_two_step_dag(monkeypatch) + env_files = [c for p, c in runner.files.items() if p.endswith("/env")] + assert env_files + # The run id is the placeholder run id (unique per run), not the snapshot + # id (which repeats across runs of the same snapshot). + assert all( + f"{ENV_ZENML_SLURM_RUN_ID}={run_id}" in content + for content in env_files + ) + + +def test_submit_uses_placeholder_run_id_for_paths(monkeypatch): + """Staging paths are namespaced by the placeholder run id.""" + runner, run_id = _submit_two_step_dag(monkeypatch) + assert all(f"/{run_id}/" in path for path in runner.files) + + +def test_submit_requires_a_placeholder_run(monkeypatch): + """Submission without a placeholder run is rejected (no run id).""" + op = _build_orchestrator() + op.get_settings = lambda _step: SlurmOrchestratorSettings() + _use_fake_client(monkeypatch, FakeRunner()) + with pytest.raises(AssertionError): + op.submit_pipeline( + snapshot=_snapshot({"load": _step([])}), + stack=_stack(), + base_environment={}, + step_environments={"load": {}}, + placeholder_run=None, + ) + + +def test_submit_rejects_scheduled_pipelines(monkeypatch): + """The orchestrator refuses scheduled pipelines with a clear error.""" + op = _build_orchestrator() + _use_fake_client(monkeypatch, FakeRunner()) + snapshot = _snapshot({"load": _step([])}) + snapshot.schedule = SimpleNamespace(name="daily") + with pytest.raises(RuntimeError, match="does not support scheduled"): + op.submit_pipeline( + snapshot=snapshot, + stack=_stack(), + base_environment={}, + step_environments={"load": {}}, + placeholder_run=_placeholder_run(), + ) + + +def test_submit_dynamic_pipeline_submits_orchestration_job(monkeypatch): + """Dynamic pipelines launch a Slurm orchestration job. + + All steps of a dynamic pipeline run inline inside that job, so its + allocation must carry the pipeline-level resource settings. + """ + op = _build_orchestrator() + op.get_settings = lambda _snapshot: SlurmOrchestratorSettings() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + placeholder = _placeholder_run() + + result = op.submit_dynamic_pipeline( + snapshot=_snapshot( + {}, pipeline_resources=ResourceSettings(gpu_count=2) + ), + stack=_stack(), + environment={"ZENML_STORE_API_KEY": SECRET_TOKEN}, + placeholder_run=placeholder, + ) + + assert result is not None + assert result.metadata == { + METADATA_ORCHESTRATOR_RUN_ID: str(placeholder.id), + SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000", + SLURM_CLEANUP_JOB_ID_METADATA_KEY: "1001", + } + sbatch = [c for c in runner.commands if c.startswith("sbatch")] + assert len(sbatch) == 2 + assert "orchestration/job.sh" in sbatch[0] + job_script = runner.files[f"/runs/{placeholder.id}/orchestration/job.sh"] + assert "#SBATCH --gres=gpu:2" in job_script + env_file = runner.files[f"/runs/{placeholder.id}/orchestration/env"] + assert f"{ENV_ZENML_SLURM_RUN_ID}={placeholder.id}" in env_file + assert SECRET_TOKEN in env_file + # The credential reaper: an `afterany` cleanup job scrubs the staged + # sensitive files even if the orchestration job is killed before its + # own EXIT trap can run. + assert "--dependency=afterany:1000" in sbatch[1] + cleanup_script = runner.files[f"/runs/{placeholder.id}/cleanup/job.sh"] + assert f"/runs/{placeholder.id}/orchestration/env" in cleanup_script + + +def test_submit_dynamic_pipeline_rejects_schedules(monkeypatch): + """Scheduled dynamic pipelines are rejected consistently.""" + op = _build_orchestrator() + _use_fake_client(monkeypatch, FakeRunner()) + snapshot = _snapshot({}) + snapshot.schedule = SimpleNamespace(name="daily") + + with pytest.raises(RuntimeError, match="does not support scheduled"): + op.submit_dynamic_pipeline( + snapshot=snapshot, + stack=_stack(), + environment={}, + placeholder_run=_placeholder_run(), + ) + + +def test_submit_keeps_secrets_off_the_command_line(monkeypatch): + """Secrets live only in the 0600 env file, never on a command line.""" + runner, _ = _submit_two_step_dag(monkeypatch) + assert all(SECRET_TOKEN not in c for c in runner.commands) + + env_files = {p: c for p, c in runner.files.items() if p.endswith("/env")} + assert env_files + for path, content in env_files.items(): + assert SECRET_TOKEN in content + assert runner.modes[path] == 0o600 + scripts = {p: m for p, m in runner.modes.items() if p.endswith("job.sh")} + assert scripts + assert all(mode == 0o700 for mode in scripts.values()) + + +def test_submit_cancels_queued_jobs_on_failure(monkeypatch): + """A mid-DAG failure cancels the jobs already queued for the run.""" + op = _build_orchestrator() + op.get_settings = lambda _step: SlurmOrchestratorSettings() + + class FailingRunner(FakeRunner): + def put_text(self, remote_path, content, mode=0o600): + # Fail while staging the second step, after the first is queued. + if remote_path.endswith("/env") and self._next_job_id > 1000: + raise RuntimeError("disk full") + super().put_text(remote_path, content, mode=mode) + + runner = FailingRunner() + _use_fake_client(monkeypatch, runner) + snapshot = _snapshot({"load": _step([]), "train": _step(["load"])}) + with pytest.raises(RuntimeError): + op.submit_pipeline( + snapshot=snapshot, + stack=_stack(), + base_environment={}, + step_environments={"load": {}, "train": {}}, + placeholder_run=_placeholder_run(), + ) + assert any(c.startswith("scancel 1000") for c in runner.commands) + assert any(c.startswith("rm -rf -- ") for c in runner.commands) + + +def test_fetch_status_reports_pre_entrypoint_failure(monkeypatch): + """A failed Slurm wrapper reconciles a detached run as failed.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + run_dir = op._run_dir(str(run_id), "load") + runner.files[f"{run_dir}/exit_code"] = "1\n" + run = _run( + run_id=run_id, + run_metadata={SLURM_JOB_IDS_METADATA_KEY: {"load": "1000"}}, + ) + + pipeline_status, step_statuses = op.fetch_status(run, include_steps=True) + + assert pipeline_status is ExecutionStatus.FAILED + assert step_statuses == {"load": ExecutionStatus.FAILED} + + +def test_fetch_status_uses_cleanup_marker_for_never_started_job(monkeypatch): + """A dependency-cancelled job is terminal after cleanup completes.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + cleanup_marker = f"/runs/{run_id}/cleanup/cleanup_complete" + runner.files[cleanup_marker] = "" + run = _run( + run_id=run_id, + run_metadata={SLURM_JOB_IDS_METADATA_KEY: {"train": "1001"}}, + ) + + pipeline_status, _ = op.fetch_status(run) + + assert pipeline_status is ExecutionStatus.FAILED + + +def test_fetch_status_continue_on_failure_keeps_siblings(monkeypatch): + """Continue-on-failure does not cancel independent sibling jobs.""" + op = _build_orchestrator() + + class StatusRunner(FakeRunner): + def run(self, command: str) -> CommandResult: + self.commands.append(command) + if command.startswith("squeue"): + return CommandResult( + exit_code=0, + stdout=("1000|FAILED\n1001|RUNNING\n1002|PENDING\n"), + stderr="", + ) + return CommandResult(exit_code=0, stdout="", stderr="") + + runner = StatusRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + run = _run( + run_id=run_id, + run_metadata={ + SLURM_JOB_IDS_METADATA_KEY: { + "load": "1000", + "train": "1001", + "evaluate": "1002", + } + }, + ) + + pipeline_status, step_statuses = op.fetch_status(run, include_steps=True) + + assert pipeline_status is ExecutionStatus.RUNNING + assert step_statuses == { + "load": ExecutionStatus.FAILED, + "train": ExecutionStatus.RUNNING, + "evaluate": ExecutionStatus.QUEUED, + } + assert [ + command for command in runner.commands if command.startswith("squeue") + ] == ["squeue --noheader --format='%i|%T' --jobs=1000,1001,1002"] + assert not any( + command.startswith("scancel") for command in runner.commands + ) + + +def test_fetch_status_stop_on_failure_cancels_siblings(monkeypatch): + """Stop-on-failure cancels unfinished sibling jobs.""" + op = _build_orchestrator() + + class StatusRunner(FakeRunner): + def run(self, command: str) -> CommandResult: + self.commands.append(command) + if command.startswith("squeue"): + return CommandResult( + exit_code=0, + stdout=("1000|FAILED\n1001|RUNNING\n1002|PENDING\n"), + stderr="", + ) + return CommandResult(exit_code=0, stdout="", stderr="") + + runner = StatusRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + run = _run( + run_id=run_id, + run_metadata={ + SLURM_JOB_IDS_METADATA_KEY: { + "load": "1000", + "train": "1001", + "evaluate": "1002", + } + }, + execution_mode=ExecutionMode.STOP_ON_FAILURE, + ) + + pipeline_status, _ = op.fetch_status(run) + + assert pipeline_status is ExecutionStatus.FAILED + assert any(command == "scancel 1001 1002" for command in runner.commands) + assert ( + runner.files[f"{op._run_dir(str(run_id), 'train')}/cancelled"] == "1\n" + ) + assert ( + runner.files[f"{op._run_dir(str(run_id), 'evaluate')}/cancelled"] + == "1\n" + ) + assert any( + command.startswith("rm -rf --") + and f"{op._run_dir(str(run_id), 'train')}/env" in command + for command in runner.commands + ) + + +def test_stop_run_cancels_submitted_jobs(monkeypatch): + """Stopping a detached Slurm run scancels all submitted step jobs.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + monkeypatch.setattr( + "zenml.orchestrators.base_orchestrator" + ".publish_pipeline_run_status_update", + lambda **kwargs: None, + ) + run_id = uuid4() + run = SimpleNamespace( + id=run_id, + orchestrator_run_id=str(run_id), + run_metadata={ + SLURM_JOB_IDS_METADATA_KEY: {"load": "1000", "train": "1001"} + }, + ) + + op.stop_run(run) + + assert any(command == "scancel 1000 1001" for command in runner.commands) + assert runner.files[f"{op._run_dir(str(run_id), 'load')}/cancelled"] == ( + "1\n" + ) + assert runner.files[f"{op._run_dir(str(run_id), 'train')}/cancelled"] == ( + "1\n" + ) + assert any( + command.startswith("rm -rf --") + and f"{op._run_dir(str(run_id), 'load')}/env" in command + for command in runner.commands + ) + + +def test_fetch_status_reconciles_dynamic_orchestration_job(monkeypatch): + """Dynamic runs reconcile through the orchestration job metadata.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + run_dir = op._orchestration_run_dir(str(run_id)) + runner.files[f"{run_dir}/exit_code"] = "1\n" + run = _run( + run_id=run_id, + run_metadata={SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000"}, + ) + + pipeline_status, step_statuses = op.fetch_status(run, include_steps=True) + + assert pipeline_status is ExecutionStatus.FAILED + assert step_statuses is None + + +def test_fetch_status_keeps_dynamic_run_unknown_without_sentinel(monkeypatch): + """A vanished job with no sentinel yet is unknown, never FAILED. + + Between a job leaving `squeue --states=all` and its exit-code + sentinel flushing to the shared filesystem, the run's outcome is + not knowable; declaring FAILED in that window would terminally + mislabel runs that actually completed. + """ + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + run = _run( + run_id=uuid4(), + run_metadata={SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000"}, + ) + + pipeline_status, step_statuses = op.fetch_status(run, include_steps=True) + + assert pipeline_status is None + assert step_statuses is None + + +def test_fetch_status_fails_dynamic_run_after_cleanup_without_sentinel( + monkeypatch, +): + """Once the cleanup job has run, a missing sentinel is conclusive.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + runner.files[f"/runs/{run_id}/cleanup/cleanup_complete"] = "1\n" + run = _run( + run_id=run_id, + run_metadata={SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000"}, + ) + + pipeline_status, _ = op.fetch_status(run, include_steps=True) + + assert pipeline_status is ExecutionStatus.FAILED + + +def test_does_not_support_isolated_steps(): + """Dynamic steps run inline in the orchestration job's allocation. + + Submitting per-step jobs from within an active Slurm job deadlocks + under per-user job limits, so the orchestrator must not advertise + isolated-step support. + """ + op = _build_orchestrator() + assert op.can_run_isolated_steps is False + assert op.can_stop_isolated_steps is False + + +def test_stop_run_cancels_dynamic_orchestration_job(monkeypatch): + """Stopping a dynamic run cancels the orchestration job.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + monkeypatch.setattr( + "zenml.orchestrators.base_orchestrator" + ".publish_pipeline_run_status_update", + lambda **kwargs: None, + ) + run_id = uuid4() + run = SimpleNamespace( + id=run_id, + orchestrator_run_id=str(run_id), + run_metadata={SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000"}, + ) + + op.stop_run(run) + + assert "scancel 1000" in runner.commands + assert ( + runner.files[f"{op._orchestration_run_dir(str(run_id))}/cancelled"] + == "1\n" + ) + assert any( + command.startswith("rm -rf --") + and f"{op._orchestration_run_dir(str(run_id))}/env" in command + for command in runner.commands + ) + + +# --- validator ---------------------------------------------------------------- + + +def test_validator_rejects_local_components(): + """The orchestrator needs a remote registry and artifact store.""" + op = _build_orchestrator() + validator = op.validator + assert validator is not None + assert validator._required_components == { + StackComponentType.CONTAINER_REGISTRY, + StackComponentType.IMAGE_BUILDER, + } + check = validator._custom_validation_function + assert check is not None + + def _stack(artifact_local: bool, registry_local: bool) -> SimpleNamespace: + return SimpleNamespace( + artifact_store=SimpleNamespace( + name="a", config=SimpleNamespace(is_local=artifact_local) + ), + container_registry=SimpleNamespace( + name="r", config=SimpleNamespace(is_local=registry_local) + ), + ) + + assert check(_stack(True, False))[0] is False + assert check(_stack(False, True))[0] is False + assert check(_stack(False, False)) == (True, "") diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py new file mode 100644 index 00000000000..60a540057d6 --- /dev/null +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -0,0 +1,619 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Unit tests for the Slurm integration. + +The Slurm CLI and SSH are never invoked: the client and step operator are +driven through an in-memory fake command runner, so the tests run without a +cluster. +""" + +from datetime import datetime +from types import SimpleNamespace +from typing import Dict, List, Optional +from uuid import UUID, uuid4 + +import pytest + +from zenml.config.resource_settings import ResourceSettings +from zenml.enums import ExecutionStatus, StackComponentType +from zenml.integrations.slurm.flavors import ( + SlurmStepOperatorConfig, + SlurmStepOperatorFlavor, + SlurmStepOperatorSettings, +) +from zenml.integrations.slurm.slurm_client import ( + CommandResult, + SlurmClient, + SlurmCommandRunner, +) +from zenml.integrations.slurm.slurm_job import ( + build_container_command, + build_registry_auth, + build_sbatch_script, + serialize_environment, +) +from zenml.integrations.slurm.step_operators import SlurmStepOperator +from zenml.integrations.slurm.step_operators.slurm_step_operator import ( + SLURM_JOB_ID_METADATA_KEY, + SLURM_STATE_METADATA_KEY, +) + +SECRET_TOKEN = "super-secret-zenml-token" +IMAGE = "registry.example.com/zenml:abc123" + + +class FakeRunner(SlurmCommandRunner): + """In-memory fake of the Slurm submission host.""" + + def __init__( + self, + queue_state: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + ) -> None: + """Initialize the fake. + + Args: + queue_state: State reported by the fake squeue, None = not queued. + files: Initial files present on the fake host. + """ + self.queue_state = queue_state + self.files = files or {} + self.modes: Dict[str, int] = {} + self.commands: List[str] = [] + + def run(self, command: str) -> CommandResult: + """Fake command execution. + + Args: + command: The command line. + + Returns: + A canned result per Slurm command. + """ + self.commands.append(command) + if command.startswith("sbatch"): + return CommandResult(exit_code=0, stdout="12345\n", stderr="") + if command.startswith("squeue"): + stdout = f"12345|{self.queue_state}\n" if self.queue_state else "" + return CommandResult(exit_code=0, stdout=stdout, stderr="") + return CommandResult(exit_code=0, stdout="", stderr="") + + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: + """Record a text upload with its mode. + + Args: + remote_path: Destination path. + content: The text content. + mode: File permissions. + """ + self.files[remote_path] = content + self.modes[remote_path] = mode + + def read_text(self, remote_path: str) -> str: + """Read a recorded file. + + Args: + remote_path: Path to read. + + Returns: + The recorded content. + + Raises: + FileNotFoundError: If nothing was recorded at that path. + """ + if remote_path not in self.files: + raise FileNotFoundError(remote_path) + return self.files[remote_path] + + +def _build_operator( + config: Optional[SlurmStepOperatorConfig] = None, +) -> SlurmStepOperator: + """Construct a Slurm step operator. + + Args: + config: The config for the operator. + + Returns: + A SlurmStepOperator instance. + """ + return SlurmStepOperator( + name="", + id=uuid4(), + config=config + or SlurmStepOperatorConfig( + transport="local", workdir="/shared/zenml-runs" + ), + flavor="slurm", + type=StackComponentType.STEP_OPERATOR, + user=uuid4(), + created=datetime.now(), + updated=datetime.now(), + ) + + +def _fake_info( + step_run_id: Optional[UUID] = None, gpu: int = 0 +) -> SimpleNamespace: + """Build a StepRunInfo stand-in. + + Args: + step_run_id: The step run id. + gpu: Number of GPUs the step requests. + + Returns: + A namespace usable where StepRunInfo is expected. + """ + return SimpleNamespace( + step_run_id=step_run_id or uuid4(), + pipeline_step_name="trainer", + config=SimpleNamespace( + resource_settings=ResourceSettings( + cpu_count=4, gpu_count=gpu or None, memory="16GB" + ) + ), + step_run=SimpleNamespace(run_metadata={}), + get_image=lambda key: IMAGE, + ) + + +# --- config validation -------------------------------------------------------- + + +def test_ssh_transport_requires_connection_details(): + """The ssh transport needs hostname and username; local does not.""" + with pytest.raises(ValueError): + SlurmStepOperatorConfig(transport="ssh", workdir="/s") + + SlurmStepOperatorConfig(transport="local", workdir="/s") + SlurmStepOperatorConfig( + transport="ssh", workdir="/s", hostname="h", username="u" + ) + + +def test_flavor_identity(): + """The flavor identifier and config class are Slurm's.""" + flavor = SlurmStepOperatorFlavor() + assert flavor.name == "slurm" + assert flavor.display_name == "Slurm" + assert flavor.config_class is SlurmStepOperatorConfig + assert flavor.implementation_class is SlurmStepOperator + + +def test_default_container_runtime_is_apptainer(): + """Apptainer is the default (rootless, HPC-safe).""" + config = SlurmStepOperatorConfig(transport="local", workdir="/s") + assert config.container_runtime.value == "apptainer" + + +def test_operator_config_is_remote(): + """Steps execute on the cluster, so the config is remote.""" + config = SlurmStepOperatorConfig(transport="local", workdir="/s") + assert config.is_remote is True + + +# --- validator ---------------------------------------------------------------- + + +def test_validator_requires_remote_registry_and_artifact_store(): + """The operator needs a remote registry, image builder, artifact store.""" + op = _build_operator() + validator = op.validator + assert validator is not None + assert validator._required_components == { + StackComponentType.CONTAINER_REGISTRY, + StackComponentType.IMAGE_BUILDER, + } + check = validator._custom_validation_function + assert check is not None + + def _stack(artifact_local: bool, registry_local: bool) -> SimpleNamespace: + return SimpleNamespace( + artifact_store=SimpleNamespace( + name="a", config=SimpleNamespace(is_local=artifact_local) + ), + container_registry=SimpleNamespace( + name="r", config=SimpleNamespace(is_local=registry_local) + ), + ) + + assert check(_stack(True, False))[0] is False + assert check(_stack(False, True))[0] is False + assert check(_stack(False, False)) == (True, "") + + +# --- slurm client ------------------------------------------------------------- + + +def test_client_submit_parses_job_id(): + """`sbatch --parsable` output is parsed into a job id.""" + runner = FakeRunner() + assert SlurmClient(runner).submit("/s/job.sh") == "12345" + assert runner.commands == ["sbatch --parsable /s/job.sh"] + + +def test_client_state_and_cancel_use_job_id(): + """Status lookup and cancellation address the exact job ID.""" + runner = FakeRunner(queue_state="RUNNING") + client = SlurmClient(runner) + assert client.get_job_state("12345") == "RUNNING" + client.cancel("12345") + assert runner.commands[-1] == "scancel 12345" + + +def test_client_batches_job_state_lookup(): + """Multiple job states are fetched with one squeue call.""" + + class BatchRunner(FakeRunner): + def run(self, command: str) -> CommandResult: + self.commands.append(command) + if ( + command.startswith("squeue --noheader") + and "--states=all" in command + ): + return CommandResult( + exit_code=0, + stdout="34567|COMPLETED\n45678|FAILED\n", + stderr="", + ) + if command.startswith("squeue"): + return CommandResult( + exit_code=0, + stdout="12345|RUNNING\n23456|PENDING\n", + stderr="", + ) + return CommandResult(exit_code=0, stdout="", stderr="") + + runner = BatchRunner() + states = SlurmClient(runner).get_job_states( + ["12345", "23456", "34567", "45678"] + ) + + assert states == { + "12345": "RUNNING", + "23456": "PENDING", + "34567": "COMPLETED", + "45678": "FAILED", + } + assert runner.commands == [ + "squeue --noheader --format='%i|%T' --jobs=12345,23456,34567,45678", + "squeue --noheader --format='%i|%T' --states=all --jobs=34567,45678", + ] + + +# --- container command rendering (per runtime) -------------------------------- + + +def _container_cmd(runtime: str, gpu: bool = False) -> str: + config = SlurmStepOperatorConfig( + transport="local", workdir="/s", container_runtime=runtime + ) + return build_container_command( + runtime=config.container_runtime, + image=IMAGE, + entrypoint_command=["python", "-m", "zenml.entrypoint"], + env_file="/s/zenml-x/env", + env_keys=["ZENML_STORE_API_KEY", "ZENML_STORE_URL"], + use_gpu=gpu, + settings=SlurmStepOperatorSettings( + container_mounts={"/scratch": "/scratch"} + ), + ) + + +def test_apptainer_command(): + """Apptainer exec pulls the docker image and passes the env file.""" + cmd = _container_cmd("apptainer", gpu=True) + assert cmd.startswith("apptainer exec --no-eval") + assert "--nv" in cmd + assert "--env-file" in cmd and "/s/zenml-x/env" in cmd + assert "docker://registry.example.com/zenml:abc123" in cmd + assert "--bind" in cmd and "/scratch:/scratch" in cmd + assert "python -m zenml.entrypoint" in cmd + + +def test_docker_command(): + """The docker runtime runs a foreground --rm container with the env file.""" + cmd = _container_cmd("docker", gpu=True) + assert cmd.startswith("docker run --rm") + assert "--gpus all" in cmd + assert "--env-file" in cmd and "/s/zenml-x/env" in cmd + assert "-v" in cmd and "/scratch:/scratch" in cmd + + +def test_pyxis_command_passes_env_names_not_values(): + """Pyxis sources the env file and passes only variable names to srun.""" + cmd = _container_cmd("pyxis") + assert "srun --container-image=" in cmd + assert "--container-env=ZENML_STORE_API_KEY,ZENML_STORE_URL" in cmd + assert "source /s/zenml-x/env" in cmd + assert "--container-mounts=/scratch:/scratch" in cmd + + +def test_pyxis_environment_is_shell_escaped(): + """Pyxis can source values without executing shell substitutions.""" + content = serialize_environment( + {"TOKEN": "value with spaces $(touch /tmp/unsafe)"}, + runtime=SlurmStepOperatorConfig( + transport="local", workdir="/s", container_runtime="pyxis" + ).container_runtime, + ) + assert content == "export TOKEN='value with spaces $(touch /tmp/unsafe)'\n" + + +@pytest.mark.parametrize( + "runtime", ["apptainer", "singularity", "pyxis", "docker"] +) +def test_registry_auth_keeps_password_out_of_shell_command(runtime): + """Runtime-specific registry auth is staged instead of passed in argv.""" + config = SlurmStepOperatorConfig( + transport="local", workdir="/s", container_runtime=runtime + ) + auth = build_registry_auth( + runtime=config.container_runtime, + run_dir="/s/run", + registry_uri="registry.example.com", + credentials=("user", SECRET_TOKEN), + ) + assert auth.files + assert SECRET_TOKEN not in auth.shell_setup + if runtime == "docker": + assert all( + SECRET_TOKEN not in content for content in auth.files.values() + ) + else: + assert any(SECRET_TOKEN in content for content in auth.files.values()) + + +# --- sbatch script ------------------------------------------------------------ + + +def test_sbatch_script_directives_and_scrub(): + """The script carries directives, fails fast, and scrubs the env file.""" + sid = uuid4() + script = build_sbatch_script( + job_name=f"zenml-{sid}", + run_dir="/s/zenml-x", + container_command="RUN_THE_CONTAINER", + resources=ResourceSettings(cpu_count=4, gpu_count=2, memory="16GB"), + settings=SlurmStepOperatorSettings( + partition="gpu", time_limit="2:00:00" + ), + ) + assert f"#SBATCH --job-name=zenml-{sid}" in script + assert "#SBATCH --partition=gpu" in script + assert "#SBATCH --gres=gpu:2" in script + assert "#SBATCH --mem=16000M" in script + assert "set -eo pipefail" in script + assert "RUN_THE_CONTAINER" in script + # the EXIT trap records the outcome AND scrubs the credential env file + assert 'echo "$ec" > /s/zenml-x/exit_code' in script + assert "rm -rf -- /s/zenml-x/env" in script + + +def test_sbatch_resources_round_up(): + """Fractional CPUs and sub-GB memory never become zero requests.""" + script = build_sbatch_script( + job_name="zenml-small", + run_dir="/s/small", + container_command="true", + resources=ResourceSettings(cpu_count=0.5, memory="500MB"), + settings=SlurmStepOperatorSettings(), + ) + assert "#SBATCH --cpus-per-task=1" in script + assert "#SBATCH --mem=500M" in script + + +# --- submit: security-sensitive handling -------------------------------------- + + +def test_submit_writes_secret_env_file_owner_only(monkeypatch): + """The env file is written 0600, the script 0700, and no secret leaks.""" + op = _build_operator() + op.get_settings = lambda _i: SlurmStepOperatorSettings() + runner = FakeRunner() + monkeypatch.setattr( + "zenml.integrations.slurm.step_operators.slurm_step_operator" + ".build_slurm_client", + lambda config: SlurmClient(runner), + ) + monkeypatch.setattr( + "zenml.integrations.slurm.step_operators.slurm_step_operator" + ".publish_step_run_metadata", + lambda **kwargs: None, + ) + fake_registry = SimpleNamespace( + config=SimpleNamespace(uri="registry.example.com"), + credentials=("registry-user", "registry-password"), + ) + monkeypatch.setattr( + "zenml.client.Client", + lambda: SimpleNamespace( + active_stack=SimpleNamespace(container_registry=fake_registry) + ), + ) + + sid = uuid4() + info = _fake_info(sid) + op.submit( + info=info, + entrypoint_command=["python", "-m", "zenml.entrypoint"], + environment={ + "ZENML_STORE_API_KEY": SECRET_TOKEN, + "ZENML_STORE_URL": "https://z.example.com", + }, + ) + run_dir = f"/shared/zenml-runs/zenml-{sid}" + + # env file present, contains the secret, and is owner-only + assert runner.modes[f"{run_dir}/env"] == 0o600 + assert SECRET_TOKEN in runner.files[f"{run_dir}/env"] + auth_path = f"{run_dir}/registry_auth.sh" + assert runner.modes[auth_path] == 0o600 + assert "registry-password" in runner.files[auth_path] + + # the run directory is created owner-only + assert any("chmod 700" in c for c in runner.commands) + # the job script is written owner-only + assert runner.modes[f"{run_dir}/job.sh"] == 0o700 + + # the secret value is NOT baked into the job script (only the env-file ref) + assert SECRET_TOKEN not in runner.files[f"{run_dir}/job.sh"] + assert "registry-password" not in runner.files[f"{run_dir}/job.sh"] + # the secret value is NOT on any command line issued to the host + assert all(SECRET_TOKEN not in c for c in runner.commands) + assert all("registry-password" not in c for c in runner.commands) + assert info.step_run.run_metadata[SLURM_JOB_ID_METADATA_KEY] == "12345" + + +# --- status mapping ----------------------------------------------------------- + + +@pytest.fixture +def op_and_runner(monkeypatch): + """A local operator whose runner is the in-memory fake.""" + op = _build_operator() + runner = FakeRunner() + monkeypatch.setattr( + "zenml.integrations.slurm.step_operators.slurm_step_operator" + ".build_slurm_client", + lambda config: SlurmClient(runner), + ) + monkeypatch.setattr( + "zenml.integrations.slurm.step_operators.slurm_step_operator" + ".publish_step_run_metadata", + lambda **kwargs: None, + ) + return op, runner + + +@pytest.mark.parametrize( + "queue_state,expected", + [ + ("PENDING", ExecutionStatus.QUEUED), + ("RUNNING", ExecutionStatus.RUNNING), + ("CANCELLED", ExecutionStatus.CANCELLED), + ("FAILED", ExecutionStatus.FAILED), + ], +) +def test_get_status_maps_queue_states(op_and_runner, queue_state, expected): + """Queue states map onto ZenML execution statuses.""" + op, runner = op_and_runner + runner.queue_state = queue_state + assert ( + op.get_status( + SimpleNamespace( + id=uuid4(), + run_metadata={SLURM_JOB_ID_METADATA_KEY: "12345"}, + ) + ) + is expected + ) + + +def test_get_status_reads_sentinel_after_queue(op_and_runner): + """After the job leaves the queue, the sentinel file decides the status.""" + op, runner = op_and_runner + step_run = SimpleNamespace( + id=uuid4(), run_metadata={SLURM_JOB_ID_METADATA_KEY: "12345"} + ) + run_dir = op._run_dir(step_run.id) + runner.files[f"{run_dir}/exit_code"] = "0\n" + assert op.get_status(step_run) is ExecutionStatus.COMPLETED + runner.files[f"{run_dir}/exit_code"] = "1\n" + assert op.get_status(step_run) is ExecutionStatus.FAILED + + +def test_get_status_without_queue_or_sentinel_fails(op_and_runner): + """A vanished job without a sentinel is treated as failed.""" + op, runner = op_and_runner + assert ( + op.get_status( + SimpleNamespace( + id=uuid4(), + run_metadata={SLURM_JOB_ID_METADATA_KEY: "12345"}, + ) + ) + is ExecutionStatus.FAILED + ) + + +def test_get_status_publishes_raw_terminal_state(op_and_runner, monkeypatch): + """Scheduler-observed terminal states land in step metadata exactly once. + + A scheduler kill (TIMEOUT, NODE_FAIL, OOM, PREEMPTED) leaves no exit-code + sentinel; the raw state is the only record of why the step failed. + """ + op, runner = op_and_runner + published = [] + monkeypatch.setattr( + "zenml.integrations.slurm.step_operators.slurm_step_operator" + ".publish_step_run_metadata", + lambda **kwargs: published.append(kwargs), + ) + step_run = SimpleNamespace( + id=uuid4(), run_metadata={SLURM_JOB_ID_METADATA_KEY: "12345"} + ) + runner.queue_state = "TIMEOUT" + + assert op.get_status(step_run) is ExecutionStatus.FAILED + assert published == [ + { + "step_run_id": step_run.id, + "step_run_metadata": { + op.id: {SLURM_STATE_METADATA_KEY: "TIMEOUT"} + }, + } + ] + + # A state already recorded on the step run is not re-published. + step_run.run_metadata[SLURM_STATE_METADATA_KEY] = "TIMEOUT" + assert op.get_status(step_run) is ExecutionStatus.FAILED + assert len(published) == 1 + + +def test_get_status_reads_cancelled_marker(op_and_runner): + """A cancelled-before-start marker produces a terminal status.""" + op, runner = op_and_runner + step_run = SimpleNamespace( + id=uuid4(), run_metadata={SLURM_JOB_ID_METADATA_KEY: "12345"} + ) + run_dir = op._run_dir(step_run.id) + runner.files[f"{run_dir}/cancelled"] = "1\n" + + assert op.get_status(step_run) is ExecutionStatus.CANCELLED + + +def test_cleanup_removes_run_dir(op_and_runner): + """Cleanup removes the whole per-run directory.""" + op, runner = op_and_runner + step_run = SimpleNamespace(id=uuid4()) + op.cleanup_step_submission(step_run) + run_dir = op._run_dir(step_run.id) + assert any(f"rm -rf -- {run_dir}" in c for c in runner.commands) + + +# --- integration registration ------------------------------------------------- + + +def test_integration_registered(): + """The Slurm integration is discoverable in the registry.""" + from zenml.integrations.registry import integration_registry + from zenml.integrations.slurm import SlurmIntegration + + integration_registry._initialize() + assert integration_registry.integrations["slurm"] is SlurmIntegration + assert {f().name for f in SlurmIntegration.flavors()} == {"slurm"}