From 84b3e39b61fb5d583fc844e6636bec2016b18010 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 11:40:08 -0700 Subject: [PATCH 01/22] Add Slurm integration scaffold (step operator over SSH/local sbatch) Adds a `slurm` integration with a step operator that runs individual pipeline steps as Slurm batch jobs: - SlurmClient wraps the Slurm CLI (sbatch --parsable, squeue, scancel) behind a transport abstraction: LocalSlurmCommandRunner (client already on a login node) and SSHSlurmCommandRunner (reusing the ssh integration's SSHClient; the config mirrors its connection fields). - The operator is container-free by design: Slurm compute nodes cannot run Docker, so jobs activate a user-provided Python environment (env_setup_command) and unpack a code archive shipped at submission time into a per-run directory on the cluster. - Submission is stateless: jobs are named after the step run id, so get_status/cancel work by job name. Queue states come from squeue; after a job leaves the queue, a sentinel exit-code file written by the job script's EXIT trap decides success vs failure, so Slurm accounting (sacct) is never required. - ResourceSettings map onto sbatch directives (--cpus-per-task, --mem, --gres=gpu:N) plus partition/time/account/qos settings and a raw directive escape hatch. All unit tests drive the client and operator through an in-memory fake runner, so they run without a cluster. Live validation on a real Slurm cluster is the next step and will be reported on the issue. Refs #5063 Co-Authored-By: Claude Fable 5 --- src/zenml/integrations/constants.py | 1 + src/zenml/integrations/slurm/__init__.py | 53 +++ .../integrations/slurm/flavors/__init__.py | 26 ++ .../flavors/slurm_step_operator_flavor.py | 245 +++++++++++++ src/zenml/integrations/slurm/slurm_client.py | 320 +++++++++++++++++ .../slurm/step_operators/__init__.py | 20 ++ .../step_operators/slurm_step_operator.py | 324 +++++++++++++++++ .../slurm/test_slurm_step_operator.py | 329 ++++++++++++++++++ 8 files changed, 1318 insertions(+) create mode 100644 src/zenml/integrations/slurm/__init__.py create mode 100644 src/zenml/integrations/slurm/flavors/__init__.py create mode 100644 src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py create mode 100644 src/zenml/integrations/slurm/slurm_client.py create mode 100644 src/zenml/integrations/slurm/step_operators/__init__.py create mode 100644 src/zenml/integrations/slurm/step_operators/slurm_step_operator.py create mode 100644 tests/unit/integrations/slurm/test_slurm_step_operator.py diff --git a/src/zenml/integrations/constants.py b/src/zenml/integrations/constants.py index 26ce5393ca9..b1a49d48fac 100644 --- a/src/zenml/integrations/constants.py +++ b/src/zenml/integrations/constants.py @@ -73,6 +73,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..8d73c15be1a --- /dev/null +++ b/src/zenml/integrations/slurm/__init__.py @@ -0,0 +1,53 @@ +# 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 provides a step operator that runs individual pipeline +steps as Slurm jobs on an HPC cluster, either over SSH to a login/controller +node or via local ``sbatch`` when the client already runs on the cluster. + +Unlike container-based step operators, the Slurm step operator is +container-free: Slurm compute nodes generally cannot run Docker, so steps +execute inside a user-provided Python environment on the cluster (typically on +a shared filesystem), and the pipeline code is shipped to the cluster as an +archive at submission time. +""" + +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" + + +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 SlurmStepOperatorFlavor + + return [SlurmStepOperatorFlavor] diff --git a/src/zenml/integrations/slurm/flavors/__init__.py b/src/zenml/integrations/slurm/flavors/__init__.py new file mode 100644 index 00000000000..ee1e7d6ccfa --- /dev/null +++ b/src/zenml/integrations/slurm/flavors/__init__.py @@ -0,0 +1,26 @@ +# 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_step_operator_flavor import ( # noqa + SlurmStepOperatorConfig, + SlurmStepOperatorFlavor, + SlurmStepOperatorSettings, +) + +__all__ = [ + "SlurmStepOperatorConfig", + "SlurmStepOperatorFlavor", + "SlurmStepOperatorSettings", +] 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..464b2defd6f --- /dev/null +++ b/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py @@ -0,0 +1,245 @@ +# 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 enum import Enum +from typing import TYPE_CHECKING, List, Optional, Type + +from pydantic import Field, model_validator + +from zenml.config.base_settings import BaseSettings +from zenml.integrations.slurm import SLURM_STEP_OPERATOR_FLAVOR +from zenml.step_operators.base_step_operator import ( + BaseStepOperatorConfig, + BaseStepOperatorFlavor, +) +from zenml.utils.secret_utils import SecretField + +if TYPE_CHECKING: + from zenml.integrations.slurm.step_operators import SlurmStepOperator + + +class SlurmTransport(str, Enum): + """How the client reaches the Slurm submission host.""" + + SSH = "ssh" + LOCAL = "local" + + +class SlurmStepOperatorSettings(BaseSettings): + """Settings for the Slurm step operator. + + These map onto ``sbatch`` directives 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']", + ) + + +class SlurmStepOperatorConfig( + BaseStepOperatorConfig, SlurmStepOperatorSettings +): + """Configuration for the Slurm step operator.""" + + 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'", + ) + workdir: str = Field( + description="Directory on the cluster used to stage job scripts, " + "code archives and job output, ideally on a filesystem shared " + "between the submission host and the compute nodes. Example: " + "'/shared/zenml-runs'", + ) + env_setup_command: str = Field( + description="Shell command executed at the start of every job to " + "activate a Python environment that has zenml and the pipeline " + "requirements installed. Examples: 'source /shared/venvs/zenml/bin/" + "activate', 'module load python && conda activate zenml'", + ) + # SSH connection settings, required for the `ssh` transport. Field names + # deliberately match the ssh integration's connection mixin so its + # SSHClient can consume this config directly. + hostname: Optional[str] = Field( + 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( + default=None, + description="SSH username on the login node, required for the 'ssh' " + "transport. Example: 'mlops'", + ) + port: int = Field( + default=22, + description="SSH port on the login node. Standard SSH port is 22", + ) + ssh_key_path: Optional[str] = Field( + default=None, + description="Path to the SSH private key file on the submitting " + "machine. Supports RSA, Ed25519, and ECDSA keys", + ) + ssh_private_key: Optional[str] = SecretField( + default=None, + description="SSH private key content, used instead of ssh_key_path " + "when the key is stored in a ZenML secret. Supports {{secret.key}} " + "references", + ) + ssh_key_passphrase: Optional[str] = SecretField( + default=None, + description="Passphrase for an encrypted SSH private key. Leave " + "unset if the key is not encrypted", + ) + verify_host_key: bool = Field( + default=True, + description="Require the remote host key to be present in " + "known_hosts. Set to False to auto-accept unknown host keys (less " + "secure, convenient for ephemeral clusters)", + ) + known_hosts_path: Optional[str] = Field( + default=None, + description="Path to a custom known_hosts file used for host key " + "verification. Uses ~/.ssh/known_hosts if unset", + ) + + @model_validator(mode="after") + def _validate_transport_requirements(self) -> "SlurmStepOperatorConfig": + """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 step operator 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 + + @property + def is_remote(self) -> bool: + """Whether this component runs remotely. + + Returns: + True, since jobs 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/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py new file mode 100644 index 00000000000..07caefc031e --- /dev/null +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -0,0 +1,320 @@ +# 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 shlex +import subprocess +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from zenml.integrations.slurm.flavors.slurm_step_operator_flavor import ( + SlurmStepOperatorConfig, + ) + +# 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_file(self, local_path: str, remote_path: str) -> None: + """Copy a local file to the submission host. + + Args: + local_path: Path of the file on the local machine. + remote_path: Destination path on the submission host. + """ + + @abstractmethod + def put_text(self, content: str, remote_path: str) -> None: + """Write text content to a file on the submission host. + + Args: + content: The text content to write. + remote_path: Destination path on the submission host. + """ + + @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_file(self, local_path: str, remote_path: str) -> None: + """Copy a local file to the destination path. + + Args: + local_path: Path of the source file. + remote_path: Destination path. + """ + import shutil + + shutil.copyfile(local_path, remote_path) + + def put_text(self, content: str, remote_path: str) -> None: + """Write text content to the destination path. + + Args: + content: The text content to write. + remote_path: Destination path. + """ + with open(remote_path, "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: "SlurmStepOperatorConfig") -> None: + """Initialize the runner with an open SSH connection. + + Args: + config: The step operator config holding SSH connection settings. + """ + from typing import cast + + from zenml.integrations.ssh.flavors.base import ( + SSHConnectionConfigMixin, + ) + from zenml.integrations.ssh.ssh_client import SSHClient + + # The Slurm config deliberately mirrors the SSH connection mixin's + # field names, so the SSH client can consume it directly. + self._ssh = SSHClient(cast(SSHConnectionConfigMixin, 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_file(self, local_path: str, remote_path: str) -> None: + """Copy a local file to the remote host via SFTP. + + Args: + local_path: Path of the file on the local machine. + remote_path: Destination path on the remote host. + """ + with self._ssh.sftp() as sftp: + sftp.put(local_path, remote_path) + + def put_text(self, content: str, remote_path: str) -> None: + """Write text content to a file on the remote host. + + Args: + content: The text content to write. + remote_path: Destination path on the remote host. + """ + self._ssh.put_text(content, remote_path) + + 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) -> str: + """Submit a batch script and return the Slurm job id. + + Args: + script_path: Path of the sbatch script on the submission host. + + Returns: + The Slurm job id. + + Raises: + RuntimeError: If the submission fails. + """ + result = self.runner.run( + f"sbatch --parsable {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]` + return result.stdout.strip().split(";")[0] + + def get_job_state(self, job_name: str) -> Optional[str]: + """Get the queue state of a job by its name. + + Args: + job_name: The Slurm job name to look up. + + Returns: + The Slurm job state (e.g. ``PENDING``, ``RUNNING``) or ``None`` + if the job is no longer in the queue. 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: + RuntimeError: If ``squeue`` fails. + """ + result = self.runner.run( + f"squeue --noheader --format=%T --name {shlex.quote(job_name)}" + ) + if result.exit_code != 0: + raise RuntimeError( + f"`squeue` failed with exit code {result.exit_code}: " + f"{result.stderr.strip()}" + ) + state = result.stdout.strip().splitlines() + return state[0].strip() if state else None + + def cancel(self, job_name: str) -> None: + """Cancel a job by its name. + + Args: + job_name: The Slurm job name to cancel. + + Raises: + RuntimeError: If ``scancel`` fails. + """ + result = self.runner.run(f"scancel --name {shlex.quote(job_name)}") + if result.exit_code != 0: + raise RuntimeError( + f"`scancel` failed with exit code {result.exit_code}: " + f"{result.stderr.strip()}" + ) 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..b25ea011930 --- /dev/null +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -0,0 +1,324 @@ +# 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 operator is container-free: Slurm compute nodes generally cannot run +Docker, so the step executes inside a user-provided Python environment on the +cluster (``env_setup_command``). Pipeline code is shipped at submission time +as an archive that the job script unpacks into a per-run directory, so no +image build and no artifact-store code upload are required. + +Submission is stateless: the Slurm job is named after the step run id, so +status lookups and cancellation work by job name without persisting job ids. +While a job is in the queue, its state comes from ``squeue``; once it leaves +the queue, a sentinel exit-code file written by the job script disambiguates +success from failure, so Slurm accounting (``sacct``) is never required. +""" + +import shlex +import tempfile +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast +from uuid import UUID + +from zenml.config.base_settings import BaseSettings +from zenml.enums import ExecutionStatus +from zenml.integrations.slurm.flavors.slurm_step_operator_flavor import ( + SlurmStepOperatorConfig, + SlurmStepOperatorSettings, + SlurmTransport, +) +from zenml.integrations.slurm.slurm_client import ( + PENDING_STATES, + RUNNING_STATES, + LocalSlurmCommandRunner, + SlurmClient, + SlurmCommandRunner, + SSHSlurmCommandRunner, +) +from zenml.logger import get_logger +from zenml.step_operators import BaseStepOperator +from zenml.utils import code_utils, source_utils + +if TYPE_CHECKING: + from zenml.config.step_run_info import StepRunInfo + from zenml.models import StepRunResponse + +logger = get_logger(__name__) + +_JOB_NAME_PREFIX = "zenml" +_EXIT_CODE_FILE = "exit_code" +_OUTPUT_FILE = "output.log" + + +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 + + def _get_client(self) -> Tuple[SlurmClient, SlurmCommandRunner]: + """Build a Slurm client for the configured transport. + + Returns: + The Slurm client and its underlying runner (which the caller is + responsible for closing). + """ + runner: SlurmCommandRunner + if self.config.transport == SlurmTransport.LOCAL: + runner = LocalSlurmCommandRunner() + else: + runner = SSHSlurmCommandRunner(self.config) + return SlurmClient(runner), runner + + @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 _build_sbatch_script( + self, + info: "StepRunInfo", + entrypoint_command: List[str], + environment: Dict[str, str], + run_dir: str, + ) -> str: + """Render the sbatch job script for a step. + + Args: + info: Information about the step run. + entrypoint_command: Command that executes the step. + environment: Environment variables to set for the step. + run_dir: The per-run staging directory on the cluster. + + Returns: + The job script content. + """ + settings = cast( + SlurmStepOperatorSettings, self.get_settings(info.step_run) + ) + resources = info.config.resource_settings + + directives = [ + f"#SBATCH --job-name={self._job_name(info.step_run_id)}", + 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={int(resources.cpu_count)}" + ) + if memory_gb := resources.get_memory(unit="GB"): + directives.append(f"#SBATCH --mem={int(memory_gb)}G") + 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}") + + env_exports = "\n".join( + f"export {key}={shlex.quote(value)}" + for key, value in sorted(environment.items()) + ) + command = " ".join(shlex.quote(part) for part in entrypoint_command) + + # The EXIT trap records the job outcome in a sentinel file that + # `get_status` reads after the job leaves the squeue output, so the + # operator never depends on Slurm accounting being configured. + return f"""#!/bin/bash +{chr(10).join(directives)} + +trap 'echo $? > {run_dir}/{_EXIT_CODE_FILE}' EXIT + +{self.config.env_setup_command} + +{env_exports} + +mkdir -p {run_dir}/code +tar -xzf {run_dir}/code.tar.gz -C {run_dir}/code +cd {run_dir}/code +export PYTHONPATH="{run_dir}/code:${{PYTHONPATH:-}}" + +{command} +""" + + 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. + """ + run_dir = self._run_dir(info.step_run_id) + client, runner = self._get_client() + try: + # The run directory holds the code archive, the job script (which + # contains credentials as env exports) and the job output, so it + # is created private to the submitting user. + result = runner.run( + f"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 " + f"cluster: {result.stderr.strip()}" + ) + + with tempfile.NamedTemporaryFile(suffix=".tar.gz") as archive: + code_archive = code_utils.CodeArchive( + root=source_utils.get_source_root() + ) + code_archive.write_archive(archive) + archive.flush() + runner.put_file(archive.name, f"{run_dir}/code.tar.gz") + + script = self._build_sbatch_script( + info=info, + entrypoint_command=entrypoint_command, + environment=environment, + run_dir=run_dir, + ) + script_path = f"{run_dir}/job.sh" + runner.put_text(script, script_path) + + job_id = client.submit(script_path) + 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), + ) + finally: + 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_name = self._job_name(step_run.id) + run_dir = self._run_dir(step_run.id) + client, runner = self._get_client() + try: + state = client.get_job_state(job_name) + if state in PENDING_STATES: + return ExecutionStatus.QUEUED + if state in RUNNING_STATES: + return ExecutionStatus.RUNNING + if state is not None: + # Cancellation and failure states can linger in the queue + # output briefly before the job is purged from it. + 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: + exit_code = runner.read_text( + f"{run_dir}/{_EXIT_CODE_FILE}" + ).strip() + except Exception: + # No queue entry and no sentinel: the job was cancelled + # before its trap could run, or never started. + return ExecutionStatus.CANCELLED + + if exit_code == "0": + return ExecutionStatus.COMPLETED + + self._log_failure_output(runner, run_dir) + return ExecutionStatus.FAILED + finally: + runner.close() + + def cancel(self, step_run: "StepRunResponse") -> None: + """Cancel the Slurm job of a step run. + + Args: + step_run: The step run to cancel. + """ + client, runner = self._get_client() + try: + client.cancel(self._job_name(step_run.id)) + finally: + 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 + ) 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..03f487ece89 --- /dev/null +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -0,0 +1,329 @@ +# 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 is never invoked: all tests drive the client and step operator +through a fake in-memory command runner, so they run without a cluster. +""" + +from datetime import datetime +from typing import Dict, Optional +from uuid import uuid4 + +import pytest + +from zenml.enums import ExecutionStatus, StackComponentType +from zenml.integrations.slurm.flavors import ( + SlurmStepOperatorConfig, + SlurmStepOperatorFlavor, +) +from zenml.integrations.slurm.slurm_client import ( + CommandResult, + SlurmClient, + SlurmCommandRunner, +) +from zenml.integrations.slurm.step_operators import SlurmStepOperator + + +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.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"{self.queue_state}\n" if self.queue_state else "" + return CommandResult(exit_code=0, stdout=stdout, stderr="") + if command.startswith("scancel"): + return CommandResult(exit_code=0, stdout="", stderr="") + return CommandResult(exit_code=0, stdout="", stderr="") + + def put_file(self, local_path: str, remote_path: str) -> None: + """Record a file upload. + + Args: + local_path: Source path. + remote_path: Destination path. + """ + self.files[remote_path] = f"" + + def put_text(self, content: str, remote_path: str) -> None: + """Record a text upload. + + Args: + content: The text content. + remote_path: Destination path. + """ + self.files[remote_path] = content + + 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: + return SlurmStepOperator( + name="", + id=uuid4(), + config=config + or SlurmStepOperatorConfig( + transport="local", + workdir="/shared/zenml-runs", + env_setup_command="source /shared/venv/bin/activate", + ), + flavor="slurm", + type=StackComponentType.STEP_OPERATOR, + user=uuid4(), + created=datetime.now(), + updated=datetime.now(), + ) + + +# --- 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="/shared/zenml-runs", + env_setup_command="true", + ) + + # local transport without connection details is fine + SlurmStepOperatorConfig( + transport="local", + workdir="/shared/zenml-runs", + env_setup_command="true", + ) + + # ssh transport with connection details is fine + SlurmStepOperatorConfig( + transport="ssh", + workdir="/shared/zenml-runs", + env_setup_command="true", + hostname="login.example.com", + username="mlops", + ) + + +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 + + +# --- slurm client ------------------------------------------------------------- + + +def test_client_submit_parses_job_id(): + """`sbatch --parsable` output is parsed into a job id.""" + runner = FakeRunner() + client = SlurmClient(runner) + assert client.submit("/shared/job.sh") == "12345" + assert runner.commands == ["sbatch --parsable /shared/job.sh"] + + +def test_client_job_state_and_cancel_use_job_name(): + """Status lookup and cancellation address the job by name.""" + runner = FakeRunner(queue_state="RUNNING") + client = SlurmClient(runner) + assert client.get_job_state("zenml-abc") == "RUNNING" + client.cancel("zenml-abc") + assert runner.commands[-1] == "scancel --name zenml-abc" + + +def test_client_job_state_returns_none_when_not_queued(): + """A job absent from squeue returns None (terminal, see sentinel).""" + runner = FakeRunner(queue_state=None) + client = SlurmClient(runner) + assert client.get_job_state("zenml-abc") is None + + +# --- step operator status mapping --------------------------------------------- + + +class _FakeStepRun: + """Minimal stand-in for a StepRunResponse.""" + + def __init__(self) -> None: + """Initialize with a random id.""" + self.id = uuid4() + + +@pytest.fixture +def operator_with_fake_runner(monkeypatch): + """A local-transport operator whose runner is the in-memory fake.""" + operator = _build_operator() + runner = FakeRunner() + monkeypatch.setattr( + operator, + "_get_client", + lambda: (SlurmClient(runner), runner), + ) + return operator, runner + + +@pytest.mark.parametrize( + "queue_state,expected", + [ + ("PENDING", ExecutionStatus.QUEUED), + ("CONFIGURING", ExecutionStatus.QUEUED), + ("RUNNING", ExecutionStatus.RUNNING), + ("COMPLETING", ExecutionStatus.RUNNING), + ("CANCELLED", ExecutionStatus.CANCELLED), + ("FAILED", ExecutionStatus.FAILED), + ], +) +def test_get_status_maps_queue_states( + operator_with_fake_runner, queue_state, expected +): + """Queue states map onto ZenML execution statuses.""" + operator, runner = operator_with_fake_runner + runner.queue_state = queue_state + assert operator.get_status(_FakeStepRun()) is expected + + +def test_get_status_reads_sentinel_after_queue(operator_with_fake_runner): + """After the job leaves the queue, the sentinel file decides the status.""" + operator, runner = operator_with_fake_runner + step_run = _FakeStepRun() + run_dir = operator._run_dir(step_run.id) + + runner.files[f"{run_dir}/exit_code"] = "0\n" + assert operator.get_status(step_run) is ExecutionStatus.COMPLETED + + runner.files[f"{run_dir}/exit_code"] = "1\n" + assert operator.get_status(step_run) is ExecutionStatus.FAILED + + +def test_get_status_without_queue_entry_or_sentinel_is_cancelled( + operator_with_fake_runner, +): + """No queue entry and no sentinel means the job was torn down early.""" + operator, runner = operator_with_fake_runner + assert operator.get_status(_FakeStepRun()) is ExecutionStatus.CANCELLED + + +# --- sbatch script rendering ---------------------------------------------------- + + +def test_sbatch_script_rendering(monkeypatch): + """The generated script carries directives, env setup, and the command.""" + from types import SimpleNamespace + + from zenml.config.resource_settings import ResourceSettings + + operator = _build_operator() + step_run_id = uuid4() + + fake_info = SimpleNamespace( + step_run_id=step_run_id, + pipeline_step_name="trainer", + config=SimpleNamespace( + resource_settings=ResourceSettings( + cpu_count=4, gpu_count=2, memory="16GB" + ) + ), + step_run=None, + ) + + from zenml.integrations.slurm.flavors import SlurmStepOperatorSettings + + monkeypatch.setattr( + operator, + "get_settings", + lambda _: SlurmStepOperatorSettings( + partition="gpu", + time_limit="2:00:00", + extra_sbatch_directives=["--constraint=a100"], + ), + ) + + script = operator._build_sbatch_script( + info=fake_info, + entrypoint_command=["python", "-m", "zenml.entrypoint"], + environment={"ZENML_STORE_URL": "https://z.example.com"}, + run_dir="/shared/zenml-runs/zenml-x", + ) + + assert f"#SBATCH --job-name=zenml-{step_run_id}" in script + assert "#SBATCH --partition=gpu" in script + assert "#SBATCH --time=2:00:00" in script + assert "#SBATCH --cpus-per-task=4" in script + assert "#SBATCH --mem=16G" in script + assert "#SBATCH --gres=gpu:2" in script + assert "#SBATCH --constraint=a100" in script + assert "source /shared/venv/bin/activate" in script + assert "export ZENML_STORE_URL=https://z.example.com" in script + assert "python -m zenml.entrypoint" in script + assert ( + "trap 'echo $? > /shared/zenml-runs/zenml-x/exit_code' EXIT" in script + ) + + +# --- 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 + flavors = {f().name for f in SlurmIntegration.flavors()} + assert flavors == {"slurm"} From b1bfc757da8e40d3fba0ed3ce1e10cbbe3a41cce Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 11:51:58 -0700 Subject: [PATCH 02/22] Fail fast in the Slurm job script on setup/extraction errors Add `set -eo pipefail` after the EXIT trap so a failing env_setup_command or a failed code-archive extraction aborts the job (with the failure code captured by the sentinel) instead of falling through to run the step in a broken environment. Verified by executing the generated script directly in bash (the #SBATCH directives are comments, so it is runnable without Slurm). Refs #5063 Co-Authored-By: Claude Fable 5 --- .../integrations/slurm/step_operators/slurm_step_operator.py | 4 ++++ tests/unit/integrations/slurm/test_slurm_step_operator.py | 1 + 2 files changed, 5 insertions(+) diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index b25ea011930..dad05d7fd4e 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -176,10 +176,14 @@ def _build_sbatch_script( # The EXIT trap records the job outcome in a sentinel file that # `get_status` reads after the job leaves the squeue output, so the # operator never depends on Slurm accounting being configured. + # `set -e` makes environment setup or code extraction failures abort + # the job (with the failing code captured by the trap) rather than + # falling through to run the step in a broken environment. return f"""#!/bin/bash {chr(10).join(directives)} trap 'echo $? > {run_dir}/{_EXIT_CODE_FILE}' EXIT +set -eo pipefail {self.config.env_setup_command} diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index 03f487ece89..37b4519bddc 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -307,6 +307,7 @@ def test_sbatch_script_rendering(monkeypatch): assert "#SBATCH --mem=16G" in script assert "#SBATCH --gres=gpu:2" in script assert "#SBATCH --constraint=a100" in script + assert "set -eo pipefail" in script assert "source /shared/venv/bin/activate" in script assert "export ZENML_STORE_URL=https://z.example.com" in script assert "python -m zenml.entrypoint" in script From 89a1ed1cf7cfebc579f8ba8d1c2a8b50dacefb9b Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 12:09:59 -0700 Subject: [PATCH 03/22] Reuse the ssh integration's connection mixin instead of copying its fields Follow the ssh step operator's pattern: SlurmStepOperatorConfig now inherits SSHConnectionConfigMixin rather than re-declaring the SSH connection fields. This fixes a latent bug: SSHClient reads config.connection_timeout and config.keepalive_interval, which the hand-copied config never declared, so the ssh transport would have crashed with AttributeError on connect. The mixin provides them (and keeps every SSH field in sync with the client we already reuse). hostname/username are relaxed to optional (the mixin makes them required) because the `local` transport uses no SSH; the ssh transport still enforces them via the existing validator. The now-unnecessary cast in the SSH runner is removed. Decision recorded: no Slurm service connector. SSH-key auth is served by SecretFields + ZenML secret references exactly as the ssh operators do; a connector (HyperAI-style) would not fit the SSH-key model well. Refs #5063 Co-Authored-By: Claude Fable 5 --- .../flavors/slurm_step_operator_flavor.py | 55 ++++++------------- src/zenml/integrations/slurm/slurm_client.py | 11 +--- 2 files changed, 19 insertions(+), 47 deletions(-) diff --git a/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py b/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py index 464b2defd6f..4f6daefeb72 100644 --- a/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py +++ b/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py @@ -20,11 +20,11 @@ from zenml.config.base_settings import BaseSettings from zenml.integrations.slurm import SLURM_STEP_OPERATOR_FLAVOR +from zenml.integrations.ssh.flavors.base import SSHConnectionConfigMixin from zenml.step_operators.base_step_operator import ( BaseStepOperatorConfig, BaseStepOperatorFlavor, ) -from zenml.utils.secret_utils import SecretField if TYPE_CHECKING: from zenml.integrations.slurm.step_operators import SlurmStepOperator @@ -76,9 +76,17 @@ class SlurmStepOperatorSettings(BaseSettings): class SlurmStepOperatorConfig( - BaseStepOperatorConfig, SlurmStepOperatorSettings + BaseStepOperatorConfig, + SSHConnectionConfigMixin, + SlurmStepOperatorSettings, ): - """Configuration for the Slurm step operator.""" + """Configuration for the Slurm step operator. + + 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, @@ -98,51 +106,20 @@ class SlurmStepOperatorConfig( "requirements installed. Examples: 'source /shared/venvs/zenml/bin/" "activate', 'module load python && conda activate zenml'", ) - # SSH connection settings, required for the `ssh` transport. Field names - # deliberately match the ssh integration's connection mixin so its - # SSHClient can consume this config directly. - hostname: Optional[str] = Field( + # 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( + username: Optional[str] = Field( # type: ignore[assignment] default=None, description="SSH username on the login node, required for the 'ssh' " "transport. Example: 'mlops'", ) - port: int = Field( - default=22, - description="SSH port on the login node. Standard SSH port is 22", - ) - ssh_key_path: Optional[str] = Field( - default=None, - description="Path to the SSH private key file on the submitting " - "machine. Supports RSA, Ed25519, and ECDSA keys", - ) - ssh_private_key: Optional[str] = SecretField( - default=None, - description="SSH private key content, used instead of ssh_key_path " - "when the key is stored in a ZenML secret. Supports {{secret.key}} " - "references", - ) - ssh_key_passphrase: Optional[str] = SecretField( - default=None, - description="Passphrase for an encrypted SSH private key. Leave " - "unset if the key is not encrypted", - ) - verify_host_key: bool = Field( - default=True, - description="Require the remote host key to be present in " - "known_hosts. Set to False to auto-accept unknown host keys (less " - "secure, convenient for ephemeral clusters)", - ) - known_hosts_path: Optional[str] = Field( - default=None, - description="Path to a custom known_hosts file used for host key " - "verification. Uses ~/.ssh/known_hosts if unset", - ) @model_validator(mode="after") def _validate_transport_requirements(self) -> "SlurmStepOperatorConfig": diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index 07caefc031e..43822162569 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -179,16 +179,11 @@ def __init__(self, config: "SlurmStepOperatorConfig") -> None: Args: config: The step operator config holding SSH connection settings. """ - from typing import cast - - from zenml.integrations.ssh.flavors.base import ( - SSHConnectionConfigMixin, - ) from zenml.integrations.ssh.ssh_client import SSHClient - # The Slurm config deliberately mirrors the SSH connection mixin's - # field names, so the SSH client can consume it directly. - self._ssh = SSHClient(cast(SSHConnectionConfigMixin, config)) + # SlurmStepOperatorConfig inherits SSHConnectionConfigMixin, so the SSH + # client consumes it directly. + self._ssh = SSHClient(config) self._ssh.__enter__() def run(self, command: str) -> CommandResult: From 8088c337b79b43c5dc195c4e0bba9de12f0125ed Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 12:24:26 -0700 Subject: [PATCH 04/22] Conform the Slurm step operator to step-operator conventions Surveyed the 11 other step operators; two gaps found and fixed: - Add a `validator` (every modern step operator has one). It rejects a local artifact store, since the cluster reads inputs / writes outputs there. Unlike the Docker-based operators it requires no container registry or image builder - the operator is container-free - matching the databricks operator, the only other container-free one. - Pass `info` (not `info.step_run`) to `get_settings` in the script builder, matching every other operator's `submit`. Verified consistent with convention (no change needed): not overriding `get_docker_builds` is correct - the StackComponent default returns [], and container-free operators (databricks) rely on it; the submit/get_status/cancel lifecycle plus the inherited `wait` matches the modern interface; and the default entrypoint config is right for the container-free path. Refs #5063 Co-Authored-By: Claude Fable 5 --- .../step_operators/slurm_step_operator.py | 36 +++++++++++++++++-- .../slurm/test_slurm_step_operator.py | 28 +++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index dad05d7fd4e..3f5bf10534a 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -47,12 +47,14 @@ SSHSlurmCommandRunner, ) from zenml.logger import get_logger +from zenml.stack import StackValidator from zenml.step_operators import BaseStepOperator from zenml.utils import code_utils, source_utils if TYPE_CHECKING: from zenml.config.step_run_info import StepRunInfo from zenml.models import StepRunResponse + from zenml.stack import Stack logger = get_logger(__name__) @@ -82,6 +84,36 @@ def settings_class(self) -> Optional[type[BaseSettings]]: """ return SlurmStepOperatorSettings + @property + def validator(self) -> Optional[StackValidator]: + """Validate that the stack meets the operator's requirements. + + Unlike Docker-based step operators, the Slurm operator is + container-free, so it needs no container registry or image builder. It + does require a remote artifact store, since the cluster reads inputs + and writes outputs there. + + Returns: + A stack validator. + """ + + def _validate_remote_artifact_store( + stack: "Stack", + ) -> Tuple[bool, str]: + if stack.artifact_store.config.is_local: + return False, ( + "The Slurm step operator runs steps on a remote cluster " + "that must read inputs and write outputs to a shared " + "artifact store, but the artifact store " + f"'{stack.artifact_store.name}' is local. Please use a " + "remote artifact store (S3, GCS, Azure Blob, etc.)." + ) + return True, "" + + return StackValidator( + custom_validation_function=_validate_remote_artifact_store, + ) + def _get_client(self) -> Tuple[SlurmClient, SlurmCommandRunner]: """Build a Slurm client for the configured transport. @@ -139,9 +171,7 @@ def _build_sbatch_script( Returns: The job script content. """ - settings = cast( - SlurmStepOperatorSettings, self.get_settings(info.step_run) - ) + settings = cast(SlurmStepOperatorSettings, self.get_settings(info)) resources = info.config.resource_settings directives = [ diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index 37b4519bddc..05ac22e6c3a 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -166,6 +166,34 @@ def test_flavor_identity(): assert flavor.implementation_class is SlurmStepOperator +def test_validator_rejects_local_artifact_store(): + """The operator requires a remote artifact store, not a container registry.""" + from types import SimpleNamespace + + op = _build_operator() + validator = op.validator + assert validator is not None + # container-free: no required components (no container registry / image builder) + assert not validator._required_components + + check = validator._custom_validation_function + assert check is not None + local_stack = SimpleNamespace( + artifact_store=SimpleNamespace( + name="local", config=SimpleNamespace(is_local=True) + ) + ) + ok, msg = check(local_stack) + assert ok is False and "local" in msg + + remote_stack = SimpleNamespace( + artifact_store=SimpleNamespace( + name="s3", config=SimpleNamespace(is_local=False) + ) + ) + assert check(remote_stack) == (True, "") + + # --- slurm client ------------------------------------------------------------- From f9a5dd6c5b28661ee0cc8f62eb3644c00e4756ba Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 12:45:07 -0700 Subject: [PATCH 05/22] Run Slurm steps from the step's Docker image instead of a shipped tarball The step operator previously shipped the step code as a tarball extracted on the cluster and put on PYTHONPATH. That fought the ZenML entrypoint (which chdirs into the image's /app) and reimplemented code delivery the framework already does well. Instead, declare the step's Docker build via get_docker_builds and run that image on the compute node with a rootless HPC container runtime (Apptainer/Singularity by default, or NVIDIA Pyxis, or the Docker daemon where available), so code delivery uses ZenML's standard, tested image path and no code is shipped by the operator. Security is handled up front: the environment file holding the step's credentials is written owner-only (0600) into a per-run directory created atomically owner-only (mkdir -m 700), passed to the container via --env-file (or, for pyxis, --container-env variable names sourced from that 0600 file) so values never appear on a command line visible to other cluster users, and scrubbed by the job's EXIT trap the moment the job ends. All interpolated paths and values go through shlex.quote. Co-Authored-By: Claude Fable 5 --- .../flavors/slurm_step_operator_flavor.py | 53 ++- src/zenml/integrations/slurm/slurm_client.py | 58 +-- .../step_operators/slurm_step_operator.py | 261 +++++++++--- .../slurm/test_slurm_step_operator.py | 378 ++++++++++-------- 4 files changed, 481 insertions(+), 269 deletions(-) diff --git a/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py b/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py index 4f6daefeb72..460cd1f3ad9 100644 --- a/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py +++ b/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py @@ -14,7 +14,7 @@ """Slurm step operator flavor.""" from enum import Enum -from typing import TYPE_CHECKING, List, Optional, Type +from typing import TYPE_CHECKING, Dict, List, Optional, Type from pydantic import Field, model_validator @@ -37,10 +37,24 @@ class SlurmTransport(str, Enum): 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 SlurmStepOperatorSettings(BaseSettings): """Settings for the Slurm step operator. - These map onto ``sbatch`` directives and can be overridden per step. + These map onto ``sbatch`` directives / container-run flags and can be + overridden per step. """ partition: Optional[str] = Field( @@ -73,6 +87,18 @@ class SlurmStepOperatorSettings(BaseSettings): "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 SlurmStepOperatorConfig( @@ -94,17 +120,20 @@ class SlurmStepOperatorConfig( "to a remote login/controller node, 'local' runs sbatch directly and " "requires the client to already run on the cluster. Example: 'ssh'", ) - workdir: str = Field( - description="Directory on the cluster used to stage job scripts, " - "code archives and job output, ideally on a filesystem shared " - "between the submission host and the compute nodes. Example: " - "'/shared/zenml-runs'", + 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; the safe HPC default), 'pyxis' (NVIDIA " + "enroot/pyxis via `srun --container-image`), or 'docker' (only where " + "a Docker daemon is available). Example: 'apptainer'", ) - env_setup_command: str = Field( - description="Shell command executed at the start of every job to " - "activate a Python environment that has zenml and the pipeline " - "requirements installed. Examples: 'source /shared/venvs/zenml/bin/" - "activate', 'module load python && conda activate zenml'", + 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 diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index 43822162569..824dd1f7a11 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -72,21 +72,17 @@ def run(self, command: str) -> CommandResult: """ @abstractmethod - def put_file(self, local_path: str, remote_path: str) -> None: - """Copy a local file to the submission host. - - Args: - local_path: Path of the file on the local machine. - remote_path: Destination path on the submission host. - """ - - @abstractmethod - def put_text(self, content: str, remote_path: str) -> None: + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: """Write text content to a file on the submission host. Args: - content: The text content to write. 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 @@ -132,26 +128,21 @@ def run(self, command: str) -> CommandResult: stderr=completed.stderr, ) - def put_file(self, local_path: str, remote_path: str) -> None: - """Copy a local file to the destination path. - - Args: - local_path: Path of the source file. - remote_path: Destination path. - """ - import shutil - - shutil.copyfile(local_path, remote_path) - - def put_text(self, content: str, remote_path: str) -> None: + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: """Write text content to the destination path. Args: - content: The text content to write. remote_path: Destination path. + content: The text content to write. + mode: File permissions to apply. """ + import os + with open(remote_path, "w") as f: f.write(content) + os.chmod(remote_path, mode) def read_text(self, remote_path: str) -> str: """Read a text file. @@ -202,24 +193,17 @@ def run(self, command: str) -> CommandResult: stderr=result.stderr, ) - def put_file(self, local_path: str, remote_path: str) -> None: - """Copy a local file to the remote host via SFTP. - - Args: - local_path: Path of the file on the local machine. - remote_path: Destination path on the remote host. - """ - with self._ssh.sftp() as sftp: - sftp.put(local_path, remote_path) - - def put_text(self, content: str, remote_path: str) -> None: + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: """Write text content to a file on the remote host. Args: - content: The text content to write. remote_path: Destination path on the remote host. + content: The text content to write. + mode: File permissions to apply. """ - self._ssh.put_text(content, remote_path) + 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. diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index 3f5bf10534a..6acb4204834 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -13,27 +13,36 @@ # permissions and limitations under the License. """Step operator that runs individual steps as Slurm jobs. -The operator is container-free: Slurm compute nodes generally cannot run -Docker, so the step executes inside a user-provided Python environment on the -cluster (``env_setup_command``). Pipeline code is shipped at submission time -as an archive that the job script unpacks into a per-run directory, so no -image build and no artifact-store code upload are required. +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. Submission is stateless: the Slurm job is named after the step run id, so status lookups and cancellation work by job name without persisting job ids. While a job is in the queue, its state comes from ``squeue``; once it leaves the queue, a sentinel exit-code file written by the job script disambiguates success from failure, so Slurm accounting (``sacct``) is never required. + +Security: the environment file holds the step's credentials (ZenML store +token, etc.). It is written owner-only (0600) into a per-run directory created +owner-only (0700), passed to the container via ``--env-file`` / +``--container-env`` so the values never appear on a command line visible to +other cluster users, and deleted by the job's EXIT trap the moment the job +ends - even if the orchestrator process dies. """ import shlex -import tempfile from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast from uuid import UUID from zenml.config.base_settings import BaseSettings -from zenml.enums import ExecutionStatus +from zenml.config.build_configuration import BuildConfiguration +from zenml.enums import ExecutionStatus, StackComponentType from zenml.integrations.slurm.flavors.slurm_step_operator_flavor import ( + SlurmContainerRuntime, SlurmStepOperatorConfig, SlurmStepOperatorSettings, SlurmTransport, @@ -46,20 +55,23 @@ SlurmCommandRunner, SSHSlurmCommandRunner, ) +from zenml.integrations.ssh.utils import serialize_env_for_docker_env_file from zenml.logger import get_logger from zenml.stack import StackValidator from zenml.step_operators import BaseStepOperator -from zenml.utils import code_utils, source_utils if TYPE_CHECKING: from zenml.config.step_run_info import StepRunInfo - from zenml.models import StepRunResponse + from zenml.models import PipelineSnapshotBase, StepRunResponse from zenml.stack import Stack logger = get_logger(__name__) +SLURM_STEP_OPERATOR_DOCKER_IMAGE_KEY = "slurm_step_operator" + _JOB_NAME_PREFIX = "zenml" _EXIT_CODE_FILE = "exit_code" +_ENV_FILE = "env" _OUTPUT_FILE = "output.log" @@ -88,16 +100,16 @@ def settings_class(self) -> Optional[type[BaseSettings]]: def validator(self) -> Optional[StackValidator]: """Validate that the stack meets the operator's requirements. - Unlike Docker-based step operators, the Slurm operator is - container-free, so it needs no container registry or image builder. It - does require a remote artifact store, since the cluster reads inputs - and writes outputs there. + The operator runs the step's Docker image on the cluster, so it needs + a remote container registry and image builder (to build and pull the + image) and a remote artifact store (the cluster reads inputs and + writes outputs there). Returns: A stack validator. """ - def _validate_remote_artifact_store( + def _validate_remote_components( stack: "Stack", ) -> Tuple[bool, str]: if stack.artifact_store.config.is_local: @@ -108,12 +120,50 @@ def _validate_remote_artifact_store( f"'{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 step operator runs the step's image on a " + "remote cluster, which must pull it from a registry, but " + f"the container registry '{container_registry.name}' is " + "local. Please use a remote container registry (ECR, GCR, " + "ACR, DockerHub, etc.)." + ) return True, "" return StackValidator( - custom_validation_function=_validate_remote_artifact_store, + required_components={ + StackComponentType.CONTAINER_REGISTRY, + StackComponentType.IMAGE_BUILDER, + }, + custom_validation_function=_validate_remote_components, ) + 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 + def _get_client(self) -> Tuple[SlurmClient, SlurmCommandRunner]: """Build a Slurm client for the configured transport. @@ -153,19 +203,92 @@ def _run_dir(self, step_run_id: UUID) -> str: f"{self.config.workdir.rstrip('/')}/{self._job_name(step_run_id)}" ) + def _build_container_command( + self, + image: str, + entrypoint_command: List[str], + env_file: str, + env_keys: List[str], + use_gpu: bool, + settings: SlurmStepOperatorSettings, + ) -> 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. + + Args: + image: Fully-qualified image reference to run. + entrypoint_command: The ZenML entrypoint command. + 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). + + 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) + runtime = self.config.container_runtime + + if runtime in ( + SlurmContainerRuntime.APPTAINER, + SlurmContainerRuntime.SINGULARITY, + ): + binary = runtime.value # "apptainer" or "singularity" + parts = [binary, "exec"] + 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"{' '.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"{' '.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"{source_env}{' '.join(srun)} {entrypoint}" + def _build_sbatch_script( self, info: "StepRunInfo", - entrypoint_command: List[str], - environment: Dict[str, str], + container_command: str, run_dir: str, ) -> str: """Render the sbatch job script for a step. Args: info: Information about the step run. - entrypoint_command: Command that executes the step. - environment: Environment variables to set for the step. + container_command: The shell snippet that runs the step container. run_dir: The per-run staging directory on the cluster. Returns: @@ -197,34 +320,19 @@ def _build_sbatch_script( for directive in settings.extra_sbatch_directives: directives.append(f"#SBATCH {directive}") - env_exports = "\n".join( - f"export {key}={shlex.quote(value)}" - for key, value in sorted(environment.items()) - ) - command = " ".join(shlex.quote(part) for part in entrypoint_command) - # The EXIT trap records the job outcome in a sentinel file that - # `get_status` reads after the job leaves the squeue output, so the - # operator never depends on Slurm accounting being configured. - # `set -e` makes environment setup or code extraction failures abort - # the job (with the failing code captured by the trap) rather than - # falling through to run the step in a broken environment. + # `get_status` 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 orchestrator dies. + # `set -e` aborts the job (with the failing code captured) if the + # container runtime itself fails. return f"""#!/bin/bash {chr(10).join(directives)} -trap 'echo $? > {run_dir}/{_EXIT_CODE_FILE}' EXIT +trap 'ec=$?; echo "$ec" > {run_dir}/{_EXIT_CODE_FILE}; rm -f {run_dir}/{_ENV_FILE}' EXIT set -eo pipefail -{self.config.env_setup_command} - -{env_exports} - -mkdir -p {run_dir}/code -tar -xzf {run_dir}/code.tar.gz -C {run_dir}/code -cd {run_dir}/code -export PYTHONPATH="{run_dir}/code:${{PYTHONPATH:-}}" - -{command} +{container_command} """ def submit( @@ -240,39 +348,46 @@ def submit( entrypoint_command: Command that executes the step. environment: Environment variables to set in the step operator environment. + + Raises: + RuntimeError: If the run directory cannot be created on the + cluster. """ + 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_env_for_docker_env_file(environment) + container_command = self._build_container_command( + image=image, + entrypoint_command=entrypoint_command, + env_file=env_file, + env_keys=sorted(environment), + use_gpu=use_gpu, + settings=settings, + ) + script = self._build_sbatch_script( + info=info, container_command=container_command, run_dir=run_dir + ) + client, runner = self._get_client() try: - # The run directory holds the code archive, the job script (which - # contains credentials as env exports) and the job output, so it - # is created private to the submitting user. - result = runner.run( - f"mkdir -p {shlex.quote(run_dir)} " - f"&& chmod 700 {shlex.quote(run_dir)}" - ) + # Create the run directory owner-only: it holds the env file (with + # credentials), the job script, and the job output. `-m 700` sets + # the mode atomically on creation, so there is no window where the + # directory is readable by other cluster users before a chmod. + result = runner.run(f"mkdir -m 700 -p {shlex.quote(run_dir)}") if result.exit_code != 0: raise RuntimeError( f"Failed to create run directory `{run_dir}` on the " f"cluster: {result.stderr.strip()}" ) - with tempfile.NamedTemporaryFile(suffix=".tar.gz") as archive: - code_archive = code_utils.CodeArchive( - root=source_utils.get_source_root() - ) - code_archive.write_archive(archive) - archive.flush() - runner.put_file(archive.name, f"{run_dir}/code.tar.gz") - - script = self._build_sbatch_script( - info=info, - entrypoint_command=entrypoint_command, - environment=environment, - run_dir=run_dir, - ) + runner.put_text(env_file, env_content, mode=0o600) script_path = f"{run_dir}/job.sh" - runner.put_text(script, script_path) + runner.put_text(script_path, script, mode=0o700) job_id = client.submit(script_path) logger.info( @@ -340,6 +455,26 @@ def cancel(self, step_run: "StepRunResponse") -> None: finally: 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, runner = self._get_client() + try: + 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: + runner.close() + def _log_failure_output( self, runner: SlurmCommandRunner, run_dir: str ) -> None: diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index 05ac22e6c3a..8ac263342d4 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -13,20 +13,24 @@ # permissions and limitations under the License. """Unit tests for the Slurm integration. -The Slurm CLI is never invoked: all tests drive the client and step operator -through a fake in-memory command runner, so they run without a cluster. +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 typing import Dict, Optional -from uuid import uuid4 +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, @@ -35,6 +39,9 @@ ) from zenml.integrations.slurm.step_operators import SlurmStepOperator +SECRET_TOKEN = "super-secret-zenml-token" +IMAGE = "registry.example.com/zenml:abc123" + class FakeRunner(SlurmCommandRunner): """In-memory fake of the Slurm submission host.""" @@ -52,7 +59,8 @@ def __init__( """ self.queue_state = queue_state self.files = files or {} - self.commands: list[str] = [] + self.modes: Dict[str, int] = {} + self.commands: List[str] = [] def run(self, command: str) -> CommandResult: """Fake command execution. @@ -69,27 +77,20 @@ def run(self, command: str) -> CommandResult: if command.startswith("squeue"): stdout = f"{self.queue_state}\n" if self.queue_state else "" return CommandResult(exit_code=0, stdout=stdout, stderr="") - if command.startswith("scancel"): - return CommandResult(exit_code=0, stdout="", stderr="") return CommandResult(exit_code=0, stdout="", stderr="") - def put_file(self, local_path: str, remote_path: str) -> None: - """Record a file upload. + def put_text( + self, remote_path: str, content: str, mode: int = 0o600 + ) -> None: + """Record a text upload with its mode. Args: - local_path: Source path. remote_path: Destination path. - """ - self.files[remote_path] = f"" - - def put_text(self, content: str, remote_path: str) -> None: - """Record a text upload. - - Args: content: The text content. - remote_path: Destination path. + 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. @@ -111,14 +112,20 @@ def read_text(self, remote_path: str) -> str: 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", - env_setup_command="source /shared/venv/bin/activate", + transport="local", workdir="/shared/zenml-runs" ), flavor="slurm", type=StackComponentType.STEP_OPERATOR, @@ -128,32 +135,42 @@ def _build_operator( ) +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=None, + 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="/shared/zenml-runs", - env_setup_command="true", - ) - - # local transport without connection details is fine - SlurmStepOperatorConfig( - transport="local", - workdir="/shared/zenml-runs", - env_setup_command="true", - ) + SlurmStepOperatorConfig(transport="ssh", workdir="/s") - # ssh transport with connection details is fine + SlurmStepOperatorConfig(transport="local", workdir="/s") SlurmStepOperatorConfig( - transport="ssh", - workdir="/shared/zenml-runs", - env_setup_command="true", - hostname="login.example.com", - username="mlops", + transport="ssh", workdir="/s", hostname="h", username="u" ) @@ -166,32 +183,40 @@ def test_flavor_identity(): assert flavor.implementation_class is SlurmStepOperator -def test_validator_rejects_local_artifact_store(): - """The operator requires a remote artifact store, not a container registry.""" - from types import SimpleNamespace +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" + +# --- 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 - # container-free: no required components (no container registry / image builder) - assert not validator._required_components - + assert validator._required_components == { + StackComponentType.CONTAINER_REGISTRY, + StackComponentType.IMAGE_BUILDER, + } check = validator._custom_validation_function assert check is not None - local_stack = SimpleNamespace( - artifact_store=SimpleNamespace( - name="local", config=SimpleNamespace(is_local=True) - ) - ) - ok, msg = check(local_stack) - assert ok is False and "local" in msg - remote_stack = SimpleNamespace( - artifact_store=SimpleNamespace( - name="s3", config=SimpleNamespace(is_local=False) + 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(remote_stack) == (True, "") + + assert check(_stack(True, False))[0] is False + assert check(_stack(False, True))[0] is False + assert check(_stack(False, False)) == (True, "") # --- slurm client ------------------------------------------------------------- @@ -200,12 +225,11 @@ def test_validator_rejects_local_artifact_store(): def test_client_submit_parses_job_id(): """`sbatch --parsable` output is parsed into a job id.""" runner = FakeRunner() - client = SlurmClient(runner) - assert client.submit("/shared/job.sh") == "12345" - assert runner.commands == ["sbatch --parsable /shared/job.sh"] + assert SlurmClient(runner).submit("/s/job.sh") == "12345" + assert runner.commands == ["sbatch --parsable /s/job.sh"] -def test_client_job_state_and_cancel_use_job_name(): +def test_client_state_and_cancel_use_job_name(): """Status lookup and cancellation address the job by name.""" runner = FakeRunner(queue_state="RUNNING") client = SlurmClient(runner) @@ -214,137 +238,178 @@ def test_client_job_state_and_cancel_use_job_name(): assert runner.commands[-1] == "scancel --name zenml-abc" -def test_client_job_state_returns_none_when_not_queued(): - """A job absent from squeue returns None (terminal, see sentinel).""" - runner = FakeRunner(queue_state=None) - client = SlurmClient(runner) - assert client.get_job_state("zenml-abc") is None +# --- container command rendering (per runtime) -------------------------------- + + +def _container_cmd(runtime: str, gpu: bool = False) -> str: + config = SlurmStepOperatorConfig( + transport="local", workdir="/s", container_runtime=runtime + ) + op = _build_operator(config) + return op._build_container_command( + 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") + 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 + + +# --- sbatch script ------------------------------------------------------------ -# --- step operator status mapping --------------------------------------------- +def test_sbatch_script_directives_and_scrub(): + """The script carries directives, fails fast, and scrubs the env file.""" + op = _build_operator( + SlurmStepOperatorConfig(transport="local", workdir="/s") + ) + op.get_settings = lambda _i: SlurmStepOperatorSettings( + partition="gpu", time_limit="2:00:00" + ) + sid = uuid4() + script = op._build_sbatch_script( + info=_fake_info(sid, gpu=2), + container_command="RUN_THE_CONTAINER", + run_dir="/s/zenml-x", + ) + 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=16G" 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 -f /s/zenml-x/env" in script + + +# --- submit: security-sensitive handling -------------------------------------- + + +def test_submit_writes_secret_env_file_owner_only(): + """The env file is written 0600, the script 0700, and no secret leaks.""" + op = _build_operator() + op.get_settings = lambda _i: SlurmStepOperatorSettings() + runner = FakeRunner() + op._get_client = lambda: (SlurmClient(runner), runner) + + sid = uuid4() + op.submit( + info=_fake_info(sid), + 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"] + + # the run directory is created owner-only, atomically + assert any("mkdir -m 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"] + # 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) -class _FakeStepRun: - """Minimal stand-in for a StepRunResponse.""" - def __init__(self) -> None: - """Initialize with a random id.""" - self.id = uuid4() +# --- status mapping ----------------------------------------------------------- @pytest.fixture -def operator_with_fake_runner(monkeypatch): - """A local-transport operator whose runner is the in-memory fake.""" - operator = _build_operator() +def op_and_runner(monkeypatch): + """A local operator whose runner is the in-memory fake.""" + op = _build_operator() runner = FakeRunner() monkeypatch.setattr( - operator, - "_get_client", - lambda: (SlurmClient(runner), runner), + op, "_get_client", lambda: (SlurmClient(runner), runner) ) - return operator, runner + return op, runner @pytest.mark.parametrize( "queue_state,expected", [ ("PENDING", ExecutionStatus.QUEUED), - ("CONFIGURING", ExecutionStatus.QUEUED), ("RUNNING", ExecutionStatus.RUNNING), - ("COMPLETING", ExecutionStatus.RUNNING), ("CANCELLED", ExecutionStatus.CANCELLED), ("FAILED", ExecutionStatus.FAILED), ], ) -def test_get_status_maps_queue_states( - operator_with_fake_runner, queue_state, expected -): +def test_get_status_maps_queue_states(op_and_runner, queue_state, expected): """Queue states map onto ZenML execution statuses.""" - operator, runner = operator_with_fake_runner + op, runner = op_and_runner runner.queue_state = queue_state - assert operator.get_status(_FakeStepRun()) is expected + assert op.get_status(SimpleNamespace(id=uuid4())) is expected -def test_get_status_reads_sentinel_after_queue(operator_with_fake_runner): +def test_get_status_reads_sentinel_after_queue(op_and_runner): """After the job leaves the queue, the sentinel file decides the status.""" - operator, runner = operator_with_fake_runner - step_run = _FakeStepRun() - run_dir = operator._run_dir(step_run.id) - + op, runner = op_and_runner + step_run = SimpleNamespace(id=uuid4()) + run_dir = op._run_dir(step_run.id) runner.files[f"{run_dir}/exit_code"] = "0\n" - assert operator.get_status(step_run) is ExecutionStatus.COMPLETED - + assert op.get_status(step_run) is ExecutionStatus.COMPLETED runner.files[f"{run_dir}/exit_code"] = "1\n" - assert operator.get_status(step_run) is ExecutionStatus.FAILED + assert op.get_status(step_run) is ExecutionStatus.FAILED -def test_get_status_without_queue_entry_or_sentinel_is_cancelled( - operator_with_fake_runner, -): +def test_get_status_without_queue_or_sentinel_is_cancelled(op_and_runner): """No queue entry and no sentinel means the job was torn down early.""" - operator, runner = operator_with_fake_runner - assert operator.get_status(_FakeStepRun()) is ExecutionStatus.CANCELLED - - -# --- sbatch script rendering ---------------------------------------------------- - - -def test_sbatch_script_rendering(monkeypatch): - """The generated script carries directives, env setup, and the command.""" - from types import SimpleNamespace - - from zenml.config.resource_settings import ResourceSettings - - operator = _build_operator() - step_run_id = uuid4() - - fake_info = SimpleNamespace( - step_run_id=step_run_id, - pipeline_step_name="trainer", - config=SimpleNamespace( - resource_settings=ResourceSettings( - cpu_count=4, gpu_count=2, memory="16GB" - ) - ), - step_run=None, - ) - - from zenml.integrations.slurm.flavors import SlurmStepOperatorSettings - - monkeypatch.setattr( - operator, - "get_settings", - lambda _: SlurmStepOperatorSettings( - partition="gpu", - time_limit="2:00:00", - extra_sbatch_directives=["--constraint=a100"], - ), + op, runner = op_and_runner + assert ( + op.get_status(SimpleNamespace(id=uuid4())) is ExecutionStatus.CANCELLED ) - script = operator._build_sbatch_script( - info=fake_info, - entrypoint_command=["python", "-m", "zenml.entrypoint"], - environment={"ZENML_STORE_URL": "https://z.example.com"}, - run_dir="/shared/zenml-runs/zenml-x", - ) - assert f"#SBATCH --job-name=zenml-{step_run_id}" in script - assert "#SBATCH --partition=gpu" in script - assert "#SBATCH --time=2:00:00" in script - assert "#SBATCH --cpus-per-task=4" in script - assert "#SBATCH --mem=16G" in script - assert "#SBATCH --gres=gpu:2" in script - assert "#SBATCH --constraint=a100" in script - assert "set -eo pipefail" in script - assert "source /shared/venv/bin/activate" in script - assert "export ZENML_STORE_URL=https://z.example.com" in script - assert "python -m zenml.entrypoint" in script - assert ( - "trap 'echo $? > /shared/zenml-runs/zenml-x/exit_code' EXIT" in script - ) +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("rm -rf" in c and run_dir in c for c in runner.commands) -# --- integration registration --------------------------------------------------- +# --- integration registration ------------------------------------------------- def test_integration_registered(): @@ -354,5 +419,4 @@ def test_integration_registered(): integration_registry._initialize() assert integration_registry.integrations["slurm"] is SlurmIntegration - flavors = {f().name for f in SlurmIntegration.flavors()} - assert flavors == {"slurm"} + assert {f().name for f in SlurmIntegration.flavors()} == {"slurm"} From f48eaff9eeb990f4c075806b2c403f169f9cc2cd Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 13:39:52 -0700 Subject: [PATCH 06/22] Add a Slurm orchestrator and factor out the shared Slurm-job machinery Introduce a Slurm orchestrator that submits a whole pipeline as a graph of Slurm jobs. Each step becomes its own `sbatch` job running the step's ZenML Docker image on a compute node, and the pipeline DAG is reproduced with job dependencies (`--dependency=afterok`, plus `--kill-on-invalid-dep` so a failed upstream propagates instead of leaving jobs pending forever). Submission is fully detached: once the jobs are queued the scheduler runs them and each step reports its own status to the ZenML server, so the orchestrator never polls. Jobs are submitted upstream-first (topological order via the existing `topsorted_layers`) so each step's dependencies reference already-submitted job ids, and a mid-submission failure best-effort cancels the jobs already queued. The orchestrator and the existing step operator run the same kind of Slurm job, so the shared machinery is extracted to avoid divergence: - `flavors/base.py` holds the transport/runtime/`sbatch` settings and the SSH connection mixin both components inherit. - `slurm_job.py` owns the container-command and sbatch-script builders, the stack validator, and `stage_and_submit`, which centralizes the security-sensitive staging sequence (atomic `mkdir -m 700`, 0600 env file, 0700 script, submit) so the file-mode invariants cannot drift between the two components. - `build_slurm_client` in `slurm_client.py` is the single transport factory both components call. `SlurmClient.submit` gains optional dependency support for the DAG edges. Co-Authored-By: Claude Fable 5 --- src/zenml/integrations/slurm/__init__.py | 30 +- .../integrations/slurm/flavors/__init__.py | 8 + src/zenml/integrations/slurm/flavors/base.py | 166 ++++++++ .../flavors/slurm_orchestrator_flavor.py | 110 ++++++ .../flavors/slurm_step_operator_flavor.py | 152 +------- .../slurm/orchestrators/__init__.py | 20 + .../slurm/orchestrators/slurm_orchestrator.py | 313 +++++++++++++++ src/zenml/integrations/slurm/slurm_client.py | 55 ++- src/zenml/integrations/slurm/slurm_job.py | 260 +++++++++++++ .../step_operators/slurm_step_operator.py | 265 ++----------- .../slurm/test_slurm_orchestrator.py | 359 ++++++++++++++++++ .../slurm/test_slurm_step_operator.py | 42 +- 12 files changed, 1372 insertions(+), 408 deletions(-) create mode 100644 src/zenml/integrations/slurm/flavors/base.py create mode 100644 src/zenml/integrations/slurm/flavors/slurm_orchestrator_flavor.py create mode 100644 src/zenml/integrations/slurm/orchestrators/__init__.py create mode 100644 src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py create mode 100644 src/zenml/integrations/slurm/slurm_job.py create mode 100644 tests/unit/integrations/slurm/test_slurm_orchestrator.py diff --git a/src/zenml/integrations/slurm/__init__.py b/src/zenml/integrations/slurm/__init__.py index 8d73c15be1a..5b8e014e544 100644 --- a/src/zenml/integrations/slurm/__init__.py +++ b/src/zenml/integrations/slurm/__init__.py @@ -13,15 +13,19 @@ # permissions and limitations under the License. """Initialization of the Slurm integration. -The Slurm integration provides a step operator that runs individual pipeline -steps as Slurm jobs on an HPC cluster, either over SSH to a login/controller -node or via local ``sbatch`` when the client already runs on the cluster. - -Unlike container-based step operators, the Slurm step operator is -container-free: Slurm compute nodes generally cannot run Docker, so steps -execute inside a user-provided Python environment on the cluster (typically on -a shared filesystem), and the pipeline code is shipped to the cluster as an -archive at submission time. +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 @@ -31,6 +35,7 @@ from zenml.stack import Flavor SLURM_STEP_OPERATOR_FLAVOR = "slurm" +SLURM_ORCHESTRATOR_FLAVOR = "slurm" class SlurmIntegration(Integration): @@ -48,6 +53,9 @@ def flavors(cls) -> List[Type[Flavor]]: Returns: List of stack component flavors for this integration. """ - from zenml.integrations.slurm.flavors import SlurmStepOperatorFlavor + from zenml.integrations.slurm.flavors import ( + SlurmOrchestratorFlavor, + SlurmStepOperatorFlavor, + ) - return [SlurmStepOperatorFlavor] + return [SlurmStepOperatorFlavor, SlurmOrchestratorFlavor] diff --git a/src/zenml/integrations/slurm/flavors/__init__.py b/src/zenml/integrations/slurm/flavors/__init__.py index ee1e7d6ccfa..58ea0f88a2d 100644 --- a/src/zenml/integrations/slurm/flavors/__init__.py +++ b/src/zenml/integrations/slurm/flavors/__init__.py @@ -13,6 +13,11 @@ # 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, @@ -20,6 +25,9 @@ ) __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..0f93503bbcf --- /dev/null +++ b/src/zenml/integrations/slurm/flavors/base.py @@ -0,0 +1,166 @@ +# 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; the safe HPC default), '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 index 460cd1f3ad9..7fbef9db2d4 100644 --- a/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py +++ b/src/zenml/integrations/slurm/flavors/slurm_step_operator_flavor.py @@ -13,14 +13,13 @@ # permissions and limitations under the License. """Slurm step operator flavor.""" -from enum import Enum -from typing import TYPE_CHECKING, Dict, List, Optional, Type +from typing import TYPE_CHECKING, Optional, Type -from pydantic import Field, model_validator - -from zenml.config.base_settings import BaseSettings from zenml.integrations.slurm import SLURM_STEP_OPERATOR_FLAVOR -from zenml.integrations.ssh.flavors.base import SSHConnectionConfigMixin +from zenml.integrations.slurm.flavors.base import ( + SlurmConnectionConfig, + SlurmJobSettings, +) from zenml.step_operators.base_step_operator import ( BaseStepOperatorConfig, BaseStepOperatorFlavor, @@ -30,154 +29,23 @@ from zenml.integrations.slurm.step_operators import SlurmStepOperator -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 SlurmStepOperatorSettings(BaseSettings): - """Settings for the Slurm step operator. - - 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 SlurmStepOperatorSettings(SlurmJobSettings): + """Settings for the Slurm step operator.""" class SlurmStepOperatorConfig( BaseStepOperatorConfig, - SSHConnectionConfigMixin, + SlurmConnectionConfig, SlurmStepOperatorSettings, ): - """Configuration for the Slurm step operator. - - 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; the safe HPC default), '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) -> "SlurmStepOperatorConfig": - """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 step operator 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 + """Configuration for the Slurm step operator.""" @property def is_remote(self) -> bool: """Whether this component runs remotely. Returns: - True, since jobs execute on the Slurm cluster. + True, since steps execute on the Slurm cluster. """ return True 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..4402a6c9e3f --- /dev/null +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -0,0 +1,313 @@ +# 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. +""" + +import os +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast + +from zenml.config.base_settings import BaseSettings +from zenml.entrypoints.step_entrypoint_configuration import ( + StepEntrypointConfiguration, +) +from zenml.integrations.slurm.flavors.slurm_orchestrator_flavor import ( + SlurmOrchestratorConfig, + SlurmOrchestratorSettings, +) +from zenml.integrations.slurm.slurm_client import ( + SlurmClient, + build_slurm_client, +) +from zenml.integrations.slurm.slurm_job import ( + ENV_FILE, + REQUIRED_COMPONENTS, + build_container_command, + build_sbatch_script, + stage_and_submit, + validate_remote_stack, +) +from zenml.integrations.ssh.utils import serialize_env_for_docker_env_file +from zenml.logger import get_logger +from zenml.orchestrators import ContainerizedOrchestrator +from zenml.orchestrators.base_orchestrator import SubmissionResult +from zenml.orchestrators.topsort import topsorted_layers +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" + + +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 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}." + ) + + @staticmethod + def _sorted_step_names( + steps: Dict[str, "Step"], + ) -> List[str]: + """Return the step invocation ids in a topological order. + + Jobs must be submitted upstream-first so that a step's Slurm + dependencies reference the already-submitted job ids of its parents. + + Args: + steps: The steps keyed by invocation id. + + Returns: + The invocation ids in a valid topological order. + """ + + def parents(name: str) -> List[str]: + return [u for u in steps[name].spec.upstream_steps if u in steps] + + def children(name: str) -> List[str]: + return [ + other + for other in steps + if name in steps[other].spec.upstream_steps + ] + + layers = topsorted_layers( + nodes=list(steps), + get_node_id_fn=lambda name: name, + get_parent_nodes=parents, + get_child_nodes=children, + ) + return [name for layer in layers for name in layer] + + 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 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. + step_environments: Environment variables per step. + placeholder_run: An optional placeholder run for the snapshot. + + Returns: + None, since the pipeline is submitted detached. + + Raises: + RuntimeError: If a step directory cannot be created on the cluster. + """ + run_id = str(snapshot.id) + # Set locally too, so get_orchestrator_run_id works if the framework + # queries it on the submitting machine; the authoritative value is + # injected into each step's job environment below. + os.environ[ENV_ZENML_SLURM_RUN_ID] = run_id + + steps = snapshot.step_configurations + client = build_slurm_client(self.config) + submitted: Dict[str, str] = {} + try: + for step_name in self._sorted_step_names(steps): + step = steps[step_name] + job_id = 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 + ], + ) + submitted[step_name] = job_id + logger.info( + "Submitted step `%s` as Slurm job `%s`.", + step_name, + job_id, + ) + 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: + client.runner.run(f"scancel {' '.join(submitted.values())}") + raise + finally: + client.runner.close() + + logger.info( + "Submitted pipeline `%s` as %d Slurm jobs.", + snapshot.pipeline_configuration.name, + len(submitted), + ) + return None + + 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], + ) -> 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. + + Returns: + The submitted Slurm job id. + + Raises: + RuntimeError: If the step directory cannot be created. + """ + image = self.get_image(snapshot=snapshot, step_name=step_name) + settings = cast(SlurmOrchestratorSettings, self.get_settings(step)) + run_dir = self._run_dir(run_id, step_name) + env_file = f"{run_dir}/{ENV_FILE}" + use_gpu = bool(step.config.resource_settings.gpu_count) + + environment = step_environment.copy() + environment[ENV_ZENML_SLURM_RUN_ID] = run_id + env_content = serialize_env_for_docker_env_file(environment) + + entrypoint_command = ( + StepEntrypointConfiguration.get_entrypoint_command() + + StepEntrypointConfiguration.get_entrypoint_arguments( + step_name=step_name, snapshot_id=snapshot.id + ) + ) + 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, + ) + script = build_sbatch_script( + job_name=self._job_name(run_id, step_name), + run_dir=run_dir, + container_command=container_command, + resources=step.config.resource_settings, + settings=settings, + ) + + return stage_and_submit( + client, + run_dir=run_dir, + env_content=env_content, + script=script, + dependencies=upstream_job_ids, + ) diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index 824dd1f7a11..758cc546c13 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -29,12 +29,10 @@ import subprocess from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, List, Optional if TYPE_CHECKING: - from zenml.integrations.slurm.flavors.slurm_step_operator_flavor import ( - SlurmStepOperatorConfig, - ) + from zenml.integrations.slurm.flavors.base import SlurmConnectionConfig # Slurm job states that mean the job is still on its way to running. PENDING_STATES = { @@ -164,15 +162,15 @@ class SSHSlurmCommandRunner(SlurmCommandRunner): provide the same connection fields (``hostname``, ``username``, ...). """ - def __init__(self, config: "SlurmStepOperatorConfig") -> None: + def __init__(self, config: "SlurmConnectionConfig") -> None: """Initialize the runner with an open SSH connection. Args: - config: The step operator config holding SSH connection settings. + config: The Slurm component config holding SSH connection settings. """ from zenml.integrations.ssh.ssh_client import SSHClient - # SlurmStepOperatorConfig inherits SSHConnectionConfigMixin, so the SSH + # SlurmConnectionConfig inherits SSHConnectionConfigMixin, so the SSH # client consumes it directly. self._ssh = SSHClient(config) self._ssh.__enter__() @@ -232,11 +230,18 @@ def __init__(self, runner: SlurmCommandRunner) -> None: """ self.runner = runner - def submit(self, script_path: str) -> str: + def submit( + self, script_path: str, dependencies: Optional[List[str]] = None + ) -> 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. Returns: The Slurm job id. @@ -244,9 +249,13 @@ def submit(self, script_path: str) -> str: Raises: RuntimeError: If the submission fails. """ - result = self.runner.run( - f"sbatch --parsable {shlex.quote(script_path)}" - ) + command = "sbatch --parsable" + if dependencies: + command += ( + f" --dependency=afterok:{':'.join(dependencies)}" + " --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}: " @@ -297,3 +306,27 @@ def cancel(self, job_name: str) -> None: 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..b7907d204fe --- /dev/null +++ b/src/zenml/integrations/slurm/slurm_job.py @@ -0,0 +1,260 @@ +# 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 shlex +from typing import TYPE_CHECKING, List, 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" + +# The required stack components shared by both Slurm components. +REQUIRED_COMPONENTS = { + StackComponentType.CONTAINER_REGISTRY, + StackComponentType.IMAGE_BUILDER, +} + + +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, +) -> 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. + + 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). + + 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) + + if runtime in ( + SlurmContainerRuntime.APPTAINER, + SlurmContainerRuntime.SINGULARITY, + ): + binary = runtime.value # "apptainer" or "singularity" + parts = [binary, "exec"] + 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"{' '.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"{' '.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"{source_env}{' '.join(srun)} {entrypoint}" + + +def build_sbatch_script( + job_name: str, + run_dir: str, + container_command: str, + resources: "ResourceSettings", + settings: SlurmJobSettings, +) -> 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, ...). + + 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={int(resources.cpu_count)}" + ) + if memory_gb := resources.get_memory(unit="GB"): + directives.append(f"#SBATCH --mem={int(memory_gb)}G") + 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. + return f"""#!/bin/bash +{chr(10).join(directives)} + +trap 'ec=$?; echo "$ec" > {run_dir}/{EXIT_CODE_FILE}; rm -f {run_dir}/{ENV_FILE}' 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, +) -> 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). + + Returns: + The submitted Slurm job id. + + Raises: + RuntimeError: If the run directory cannot be created on the cluster. + """ + runner = client.runner + result = runner.run(f"mkdir -m 700 -p {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()}" + ) + + runner.put_text(f"{run_dir}/{ENV_FILE}", env_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) diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index 6acb4204834..da98d3514c6 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -25,13 +25,6 @@ While a job is in the queue, its state comes from ``squeue``; once it leaves the queue, a sentinel exit-code file written by the job script disambiguates success from failure, so Slurm accounting (``sacct``) is never required. - -Security: the environment file holds the step's credentials (ZenML store -token, etc.). It is written owner-only (0600) into a per-run directory created -owner-only (0700), passed to the container via ``--env-file`` / -``--container-env`` so the values never appear on a command line visible to -other cluster users, and deleted by the job's EXIT trap the moment the job -ends - even if the orchestrator process dies. """ import shlex @@ -40,20 +33,26 @@ from zenml.config.base_settings import BaseSettings from zenml.config.build_configuration import BuildConfiguration -from zenml.enums import ExecutionStatus, StackComponentType +from zenml.enums import ExecutionStatus from zenml.integrations.slurm.flavors.slurm_step_operator_flavor import ( - SlurmContainerRuntime, SlurmStepOperatorConfig, SlurmStepOperatorSettings, - SlurmTransport, ) from zenml.integrations.slurm.slurm_client import ( PENDING_STATES, RUNNING_STATES, - LocalSlurmCommandRunner, - SlurmClient, SlurmCommandRunner, - SSHSlurmCommandRunner, + build_slurm_client, +) +from zenml.integrations.slurm.slurm_job import ( + ENV_FILE, + EXIT_CODE_FILE, + OUTPUT_FILE, + REQUIRED_COMPONENTS, + build_container_command, + build_sbatch_script, + stage_and_submit, + validate_remote_stack, ) from zenml.integrations.ssh.utils import serialize_env_for_docker_env_file from zenml.logger import get_logger @@ -63,16 +62,12 @@ if TYPE_CHECKING: from zenml.config.step_run_info import StepRunInfo from zenml.models import PipelineSnapshotBase, StepRunResponse - from zenml.stack import Stack logger = get_logger(__name__) SLURM_STEP_OPERATOR_DOCKER_IMAGE_KEY = "slurm_step_operator" _JOB_NAME_PREFIX = "zenml" -_EXIT_CODE_FILE = "exit_code" -_ENV_FILE = "env" -_OUTPUT_FILE = "output.log" class SlurmStepOperator(BaseStepOperator): @@ -100,45 +95,12 @@ def settings_class(self) -> Optional[type[BaseSettings]]: def validator(self) -> Optional[StackValidator]: """Validate that the stack meets the operator's requirements. - The operator runs the step's Docker image on the cluster, so it needs - a remote container registry and image builder (to build and pull the - image) and a remote artifact store (the cluster reads inputs and - writes outputs there). - Returns: A stack validator. """ - - def _validate_remote_components( - stack: "Stack", - ) -> Tuple[bool, str]: - if stack.artifact_store.config.is_local: - return False, ( - "The Slurm step operator runs steps on a remote cluster " - "that must read inputs and write outputs to a shared " - "artifact store, but the artifact store " - f"'{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 step operator runs the step's image on a " - "remote cluster, which must pull it from a registry, but " - f"the container registry '{container_registry.name}' is " - "local. Please use a remote container registry (ECR, GCR, " - "ACR, DockerHub, etc.)." - ) - return True, "" - return StackValidator( - required_components={ - StackComponentType.CONTAINER_REGISTRY, - StackComponentType.IMAGE_BUILDER, - }, - custom_validation_function=_validate_remote_components, + required_components=REQUIRED_COMPONENTS, + custom_validation_function=validate_remote_stack, ) def get_docker_builds( @@ -164,20 +126,6 @@ def get_docker_builds( ) return builds - def _get_client(self) -> Tuple[SlurmClient, SlurmCommandRunner]: - """Build a Slurm client for the configured transport. - - Returns: - The Slurm client and its underlying runner (which the caller is - responsible for closing). - """ - runner: SlurmCommandRunner - if self.config.transport == SlurmTransport.LOCAL: - runner = LocalSlurmCommandRunner() - else: - runner = SSHSlurmCommandRunner(self.config) - return SlurmClient(runner), runner - @staticmethod def _job_name(step_run_id: UUID) -> str: """Deterministic Slurm job name for a step run. @@ -203,138 +151,6 @@ def _run_dir(self, step_run_id: UUID) -> str: f"{self.config.workdir.rstrip('/')}/{self._job_name(step_run_id)}" ) - def _build_container_command( - self, - image: str, - entrypoint_command: List[str], - env_file: str, - env_keys: List[str], - use_gpu: bool, - settings: SlurmStepOperatorSettings, - ) -> 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. - - Args: - image: Fully-qualified image reference to run. - entrypoint_command: The ZenML entrypoint command. - 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). - - 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) - runtime = self.config.container_runtime - - if runtime in ( - SlurmContainerRuntime.APPTAINER, - SlurmContainerRuntime.SINGULARITY, - ): - binary = runtime.value # "apptainer" or "singularity" - parts = [binary, "exec"] - 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"{' '.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"{' '.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"{source_env}{' '.join(srun)} {entrypoint}" - - def _build_sbatch_script( - self, - info: "StepRunInfo", - container_command: str, - run_dir: str, - ) -> str: - """Render the sbatch job script for a step. - - Args: - info: Information about the step run. - container_command: The shell snippet that runs the step container. - run_dir: The per-run staging directory on the cluster. - - Returns: - The job script content. - """ - settings = cast(SlurmStepOperatorSettings, self.get_settings(info)) - resources = info.config.resource_settings - - directives = [ - f"#SBATCH --job-name={self._job_name(info.step_run_id)}", - 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={int(resources.cpu_count)}" - ) - if memory_gb := resources.get_memory(unit="GB"): - directives.append(f"#SBATCH --mem={int(memory_gb)}G") - 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 - # `get_status` 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 orchestrator dies. - # `set -e` aborts the job (with the failing code captured) if the - # container runtime itself fails. - return f"""#!/bin/bash -{chr(10).join(directives)} - -trap 'ec=$?; echo "$ec" > {run_dir}/{_EXIT_CODE_FILE}; rm -f {run_dir}/{_ENV_FILE}' EXIT -set -eo pipefail - -{container_command} -""" - def submit( self, info: "StepRunInfo", @@ -355,12 +171,13 @@ def submit( """ 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}" + 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_env_for_docker_env_file(environment) - container_command = self._build_container_command( + container_command = build_container_command( + runtime=self.config.container_runtime, image=image, entrypoint_command=entrypoint_command, env_file=env_file, @@ -368,28 +185,17 @@ def submit( use_gpu=use_gpu, settings=settings, ) - script = self._build_sbatch_script( - info=info, container_command=container_command, run_dir=run_dir + 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, ) - client, runner = self._get_client() + client = build_slurm_client(self.config) try: - # Create the run directory owner-only: it holds the env file (with - # credentials), the job script, and the job output. `-m 700` sets - # the mode atomically on creation, so there is no window where the - # directory is readable by other cluster users before a chmod. - result = runner.run(f"mkdir -m 700 -p {shlex.quote(run_dir)}") - if result.exit_code != 0: - raise RuntimeError( - f"Failed to create run directory `{run_dir}` on the " - f"cluster: {result.stderr.strip()}" - ) - - runner.put_text(env_file, env_content, mode=0o600) - script_path = f"{run_dir}/job.sh" - runner.put_text(script_path, script, mode=0o700) - - job_id = client.submit(script_path) + job_id = stage_and_submit(client, run_dir, env_content, script) logger.info( "Submitted step `%s` as Slurm job `%s` (job name `%s`).", info.pipeline_step_name, @@ -397,7 +203,7 @@ def submit( self._job_name(info.step_run_id), ) finally: - runner.close() + client.runner.close() def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: """Get the status of a step run from the Slurm queue. @@ -410,7 +216,8 @@ def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: """ job_name = self._job_name(step_run.id) run_dir = self._run_dir(step_run.id) - client, runner = self._get_client() + client = build_slurm_client(self.config) + runner = client.runner try: state = client.get_job_state(job_name) if state in PENDING_STATES: @@ -428,7 +235,7 @@ def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: # job script's EXIT trap holds the outcome. try: exit_code = runner.read_text( - f"{run_dir}/{_EXIT_CODE_FILE}" + f"{run_dir}/{EXIT_CODE_FILE}" ).strip() except Exception: # No queue entry and no sentinel: the job was cancelled @@ -449,11 +256,11 @@ def cancel(self, step_run: "StepRunResponse") -> None: Args: step_run: The step run to cancel. """ - client, runner = self._get_client() + client = build_slurm_client(self.config) try: client.cancel(self._job_name(step_run.id)) finally: - runner.close() + client.runner.close() def cleanup_step_submission(self, step_run: "StepRunResponse") -> None: """Remove the per-run directory after the step has finished. @@ -465,15 +272,15 @@ def cleanup_step_submission(self, step_run: "StepRunResponse") -> None: step_run: The finished step run. """ run_dir = self._run_dir(step_run.id) - client, runner = self._get_client() + client = build_slurm_client(self.config) try: - runner.run(f"rm -rf {shlex.quote(run_dir)}") + 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: - runner.close() + client.runner.close() def _log_failure_output( self, runner: SlurmCommandRunner, run_dir: str @@ -484,9 +291,7 @@ def _log_failure_output( 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}" - ) + 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 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..61a630d80c2 --- /dev/null +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -0,0 +1,359 @@ +# 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 UUID, uuid4 + +import pytest + +from zenml.config.resource_settings import ResourceSettings +from zenml.enums import 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, +) +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]) -> SimpleNamespace: + """Build a PipelineSnapshotResponse stand-in. + + Args: + steps: The step configurations keyed by invocation id. + + 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"), + schedule=None, + ) + + +# --- 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 + + +# --- 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() + + +# --- topological ordering ----------------------------------------------------- + + +def test_sorted_steps_is_topological(): + """Upstream steps are submitted before their dependents.""" + # Declared out of order: a dependent appears before its parent. + steps = { + "evaluate": _step(["train"]), + "train": _step(["load"]), + "load": _step([]), + } + order = SlurmOrchestrator._sorted_step_names(steps) + assert order.index("load") < order.index("train") + assert order.index("train") < order.index("evaluate") + + +# --- 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 _submit_two_step_dag(monkeypatch) -> FakeRunner: + """Submit a load -> train DAG and return the fake host. + + Args: + monkeypatch: The pytest monkeypatch fixture. + + Returns: + The fake runner recording everything the submission did. + """ + 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)}) + op.submit_pipeline( + snapshot=snapshot, + stack=None, + base_environment={}, + step_environments={ + "load": {"ZENML_STORE_API_KEY": SECRET_TOKEN}, + "train": {"ZENML_STORE_API_KEY": SECRET_TOKEN}, + }, + ) + return runner + + +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) == 2 + # 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] + + +def test_submit_injects_run_id_into_step_env(monkeypatch): + """Each step's env file carries the orchestrator run id.""" + runner = _submit_two_step_dag(monkeypatch) + env_files = [c for p, c in runner.files.items() if p.endswith("/env")] + assert env_files + assert all(ENV_ZENML_SLURM_RUN_ID in content for content in env_files) + + +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=None, + base_environment={}, + step_environments={"load": {}, "train": {}}, + ) + assert any(c.startswith("scancel 1000") for c 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 index 8ac263342d4..37595b1d599 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -37,6 +37,10 @@ SlurmClient, SlurmCommandRunner, ) +from zenml.integrations.slurm.slurm_job import ( + build_container_command, + build_sbatch_script, +) from zenml.integrations.slurm.step_operators import SlurmStepOperator SECRET_TOKEN = "super-secret-zenml-token" @@ -189,6 +193,12 @@ def test_default_container_runtime_is_apptainer(): 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 ---------------------------------------------------------------- @@ -245,8 +255,8 @@ def _container_cmd(runtime: str, gpu: bool = False) -> str: config = SlurmStepOperatorConfig( transport="local", workdir="/s", container_runtime=runtime ) - op = _build_operator(config) - return op._build_container_command( + return build_container_command( + runtime=config.container_runtime, image=IMAGE, entrypoint_command=["python", "-m", "zenml.entrypoint"], env_file="/s/zenml-x/env", @@ -292,17 +302,15 @@ def test_pyxis_command_passes_env_names_not_values(): def test_sbatch_script_directives_and_scrub(): """The script carries directives, fails fast, and scrubs the env file.""" - op = _build_operator( - SlurmStepOperatorConfig(transport="local", workdir="/s") - ) - op.get_settings = lambda _i: SlurmStepOperatorSettings( - partition="gpu", time_limit="2:00:00" - ) sid = uuid4() - script = op._build_sbatch_script( - info=_fake_info(sid, gpu=2), - container_command="RUN_THE_CONTAINER", + 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 @@ -318,12 +326,16 @@ def test_sbatch_script_directives_and_scrub(): # --- submit: security-sensitive handling -------------------------------------- -def test_submit_writes_secret_env_file_owner_only(): +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() - op._get_client = lambda: (SlurmClient(runner), runner) + monkeypatch.setattr( + "zenml.integrations.slurm.step_operators.slurm_step_operator" + ".build_slurm_client", + lambda config: SlurmClient(runner), + ) sid = uuid4() op.submit( @@ -360,7 +372,9 @@ def op_and_runner(monkeypatch): op = _build_operator() runner = FakeRunner() monkeypatch.setattr( - op, "_get_client", lambda: (SlurmClient(runner), runner) + "zenml.integrations.slurm.step_operators.slurm_step_operator" + ".build_slurm_client", + lambda config: SlurmClient(runner), ) return op, runner From 4f9c92961ed07ac73f7dc27a9f030a338e043c91 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 14:56:53 -0700 Subject: [PATCH 07/22] Derive the orchestrator run id from the placeholder run, not the snapshot The Slurm orchestrator was modelled on the HyperAI orchestrator, which is deprecated. Re-checked against the current orchestrators (the generic `ssh` orchestrator that HyperAI defers to, plus Kubernetes/Vertex/SageMaker/AzureML) and corrected the run-id handling to match them: - Use `placeholder_run.id` as the orchestrator run id, not `snapshot.id`. A snapshot can be executed many times, so `snapshot.id` is not unique per run: two runs of the same pipeline would collide on `get_orchestrator_run_id`, the staging directory, and the job names. Every current orchestrator derives the run id from the placeholder run; only the deprecated HyperAI used the snapshot id. - Drop the local `os.environ[...] = run_id` write. No current orchestrator does this; the run id only needs to reach `get_orchestrator_run_id` inside each step's job, which happens via the per-step environment injection that was already in place. (Another HyperAI-ism.) - Reject scheduled pipelines with a clear error, matching the `ssh` orchestrator, instead of silently ignoring the schedule. Also drop now-unused `Tuple` imports left over from the earlier client-factory refactor. Co-Authored-By: Claude Fable 5 --- .../slurm/orchestrators/slurm_orchestrator.py | 32 +++++--- .../step_operators/slurm_step_operator.py | 2 +- .../slurm/test_slurm_orchestrator.py | 75 ++++++++++++++++--- 3 files changed, 87 insertions(+), 22 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 4402a6c9e3f..18cad0d3345 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -25,7 +25,7 @@ """ import os -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast +from typing import TYPE_CHECKING, Dict, List, Optional, cast from zenml.config.base_settings import BaseSettings from zenml.entrypoints.step_entrypoint_configuration import ( @@ -49,8 +49,7 @@ ) from zenml.integrations.ssh.utils import serialize_env_for_docker_env_file from zenml.logger import get_logger -from zenml.orchestrators import ContainerizedOrchestrator -from zenml.orchestrators.base_orchestrator import SubmissionResult +from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult from zenml.orchestrators.topsort import topsorted_layers from zenml.stack import StackValidator @@ -187,21 +186,32 @@ def submit_pipeline( Args: snapshot: The pipeline snapshot to submit. stack: The stack the pipeline will run on. - base_environment: Base environment shared by all steps. + base_environment: Base environment shared by all steps. Unused: the + per-step environments already include it. step_environments: Environment variables per step. - placeholder_run: An optional placeholder run for the snapshot. + placeholder_run: The placeholder run for the pipeline. Returns: None, since the pipeline is submitted detached. Raises: - RuntimeError: If a step directory cannot be created on the cluster. + RuntimeError: If the pipeline has a schedule (not supported), or if + a step directory cannot be created on the cluster. """ - run_id = str(snapshot.id) - # Set locally too, so get_orchestrator_run_id works if the framework - # queries it on the submitting machine; the authoritative value is - # injected into each step's job environment below. - os.environ[ENV_ZENML_SLURM_RUN_ID] = run_id + 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) diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index da98d3514c6..7fd7535e57e 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -28,7 +28,7 @@ """ import shlex -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast +from typing import TYPE_CHECKING, Dict, List, Optional, cast from uuid import UUID from zenml.config.base_settings import BaseSettings diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index 61a630d80c2..debb4d74adf 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -243,14 +243,24 @@ def _use_fake_client(monkeypatch, runner: FakeRunner) -> None: ) -def _submit_two_step_dag(monkeypatch) -> FakeRunner: - """Submit a load -> train DAG and return the fake host. +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 _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. + 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() @@ -258,6 +268,7 @@ def _submit_two_step_dag(monkeypatch) -> FakeRunner: _use_fake_client(monkeypatch, runner) snapshot = _snapshot({"load": _step([]), "train": _step(["load"], gpu=1)}) + placeholder = _placeholder_run() op.submit_pipeline( snapshot=snapshot, stack=None, @@ -266,13 +277,14 @@ def _submit_two_step_dag(monkeypatch) -> FakeRunner: "load": {"ZENML_STORE_API_KEY": SECRET_TOKEN}, "train": {"ZENML_STORE_API_KEY": SECRET_TOKEN}, }, + placeholder_run=placeholder, ) - return runner + return runner, str(placeholder.id) 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) + runner, _ = _submit_two_step_dag(monkeypatch) sbatch = [c for c in runner.commands if c.startswith("sbatch")] assert len(sbatch) == 2 # load submitted first (job id 1000, no dependency), train second. @@ -281,17 +293,59 @@ def test_submit_wires_dependencies_in_topological_order(monkeypatch): assert "--kill-on-invalid-dep=yes" in sbatch[1] -def test_submit_injects_run_id_into_step_env(monkeypatch): - """Each step's env file carries the orchestrator run id.""" - runner = _submit_two_step_dag(monkeypatch) +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 - assert all(ENV_ZENML_SLURM_RUN_ID in content for content in 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=None, + 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=None, + base_environment={}, + step_environments={"load": {}}, + 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) + 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")} @@ -325,6 +379,7 @@ def put_text(self, remote_path, content, mode=0o600): stack=None, base_environment={}, step_environments={"load": {}, "train": {}}, + placeholder_run=_placeholder_run(), ) assert any(c.startswith("scancel 1000") for c in runner.commands) From fee6f21f8e1d0fee239437638b20b2fe6ebe1816 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 15:50:26 -0700 Subject: [PATCH 08/22] Harden Slurm failure reporting, credential handling, and registry auth Addresses an external review of the detached failure/status and credential lifecycle, which were not production-safe. Status and cancellation: - Persist the Slurm job id as step-run metadata and query/cancel by exact id instead of reconstructing the job name. - Add orchestrator `fetch_status` so a detached run is reconciled from its Slurm jobs, catching failures that happen before the ZenML entrypoint starts (bad image, container runtime error) which the steps can never self-report. - Reconcile job state by id via `squeue --jobs` with a `scontrol show job` fallback (controller memory), closing the window where a finished job has left the queue before its shared-filesystem sentinel is visible. Treat that gap as still-running rather than terminally cancelled. Credentials: - Add per-runtime private registry authentication (Apptainer/Singularity `*_DOCKER_USERNAME/PASSWORD`, enroot `.credentials`, Docker `config.json`) written as owner-only files, so private ECR/GCR/ACR/Docker Hub images pull without exposing passwords on any command line. - Roll back staged files (`rm -rf` the run dir) on any staging/submission failure, and submit an `afterany` cleanup job depending on all step jobs to scrub credential files even for jobs cancelled before their EXIT trap runs. Runtime safety: - Shell-escape and validate environment variable names; the pyxis path now `export`s quoted values instead of sourcing a Docker-format file, and Apptainer/Singularity run with `--no-eval`. - Round CPU requests up and emit memory in MB so fractional/sub-GB resource settings produce valid sbatch directives. Co-Authored-By: Claude Opus 4.8 --- src/zenml/integrations/slurm/flavors/base.py | 3 +- .../slurm/orchestrators/slurm_orchestrator.py | 269 +++++++++++++++++- src/zenml/integrations/slurm/slurm_client.py | 73 +++-- src/zenml/integrations/slurm/slurm_job.py | 202 ++++++++++++- .../step_operators/slurm_step_operator.py | 122 +++++++- .../slurm/test_slurm_orchestrator.py | 81 +++++- .../slurm/test_slurm_step_operator.py | 137 +++++++-- 7 files changed, 809 insertions(+), 78 deletions(-) diff --git a/src/zenml/integrations/slurm/flavors/base.py b/src/zenml/integrations/slurm/flavors/base.py index 0f93503bbcf..9a89e43e404 100644 --- a/src/zenml/integrations/slurm/flavors/base.py +++ b/src/zenml/integrations/slurm/flavors/base.py @@ -117,7 +117,8 @@ class SlurmConnectionConfig(SSHConnectionConfigMixin): default=SlurmContainerRuntime.APPTAINER, description="Runtime used to run the step's Docker image on the " "compute node: 'apptainer' or 'singularity' (rootless, pulls " - "`docker://` images; the safe HPC default), 'pyxis' (NVIDIA " + "`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'", ) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 18cad0d3345..6c08cd0eb68 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -25,29 +25,38 @@ """ import os -from typing import TYPE_CHECKING, Dict, List, Optional, cast +import shlex +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast from zenml.config.base_settings import BaseSettings from zenml.entrypoints.step_entrypoint_configuration import ( StepEntrypointConfiguration, ) +from zenml.enums import 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, ENV_FILE, + EXIT_CODE_FILE, + OUTPUT_FILE, REQUIRED_COMPONENTS, + SCRIPT_FILE, build_container_command, + build_registry_auth, build_sbatch_script, + serialize_environment, stage_and_submit, validate_remote_stack, ) -from zenml.integrations.ssh.utils import serialize_env_for_docker_env_file from zenml.logger import get_logger from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult from zenml.orchestrators.topsort import topsorted_layers @@ -61,6 +70,9 @@ 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" +_CLEANUP_COMPLETE_FILE = "cleanup_complete" class SlurmOrchestrator(ContainerizedOrchestrator): @@ -216,10 +228,17 @@ def submit_pipeline( 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: for step_name in self._sorted_step_names(steps): step = steps[step_name] - job_id = self._submit_step( + job_id, step_sensitive_paths = self._submit_step( client=client, snapshot=snapshot, run_id=run_id, @@ -231,18 +250,43 @@ def submit_pipeline( 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: - client.runner.run(f"scancel {' '.join(submitted.values())}") + 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() @@ -252,7 +296,13 @@ def submit_pipeline( snapshot.pipeline_configuration.name, len(submitted), ) - return None + assert cleanup_job_id is not None + return SubmissionResult( + metadata={ + SLURM_JOB_IDS_METADATA_KEY: submitted, + SLURM_CLEANUP_JOB_ID_METADATA_KEY: cleanup_job_id, + } + ) def _submit_step( self, @@ -263,7 +313,9 @@ def _submit_step( step: "Step", step_environment: Dict[str, str], upstream_job_ids: List[str], - ) -> 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: @@ -274,9 +326,11 @@ def _submit_step( 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. + The submitted Slurm job id and credential-bearing paths. Raises: RuntimeError: If the step directory cannot be created. @@ -289,7 +343,15 @@ def _submit_step( environment = step_environment.copy() environment[ENV_ZENML_SLURM_RUN_ID] = run_id - env_content = serialize_env_for_docker_env_file(environment) + env_content = serialize_environment( + 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, + ) entrypoint_command = ( StepEntrypointConfiguration.get_entrypoint_command() @@ -305,6 +367,7 @@ def _submit_step( env_keys=sorted(environment), use_gpu=use_gpu, settings=settings, + registry_auth=registry_auth, ) script = build_sbatch_script( job_name=self._job_name(run_id, step_name), @@ -312,12 +375,200 @@ def _submit_step( container_command=container_command, resources=step.config.resource_settings, settings=settings, + sensitive_paths=registry_auth.sensitive_paths, ) - return stage_and_submit( + job_id = stage_and_submit( client, run_dir=run_dir, env_content=env_content, script=script, dependencies=upstream_job_ids, + extra_files=registry_auth.files, + ) + return job_id, [env_file, *registry_auth.sensitive_paths] + + 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}") + + 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): + logger.warning("No Slurm job metadata found for run `%s`.", run.id) + return None, None + + 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 + + statuses = { + step_name: self._get_job_status( + client=client, + job_id=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() + } + finally: + client.runner.close() + + known_statuses = [status for status in statuses.values() if status] + pipeline_status: Optional[ExecutionStatus] = None + if not run.status.is_finished: + if any( + status in {ExecutionStatus.FAILED, ExecutionStatus.CANCELLED} + for status in known_statuses + ): + pipeline_status = ExecutionStatus.FAILED + 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 + + 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 _get_job_status( + client: SlurmClient, + job_id: 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. + job_id: Slurm job ID. + 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. + """ + state = client.get_job_state(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: + 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 ) diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index 758cc546c13..0a3178922c9 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -25,11 +25,12 @@ 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, List, Optional +from typing import TYPE_CHECKING, List, Literal, Optional if TYPE_CHECKING: from zenml.integrations.slurm.flavors.base import SlurmConnectionConfig @@ -138,9 +139,12 @@ def put_text( """ import os - with open(remote_path, "w") as f: + 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) - os.chmod(remote_path, mode) def read_text(self, remote_path: str) -> str: """Read a text file. @@ -231,7 +235,10 @@ def __init__(self, runner: SlurmCommandRunner) -> None: self.runner = runner def submit( - self, script_path: str, dependencies: Optional[List[str]] = None + 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. @@ -242,6 +249,7 @@ def submit( 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. @@ -251,10 +259,15 @@ def submit( """ 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=afterok:{':'.join(dependencies)}" - " --kill-on-invalid-dep=yes" + 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( @@ -262,13 +275,18 @@ def submit( f"{result.stderr.strip() or result.stdout.strip()}" ) # --parsable prints `jobid[;cluster]` - return result.stdout.strip().split(";")[0] + 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_state(self, job_name: str) -> Optional[str]: - """Get the queue state of a job by its name. + def get_job_state(self, job_id: str) -> Optional[str]: + """Get the queue state of a job by its ID. Args: - job_name: The Slurm job name to look up. + job_id: The Slurm job ID to look up. Returns: The Slurm job state (e.g. ``PENDING``, ``RUNNING``) or ``None`` @@ -280,27 +298,50 @@ def get_job_state(self, job_name: str) -> Optional[str]: Raises: RuntimeError: If ``squeue`` fails. """ + if not re.fullmatch(r"[0-9]+", job_id): + raise ValueError("Slurm job IDs must be numeric.") result = self.runner.run( - f"squeue --noheader --format=%T --name {shlex.quote(job_name)}" + f"squeue --noheader --format=%T --jobs {shlex.quote(job_id)}" ) if result.exit_code != 0: + if "invalid job id" in result.stderr.lower(): + return None raise RuntimeError( f"`squeue` failed with exit code {result.exit_code}: " f"{result.stderr.strip()}" ) state = result.stdout.strip().splitlines() - return state[0].strip() if state else None + if state: + return state[0].strip() + + # Completed jobs can disappear from squeue before a shared-filesystem + # sentinel becomes visible. Slurm keeps them in controller memory for + # MinJobAge, so scontrol closes that race without requiring sacct. + result = self.runner.run( + f"scontrol show job --oneliner {shlex.quote(job_id)}" + ) + if result.exit_code != 0: + if "invalid job id" in result.stderr.lower(): + return None + raise RuntimeError( + f"`scontrol show job` failed with exit code " + f"{result.exit_code}: {result.stderr.strip()}" + ) + match = re.search(r"(?:^|\s)JobState=([^\s]+)", result.stdout) + return match.group(1) if match else None - def cancel(self, job_name: str) -> None: - """Cancel a job by its name. + def cancel(self, job_id: str) -> None: + """Cancel a job by its ID. Args: - job_name: The Slurm job name to cancel. + job_id: The Slurm job ID to cancel. Raises: RuntimeError: If ``scancel`` fails. """ - result = self.runner.run(f"scancel --name {shlex.quote(job_name)}") + if not re.fullmatch(r"[0-9]+", job_id): + raise ValueError("Slurm job IDs must be numeric.") + result = self.runner.run(f"scancel {shlex.quote(job_id)}") if result.exit_code != 0: raise RuntimeError( f"`scancel` failed with exit code {result.exit_code}: " diff --git a/src/zenml/integrations/slurm/slurm_job.py b/src/zenml/integrations/slurm/slurm_job.py index b7907d204fe..6d3293dd65f 100644 --- a/src/zenml/integrations/slurm/slurm_job.py +++ b/src/zenml/integrations/slurm/slurm_job.py @@ -26,8 +26,14 @@ go through ``shlex.quote``. """ +import base64 +import json +import math +import posixpath +import re import shlex -from typing import TYPE_CHECKING, List, Optional, Tuple +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 ( @@ -45,6 +51,12 @@ 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 = { @@ -53,6 +65,126 @@ } +@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. + """ + 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. @@ -93,6 +225,7 @@ def build_container_command( 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. @@ -108,19 +241,21 @@ def build_container_command( 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"] + parts = [binary, "exec", "--no-eval"] if use_gpu: parts.append("--nv") parts += ["--env-file", shlex.quote(env_file)] @@ -129,7 +264,7 @@ def build_container_command( if extra: parts.append(extra) parts.append(shlex.quote(f"docker://{image}")) - return f"{' '.join(parts)} {entrypoint}" + return f"{shell_setup}{' '.join(parts)} {entrypoint}" if runtime == SlurmContainerRuntime.DOCKER: parts = ["docker", "run", "--rm"] @@ -141,7 +276,7 @@ def build_container_command( if extra: parts.append(extra) parts.append(shlex.quote(image)) - return f"{' '.join(parts)} {entrypoint}" + 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 @@ -158,7 +293,7 @@ def build_container_command( if extra: srun.append(extra) source_env = f"set -a\nsource {shlex.quote(env_file)}\nset +a\n" - return f"{source_env}{' '.join(srun)} {entrypoint}" + return f"{shell_setup}{source_env}{' '.join(srun)} {entrypoint}" def build_sbatch_script( @@ -167,6 +302,7 @@ def build_sbatch_script( container_command: str, resources: "ResourceSettings", settings: SlurmJobSettings, + sensitive_paths: Optional[List[str]] = None, ) -> str: """Render the sbatch job script for a step. @@ -176,6 +312,7 @@ def build_sbatch_script( 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. @@ -194,10 +331,10 @@ def build_sbatch_script( directives.append(f"#SBATCH --qos={settings.qos}") if resources.cpu_count: directives.append( - f"#SBATCH --cpus-per-task={int(resources.cpu_count)}" + f"#SBATCH --cpus-per-task={math.ceil(resources.cpu_count)}" ) - if memory_gb := resources.get_memory(unit="GB"): - directives.append(f"#SBATCH --mem={int(memory_gb)}G") + 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: @@ -208,10 +345,23 @@ def build_sbatch_script( # 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)} -trap 'ec=$?; echo "$ec" > {run_dir}/{EXIT_CODE_FILE}; rm -f {run_dir}/{ENV_FILE}' EXIT +cleanup() {{ + ec=$? + trap - EXIT + echo "$ec" > {exit_code_path} + {cleanup_command} + exit "$ec" +}} +trap cleanup EXIT set -eo pipefail {container_command} @@ -224,6 +374,7 @@ def stage_and_submit( 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. @@ -238,6 +389,7 @@ def stage_and_submit( 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. @@ -246,15 +398,35 @@ def stage_and_submit( RuntimeError: If the run directory cannot be created on the cluster. """ runner = client.runner - result = runner.run(f"mkdir -m 700 -p {shlex.quote(run_dir)}") + 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()}" ) - runner.put_text(f"{run_dir}/{ENV_FILE}", env_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) + 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/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index 7fd7535e57e..1bd16d422c4 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -20,11 +20,10 @@ login node, or via a local ``sbatch`` when the client already runs on the cluster. -Submission is stateless: the Slurm job is named after the step run id, so -status lookups and cancellation work by job name without persisting job ids. -While a job is in the queue, its state comes from ``squeue``; once it leaves -the queue, a sentinel exit-code file written by the job script disambiguates -success from failure, so Slurm accounting (``sacct``) is never required. +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 @@ -45,27 +44,35 @@ 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.integrations.ssh.utils import serialize_env_for_docker_env_file 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" _JOB_NAME_PREFIX = "zenml" @@ -175,7 +182,19 @@ def submit( settings = cast(SlurmStepOperatorSettings, self.get_settings(info)) use_gpu = bool(info.config.resource_settings.gpu_count) - env_content = serialize_env_for_docker_env_file(environment) + 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, @@ -184,6 +203,7 @@ def submit( 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), @@ -191,17 +211,45 @@ def submit( 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) + 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() @@ -214,16 +262,23 @@ def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: Returns: The execution status of the step run. """ - job_name = self._job_name(step_run.id) + 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(job_name) + 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. @@ -233,14 +288,22 @@ def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: # 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: - # No queue entry and no sentinel: the job was cancelled - # before its trap could run, or never started. - return ExecutionStatus.CANCELLED + # Slurm queue updates and shared-filesystem writes are not + # atomic. Keep polling until either source proves a terminal + # outcome instead of misreporting a transient gap as cancelled. + return ExecutionStatus.QUEUED if exit_code == "0": return ExecutionStatus.COMPLETED @@ -258,7 +321,18 @@ def cancel(self, step_run: "StepRunResponse") -> None: """ client = build_slurm_client(self.config) try: - client.cancel(self._job_name(step_run.id)) + 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() @@ -296,3 +370,23 @@ def _log_failure_output( 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 index debb4d74adf..19cde6b3e9f 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -25,7 +25,7 @@ import pytest from zenml.config.resource_settings import ResourceSettings -from zenml.enums import StackComponentType +from zenml.enums import ExecutionStatus, StackComponentType from zenml.integrations.slurm.flavors import ( SlurmOrchestratorConfig, SlurmOrchestratorFlavor, @@ -34,6 +34,7 @@ from zenml.integrations.slurm.orchestrators import SlurmOrchestrator from zenml.integrations.slurm.orchestrators.slurm_orchestrator import ( ENV_ZENML_SLURM_RUN_ID, + SLURM_JOB_IDS_METADATA_KEY, ) from zenml.integrations.slurm.slurm_client import ( CommandResult, @@ -252,6 +253,25 @@ def _placeholder_run() -> SimpleNamespace: 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 _submit_two_step_dag(monkeypatch) -> "tuple[FakeRunner, str]": """Submit a load -> train DAG and return the fake host and run id. @@ -271,7 +291,7 @@ def _submit_two_step_dag(monkeypatch) -> "tuple[FakeRunner, str]": placeholder = _placeholder_run() op.submit_pipeline( snapshot=snapshot, - stack=None, + stack=_stack(), base_environment={}, step_environments={ "load": {"ZENML_STORE_API_KEY": SECRET_TOKEN}, @@ -286,11 +306,20 @@ 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) == 2 + 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_submit_injects_placeholder_run_id_into_step_env(monkeypatch): @@ -320,7 +349,7 @@ def test_submit_requires_a_placeholder_run(monkeypatch): with pytest.raises(AssertionError): op.submit_pipeline( snapshot=_snapshot({"load": _step([])}), - stack=None, + stack=_stack(), base_environment={}, step_environments={"load": {}}, placeholder_run=None, @@ -336,7 +365,7 @@ def test_submit_rejects_scheduled_pipelines(monkeypatch): with pytest.raises(RuntimeError, match="does not support scheduled"): op.submit_pipeline( snapshot=snapshot, - stack=None, + stack=_stack(), base_environment={}, step_environments={"load": {}}, placeholder_run=_placeholder_run(), @@ -376,12 +405,52 @@ def put_text(self, remote_path, content, mode=0o600): with pytest.raises(RuntimeError): op.submit_pipeline( snapshot=snapshot, - stack=None, + 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 = SimpleNamespace( + id=run_id, + status=ExecutionStatus.PROVISIONING, + 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 = SimpleNamespace( + id=run_id, + status=ExecutionStatus.PROVISIONING, + run_metadata={SLURM_JOB_IDS_METADATA_KEY: {"train": "1001"}}, + ) + + pipeline_status, _ = op.fetch_status(run) + + assert pipeline_status is ExecutionStatus.FAILED # --- validator ---------------------------------------------------------------- diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index 37595b1d599..83db722b28d 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -39,9 +39,14 @@ ) 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, +) SECRET_TOKEN = "super-secret-zenml-token" IMAGE = "registry.example.com/zenml:abc123" @@ -159,7 +164,7 @@ def _fake_info( cpu_count=4, gpu_count=gpu or None, memory="16GB" ) ), - step_run=None, + step_run=SimpleNamespace(run_metadata={}), get_image=lambda key: IMAGE, ) @@ -239,13 +244,13 @@ def test_client_submit_parses_job_id(): assert runner.commands == ["sbatch --parsable /s/job.sh"] -def test_client_state_and_cancel_use_job_name(): - """Status lookup and cancellation address the job by name.""" +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("zenml-abc") == "RUNNING" - client.cancel("zenml-abc") - assert runner.commands[-1] == "scancel --name zenml-abc" + assert client.get_job_state("12345") == "RUNNING" + client.cancel("12345") + assert runner.commands[-1] == "scancel 12345" # --- container command rendering (per runtime) -------------------------------- @@ -271,7 +276,7 @@ def _container_cmd(runtime: str, gpu: bool = False) -> str: 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") + 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 @@ -297,6 +302,41 @@ def test_pyxis_command_passes_env_names_not_values(): 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 ------------------------------------------------------------ @@ -315,12 +355,25 @@ def test_sbatch_script_directives_and_scrub(): 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=16G" 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 -f /s/zenml-x/env" 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 -------------------------------------- @@ -336,10 +389,26 @@ def test_submit_writes_secret_env_file_owner_only(monkeypatch): ".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=_fake_info(sid), + info=info, entrypoint_command=["python", "-m", "zenml.entrypoint"], environment={ "ZENML_STORE_API_KEY": SECRET_TOKEN, @@ -351,16 +420,22 @@ def test_submit_writes_secret_env_file_owner_only(monkeypatch): # 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, atomically - assert any("mkdir -m 700" in c for c in runner.commands) + # 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 ----------------------------------------------------------- @@ -392,13 +467,23 @@ 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())) is expected + 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()) + 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 @@ -406,14 +491,32 @@ def test_get_status_reads_sentinel_after_queue(op_and_runner): assert op.get_status(step_run) is ExecutionStatus.FAILED -def test_get_status_without_queue_or_sentinel_is_cancelled(op_and_runner): - """No queue entry and no sentinel means the job was torn down early.""" +def test_get_status_without_queue_or_sentinel_is_non_terminal(op_and_runner): + """A queue/filesystem visibility gap remains non-terminal.""" op, runner = op_and_runner assert ( - op.get_status(SimpleNamespace(id=uuid4())) is ExecutionStatus.CANCELLED + op.get_status( + SimpleNamespace( + id=uuid4(), + run_metadata={SLURM_JOB_ID_METADATA_KEY: "12345"}, + ) + ) + is ExecutionStatus.QUEUED ) +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 From 8e514ff3ae6805af168cbb46c3f0ab10e03bacaf Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 16:52:00 -0700 Subject: [PATCH 09/22] Fix pydoclint Raises sections in the Slurm integration The review-hardening changes added `ValueError` input guards and `except Exception: ... raise` cleanup re-raises without updating the docstring `Raises:` sections, and moved a `raise` out of `_submit_step` into `stage_and_submit`, which the CI `pydoclint` check flagged (DOC501/DOC502/ DOC503). Align every `Raises:` section with its function body: - Document the numeric-id `ValueError` on `SlurmClient.submit`/`get_job_state`/ `cancel` and the enroot `ValueError` on `build_registry_auth`. - Document the re-raised `Exception` on `stage_and_submit`, `SlurmOrchestrator.submit_pipeline`, and `SlurmStepOperator.submit`. - Drop the stale `Raises:` from `_submit_step`, whose raise now lives in `stage_and_submit`. Co-Authored-By: Claude Opus 4.8 --- .../slurm/orchestrators/slurm_orchestrator.py | 9 ++++----- src/zenml/integrations/slurm/slurm_client.py | 3 +++ src/zenml/integrations/slurm/slurm_job.py | 6 ++++++ .../slurm/step_operators/slurm_step_operator.py | 4 ++-- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 6c08cd0eb68..bad1761379c 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -207,8 +207,10 @@ def submit_pipeline( None, since the pipeline is submitted detached. Raises: - RuntimeError: If the pipeline has a schedule (not supported), or if - a step directory cannot be created on the cluster. + 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( @@ -331,9 +333,6 @@ def _submit_step( Returns: The submitted Slurm job id and credential-bearing paths. - - Raises: - RuntimeError: If the step directory cannot be created. """ image = self.get_image(snapshot=snapshot, step_name=step_name) settings = cast(SlurmOrchestratorSettings, self.get_settings(step)) diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index 0a3178922c9..4c131e07c24 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -255,6 +255,7 @@ def submit( The Slurm job id. Raises: + ValueError: If any dependency job id is not numeric. RuntimeError: If the submission fails. """ command = "sbatch --parsable" @@ -296,6 +297,7 @@ def get_job_state(self, job_id: str) -> Optional[str]: via the exit-code sentinel file. Raises: + ValueError: If the job id is not numeric. RuntimeError: If ``squeue`` fails. """ if not re.fullmatch(r"[0-9]+", job_id): @@ -337,6 +339,7 @@ def cancel(self, job_id: str) -> None: job_id: The Slurm job ID to cancel. Raises: + ValueError: If the job id is not numeric. RuntimeError: If ``scancel`` fails. """ if not re.fullmatch(r"[0-9]+", job_id): diff --git a/src/zenml/integrations/slurm/slurm_job.py b/src/zenml/integrations/slurm/slurm_job.py index 6d3293dd65f..3f087152422 100644 --- a/src/zenml/integrations/slurm/slurm_job.py +++ b/src/zenml/integrations/slurm/slurm_job.py @@ -120,6 +120,10 @@ def build_registry_auth( 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() @@ -396,6 +400,8 @@ def stage_and_submit( 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( diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index 1bd16d422c4..51fb863e385 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -173,8 +173,8 @@ def submit( environment. Raises: - RuntimeError: If the run directory cannot be created on the - cluster. + 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) From 27cfa6412a72ebf08ecd6d8aa8f7381824763294 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 10 Jul 2026 17:39:23 -0700 Subject: [PATCH 10/22] Remove unused UUID import in the Slurm orchestrator tests The CI lint runs ruff's F401/F841 check over the tests too; UUID was imported but only uuid4 is used. Co-Authored-By: Claude Opus 4.8 --- tests/unit/integrations/slurm/test_slurm_orchestrator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index 19cde6b3e9f..75e11b28f96 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -20,7 +20,7 @@ from datetime import datetime from types import SimpleNamespace from typing import Dict, List, Optional -from uuid import UUID, uuid4 +from uuid import uuid4 import pytest From d503952bf30d18f7a7664d5737486b0aea0596f1 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Mon, 13 Jul 2026 10:55:22 +0200 Subject: [PATCH 11/22] Improve Slurm lifecycle reconciliation --- .../slurm/orchestrators/slurm_orchestrator.py | 191 ++++++++++++++++-- src/zenml/integrations/slurm/slurm_client.py | 96 ++++++--- .../step_operators/slurm_step_operator.py | 11 +- .../slurm/test_slurm_orchestrator.py | 161 +++++++++++++++ .../slurm/test_slurm_step_operator.py | 49 ++++- 5 files changed, 455 insertions(+), 53 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index bad1761379c..58c28528915 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast from zenml.config.base_settings import BaseSettings +from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID from zenml.entrypoints.step_entrypoint_configuration import ( StepEntrypointConfiguration, ) @@ -301,11 +302,41 @@ def submit_pipeline( 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]: + """Reject dynamic pipelines explicitly. + + 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: + Never returns successfully. + + Raises: + RuntimeError: Always, since dynamic pipelines need an + orchestration-container mode that the Slurm orchestrator does + not implement yet. + """ + _ = (snapshot, stack, environment, placeholder_run) + raise RuntimeError( + "The Slurm orchestrator does not support dynamic pipelines yet." + ) + def _submit_step( self, client: SlurmClient, @@ -426,6 +457,8 @@ def _submit_cleanup_job( 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)} @@ -486,36 +519,46 @@ def fetch_status( else: cleanup_complete = True + job_states = client.get_job_states(list(job_ids.values())) statuses = { step_name: self._get_job_status( client=client, - job_id=job_id, + 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: + if any( + status + in {ExecutionStatus.FAILED, ExecutionStatus.CANCELLED} + for status in known_statuses + ): + pipeline_status = ExecutionStatus.FAILED + self._cancel_unfinished_jobs( + client=client, + run_id=str(run.id), + job_ids=job_ids, + statuses=statuses, + ) + 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() - known_statuses = [status for status in statuses.values() if status] - pipeline_status: Optional[ExecutionStatus] = None - if not run.status.is_finished: - if any( - status in {ExecutionStatus.FAILED, ExecutionStatus.CANCELLED} - for status in known_statuses - ): - pipeline_status = ExecutionStatus.FAILED - 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 - step_statuses = None if include_steps: step_statuses = { @@ -526,7 +569,7 @@ def fetch_status( @staticmethod def _get_job_status( client: SlurmClient, - job_id: str, + state: Optional[str], run_dir: str, cleanup_complete: bool, ) -> Optional[ExecutionStatus]: @@ -534,14 +577,13 @@ def _get_job_status( Args: client: Slurm client used to inspect the job. - job_id: Slurm job ID. + 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. """ - state = client.get_job_state(job_id) if state in PENDING_STATES: return ExecutionStatus.QUEUED if state in RUNNING_STATES: @@ -571,3 +613,108 @@ def _get_job_status( 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) + if not isinstance(raw_job_ids, dict): + logger.warning("No Slurm job metadata found for run `%s`.", run.id) + return + + job_ids = { + str(name): str(job_id) for name, job_id in raw_job_ids.items() + } + client = build_slurm_client(self.config) + try: + self._cancel_jobs( + client=client, + run_id=str(run.id), + job_ids=job_ids, + ) + 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, + run_id=run_id, + job_ids=jobs_to_cancel, + ) + + def _cancel_jobs( + self, + client: SlurmClient, + run_id: str, + job_ids: Dict[str, str], + ) -> None: + """Cancel Slurm jobs and persist cancellation sentinels. + + Args: + client: Slurm client used to cancel jobs. + run_id: Orchestrator run ID. + job_ids: Mapping of step names to Slurm job IDs. + """ + 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 step_name in cancelled_steps: + try: + client.runner.put_text( + f"{self._run_dir(run_id, step_name)}/{CANCELLED_FILE}", + "1\n", + mode=0o600, + ) + except Exception: + logger.warning( + "Failed to write Slurm cancellation marker for step `%s`.", + step_name, + ) diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index 4c131e07c24..049e1e08d69 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -30,7 +30,7 @@ import subprocess from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, List, Literal, Optional +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Sequence if TYPE_CHECKING: from zenml.integrations.slurm.flavors.base import SlurmConnectionConfig @@ -283,68 +283,118 @@ def submit( ) return job_id - def get_job_state(self, job_id: str) -> Optional[str]: - """Get the queue state of a job by its 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_id: The Slurm job ID to look up. + job_ids: The Slurm job IDs to look up. Returns: - The Slurm job state (e.g. ``PENDING``, ``RUNNING``) or ``None`` - if the job is no longer in the queue. 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. + 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 the job id is not numeric. + ValueError: If any job id is not numeric. RuntimeError: If ``squeue`` fails. """ - if not re.fullmatch(r"[0-9]+", job_id): + 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=%T --jobs {shlex.quote(job_id)}" + 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 None + return states raise RuntimeError( f"`squeue` failed with exit code {result.exit_code}: " f"{result.stderr.strip()}" ) - state = result.stdout.strip().splitlines() - if state: - return state[0].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 squeue before a shared-filesystem # sentinel becomes visible. Slurm keeps them in controller memory for # MinJobAge, so scontrol closes that race without requiring sacct. + missing_jobs_arg = ",".join(missing_job_ids) result = self.runner.run( - f"scontrol show job --oneliner {shlex.quote(job_id)}" + "scontrol show job --oneliner " + f"{shlex.quote(missing_jobs_arg)}" ) if result.exit_code != 0: if "invalid job id" in result.stderr.lower(): - return None + return states raise RuntimeError( f"`scontrol show job` failed with exit code " f"{result.exit_code}: {result.stderr.strip()}" ) - match = re.search(r"(?:^|\s)JobState=([^\s]+)", result.stdout) - return match.group(1) if match else None + for line in result.stdout.strip().splitlines(): + job_id_match = re.search(r"(?:^|\s)JobId=([0-9]+)", line) + state_match = re.search(r"(?:^|\s)JobState=([^\s]+)", line) + if job_id_match and state_match: + job_id = job_id_match.group(1) + if job_id in states: + states[job_id] = state_match.group(1) + + 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 the job id is not numeric. + ValueError: If any job id is not numeric. RuntimeError: If ``scancel`` fails. """ - if not re.fullmatch(r"[0-9]+", job_id): + 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.") - result = self.runner.run(f"scancel {shlex.quote(job_id)}") + 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}: " diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index 51fb863e385..91a05c3afa0 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -300,10 +300,13 @@ def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: f"{run_dir}/{EXIT_CODE_FILE}" ).strip() except Exception: - # Slurm queue updates and shared-filesystem writes are not - # atomic. Keep polling until either source proves a terminal - # outcome instead of misreporting a transient gap as cancelled. - return ExecutionStatus.QUEUED + 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 diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index 75e11b28f96..4c61a8f47e0 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -25,6 +25,7 @@ import pytest from zenml.config.resource_settings import ResourceSettings +from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID from zenml.enums import ExecutionStatus, StackComponentType from zenml.integrations.slurm.flavors import ( SlurmOrchestratorConfig, @@ -34,6 +35,7 @@ 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, ) from zenml.integrations.slurm.slurm_client import ( @@ -302,6 +304,31 @@ def _submit_two_step_dag(monkeypatch) -> "tuple[FakeRunner, str]": 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) @@ -322,6 +349,36 @@ def test_submit_wires_dependencies_in_topological_order(monkeypatch): 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) @@ -372,6 +429,20 @@ def test_submit_rejects_scheduled_pipelines(monkeypatch): ) +def test_submit_dynamic_pipeline_raises(monkeypatch): + """Dynamic pipelines fail loudly instead of being silently ignored.""" + op = _build_orchestrator() + _use_fake_client(monkeypatch, FakeRunner()) + + with pytest.raises(RuntimeError, match="dynamic pipelines"): + op.submit_dynamic_pipeline( + snapshot=_snapshot({"load": _step([])}), + 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) @@ -453,6 +524,96 @@ def test_fetch_status_uses_cleanup_marker_for_never_started_job(monkeypatch): assert pipeline_status is ExecutionStatus.FAILED +def test_fetch_status_batches_and_cancels_siblings(monkeypatch): + """A failed job cancels still-active sibling jobs in one refresh.""" + 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\n" + "1001|RUNNING\n" + "1002|PENDING\n" + ), + stderr="", + ) + return CommandResult(exit_code=0, stdout="", stderr="") + + runner = StatusRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + run = SimpleNamespace( + id=run_id, + status=ExecutionStatus.PROVISIONING, + 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.FAILED + 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 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" + ) + + +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" + ) + + # --- validator ---------------------------------------------------------------- diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index 83db722b28d..6b80ae8bff3 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -84,7 +84,9 @@ def run(self, command: str) -> CommandResult: if command.startswith("sbatch"): return CommandResult(exit_code=0, stdout="12345\n", stderr="") if command.startswith("squeue"): - stdout = f"{self.queue_state}\n" if self.queue_state else "" + 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="") @@ -253,6 +255,45 @@ def test_client_state_and_cancel_use_job_id(): 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"): + return CommandResult( + exit_code=0, + stdout="12345|RUNNING\n23456|PENDING\n", + stderr="", + ) + if command.startswith("scontrol"): + return CommandResult( + exit_code=0, + stdout=( + "JobId=34567 JobState=COMPLETED " + "Command=/s/job.sh\n" + ), + stderr="", + ) + return CommandResult(exit_code=0, stdout="", stderr="") + + runner = BatchRunner() + states = SlurmClient(runner).get_job_states( + ["12345", "23456", "34567"] + ) + + assert states == { + "12345": "RUNNING", + "23456": "PENDING", + "34567": "COMPLETED", + } + assert runner.commands == [ + "squeue --noheader --format='%i|%T' --jobs=12345,23456,34567", + "scontrol show job --oneliner 34567", + ] + + # --- container command rendering (per runtime) -------------------------------- @@ -491,8 +532,8 @@ def test_get_status_reads_sentinel_after_queue(op_and_runner): assert op.get_status(step_run) is ExecutionStatus.FAILED -def test_get_status_without_queue_or_sentinel_is_non_terminal(op_and_runner): - """A queue/filesystem visibility gap remains non-terminal.""" +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( @@ -501,7 +542,7 @@ def test_get_status_without_queue_or_sentinel_is_non_terminal(op_and_runner): run_metadata={SLURM_JOB_ID_METADATA_KEY: "12345"}, ) ) - is ExecutionStatus.QUEUED + is ExecutionStatus.FAILED ) From 404072f91355cc2e6660138242b6fc1b8358aa7f Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Mon, 13 Jul 2026 11:08:32 +0200 Subject: [PATCH 12/22] Add dynamic pipeline support for Slurm --- .../slurm/orchestrators/slurm_orchestrator.py | 562 ++++++++++++++++-- .../slurm/test_slurm_orchestrator.py | 217 ++++++- 2 files changed, 711 insertions(+), 68 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 58c28528915..e22363423fa 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -26,14 +26,19 @@ import os import shlex -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, cast +from uuid import UUID from zenml.config.base_settings import BaseSettings -from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID +from zenml.config.resource_settings import ResourceSettings +from zenml.constants import ( + METADATA_ORCHESTRATOR_RUN_ID, + ORCHESTRATOR_DOCKER_IMAGE_KEY, +) from zenml.entrypoints.step_entrypoint_configuration import ( StepEntrypointConfiguration, ) -from zenml.enums import ExecutionStatus +from zenml.enums import ExecutionStatus, MetadataResourceTypes from zenml.integrations.slurm.flavors.slurm_orchestrator_flavor import ( SlurmOrchestratorConfig, SlurmOrchestratorSettings, @@ -59,20 +64,33 @@ validate_remote_stack, ) from zenml.logger import get_logger +from zenml.models.v2.misc.run_metadata import RunMetadataResource from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult +from zenml.orchestrators import utils as orchestrator_utils +from zenml.orchestrators.publish_utils import publish_step_run_metadata from zenml.orchestrators.topsort import topsorted_layers -from zenml.stack import StackValidator +from zenml.stack import Stack, StackValidator +from zenml.step_operators.step_operator_entrypoint_configuration import ( + StepOperatorEntrypointConfiguration, +) if TYPE_CHECKING: from zenml.config.step_configurations import Step - from zenml.models import PipelineRunResponse, PipelineSnapshotResponse - from zenml.stack import Stack + from zenml.config.step_run_info import StepRunInfo + from zenml.models import ( + PipelineRunResponse, + PipelineSnapshotResponse, + StepRunResponse, + ) 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" +SLURM_ISOLATED_JOB_IDS_METADATA_KEY = "slurm_isolated_job_ids" +SLURM_ISOLATED_JOB_ID_METADATA_KEY = "slurm_job_id" _CLEANUP_COMPLETE_FILE = "cleanup_complete" @@ -186,6 +204,31 @@ def _run_dir(self, run_id: str, step_name: str) -> str: 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 _isolated_run_dir(self, run_id: str, step_run_id: str) -> str: + """Staging directory for a dynamic isolated step job. + + Args: + run_id: The pipeline run id. + step_run_id: The isolated step run id. + + Returns: + The absolute path of the isolated step's run directory. + """ + workdir = self.config.workdir.rstrip("/") + return f"{workdir}/{run_id}/isolated/{step_run_id}" + def submit_pipeline( self, snapshot: "PipelineSnapshotResponse", @@ -315,7 +358,7 @@ def submit_dynamic_pipeline( environment: Dict[str, str], placeholder_run: Optional["PipelineRunResponse"] = None, ) -> Optional[SubmissionResult]: - """Reject dynamic pipelines explicitly. + """Submit a dynamic pipeline as a Slurm orchestration job. Args: snapshot: The pipeline snapshot to submit. @@ -325,56 +368,109 @@ def submit_dynamic_pipeline( placeholder_run: The placeholder run for the pipeline. Returns: - Never returns successfully. + Submission metadata containing the orchestration job id. Raises: - RuntimeError: Always, since dynamic pipelines need an - orchestration-container mode that the Slurm orchestrator does - not implement yet. + RuntimeError: If the pipeline has a schedule, which is not + supported. """ - _ = (snapshot, stack, environment, placeholder_run) - raise RuntimeError( - "The Slurm orchestrator does not support dynamic pipelines yet." + from zenml.pipelines.dynamic.entrypoint_configuration import ( + DynamicPipelineEntrypointConfiguration, ) - def _submit_step( + 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, _ = 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, + resources=ResourceSettings(), + settings=settings, + registry_uri=container_registry.config.uri, + registry_credentials=container_registry.credentials, + ) + 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, + } + ) + + def _submit_container_job( self, client: SlurmClient, - snapshot: "PipelineSnapshotResponse", run_id: str, - step_name: str, - step: "Step", - step_environment: Dict[str, str], - upstream_job_ids: List[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 a single step as a Slurm job. + """Stage and submit one containerized 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. + 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. """ - image = self.get_image(snapshot=snapshot, step_name=step_name) - settings = cast(SlurmOrchestratorSettings, self.get_settings(step)) - run_dir = self._run_dir(run_id, step_name) env_file = f"{run_dir}/{ENV_FILE}" - use_gpu = bool(step.config.resource_settings.gpu_count) - - environment = step_environment.copy() - environment[ENV_ZENML_SLURM_RUN_ID] = run_id + runtime_environment = environment.copy() + runtime_environment[ENV_ZENML_SLURM_RUN_ID] = run_id env_content = serialize_environment( - environment, runtime=self.config.container_runtime + runtime_environment, runtime=self.config.container_runtime ) registry_auth = build_registry_auth( runtime=self.config.container_runtime, @@ -382,42 +478,84 @@ def _submit_step( registry_uri=registry_uri, credentials=registry_credentials, ) - - entrypoint_command = ( - StepEntrypointConfiguration.get_entrypoint_command() - + StepEntrypointConfiguration.get_entrypoint_arguments( - step_name=step_name, snapshot_id=snapshot.id - ) - ) 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, + env_keys=sorted(runtime_environment), + use_gpu=bool(resources.gpu_count), settings=settings, registry_auth=registry_auth, ) script = build_sbatch_script( - job_name=self._job_name(run_id, step_name), + job_name=job_name, run_dir=run_dir, container_command=container_command, - resources=step.config.resource_settings, + 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=upstream_job_ids, + 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, @@ -484,6 +622,220 @@ def _submit_cleanup_job( dependency_type="afterany", ) + def submit_isolated_step( + self, step_run_info: "StepRunInfo", environment: Dict[str, str] + ) -> None: + """Submit a dynamic isolated step as its own Slurm job. + + Args: + step_run_info: Information about the isolated step run. + environment: Environment variables for the step job. + + Raises: + RuntimeError: If the snapshot has no associated stack. + Exception: Re-raised after cancelling the job if submission or + metadata publication fails. + """ + settings = cast( + SlurmOrchestratorSettings, self.get_settings(step_run_info) + ) + command, args = orchestrator_utils.get_step_entrypoint_command( + invocation_id=step_run_info.pipeline_step_name, + config=step_run_info.config, + entrypoint_config_class=StepOperatorEntrypointConfiguration, + snapshot_id=step_run_info.snapshot.id, + step_run_id=str(step_run_info.step_run_id), + ) + if not step_run_info.snapshot.stack: + raise RuntimeError( + f"Missing stack for snapshot {step_run_info.snapshot.id}." + ) + stack = Stack.from_model(step_run_info.snapshot.stack) + container_registry = stack.container_registry + assert container_registry is not None + + run_id = str(step_run_info.run_id) + run_dir = self._isolated_run_dir( + run_id=run_id, step_run_id=str(step_run_info.step_run_id) + ) + client = build_slurm_client(self.config) + job_id: Optional[str] = None + sensitive_paths: List[str] = [] + try: + job_id, sensitive_paths = self._submit_container_job( + client=client, + run_id=run_id, + run_dir=run_dir, + job_name=self._job_name( + run_id, str(step_run_info.step_run_id) + ), + image=step_run_info.get_image( + key=ORCHESTRATOR_DOCKER_IMAGE_KEY + ), + entrypoint_command=command + args, + environment=environment, + resources=step_run_info.config.resource_settings, + settings=settings, + registry_uri=container_registry.config.uri, + registry_credentials=container_registry.credentials, + ) + metadata = {SLURM_ISOLATED_JOB_ID_METADATA_KEY: job_id} + publish_step_run_metadata( + step_run_id=step_run_info.step_run_id, + step_run_metadata={self.id: metadata}, + ) + step_run_info.step_run.run_metadata.update(metadata) + except Exception: + if job_id is not None: + try: + client.cancel(job_id) + except Exception: + logger.warning( + "Failed to cancel Slurm job `%s` after isolated step " + "metadata publication failed.", + job_id, + ) + if sensitive_paths: + client.runner.run( + "rm -rf -- " + + " ".join( + shlex.quote(path) for path in sensitive_paths + ) + ) + raise + finally: + client.runner.close() + + assert job_id is not None + self._publish_isolated_job_run_metadata( + run_id=step_run_info.run_id, + step_run_id=step_run_info.step_run_id, + job_id=job_id, + ) + logger.info( + "Submitted dynamic step `%s` as Slurm job `%s`.", + step_run_info.pipeline_step_name, + job_id, + ) + + def get_isolated_step_status( + self, step_run: "StepRunResponse" + ) -> ExecutionStatus: + """Get the status of a dynamic isolated Slurm step. + + Args: + step_run: The step run to inspect. + + Returns: + The execution status. + """ + job_id = step_run.run_metadata.get(SLURM_ISOLATED_JOB_ID_METADATA_KEY) + if job_id is None: + logger.warning( + "No Slurm job ID recorded for isolated step `%s`.", + step_run.id, + ) + return ExecutionStatus.FAILED + + run_dir = self._isolated_run_dir( + run_id=str(step_run.pipeline_run_id), step_run_id=str(step_run.id) + ) + client = build_slurm_client(self.config) + try: + state = client.get_job_state(str(job_id)) + status = self._get_job_status( + client=client, + state=state, + run_dir=run_dir, + cleanup_complete=True, + ) + finally: + client.runner.close() + + return status or ExecutionStatus.FAILED + + def stop_isolated_step(self, step_run: "StepRunResponse") -> None: + """Cancel a dynamic isolated Slurm step. + + Args: + step_run: The step run to cancel. + """ + job_id = step_run.run_metadata.get(SLURM_ISOLATED_JOB_ID_METADATA_KEY) + if job_id is None: + logger.warning( + "No Slurm job ID recorded for isolated step `%s`.", + step_run.id, + ) + return + + run_dir = self._isolated_run_dir( + run_id=str(step_run.pipeline_run_id), step_run_id=str(step_run.id) + ) + client = build_slurm_client(self.config) + try: + client.cancel(str(job_id)) + client.runner.put_text( + f"{run_dir}/{CANCELLED_FILE}", "1\n", mode=0o600 + ) + finally: + client.runner.close() + + def cleanup_isolated_step(self, step_run: "StepRunResponse") -> None: + """Remove the staging directory for a dynamic isolated step. + + Args: + step_run: The finished step run. + """ + run_dir = self._isolated_run_dir( + run_id=str(step_run.pipeline_run_id), step_run_id=str(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 isolated step directory `%s`: %s", + run_dir, + e, + ) + finally: + client.runner.close() + + def _publish_isolated_job_run_metadata( + self, run_id: UUID, step_run_id: UUID, job_id: str + ) -> None: + """Publish run-level metadata for isolated Slurm jobs. + + Args: + run_id: Pipeline run ID. + step_run_id: Step run ID. + job_id: Slurm job ID. + """ + from zenml.client import Client + + try: + Client().create_run_metadata( + metadata={ + SLURM_ISOLATED_JOB_IDS_METADATA_KEY: { + str(step_run_id): job_id + } + }, + resources=[ + RunMetadataResource( + id=run_id, + type=MetadataResourceTypes.PIPELINE_RUN, + ) + ], + stack_component_id=self.id, + ) + except Exception as e: + logger.warning( + "Failed to publish Slurm run metadata for isolated step `%s`: " + "%s", + step_run_id, + e, + ) + def fetch_status( self, run: "PipelineRunResponse", include_steps: bool = False ) -> Tuple[ @@ -500,8 +852,7 @@ def fetch_status( """ raw_job_ids = run.run_metadata.get(SLURM_JOB_IDS_METADATA_KEY) if not isinstance(raw_job_ids, dict): - logger.warning("No Slurm job metadata found for run `%s`.", run.id) - return None, None + return self._fetch_dynamic_status(run) job_ids = { str(name): str(job_id) for name, job_id in raw_job_ids.items() @@ -566,6 +917,47 @@ def fetch_status( } return pipeline_status, step_statuses + 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 + + client = build_slurm_client(self.config) + try: + 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)), + cleanup_complete=True, + ) + 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, @@ -626,19 +1018,58 @@ def _stop_run( """ _ = graceful raw_job_ids = run.run_metadata.get(SLURM_JOB_IDS_METADATA_KEY) - if not isinstance(raw_job_ids, dict): + static_job_ids = ( + { + str(name): str(job_id) + for name, job_id in raw_job_ids.items() + } + if isinstance(raw_job_ids, dict) + else {} + ) + raw_isolated_job_ids = run.run_metadata.get( + SLURM_ISOLATED_JOB_IDS_METADATA_KEY + ) + isolated_job_ids = ( + { + str(step_run_id): str(job_id) + for step_run_id, job_id in raw_isolated_job_ids.items() + } + if isinstance(raw_isolated_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 isolated_job_ids or orchestration_job_ids): logger.warning("No Slurm job metadata found for run `%s`.", run.id) return - job_ids = { - str(name): str(job_id) for name, job_id in raw_job_ids.items() - } + run_id = str(run.id) client = build_slurm_client(self.config) try: self._cancel_jobs( client=client, - run_id=str(run.id), - job_ids=job_ids, + 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=isolated_job_ids, + run_dir_for_key=lambda step_run_id: self._isolated_run_dir( + run_id, step_run_id + ), + ) + 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() @@ -665,22 +1096,25 @@ def _cancel_unfinished_jobs( jobs_to_cancel[step_name] = job_id self._cancel_jobs( client=client, - run_id=run_id, 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, - run_id: str, 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. - run_id: Orchestrator run ID. - job_ids: Mapping of step names to Slurm job IDs. + 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 @@ -706,15 +1140,15 @@ def _cancel_jobs( else: cancelled_steps.add(step_name) - for step_name in cancelled_steps: + for job_key in cancelled_steps: try: client.runner.put_text( - f"{self._run_dir(run_id, step_name)}/{CANCELLED_FILE}", + f"{run_dir_for_key(job_key)}/{CANCELLED_FILE}", "1\n", mode=0o600, ) except Exception: logger.warning( - "Failed to write Slurm cancellation marker for step `%s`.", - step_name, + "Failed to write Slurm cancellation marker for job `%s`.", + job_key, ) diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index 4c61a8f47e0..68f05753617 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -36,7 +36,10 @@ from zenml.integrations.slurm.orchestrators.slurm_orchestrator import ( ENV_ZENML_SLURM_RUN_ID, SLURM_CLEANUP_JOB_ID_METADATA_KEY, + SLURM_ISOLATED_JOB_ID_METADATA_KEY, + SLURM_ISOLATED_JOB_IDS_METADATA_KEY, SLURM_JOB_IDS_METADATA_KEY, + SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY, ) from zenml.integrations.slurm.slurm_client import ( CommandResult, @@ -163,6 +166,7 @@ def _snapshot(steps: Dict[str, SimpleNamespace]) -> SimpleNamespace: step_configurations=steps, pipeline_configuration=SimpleNamespace(name="p"), schedule=None, + stack=SimpleNamespace(id=uuid4()), ) @@ -274,6 +278,28 @@ def _stack( ) +def _isolated_step_info( + snapshot: Optional[SimpleNamespace] = None, +) -> SimpleNamespace: + """Build a dynamic isolated StepRunInfo stand-in. + + Args: + snapshot: Optional snapshot to attach to the step run info. + + Returns: + A namespace usable where StepRunInfo is expected. + """ + return SimpleNamespace( + step_run_id=uuid4(), + run_id=uuid4(), + pipeline_step_name="dynamic_train", + config=SimpleNamespace(resource_settings=ResourceSettings(cpu_count=2)), + snapshot=snapshot or _snapshot({}), + step_run=SimpleNamespace(run_metadata={}), + get_image=lambda key: IMAGE, + ) + + def _submit_two_step_dag(monkeypatch) -> "tuple[FakeRunner, str]": """Submit a load -> train DAG and return the fake host and run id. @@ -429,14 +455,44 @@ def test_submit_rejects_scheduled_pipelines(monkeypatch): ) -def test_submit_dynamic_pipeline_raises(monkeypatch): - """Dynamic pipelines fail loudly instead of being silently ignored.""" +def test_submit_dynamic_pipeline_submits_orchestration_job(monkeypatch): + """Dynamic pipelines launch a Slurm orchestration job.""" + 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({}), + 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", + } + sbatch = [c for c in runner.commands if c.startswith("sbatch")] + assert len(sbatch) == 1 + assert "orchestration/job.sh" in sbatch[0] + 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 + + +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="dynamic pipelines"): + with pytest.raises(RuntimeError, match="does not support scheduled"): op.submit_dynamic_pipeline( - snapshot=_snapshot({"load": _step([])}), + snapshot=snapshot, stack=_stack(), environment={}, placeholder_run=_placeholder_run(), @@ -614,6 +670,159 @@ def test_stop_run_cancels_submitted_jobs(monkeypatch): ) +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 = SimpleNamespace( + id=run_id, + status=ExecutionStatus.PROVISIONING, + 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_submit_isolated_step_publishes_job_metadata(monkeypatch): + """Dynamic isolated steps are submitted as separate Slurm jobs.""" + op = _build_orchestrator() + op.get_settings = lambda _info: SlurmOrchestratorSettings() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + monkeypatch.setattr( + "zenml.integrations.slurm.orchestrators.slurm_orchestrator" + ".orchestrator_utils.get_step_entrypoint_command", + lambda **kwargs: (["python"], ["-m", "zenml.entrypoint"]), + ) + monkeypatch.setattr( + "zenml.integrations.slurm.orchestrators.slurm_orchestrator" + ".Stack.from_model", + lambda stack: _stack(), + ) + published_step_metadata = {} + monkeypatch.setattr( + "zenml.integrations.slurm.orchestrators.slurm_orchestrator" + ".publish_step_run_metadata", + lambda **kwargs: published_step_metadata.update(kwargs), + ) + published_run_metadata = {} + + class FakeClient: + def create_run_metadata(self, **kwargs): + published_run_metadata.update(kwargs) + + monkeypatch.setattr("zenml.client.Client", lambda: FakeClient()) + info = _isolated_step_info() + + op.submit_isolated_step( + step_run_info=info, + environment={"ZENML_STORE_API_KEY": SECRET_TOKEN}, + ) + + run_dir = f"/runs/{info.run_id}/isolated/{info.step_run_id}" + assert any(command.startswith("sbatch") for command in runner.commands) + assert runner.modes[f"{run_dir}/env"] == 0o600 + assert SECRET_TOKEN in runner.files[f"{run_dir}/env"] + assert info.step_run.run_metadata[SLURM_ISOLATED_JOB_ID_METADATA_KEY] == ( + "1000" + ) + assert published_step_metadata == { + "step_run_id": info.step_run_id, + "step_run_metadata": { + op.id: {SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"} + }, + } + assert published_run_metadata["metadata"] == { + SLURM_ISOLATED_JOB_IDS_METADATA_KEY: {str(info.step_run_id): "1000"} + } + + +def test_get_isolated_step_status_reads_sentinel(monkeypatch): + """Isolated step status uses the same Slurm sentinel mapping.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + step_run_id = uuid4() + run_dir = op._isolated_run_dir(str(run_id), str(step_run_id)) + runner.files[f"{run_dir}/exit_code"] = "0\n" + step_run = SimpleNamespace( + id=step_run_id, + pipeline_run_id=run_id, + run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, + ) + + assert op.get_isolated_step_status(step_run) is ExecutionStatus.COMPLETED + + +def test_stop_and_cleanup_isolated_step(monkeypatch): + """Isolated step stop and cleanup address the step staging directory.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + step_run_id = uuid4() + step_run = SimpleNamespace( + id=step_run_id, + pipeline_run_id=run_id, + run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, + ) + + op.stop_isolated_step(step_run) + op.cleanup_isolated_step(step_run) + + run_dir = op._isolated_run_dir(str(run_id), str(step_run_id)) + assert "scancel 1000" in runner.commands + assert runner.files[f"{run_dir}/cancelled"] == "1\n" + assert any( + command == f"rm -rf -- {run_dir}" for command in runner.commands + ) + + +def test_stop_run_cancels_dynamic_jobs(monkeypatch): + """Stopping a dynamic run cancels orchestration and isolated 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() + step_run_id = uuid4() + run = SimpleNamespace( + id=run_id, + orchestrator_run_id=str(run_id), + run_metadata={ + SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000", + SLURM_ISOLATED_JOB_IDS_METADATA_KEY: {str(step_run_id): "1001"}, + }, + ) + + op.stop_run(run) + + assert "scancel 1001" in runner.commands + assert "scancel 1000" in runner.commands + assert ( + runner.files[ + f"{op._isolated_run_dir(str(run_id), str(step_run_id))}/cancelled" + ] + == "1\n" + ) + assert ( + runner.files[f"{op._orchestration_run_dir(str(run_id))}/cancelled"] + == "1\n" + ) + + # --- validator ---------------------------------------------------------------- From 7f2de6a9b941e8a01351765325f1e981f9bffc58 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Mon, 13 Jul 2026 11:25:26 +0200 Subject: [PATCH 13/22] Fix Slurm failure modes and cancel cleanup --- .../slurm/orchestrators/slurm_orchestrator.py | 134 +++++++++++++++-- .../step_operators/slurm_step_operator.py | 2 +- .../slurm/test_slurm_orchestrator.py | 141 ++++++++++++++++-- .../slurm/test_slurm_step_operator.py | 2 +- 4 files changed, 245 insertions(+), 34 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index e22363423fa..12a69a76be7 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -26,7 +26,16 @@ import os import shlex -from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, cast +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + cast, +) from uuid import UUID from zenml.config.base_settings import BaseSettings @@ -38,7 +47,7 @@ from zenml.entrypoints.step_entrypoint_configuration import ( StepEntrypointConfiguration, ) -from zenml.enums import ExecutionStatus, MetadataResourceTypes +from zenml.enums import ExecutionMode, ExecutionStatus, MetadataResourceTypes from zenml.integrations.slurm.flavors.slurm_orchestrator_flavor import ( SlurmOrchestratorConfig, SlurmOrchestratorSettings, @@ -51,9 +60,12 @@ ) 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, @@ -92,11 +104,22 @@ SLURM_ISOLATED_JOB_IDS_METADATA_KEY = "slurm_isolated_job_ids" SLURM_ISOLATED_JOB_ID_METADATA_KEY = "slurm_job_id" _CLEANUP_COMPLETE_FILE = "cleanup_complete" +_ISOLATED_MISSING_SENTINEL_FAILURE_THRESHOLD = 3 class SlurmOrchestrator(ContainerizedOrchestrator): """Orchestrator that submits a pipeline as a graph of Slurm jobs.""" + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Initialize the Slurm orchestrator. + + Args: + *args: Forwarded to the base orchestrator. + **kwargs: Forwarded to the base orchestrator. + """ + super().__init__(*args, **kwargs) + self._missing_isolated_sentinel_counts: Dict[str, int] = {} + @property def config(self) -> SlurmOrchestratorConfig: """Returns the config of this orchestrator. @@ -106,6 +129,19 @@ def config(self) -> SlurmOrchestratorConfig: """ 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. @@ -747,12 +783,33 @@ def get_isolated_step_status( client=client, state=state, run_dir=run_dir, - cleanup_complete=True, + cleanup_complete=False, ) finally: client.runner.close() - return status or ExecutionStatus.FAILED + if status: + self._missing_isolated_sentinel_counts.pop(str(step_run.id), None) + return status + + missing_count = self._missing_isolated_sentinel_counts.get( + str(step_run.id), 0 + ) + missing_count += 1 + self._missing_isolated_sentinel_counts[str(step_run.id)] = ( + missing_count + ) + if missing_count >= _ISOLATED_MISSING_SENTINEL_FAILURE_THRESHOLD: + logger.warning( + "Slurm job `%s` for isolated step `%s` is no longer known to " + "Slurm and did not write an exit-code sentinel after %d " + "status checks.", + job_id, + step_run.id, + missing_count, + ) + return ExecutionStatus.FAILED + return ExecutionStatus.QUEUED def stop_isolated_step(self, step_run: "StepRunResponse") -> None: """Cancel a dynamic isolated Slurm step. @@ -777,6 +834,7 @@ def stop_isolated_step(self, step_run: "StepRunResponse") -> None: client.runner.put_text( f"{run_dir}/{CANCELLED_FILE}", "1\n", mode=0o600 ) + self._cleanup_sensitive_files(client, run_dir) finally: client.runner.close() @@ -886,18 +944,29 @@ def fetch_status( ] pipeline_status: Optional[ExecutionStatus] = None if not run.status.is_finished: - if any( + failed_or_cancelled = any( status in {ExecutionStatus.FAILED, ExecutionStatus.CANCELLED} for status in known_statuses - ): - pipeline_status = ExecutionStatus.FAILED - self._cancel_unfinished_jobs( - client=client, - run_id=str(run.id), - job_ids=job_ids, - statuses=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 @@ -917,6 +986,21 @@ def fetch_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]: @@ -1142,13 +1226,31 @@ def _cancel_jobs( for job_key in cancelled_steps: try: + run_dir = run_dir_for_key(job_key) client.runner.put_text( - f"{run_dir_for_key(job_key)}/{CANCELLED_FILE}", - "1\n", - mode=0o600, + 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/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index 91a05c3afa0..d56db848b19 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -351,7 +351,7 @@ def cleanup_step_submission(self, step_run: "StepRunResponse") -> None: 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)}") + 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 diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index 68f05753617..eed73914219 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -26,7 +26,7 @@ from zenml.config.resource_settings import ResourceSettings from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID -from zenml.enums import ExecutionStatus, StackComponentType +from zenml.enums import ExecutionMode, ExecutionStatus, StackComponentType from zenml.integrations.slurm.flavors import ( SlurmOrchestratorConfig, SlurmOrchestratorFlavor, @@ -199,6 +199,16 @@ def test_orchestrator_is_remote_and_detached(): 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 ------------------------------------------------------------------- @@ -300,6 +310,29 @@ def _isolated_step_info( ) +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. @@ -549,9 +582,8 @@ def test_fetch_status_reports_pre_entrypoint_failure(monkeypatch): run_id = uuid4() run_dir = op._run_dir(str(run_id), "load") runner.files[f"{run_dir}/exit_code"] = "1\n" - run = SimpleNamespace( - id=run_id, - status=ExecutionStatus.PROVISIONING, + run = _run( + run_id=run_id, run_metadata={SLURM_JOB_IDS_METADATA_KEY: {"load": "1000"}}, ) @@ -569,9 +601,8 @@ def test_fetch_status_uses_cleanup_marker_for_never_started_job(monkeypatch): run_id = uuid4() cleanup_marker = f"/runs/{run_id}/cleanup/cleanup_complete" runner.files[cleanup_marker] = "" - run = SimpleNamespace( - id=run_id, - status=ExecutionStatus.PROVISIONING, + run = _run( + run_id=run_id, run_metadata={SLURM_JOB_IDS_METADATA_KEY: {"train": "1001"}}, ) @@ -580,8 +611,8 @@ def test_fetch_status_uses_cleanup_marker_for_never_started_job(monkeypatch): assert pipeline_status is ExecutionStatus.FAILED -def test_fetch_status_batches_and_cancels_siblings(monkeypatch): - """A failed job cancels still-active sibling jobs in one refresh.""" +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): @@ -602,9 +633,8 @@ def run(self, command: str) -> CommandResult: runner = StatusRunner() _use_fake_client(monkeypatch, runner) run_id = uuid4() - run = SimpleNamespace( - id=run_id, - status=ExecutionStatus.PROVISIONING, + run = _run( + run_id=run_id, run_metadata={ SLURM_JOB_IDS_METADATA_KEY: { "load": "1000", @@ -616,7 +646,7 @@ def run(self, command: str) -> CommandResult: pipeline_status, step_statuses = op.fetch_status(run, include_steps=True) - assert pipeline_status is ExecutionStatus.FAILED + assert pipeline_status is ExecutionStatus.RUNNING assert step_statuses == { "load": ExecutionStatus.FAILED, "train": ExecutionStatus.RUNNING, @@ -625,6 +655,48 @@ def run(self, command: str) -> CommandResult: 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\n" + "1001|RUNNING\n" + "1002|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 ) @@ -636,6 +708,11 @@ def run(self, command: str) -> CommandResult: 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): @@ -668,6 +745,11 @@ def test_stop_run_cancels_submitted_jobs(monkeypatch): 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): @@ -678,9 +760,8 @@ def test_fetch_status_reconciles_dynamic_orchestration_job(monkeypatch): run_id = uuid4() run_dir = op._orchestration_run_dir(str(run_id)) runner.files[f"{run_dir}/exit_code"] = "1\n" - run = SimpleNamespace( - id=run_id, - status=ExecutionStatus.PROVISIONING, + run = _run( + run_id=run_id, run_metadata={SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000"}, ) @@ -762,6 +843,24 @@ def test_get_isolated_step_status_reads_sentinel(monkeypatch): assert op.get_isolated_step_status(step_run) is ExecutionStatus.COMPLETED +def test_get_isolated_step_status_waits_for_missing_sentinel(monkeypatch): + """A missing isolated sentinel gets a bounded retry window.""" + op = _build_orchestrator() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + run_id = uuid4() + step_run_id = uuid4() + step_run = SimpleNamespace( + id=step_run_id, + pipeline_run_id=run_id, + run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, + ) + + assert op.get_isolated_step_status(step_run) is ExecutionStatus.QUEUED + assert op.get_isolated_step_status(step_run) is ExecutionStatus.QUEUED + assert op.get_isolated_step_status(step_run) is ExecutionStatus.FAILED + + def test_stop_and_cleanup_isolated_step(monkeypatch): """Isolated step stop and cleanup address the step staging directory.""" op = _build_orchestrator() @@ -781,6 +880,11 @@ def test_stop_and_cleanup_isolated_step(monkeypatch): run_dir = op._isolated_run_dir(str(run_id), str(step_run_id)) assert "scancel 1000" in runner.commands assert runner.files[f"{run_dir}/cancelled"] == "1\n" + assert any( + command.startswith("rm -rf --") + and f"{run_dir}/env" in command + for command in runner.commands + ) assert any( command == f"rm -rf -- {run_dir}" for command in runner.commands ) @@ -821,6 +925,11 @@ def test_stop_run_cancels_dynamic_jobs(monkeypatch): 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 ---------------------------------------------------------------- diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index 6b80ae8bff3..7023904039d 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -564,7 +564,7 @@ def test_cleanup_removes_run_dir(op_and_runner): step_run = SimpleNamespace(id=uuid4()) op.cleanup_step_submission(step_run) run_dir = op._run_dir(step_run.id) - assert any("rm -rf" in c and run_dir in c for c in runner.commands) + assert any(f"rm -rf -- {run_dir}" in c for c in runner.commands) # --- integration registration ------------------------------------------------- From 6da70f46c88454f99e4296330c8ec89a1107240b Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Mon, 13 Jul 2026 11:32:56 +0200 Subject: [PATCH 14/22] Fix Slurm isolated metadata transaction --- .../slurm/orchestrators/slurm_orchestrator.py | 23 +++++----- src/zenml/integrations/slurm/slurm_client.py | 36 +++++++-------- .../slurm/test_slurm_orchestrator.py | 44 +++++++++++++++++++ .../slurm/test_slurm_step_operator.py | 14 ++++-- 4 files changed, 85 insertions(+), 32 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 12a69a76be7..17bb778c501 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -720,6 +720,11 @@ def submit_isolated_step( step_run_id=step_run_info.step_run_id, step_run_metadata={self.id: metadata}, ) + self._publish_isolated_job_run_metadata( + run_id=step_run_info.run_id, + step_run_id=step_run_info.step_run_id, + job_id=job_id, + ) step_run_info.step_run.run_metadata.update(metadata) except Exception: if job_id is not None: @@ -743,11 +748,6 @@ def submit_isolated_step( client.runner.close() assert job_id is not None - self._publish_isolated_job_run_metadata( - run_id=step_run_info.run_id, - step_run_id=step_run_info.step_run_id, - job_id=job_id, - ) logger.info( "Submitted dynamic step `%s` as Slurm job `%s`.", step_run_info.pipeline_step_name, @@ -868,6 +868,9 @@ def _publish_isolated_job_run_metadata( run_id: Pipeline run ID. step_run_id: Step run ID. job_id: Slurm job ID. + + Raises: + RuntimeError: If the metadata could not be published. """ from zenml.client import Client @@ -887,12 +890,10 @@ def _publish_isolated_job_run_metadata( stack_component_id=self.id, ) except Exception as e: - logger.warning( - "Failed to publish Slurm run metadata for isolated step `%s`: " - "%s", - step_run_id, - e, - ) + raise RuntimeError( + "Failed to publish Slurm run metadata for isolated step " + f"`{step_run_id}`." + ) from e def fetch_status( self, run: "PipelineRunResponse", include_steps: bool = False diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index 049e1e08d69..21fb9462c98 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -334,25 +334,25 @@ def get_job_states(self, job_ids: Sequence[str]) -> Dict[str, Optional[str]]: # Completed jobs can disappear from squeue before a shared-filesystem # sentinel becomes visible. Slurm keeps them in controller memory for # MinJobAge, so scontrol closes that race without requiring sacct. - missing_jobs_arg = ",".join(missing_job_ids) - result = self.runner.run( - "scontrol show job --oneliner " - f"{shlex.quote(missing_jobs_arg)}" - ) - if result.exit_code != 0: - if "invalid job id" in result.stderr.lower(): - return states - raise RuntimeError( - f"`scontrol show job` failed with exit code " - f"{result.exit_code}: {result.stderr.strip()}" + for missing_job_id in missing_job_ids: + result = self.runner.run( + "scontrol show job --oneliner " + f"{shlex.quote(missing_job_id)}" ) - for line in result.stdout.strip().splitlines(): - job_id_match = re.search(r"(?:^|\s)JobId=([0-9]+)", line) - state_match = re.search(r"(?:^|\s)JobState=([^\s]+)", line) - if job_id_match and state_match: - job_id = job_id_match.group(1) - if job_id in states: - states[job_id] = state_match.group(1) + if result.exit_code != 0: + if "invalid job id" in result.stderr.lower(): + continue + raise RuntimeError( + f"`scontrol show job` failed with exit code " + f"{result.exit_code}: {result.stderr.strip()}" + ) + for line in result.stdout.strip().splitlines(): + job_id_match = re.search(r"(?:^|\s)JobId=([0-9]+)", line) + state_match = re.search(r"(?:^|\s)JobState=([^\s]+)", line) + if job_id_match and state_match: + job_id = job_id_match.group(1) + if job_id in states: + states[job_id] = state_match.group(1) return states diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index eed73914219..dd0b2700b9b 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -825,6 +825,50 @@ def create_run_metadata(self, **kwargs): } +def test_submit_isolated_step_cancels_on_run_metadata_failure(monkeypatch): + """Run-level metadata publication is part of the submission transaction.""" + op = _build_orchestrator() + op.get_settings = lambda _info: SlurmOrchestratorSettings() + runner = FakeRunner() + _use_fake_client(monkeypatch, runner) + monkeypatch.setattr( + "zenml.integrations.slurm.orchestrators.slurm_orchestrator" + ".orchestrator_utils.get_step_entrypoint_command", + lambda **kwargs: (["python"], ["-m", "zenml.entrypoint"]), + ) + monkeypatch.setattr( + "zenml.integrations.slurm.orchestrators.slurm_orchestrator" + ".Stack.from_model", + lambda stack: _stack(), + ) + monkeypatch.setattr( + "zenml.integrations.slurm.orchestrators.slurm_orchestrator" + ".publish_step_run_metadata", + lambda **kwargs: None, + ) + + class FailingClient: + def create_run_metadata(self, **kwargs): + raise RuntimeError("metadata store unavailable") + + monkeypatch.setattr("zenml.client.Client", lambda: FailingClient()) + info = _isolated_step_info() + + with pytest.raises(RuntimeError, match="Failed to publish Slurm run"): + op.submit_isolated_step( + step_run_info=info, + environment={"ZENML_STORE_API_KEY": SECRET_TOKEN}, + ) + + run_dir = f"/runs/{info.run_id}/isolated/{info.step_run_id}" + assert "scancel 1000" in runner.commands + assert any( + command.startswith("rm -rf --") + and f"{run_dir}/env" in command + for command in runner.commands + ) + + def test_get_isolated_step_status_reads_sentinel(monkeypatch): """Isolated step status uses the same Slurm sentinel mapping.""" op = _build_orchestrator() diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index 7023904039d..8b568e28bbd 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -267,7 +267,7 @@ def run(self, command: str) -> CommandResult: stdout="12345|RUNNING\n23456|PENDING\n", stderr="", ) - if command.startswith("scontrol"): + if command == "scontrol show job --oneliner 34567": return CommandResult( exit_code=0, stdout=( @@ -276,21 +276,29 @@ def run(self, command: str) -> CommandResult: ), stderr="", ) + if command == "scontrol show job --oneliner 45678": + return CommandResult( + exit_code=0, + stdout="JobId=45678 JobState=FAILED Command=/s/job.sh\n", + stderr="", + ) return CommandResult(exit_code=0, stdout="", stderr="") runner = BatchRunner() states = SlurmClient(runner).get_job_states( - ["12345", "23456", "34567"] + ["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", + "squeue --noheader --format='%i|%T' --jobs=12345,23456,34567,45678", "scontrol show job --oneliner 34567", + "scontrol show job --oneliner 45678", ] From 04483ced17616e35e72f0f20775bd795b9e95a1e Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Mon, 13 Jul 2026 11:39:05 +0200 Subject: [PATCH 15/22] Add Slurm integration docs --- .../component-guide/orchestrators/README.md | 1 + .../component-guide/orchestrators/slurm.md | 236 ++++++++++++++++++ .../component-guide/step-operators/README.md | 1 + .../component-guide/step-operators/slurm.md | 213 ++++++++++++++++ docs/book/component-guide/toc.md | 2 + 5 files changed, 453 insertions(+) create mode 100644 docs/book/component-guide/orchestrators/slurm.md create mode 100644 docs/book/component-guide/step-operators/slurm.md 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..4e352e99f85 --- /dev/null +++ b/docs/book/component-guide/orchestrators/slurm.md @@ -0,0 +1,236 @@ +--- +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 first submits an orchestration job. That job runs +ZenML's dynamic runner inside the orchestrator image. The runner then submits +newly generated isolated steps as separate Slurm jobs, stores their Slurm job +IDs in run metadata, monitors them, and cancels them when the run is stopped. + +{% hint style="warning" %} +Dynamic pipelines require the orchestration job to be able to submit child +Slurm jobs. With `transport=ssh`, the orchestration container must be able to +SSH back to the login node, so configure `ssh_private_key` rather than a local +`ssh_key_path`. With `transport=local`, the container running on the compute +node must have working Slurm CLI access. Many clusters restrict one or both of +these patterns, so validate this with your HPC administrators. +{% endhint %} + +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, especially for dynamic pipelines. | +| `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 use a long-running orchestration job plus child Slurm jobs for +the generated isolated steps. Keep these operational details in mind: + +* Configure `time_limit`, `partition`, `account`, `qos`, and any site-required + `extra_sbatch_directives` so the orchestration job itself is allowed to run + for the expected duration. +* With SSH transport, prefer `ssh_private_key` stored as a ZenML secret. A local + `ssh_key_path` on your laptop is not available inside the orchestration + container. +* The orchestration container must be able to reach the ZenML server, artifact + store, container registry, and Slurm submission path. +* Some clusters block outbound SSH from compute nodes or disallow nested + `sbatch` calls from jobs. In that case, use static pipelines, the Slurm step + operator, or a cluster-approved submission pattern. + +### 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. Dynamic and isolated job 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..5c3b968709c --- /dev/null +++ b/docs/book/component-guide/step-operators/slurm.md @@ -0,0 +1,213 @@ +--- +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:`. + +### 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 91ac16c6b36..ab605313986 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) @@ -64,6 +65,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) From a929e654b357aa56c3cbcb3d2769c9493749dc1a Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Mon, 13 Jul 2026 11:49:58 +0200 Subject: [PATCH 16/22] Fix Slurm dynamic status polling --- .../slurm/orchestrators/slurm_orchestrator.py | 147 ++++++++++++-- src/zenml/integrations/slurm/slurm_client.py | 39 ++-- .../slurm/test_slurm_orchestrator.py | 184 +++++++++++++++++- .../slurm/test_slurm_step_operator.py | 23 +-- 4 files changed, 342 insertions(+), 51 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 17bb778c501..0d68bc0eafa 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -26,6 +26,7 @@ import os import shlex +import time from typing import ( TYPE_CHECKING, Any, @@ -105,6 +106,7 @@ SLURM_ISOLATED_JOB_ID_METADATA_KEY = "slurm_job_id" _CLEANUP_COMPLETE_FILE = "cleanup_complete" _ISOLATED_MISSING_SENTINEL_FAILURE_THRESHOLD = 3 +_ISOLATED_STATUS_CACHE_TTL_SECONDS = 1.0 class SlurmOrchestrator(ContainerizedOrchestrator): @@ -119,6 +121,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: """ super().__init__(*args, **kwargs) self._missing_isolated_sentinel_counts: Dict[str, int] = {} + self._isolated_status_cache: Dict[ + str, Tuple[float, Dict[str, Optional[str]]] + ] = {} @property def config(self) -> SlurmOrchestratorConfig: @@ -776,21 +781,32 @@ def get_isolated_step_status( run_dir = self._isolated_run_dir( run_id=str(step_run.pipeline_run_id), step_run_id=str(step_run.id) ) - client = build_slurm_client(self.config) - try: - state = client.get_job_state(str(job_id)) - status = self._get_job_status( - client=client, - state=state, - run_dir=run_dir, - cleanup_complete=False, - ) - finally: - client.runner.close() + cache_hit, state = self._get_cached_isolated_job_state( + run_id=str(step_run.pipeline_run_id), + job_id=str(job_id), + ) + status = self._status_from_slurm_state(state) if cache_hit else None + if status is None: + client = build_slurm_client(self.config) + try: + if not cache_hit: + state = self._refresh_isolated_job_state_cache( + client=client, + run_id=str(step_run.pipeline_run_id), + job_id=str(job_id), + ) + status = self._get_job_status( + client=client, + state=state, + run_dir=run_dir, + cleanup_complete=False, + ) + finally: + client.runner.close() if status: self._missing_isolated_sentinel_counts.pop(str(step_run.id), None) - return status + return self._normalize_isolated_step_status(status) missing_count = self._missing_isolated_sentinel_counts.get( str(step_run.id), 0 @@ -809,7 +825,112 @@ def get_isolated_step_status( missing_count, ) return ExecutionStatus.FAILED - return ExecutionStatus.QUEUED + return ExecutionStatus.PROVISIONING + + def _get_cached_isolated_job_state( + self, run_id: str, job_id: str + ) -> Tuple[bool, Optional[str]]: + """Get a cached isolated job state if it is still fresh. + + Args: + run_id: Pipeline run ID. + job_id: Slurm job ID to inspect. + + Returns: + A tuple with whether the cache contained the job and the cached + Slurm state. + """ + cached = self._isolated_status_cache.get(run_id) + if not cached: + return False, None + + cached_at, states = cached + if ( + time.monotonic() - cached_at <= _ISOLATED_STATUS_CACHE_TTL_SECONDS + and job_id in states + ): + return True, states[job_id] + return False, None + + def _refresh_isolated_job_state_cache( + self, client: SlurmClient, run_id: str, job_id: str + ) -> Optional[str]: + """Refresh the isolated job state cache for a pipeline run. + + Args: + client: Slurm client used to query job states. + run_id: Pipeline run ID. + job_id: Slurm job ID to inspect. + + Returns: + The refreshed Slurm job state for the requested job, or None if the + job is not known to Slurm. + """ + job_ids = [job_id] + try: + from zenml.client import Client + + run = Client().get_pipeline_run(run_id, hydrate=False) + raw_job_ids = run.run_metadata.get( + SLURM_ISOLATED_JOB_IDS_METADATA_KEY + ) + if isinstance(raw_job_ids, dict): + job_ids = [str(value) for value in raw_job_ids.values()] + if job_id not in job_ids: + job_ids.append(job_id) + except Exception: + logger.debug( + "Failed to load Slurm isolated job metadata for run `%s`; " + "falling back to a single-job status lookup.", + run_id, + exc_info=True, + ) + + states = client.get_job_states(job_ids) + self._isolated_status_cache[run_id] = (time.monotonic(), states) + return states.get(job_id) + + @staticmethod + def _normalize_isolated_step_status( + status: ExecutionStatus, + ) -> ExecutionStatus: + """Map Slurm statuses to the dynamic runner's isolated-step contract. + + Args: + status: Status produced by Slurm reconciliation. + + Returns: + Status understood by the dynamic pipeline monitor. + """ + if status == ExecutionStatus.QUEUED: + return ExecutionStatus.PROVISIONING + if status == ExecutionStatus.CANCELLED: + return ExecutionStatus.STOPPED + return status + + @staticmethod + def _status_from_slurm_state( + state: Optional[str], + ) -> Optional[ExecutionStatus]: + """Map a known Slurm state to ZenML without reading sentinels. + + Args: + state: Slurm job state. + + Returns: + The mapped status, or None if sentinels are needed. + """ + 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 + return None def stop_isolated_step(self, step_run: "StepRunResponse") -> None: """Cancel a dynamic isolated Slurm step. diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index 21fb9462c98..af64db1c81b 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -331,28 +331,25 @@ def get_job_states(self, job_ids: Sequence[str]) -> Dict[str, Optional[str]]: if not missing_job_ids: return states - # Completed jobs can disappear from squeue before a shared-filesystem - # sentinel becomes visible. Slurm keeps them in controller memory for - # MinJobAge, so scontrol closes that race without requiring sacct. - for missing_job_id in missing_job_ids: - result = self.runner.run( - "scontrol show job --oneliner " - f"{shlex.quote(missing_job_id)}" + # 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()}" ) - if result.exit_code != 0: - if "invalid job id" in result.stderr.lower(): - continue - raise RuntimeError( - f"`scontrol show job` failed with exit code " - f"{result.exit_code}: {result.stderr.strip()}" - ) - for line in result.stdout.strip().splitlines(): - job_id_match = re.search(r"(?:^|\s)JobId=([0-9]+)", line) - state_match = re.search(r"(?:^|\s)JobState=([^\s]+)", line) - if job_id_match and state_match: - job_id = job_id_match.group(1) - if job_id in states: - states[job_id] = state_match.group(1) + 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 diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index dd0b2700b9b..fe2e36b2c3f 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -27,6 +27,7 @@ from zenml.config.resource_settings import ResourceSettings from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID from zenml.enums import ExecutionMode, ExecutionStatus, StackComponentType +from zenml.execution.pipeline.dynamic.runner import DynamicPipelineRunner from zenml.integrations.slurm.flavors import ( SlurmOrchestratorConfig, SlurmOrchestratorFlavor, @@ -900,11 +901,190 @@ def test_get_isolated_step_status_waits_for_missing_sentinel(monkeypatch): run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, ) - assert op.get_isolated_step_status(step_run) is ExecutionStatus.QUEUED - assert op.get_isolated_step_status(step_run) is ExecutionStatus.QUEUED + assert ( + op.get_isolated_step_status(step_run) + is ExecutionStatus.PROVISIONING + ) + assert ( + op.get_isolated_step_status(step_run) + is ExecutionStatus.PROVISIONING + ) assert op.get_isolated_step_status(step_run) is ExecutionStatus.FAILED +def test_get_isolated_step_status_normalizes_dynamic_runner_statuses( + monkeypatch, +): + """Isolated Slurm statuses use values supported by the dynamic monitor.""" + op = _build_orchestrator() + + class StatusRunner(FakeRunner): + def __init__(self, state: str) -> None: + """Initialize the fake. + + Args: + state: Slurm state to return for the isolated job. + """ + super().__init__() + self.state = state + + def run(self, command: str) -> CommandResult: + self.commands.append(command) + if command.startswith("squeue"): + return CommandResult( + exit_code=0, + stdout=f"1000|{self.state}\n", + stderr="", + ) + return CommandResult(exit_code=0, stdout="", stderr="") + + run_id = uuid4() + step_run_id = uuid4() + step_run = SimpleNamespace( + id=step_run_id, + pipeline_run_id=run_id, + run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, + ) + + runner = StatusRunner("PENDING") + _use_fake_client(monkeypatch, runner) + assert ( + op.get_isolated_step_status(step_run) + is ExecutionStatus.PROVISIONING + ) + + op._isolated_status_cache.clear() + runner = StatusRunner("CANCELLED") + _use_fake_client(monkeypatch, runner) + assert op.get_isolated_step_status(step_run) is ExecutionStatus.STOPPED + + +def test_dynamic_monitor_keeps_pending_slurm_step(monkeypatch): + """Pending Slurm jobs stay monitored and are not cleaned up early.""" + 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|PENDING\n", stderr="" + ) + return CommandResult(exit_code=0, stdout="", stderr="") + + runner = StatusRunner() + _use_fake_client(monkeypatch, runner) + step_run = SimpleNamespace( + id=uuid4(), + name="dynamic_train", + pipeline_run_id=uuid4(), + config=SimpleNamespace(step_operator=False), + run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, + ) + + dynamic_runner = DynamicPipelineRunner.__new__(DynamicPipelineRunner) + dynamic_runner._shutdown_requested = False + dynamic_runner._steps_to_monitor = {"dynamic_train": step_run} + dynamic_runner._orchestrator = op + dynamic_runner._step_operator = None + cleaned_steps: List[object] = [] + dynamic_runner._cleanup_isolated_step = cleaned_steps.append + + class StopEvent: + def wait(self, timeout: Optional[float] = None) -> None: + """Stop the monitor after one pass. + + Args: + timeout: Unused wait timeout. + """ + _ = timeout + dynamic_runner._shutdown_requested = True + + def clear(self) -> None: + """No-op event clear.""" + + dynamic_runner._monitoring_event = StopEvent() + + dynamic_runner._monitoring_loop() + + assert dynamic_runner._steps_to_monitor == {"dynamic_train": step_run} + assert cleaned_steps == [] + + +def test_get_isolated_step_status_batches_cached_run_jobs(monkeypatch): + """Isolated step polling shares one Slurm state query per cache window.""" + 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|PENDING\n1001|RUNNING\n", + stderr="", + ) + return CommandResult(exit_code=0, stdout="", stderr="") + + runner = StatusRunner() + client_builds: List[SlurmClient] = [] + + def build_client(config: SlurmOrchestratorConfig) -> SlurmClient: + """Build a fake Slurm client and track connection churn. + + Args: + config: Unused Slurm orchestrator config. + + Returns: + A Slurm client backed by the fake runner. + """ + _ = config + client = SlurmClient(runner) + client_builds.append(client) + return client + + monkeypatch.setattr( + "zenml.integrations.slurm.orchestrators.slurm_orchestrator" + ".build_slurm_client", + build_client, + ) + run_id = uuid4() + first_step = SimpleNamespace( + id=uuid4(), + pipeline_run_id=run_id, + run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, + ) + second_step = SimpleNamespace( + id=uuid4(), + pipeline_run_id=run_id, + run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1001"}, + ) + + class FakeClient: + def get_pipeline_run(self, *args, **kwargs): + return _run( + run_id=run_id, + run_metadata={ + SLURM_ISOLATED_JOB_IDS_METADATA_KEY: { + str(first_step.id): "1000", + str(second_step.id): "1001", + } + }, + ) + + monkeypatch.setattr("zenml.client.Client", lambda: FakeClient()) + + assert ( + op.get_isolated_step_status(first_step) + is ExecutionStatus.PROVISIONING + ) + assert op.get_isolated_step_status(second_step) is ExecutionStatus.RUNNING + assert [ + command for command in runner.commands if command.startswith("squeue") + ] == ["squeue --noheader --format='%i|%T' --jobs=1000,1001"] + assert len(client_builds) == 1 + + def test_stop_and_cleanup_isolated_step(monkeypatch): """Isolated step stop and cleanup address the step staging directory.""" op = _build_orchestrator() diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index 8b568e28bbd..ffce853a6c6 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -261,25 +261,19 @@ def test_client_batches_job_state_lookup(): class BatchRunner(FakeRunner): def run(self, command: str) -> CommandResult: self.commands.append(command) - if command.startswith("squeue"): - return CommandResult( - exit_code=0, - stdout="12345|RUNNING\n23456|PENDING\n", - stderr="", - ) - if command == "scontrol show job --oneliner 34567": + if ( + command.startswith("squeue --noheader") + and "--states=all" in command + ): return CommandResult( exit_code=0, - stdout=( - "JobId=34567 JobState=COMPLETED " - "Command=/s/job.sh\n" - ), + stdout="34567|COMPLETED\n45678|FAILED\n", stderr="", ) - if command == "scontrol show job --oneliner 45678": + if command.startswith("squeue"): return CommandResult( exit_code=0, - stdout="JobId=45678 JobState=FAILED Command=/s/job.sh\n", + stdout="12345|RUNNING\n23456|PENDING\n", stderr="", ) return CommandResult(exit_code=0, stdout="", stderr="") @@ -297,8 +291,7 @@ def run(self, command: str) -> CommandResult: } assert runner.commands == [ "squeue --noheader --format='%i|%T' --jobs=12345,23456,34567,45678", - "scontrol show job --oneliner 34567", - "scontrol show job --oneliner 45678", + "squeue --noheader --format='%i|%T' --states=all --jobs=34567,45678", ] From d86145a75ddabb0b43d9fd40ef89464eada7bb65 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 17 Jul 2026 10:25:55 +0200 Subject: [PATCH 17/22] Run dynamic pipelines inline in one Slurm allocation Address review: submitting per-step Slurm jobs from inside an active orchestration job starves under per-user job limits and blocks the parent allocation while children queue. Dynamic pipelines now run as a single orchestration job sized by the pipeline-level resource settings, with every step executing inline - the isolated-step machinery is removed, and no sbatch ever happens from within the cluster. Steps that need their own allocation go through the Slurm step operator instead. Also per review, drop the redundant topological sort: snapshot step configurations are already topologically ordered by contract. RL-training readiness, kept additive: - Record the raw Slurm terminal state (TIMEOUT, NODE_FAIL, OUT_OF_MEMORY, PREEMPTED, ...) as step-run metadata in the step operator, so an infrastructure kill is distinguishable from a step that ran and exited non-zero (no exit-code sentinel exists in either case). - Document the srun-wrappable contract of build_container_command for a future multi-node command-step mode. - Document container_mounts as the supported shared-filesystem bridge for tools that manage their own environment (e.g. a prime-rl checkout with its uv venv). Co-Authored-By: Claude Fable 5 --- .../component-guide/orchestrators/slurm.md | 49 +- .../component-guide/step-operators/slurm.md | 32 ++ .../slurm/orchestrators/slurm_orchestrator.py | 487 +----------------- src/zenml/integrations/slurm/slurm_job.py | 6 + .../step_operators/slurm_step_operator.py | 33 ++ .../slurm/test_slurm_orchestrator.py | 450 ++-------------- .../slurm/test_slurm_step_operator.py | 44 +- 7 files changed, 190 insertions(+), 911 deletions(-) diff --git a/docs/book/component-guide/orchestrators/slurm.md b/docs/book/component-guide/orchestrators/slurm.md index 4e352e99f85..ce5a7d9b3ea 100644 --- a/docs/book/component-guide/orchestrators/slurm.md +++ b/docs/book/component-guide/orchestrators/slurm.md @@ -73,19 +73,15 @@ 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 first submits an orchestration job. That job runs -ZenML's dynamic runner inside the orchestrator image. The runner then submits -newly generated isolated steps as separate Slurm jobs, stores their Slurm job -IDs in run metadata, monitors them, and cancels them when the run is stopped. - -{% hint style="warning" %} -Dynamic pipelines require the orchestration job to be able to submit child -Slurm jobs. With `transport=ssh`, the orchestration container must be able to -SSH back to the login node, so configure `ssh_private_key` rather than a local -`ssh_key_path`. With `transport=local`, the container running on the compute -node must have working Slurm CLI access. Many clusters restrict one or both of -these patterns, so validate this with your HPC administrators. -{% endhint %} +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 @@ -152,7 +148,7 @@ Key registration-time options: | `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, especially for dynamic pipelines. | +| `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`. | @@ -191,20 +187,17 @@ ZenML maps standard step resource settings to Slurm directives: ### Dynamic pipeline notes -Dynamic pipelines use a long-running orchestration job plus child Slurm jobs for -the generated isolated steps. Keep these operational details in mind: +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 itself is allowed to run - for the expected duration. -* With SSH transport, prefer `ssh_private_key` stored as a ZenML secret. A local - `ssh_key_path` on your laptop is not available inside the orchestration - container. -* The orchestration container must be able to reach the ZenML server, artifact - store, container registry, and Slurm submission path. -* Some clusters block outbound SSH from compute nodes or disallow nested - `sbatch` calls from jobs. In that case, use static pipelines, the Slurm step - operator, or a cluster-approved submission pattern. + `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 @@ -214,8 +207,8 @@ 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. Dynamic and isolated job cancellation paths explicitly scrub the known -sensitive files when a pending job is cancelled before its job script starts. +files. Cancellation paths explicitly scrub the known sensitive files when a +pending job is cancelled before its job script starts. ### Limitations diff --git a/docs/book/component-guide/step-operators/slurm.md b/docs/book/component-guide/step-operators/slurm.md index 5c3b968709c..7dd40b123fc 100644 --- a/docs/book/component-guide/step-operators/slurm.md +++ b/docs/book/component-guide/step-operators/slurm.md @@ -168,6 +168,38 @@ ZenML maps standard step resource settings to Slurm directives: * 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 diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 0d68bc0eafa..4890243e1ae 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -22,11 +22,17 @@ 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 -import time from typing import ( TYPE_CHECKING, Any, @@ -37,7 +43,6 @@ Tuple, cast, ) -from uuid import UUID from zenml.config.base_settings import BaseSettings from zenml.config.resource_settings import ResourceSettings @@ -48,7 +53,7 @@ from zenml.entrypoints.step_entrypoint_configuration import ( StepEntrypointConfiguration, ) -from zenml.enums import ExecutionMode, ExecutionStatus, MetadataResourceTypes +from zenml.enums import ExecutionMode, ExecutionStatus from zenml.integrations.slurm.flavors.slurm_orchestrator_flavor import ( SlurmOrchestratorConfig, SlurmOrchestratorSettings, @@ -77,24 +82,16 @@ validate_remote_stack, ) from zenml.logger import get_logger -from zenml.models.v2.misc.run_metadata import RunMetadataResource from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult -from zenml.orchestrators import utils as orchestrator_utils -from zenml.orchestrators.publish_utils import publish_step_run_metadata -from zenml.orchestrators.topsort import topsorted_layers -from zenml.stack import Stack, StackValidator -from zenml.step_operators.step_operator_entrypoint_configuration import ( - StepOperatorEntrypointConfiguration, -) +from zenml.stack import StackValidator if TYPE_CHECKING: from zenml.config.step_configurations import Step - from zenml.config.step_run_info import StepRunInfo from zenml.models import ( PipelineRunResponse, PipelineSnapshotResponse, - StepRunResponse, ) + from zenml.stack import Stack logger = get_logger(__name__) @@ -102,29 +99,12 @@ 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" -SLURM_ISOLATED_JOB_IDS_METADATA_KEY = "slurm_isolated_job_ids" -SLURM_ISOLATED_JOB_ID_METADATA_KEY = "slurm_job_id" _CLEANUP_COMPLETE_FILE = "cleanup_complete" -_ISOLATED_MISSING_SENTINEL_FAILURE_THRESHOLD = 3 -_ISOLATED_STATUS_CACHE_TTL_SECONDS = 1.0 class SlurmOrchestrator(ContainerizedOrchestrator): """Orchestrator that submits a pipeline as a graph of Slurm jobs.""" - def __init__(self, *args: Any, **kwargs: Any) -> None: - """Initialize the Slurm orchestrator. - - Args: - *args: Forwarded to the base orchestrator. - **kwargs: Forwarded to the base orchestrator. - """ - super().__init__(*args, **kwargs) - self._missing_isolated_sentinel_counts: Dict[str, int] = {} - self._isolated_status_cache: Dict[ - str, Tuple[float, Dict[str, Optional[str]]] - ] = {} - @property def config(self) -> SlurmOrchestratorConfig: """Returns the config of this orchestrator. @@ -186,40 +166,6 @@ def get_orchestrator_run_id(self) -> str: f"variable {ENV_ZENML_SLURM_RUN_ID}." ) - @staticmethod - def _sorted_step_names( - steps: Dict[str, "Step"], - ) -> List[str]: - """Return the step invocation ids in a topological order. - - Jobs must be submitted upstream-first so that a step's Slurm - dependencies reference the already-submitted job ids of its parents. - - Args: - steps: The steps keyed by invocation id. - - Returns: - The invocation ids in a valid topological order. - """ - - def parents(name: str) -> List[str]: - return [u for u in steps[name].spec.upstream_steps if u in steps] - - def children(name: str) -> List[str]: - return [ - other - for other in steps - if name in steps[other].spec.upstream_steps - ] - - layers = topsorted_layers( - nodes=list(steps), - get_node_id_fn=lambda name: name, - get_parent_nodes=parents, - get_child_nodes=children, - ) - return [name for layer in layers for name in layer] - def _job_name(self, run_id: str, step_name: str) -> str: """Slurm job name for a step in a run. @@ -257,19 +203,6 @@ def _orchestration_run_dir(self, run_id: str) -> str: workdir = self.config.workdir.rstrip("/") return f"{workdir}/{run_id}/orchestration" - def _isolated_run_dir(self, run_id: str, step_run_id: str) -> str: - """Staging directory for a dynamic isolated step job. - - Args: - run_id: The pipeline run id. - step_run_id: The isolated step run id. - - Returns: - The absolute path of the isolated step's run directory. - """ - workdir = self.config.workdir.rstrip("/") - return f"{workdir}/{run_id}/isolated/{step_run_id}" - def submit_pipeline( self, snapshot: "PipelineSnapshotResponse", @@ -323,8 +256,10 @@ def submit_pipeline( container_registry = stack.container_registry assert container_registry is not None try: - for step_name in self._sorted_step_names(steps): - step = steps[step_name] + # `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, @@ -429,9 +364,7 @@ def submit_dynamic_pipeline( assert placeholder_run is not None run_id = str(placeholder_run.id) - settings = cast( - SlurmOrchestratorSettings, self.get_settings(snapshot) - ) + settings = cast(SlurmOrchestratorSettings, self.get_settings(snapshot)) container_registry = stack.container_registry assert container_registry is not None command = ( @@ -452,7 +385,9 @@ def submit_dynamic_pipeline( image=self.get_image(snapshot=snapshot), entrypoint_command=command, environment=environment, - resources=ResourceSettings(), + # 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, @@ -663,359 +598,6 @@ def _submit_cleanup_job( dependency_type="afterany", ) - def submit_isolated_step( - self, step_run_info: "StepRunInfo", environment: Dict[str, str] - ) -> None: - """Submit a dynamic isolated step as its own Slurm job. - - Args: - step_run_info: Information about the isolated step run. - environment: Environment variables for the step job. - - Raises: - RuntimeError: If the snapshot has no associated stack. - Exception: Re-raised after cancelling the job if submission or - metadata publication fails. - """ - settings = cast( - SlurmOrchestratorSettings, self.get_settings(step_run_info) - ) - command, args = orchestrator_utils.get_step_entrypoint_command( - invocation_id=step_run_info.pipeline_step_name, - config=step_run_info.config, - entrypoint_config_class=StepOperatorEntrypointConfiguration, - snapshot_id=step_run_info.snapshot.id, - step_run_id=str(step_run_info.step_run_id), - ) - if not step_run_info.snapshot.stack: - raise RuntimeError( - f"Missing stack for snapshot {step_run_info.snapshot.id}." - ) - stack = Stack.from_model(step_run_info.snapshot.stack) - container_registry = stack.container_registry - assert container_registry is not None - - run_id = str(step_run_info.run_id) - run_dir = self._isolated_run_dir( - run_id=run_id, step_run_id=str(step_run_info.step_run_id) - ) - client = build_slurm_client(self.config) - job_id: Optional[str] = None - sensitive_paths: List[str] = [] - try: - job_id, sensitive_paths = self._submit_container_job( - client=client, - run_id=run_id, - run_dir=run_dir, - job_name=self._job_name( - run_id, str(step_run_info.step_run_id) - ), - image=step_run_info.get_image( - key=ORCHESTRATOR_DOCKER_IMAGE_KEY - ), - entrypoint_command=command + args, - environment=environment, - resources=step_run_info.config.resource_settings, - settings=settings, - registry_uri=container_registry.config.uri, - registry_credentials=container_registry.credentials, - ) - metadata = {SLURM_ISOLATED_JOB_ID_METADATA_KEY: job_id} - publish_step_run_metadata( - step_run_id=step_run_info.step_run_id, - step_run_metadata={self.id: metadata}, - ) - self._publish_isolated_job_run_metadata( - run_id=step_run_info.run_id, - step_run_id=step_run_info.step_run_id, - job_id=job_id, - ) - step_run_info.step_run.run_metadata.update(metadata) - except Exception: - if job_id is not None: - try: - client.cancel(job_id) - except Exception: - logger.warning( - "Failed to cancel Slurm job `%s` after isolated step " - "metadata publication failed.", - job_id, - ) - if sensitive_paths: - client.runner.run( - "rm -rf -- " - + " ".join( - shlex.quote(path) for path in sensitive_paths - ) - ) - raise - finally: - client.runner.close() - - assert job_id is not None - logger.info( - "Submitted dynamic step `%s` as Slurm job `%s`.", - step_run_info.pipeline_step_name, - job_id, - ) - - def get_isolated_step_status( - self, step_run: "StepRunResponse" - ) -> ExecutionStatus: - """Get the status of a dynamic isolated Slurm step. - - Args: - step_run: The step run to inspect. - - Returns: - The execution status. - """ - job_id = step_run.run_metadata.get(SLURM_ISOLATED_JOB_ID_METADATA_KEY) - if job_id is None: - logger.warning( - "No Slurm job ID recorded for isolated step `%s`.", - step_run.id, - ) - return ExecutionStatus.FAILED - - run_dir = self._isolated_run_dir( - run_id=str(step_run.pipeline_run_id), step_run_id=str(step_run.id) - ) - cache_hit, state = self._get_cached_isolated_job_state( - run_id=str(step_run.pipeline_run_id), - job_id=str(job_id), - ) - status = self._status_from_slurm_state(state) if cache_hit else None - if status is None: - client = build_slurm_client(self.config) - try: - if not cache_hit: - state = self._refresh_isolated_job_state_cache( - client=client, - run_id=str(step_run.pipeline_run_id), - job_id=str(job_id), - ) - status = self._get_job_status( - client=client, - state=state, - run_dir=run_dir, - cleanup_complete=False, - ) - finally: - client.runner.close() - - if status: - self._missing_isolated_sentinel_counts.pop(str(step_run.id), None) - return self._normalize_isolated_step_status(status) - - missing_count = self._missing_isolated_sentinel_counts.get( - str(step_run.id), 0 - ) - missing_count += 1 - self._missing_isolated_sentinel_counts[str(step_run.id)] = ( - missing_count - ) - if missing_count >= _ISOLATED_MISSING_SENTINEL_FAILURE_THRESHOLD: - logger.warning( - "Slurm job `%s` for isolated step `%s` is no longer known to " - "Slurm and did not write an exit-code sentinel after %d " - "status checks.", - job_id, - step_run.id, - missing_count, - ) - return ExecutionStatus.FAILED - return ExecutionStatus.PROVISIONING - - def _get_cached_isolated_job_state( - self, run_id: str, job_id: str - ) -> Tuple[bool, Optional[str]]: - """Get a cached isolated job state if it is still fresh. - - Args: - run_id: Pipeline run ID. - job_id: Slurm job ID to inspect. - - Returns: - A tuple with whether the cache contained the job and the cached - Slurm state. - """ - cached = self._isolated_status_cache.get(run_id) - if not cached: - return False, None - - cached_at, states = cached - if ( - time.monotonic() - cached_at <= _ISOLATED_STATUS_CACHE_TTL_SECONDS - and job_id in states - ): - return True, states[job_id] - return False, None - - def _refresh_isolated_job_state_cache( - self, client: SlurmClient, run_id: str, job_id: str - ) -> Optional[str]: - """Refresh the isolated job state cache for a pipeline run. - - Args: - client: Slurm client used to query job states. - run_id: Pipeline run ID. - job_id: Slurm job ID to inspect. - - Returns: - The refreshed Slurm job state for the requested job, or None if the - job is not known to Slurm. - """ - job_ids = [job_id] - try: - from zenml.client import Client - - run = Client().get_pipeline_run(run_id, hydrate=False) - raw_job_ids = run.run_metadata.get( - SLURM_ISOLATED_JOB_IDS_METADATA_KEY - ) - if isinstance(raw_job_ids, dict): - job_ids = [str(value) for value in raw_job_ids.values()] - if job_id not in job_ids: - job_ids.append(job_id) - except Exception: - logger.debug( - "Failed to load Slurm isolated job metadata for run `%s`; " - "falling back to a single-job status lookup.", - run_id, - exc_info=True, - ) - - states = client.get_job_states(job_ids) - self._isolated_status_cache[run_id] = (time.monotonic(), states) - return states.get(job_id) - - @staticmethod - def _normalize_isolated_step_status( - status: ExecutionStatus, - ) -> ExecutionStatus: - """Map Slurm statuses to the dynamic runner's isolated-step contract. - - Args: - status: Status produced by Slurm reconciliation. - - Returns: - Status understood by the dynamic pipeline monitor. - """ - if status == ExecutionStatus.QUEUED: - return ExecutionStatus.PROVISIONING - if status == ExecutionStatus.CANCELLED: - return ExecutionStatus.STOPPED - return status - - @staticmethod - def _status_from_slurm_state( - state: Optional[str], - ) -> Optional[ExecutionStatus]: - """Map a known Slurm state to ZenML without reading sentinels. - - Args: - state: Slurm job state. - - Returns: - The mapped status, or None if sentinels are needed. - """ - 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 - return None - - def stop_isolated_step(self, step_run: "StepRunResponse") -> None: - """Cancel a dynamic isolated Slurm step. - - Args: - step_run: The step run to cancel. - """ - job_id = step_run.run_metadata.get(SLURM_ISOLATED_JOB_ID_METADATA_KEY) - if job_id is None: - logger.warning( - "No Slurm job ID recorded for isolated step `%s`.", - step_run.id, - ) - return - - run_dir = self._isolated_run_dir( - run_id=str(step_run.pipeline_run_id), step_run_id=str(step_run.id) - ) - client = build_slurm_client(self.config) - try: - client.cancel(str(job_id)) - client.runner.put_text( - f"{run_dir}/{CANCELLED_FILE}", "1\n", mode=0o600 - ) - self._cleanup_sensitive_files(client, run_dir) - finally: - client.runner.close() - - def cleanup_isolated_step(self, step_run: "StepRunResponse") -> None: - """Remove the staging directory for a dynamic isolated step. - - Args: - step_run: The finished step run. - """ - run_dir = self._isolated_run_dir( - run_id=str(step_run.pipeline_run_id), step_run_id=str(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 isolated step directory `%s`: %s", - run_dir, - e, - ) - finally: - client.runner.close() - - def _publish_isolated_job_run_metadata( - self, run_id: UUID, step_run_id: UUID, job_id: str - ) -> None: - """Publish run-level metadata for isolated Slurm jobs. - - Args: - run_id: Pipeline run ID. - step_run_id: Step run ID. - job_id: Slurm job ID. - - Raises: - RuntimeError: If the metadata could not be published. - """ - from zenml.client import Client - - try: - Client().create_run_metadata( - metadata={ - SLURM_ISOLATED_JOB_IDS_METADATA_KEY: { - str(step_run_id): job_id - } - }, - resources=[ - RunMetadataResource( - id=run_id, - type=MetadataResourceTypes.PIPELINE_RUN, - ) - ], - stack_component_id=self.id, - ) - except Exception as e: - raise RuntimeError( - "Failed to publish Slurm run metadata for isolated step " - f"`{step_run_id}`." - ) from e - def fetch_status( self, run: "PipelineRunResponse", include_steps: bool = False ) -> Tuple[ @@ -1061,9 +643,7 @@ def fetch_status( for step_name, job_id in job_ids.items() } - known_statuses = [ - status for status in statuses.values() if status - ] + 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( @@ -1225,24 +805,10 @@ def _stop_run( _ = 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() - } + {str(name): str(job_id) for name, job_id in raw_job_ids.items()} if isinstance(raw_job_ids, dict) else {} ) - raw_isolated_job_ids = run.run_metadata.get( - SLURM_ISOLATED_JOB_IDS_METADATA_KEY - ) - isolated_job_ids = ( - { - str(step_run_id): str(job_id) - for step_run_id, job_id in raw_isolated_job_ids.items() - } - if isinstance(raw_isolated_job_ids, dict) - else {} - ) orchestration_job_id = run.run_metadata.get( SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY ) @@ -1251,7 +817,7 @@ def _stop_run( if orchestration_job_id else {} ) - if not (static_job_ids or isolated_job_ids or orchestration_job_ids): + if not (static_job_ids or orchestration_job_ids): logger.warning("No Slurm job metadata found for run `%s`.", run.id) return @@ -1265,13 +831,6 @@ def _stop_run( run_id, step_name ), ) - self._cancel_jobs( - client=client, - job_ids=isolated_job_ids, - run_dir_for_key=lambda step_run_id: self._isolated_run_dir( - run_id, step_run_id - ), - ) self._cancel_jobs( client=client, job_ids=orchestration_job_ids, @@ -1303,9 +862,7 @@ def _cancel_unfinished_jobs( self._cancel_jobs( client=client, job_ids=jobs_to_cancel, - run_dir_for_key=lambda step_name: self._run_dir( - run_id, step_name - ), + run_dir_for_key=lambda step_name: self._run_dir(run_id, step_name), ) def _cancel_jobs( diff --git a/src/zenml/integrations/slurm/slurm_job.py b/src/zenml/integrations/slurm/slurm_job.py index 3f087152422..15ff51b93f1 100644 --- a/src/zenml/integrations/slurm/slurm_job.py +++ b/src/zenml/integrations/slurm/slurm_job.py @@ -237,6 +237,12 @@ def build_container_command( ``--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. diff --git a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py index d56db848b19..a5d2bb93f6e 100644 --- a/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py +++ b/src/zenml/integrations/slurm/step_operators/slurm_step_operator.py @@ -73,6 +73,7 @@ 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" @@ -282,6 +283,7 @@ def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: 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 @@ -316,6 +318,37 @@ def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: 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. diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index fe2e36b2c3f..2f517ebbf0c 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -27,7 +27,6 @@ from zenml.config.resource_settings import ResourceSettings from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID from zenml.enums import ExecutionMode, ExecutionStatus, StackComponentType -from zenml.execution.pipeline.dynamic.runner import DynamicPipelineRunner from zenml.integrations.slurm.flavors import ( SlurmOrchestratorConfig, SlurmOrchestratorFlavor, @@ -37,8 +36,6 @@ from zenml.integrations.slurm.orchestrators.slurm_orchestrator import ( ENV_ZENML_SLURM_RUN_ID, SLURM_CLEANUP_JOB_ID_METADATA_KEY, - SLURM_ISOLATED_JOB_ID_METADATA_KEY, - SLURM_ISOLATED_JOB_IDS_METADATA_KEY, SLURM_JOB_IDS_METADATA_KEY, SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY, ) @@ -150,11 +147,15 @@ def _step(upstream: List[str], gpu: int = 0) -> SimpleNamespace: ) -def _snapshot(steps: Dict[str, SimpleNamespace]) -> SimpleNamespace: +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. @@ -165,7 +166,10 @@ def _snapshot(steps: Dict[str, SimpleNamespace]) -> SimpleNamespace: get_image=lambda component_key, step: IMAGE, ), step_configurations=steps, - pipeline_configuration=SimpleNamespace(name="p"), + pipeline_configuration=SimpleNamespace( + name="p", + resource_settings=pipeline_resources or ResourceSettings(), + ), schedule=None, stack=SimpleNamespace(id=uuid4()), ) @@ -228,22 +232,6 @@ def test_get_orchestrator_run_id_raises_when_unset(monkeypatch): op.get_orchestrator_run_id() -# --- topological ordering ----------------------------------------------------- - - -def test_sorted_steps_is_topological(): - """Upstream steps are submitted before their dependents.""" - # Declared out of order: a dependent appears before its parent. - steps = { - "evaluate": _step(["train"]), - "train": _step(["load"]), - "load": _step([]), - } - order = SlurmOrchestrator._sorted_step_names(steps) - assert order.index("load") < order.index("train") - assert order.index("train") < order.index("evaluate") - - # --- submit ------------------------------------------------------------------- @@ -289,28 +277,6 @@ def _stack( ) -def _isolated_step_info( - snapshot: Optional[SimpleNamespace] = None, -) -> SimpleNamespace: - """Build a dynamic isolated StepRunInfo stand-in. - - Args: - snapshot: Optional snapshot to attach to the step run info. - - Returns: - A namespace usable where StepRunInfo is expected. - """ - return SimpleNamespace( - step_run_id=uuid4(), - run_id=uuid4(), - pipeline_step_name="dynamic_train", - config=SimpleNamespace(resource_settings=ResourceSettings(cpu_count=2)), - snapshot=snapshot or _snapshot({}), - step_run=SimpleNamespace(run_metadata={}), - get_image=lambda key: IMAGE, - ) - - def _run( run_id, run_metadata: Dict[str, object], @@ -490,7 +456,11 @@ def test_submit_rejects_scheduled_pipelines(monkeypatch): def test_submit_dynamic_pipeline_submits_orchestration_job(monkeypatch): - """Dynamic pipelines launch a Slurm orchestration job.""" + """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() @@ -498,7 +468,9 @@ def test_submit_dynamic_pipeline_submits_orchestration_job(monkeypatch): placeholder = _placeholder_run() result = op.submit_dynamic_pipeline( - snapshot=_snapshot({}), + snapshot=_snapshot( + {}, pipeline_resources=ResourceSettings(gpu_count=2) + ), stack=_stack(), environment={"ZENML_STORE_API_KEY": SECRET_TOKEN}, placeholder_run=placeholder, @@ -512,6 +484,8 @@ def test_submit_dynamic_pipeline_submits_orchestration_job(monkeypatch): sbatch = [c for c in runner.commands if c.startswith("sbatch")] assert len(sbatch) == 1 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 @@ -622,11 +596,7 @@ def run(self, command: str) -> CommandResult: if command.startswith("squeue"): return CommandResult( exit_code=0, - stdout=( - "1000|FAILED\n" - "1001|RUNNING\n" - "1002|PENDING\n" - ), + stdout=("1000|FAILED\n1001|RUNNING\n1002|PENDING\n"), stderr="", ) return CommandResult(exit_code=0, stdout="", stderr="") @@ -671,11 +641,7 @@ def run(self, command: str) -> CommandResult: if command.startswith("squeue"): return CommandResult( exit_code=0, - stdout=( - "1000|FAILED\n" - "1001|RUNNING\n" - "1002|PENDING\n" - ), + stdout=("1000|FAILED\n1001|RUNNING\n1002|PENDING\n"), stderr="", ) return CommandResult(exit_code=0, stdout="", stderr="") @@ -698,12 +664,9 @@ def run(self, command: str) -> CommandResult: pipeline_status, _ = op.fetch_status(run) assert pipeline_status is ExecutionStatus.FAILED - assert any( - command == "scancel 1001 1002" for command in runner.commands - ) + 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" + 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"] @@ -737,9 +700,7 @@ def test_stop_run_cancels_submitted_jobs(monkeypatch): op.stop_run(run) - assert any( - command == "scancel 1000 1001" for command in runner.commands - ) + 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" ) @@ -772,350 +733,20 @@ def test_fetch_status_reconciles_dynamic_orchestration_job(monkeypatch): assert step_statuses is None -def test_submit_isolated_step_publishes_job_metadata(monkeypatch): - """Dynamic isolated steps are submitted as separate Slurm jobs.""" - op = _build_orchestrator() - op.get_settings = lambda _info: SlurmOrchestratorSettings() - runner = FakeRunner() - _use_fake_client(monkeypatch, runner) - monkeypatch.setattr( - "zenml.integrations.slurm.orchestrators.slurm_orchestrator" - ".orchestrator_utils.get_step_entrypoint_command", - lambda **kwargs: (["python"], ["-m", "zenml.entrypoint"]), - ) - monkeypatch.setattr( - "zenml.integrations.slurm.orchestrators.slurm_orchestrator" - ".Stack.from_model", - lambda stack: _stack(), - ) - published_step_metadata = {} - monkeypatch.setattr( - "zenml.integrations.slurm.orchestrators.slurm_orchestrator" - ".publish_step_run_metadata", - lambda **kwargs: published_step_metadata.update(kwargs), - ) - published_run_metadata = {} - - class FakeClient: - def create_run_metadata(self, **kwargs): - published_run_metadata.update(kwargs) - - monkeypatch.setattr("zenml.client.Client", lambda: FakeClient()) - info = _isolated_step_info() - - op.submit_isolated_step( - step_run_info=info, - environment={"ZENML_STORE_API_KEY": SECRET_TOKEN}, - ) - - run_dir = f"/runs/{info.run_id}/isolated/{info.step_run_id}" - assert any(command.startswith("sbatch") for command in runner.commands) - assert runner.modes[f"{run_dir}/env"] == 0o600 - assert SECRET_TOKEN in runner.files[f"{run_dir}/env"] - assert info.step_run.run_metadata[SLURM_ISOLATED_JOB_ID_METADATA_KEY] == ( - "1000" - ) - assert published_step_metadata == { - "step_run_id": info.step_run_id, - "step_run_metadata": { - op.id: {SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"} - }, - } - assert published_run_metadata["metadata"] == { - SLURM_ISOLATED_JOB_IDS_METADATA_KEY: {str(info.step_run_id): "1000"} - } - - -def test_submit_isolated_step_cancels_on_run_metadata_failure(monkeypatch): - """Run-level metadata publication is part of the submission transaction.""" - op = _build_orchestrator() - op.get_settings = lambda _info: SlurmOrchestratorSettings() - runner = FakeRunner() - _use_fake_client(monkeypatch, runner) - monkeypatch.setattr( - "zenml.integrations.slurm.orchestrators.slurm_orchestrator" - ".orchestrator_utils.get_step_entrypoint_command", - lambda **kwargs: (["python"], ["-m", "zenml.entrypoint"]), - ) - monkeypatch.setattr( - "zenml.integrations.slurm.orchestrators.slurm_orchestrator" - ".Stack.from_model", - lambda stack: _stack(), - ) - monkeypatch.setattr( - "zenml.integrations.slurm.orchestrators.slurm_orchestrator" - ".publish_step_run_metadata", - lambda **kwargs: None, - ) - - class FailingClient: - def create_run_metadata(self, **kwargs): - raise RuntimeError("metadata store unavailable") - - monkeypatch.setattr("zenml.client.Client", lambda: FailingClient()) - info = _isolated_step_info() - - with pytest.raises(RuntimeError, match="Failed to publish Slurm run"): - op.submit_isolated_step( - step_run_info=info, - environment={"ZENML_STORE_API_KEY": SECRET_TOKEN}, - ) - - run_dir = f"/runs/{info.run_id}/isolated/{info.step_run_id}" - assert "scancel 1000" in runner.commands - assert any( - command.startswith("rm -rf --") - and f"{run_dir}/env" in command - for command in runner.commands - ) - - -def test_get_isolated_step_status_reads_sentinel(monkeypatch): - """Isolated step status uses the same Slurm sentinel mapping.""" - op = _build_orchestrator() - runner = FakeRunner() - _use_fake_client(monkeypatch, runner) - run_id = uuid4() - step_run_id = uuid4() - run_dir = op._isolated_run_dir(str(run_id), str(step_run_id)) - runner.files[f"{run_dir}/exit_code"] = "0\n" - step_run = SimpleNamespace( - id=step_run_id, - pipeline_run_id=run_id, - run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, - ) - - assert op.get_isolated_step_status(step_run) is ExecutionStatus.COMPLETED - +def test_does_not_support_isolated_steps(): + """Dynamic steps run inline in the orchestration job's allocation. -def test_get_isolated_step_status_waits_for_missing_sentinel(monkeypatch): - """A missing isolated sentinel gets a bounded retry window.""" - op = _build_orchestrator() - runner = FakeRunner() - _use_fake_client(monkeypatch, runner) - run_id = uuid4() - step_run_id = uuid4() - step_run = SimpleNamespace( - id=step_run_id, - pipeline_run_id=run_id, - run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, - ) - - assert ( - op.get_isolated_step_status(step_run) - is ExecutionStatus.PROVISIONING - ) - assert ( - op.get_isolated_step_status(step_run) - is ExecutionStatus.PROVISIONING - ) - assert op.get_isolated_step_status(step_run) is ExecutionStatus.FAILED - - -def test_get_isolated_step_status_normalizes_dynamic_runner_statuses( - monkeypatch, -): - """Isolated Slurm statuses use values supported by the dynamic monitor.""" - op = _build_orchestrator() - - class StatusRunner(FakeRunner): - def __init__(self, state: str) -> None: - """Initialize the fake. - - Args: - state: Slurm state to return for the isolated job. - """ - super().__init__() - self.state = state - - def run(self, command: str) -> CommandResult: - self.commands.append(command) - if command.startswith("squeue"): - return CommandResult( - exit_code=0, - stdout=f"1000|{self.state}\n", - stderr="", - ) - return CommandResult(exit_code=0, stdout="", stderr="") - - run_id = uuid4() - step_run_id = uuid4() - step_run = SimpleNamespace( - id=step_run_id, - pipeline_run_id=run_id, - run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, - ) - - runner = StatusRunner("PENDING") - _use_fake_client(monkeypatch, runner) - assert ( - op.get_isolated_step_status(step_run) - is ExecutionStatus.PROVISIONING - ) - - op._isolated_status_cache.clear() - runner = StatusRunner("CANCELLED") - _use_fake_client(monkeypatch, runner) - assert op.get_isolated_step_status(step_run) is ExecutionStatus.STOPPED - - -def test_dynamic_monitor_keeps_pending_slurm_step(monkeypatch): - """Pending Slurm jobs stay monitored and are not cleaned up early.""" - 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|PENDING\n", stderr="" - ) - return CommandResult(exit_code=0, stdout="", stderr="") - - runner = StatusRunner() - _use_fake_client(monkeypatch, runner) - step_run = SimpleNamespace( - id=uuid4(), - name="dynamic_train", - pipeline_run_id=uuid4(), - config=SimpleNamespace(step_operator=False), - run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, - ) - - dynamic_runner = DynamicPipelineRunner.__new__(DynamicPipelineRunner) - dynamic_runner._shutdown_requested = False - dynamic_runner._steps_to_monitor = {"dynamic_train": step_run} - dynamic_runner._orchestrator = op - dynamic_runner._step_operator = None - cleaned_steps: List[object] = [] - dynamic_runner._cleanup_isolated_step = cleaned_steps.append - - class StopEvent: - def wait(self, timeout: Optional[float] = None) -> None: - """Stop the monitor after one pass. - - Args: - timeout: Unused wait timeout. - """ - _ = timeout - dynamic_runner._shutdown_requested = True - - def clear(self) -> None: - """No-op event clear.""" - - dynamic_runner._monitoring_event = StopEvent() - - dynamic_runner._monitoring_loop() - - assert dynamic_runner._steps_to_monitor == {"dynamic_train": step_run} - assert cleaned_steps == [] - - -def test_get_isolated_step_status_batches_cached_run_jobs(monkeypatch): - """Isolated step polling shares one Slurm state query per cache window.""" - 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|PENDING\n1001|RUNNING\n", - stderr="", - ) - return CommandResult(exit_code=0, stdout="", stderr="") - - runner = StatusRunner() - client_builds: List[SlurmClient] = [] - - def build_client(config: SlurmOrchestratorConfig) -> SlurmClient: - """Build a fake Slurm client and track connection churn. - - Args: - config: Unused Slurm orchestrator config. - - Returns: - A Slurm client backed by the fake runner. - """ - _ = config - client = SlurmClient(runner) - client_builds.append(client) - return client - - monkeypatch.setattr( - "zenml.integrations.slurm.orchestrators.slurm_orchestrator" - ".build_slurm_client", - build_client, - ) - run_id = uuid4() - first_step = SimpleNamespace( - id=uuid4(), - pipeline_run_id=run_id, - run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, - ) - second_step = SimpleNamespace( - id=uuid4(), - pipeline_run_id=run_id, - run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1001"}, - ) - - class FakeClient: - def get_pipeline_run(self, *args, **kwargs): - return _run( - run_id=run_id, - run_metadata={ - SLURM_ISOLATED_JOB_IDS_METADATA_KEY: { - str(first_step.id): "1000", - str(second_step.id): "1001", - } - }, - ) - - monkeypatch.setattr("zenml.client.Client", lambda: FakeClient()) - - assert ( - op.get_isolated_step_status(first_step) - is ExecutionStatus.PROVISIONING - ) - assert op.get_isolated_step_status(second_step) is ExecutionStatus.RUNNING - assert [ - command for command in runner.commands if command.startswith("squeue") - ] == ["squeue --noheader --format='%i|%T' --jobs=1000,1001"] - assert len(client_builds) == 1 - - -def test_stop_and_cleanup_isolated_step(monkeypatch): - """Isolated step stop and cleanup address the step staging directory.""" + 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() - runner = FakeRunner() - _use_fake_client(monkeypatch, runner) - run_id = uuid4() - step_run_id = uuid4() - step_run = SimpleNamespace( - id=step_run_id, - pipeline_run_id=run_id, - run_metadata={SLURM_ISOLATED_JOB_ID_METADATA_KEY: "1000"}, - ) - - op.stop_isolated_step(step_run) - op.cleanup_isolated_step(step_run) - - run_dir = op._isolated_run_dir(str(run_id), str(step_run_id)) - assert "scancel 1000" in runner.commands - assert runner.files[f"{run_dir}/cancelled"] == "1\n" - assert any( - command.startswith("rm -rf --") - and f"{run_dir}/env" in command - for command in runner.commands - ) - assert any( - command == f"rm -rf -- {run_dir}" for command in runner.commands - ) + assert op.can_run_isolated_steps is False + assert op.can_stop_isolated_steps is False -def test_stop_run_cancels_dynamic_jobs(monkeypatch): - """Stopping a dynamic run cancels orchestration and isolated jobs.""" +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) @@ -1125,26 +756,15 @@ def test_stop_run_cancels_dynamic_jobs(monkeypatch): lambda **kwargs: None, ) run_id = uuid4() - step_run_id = uuid4() run = SimpleNamespace( id=run_id, orchestrator_run_id=str(run_id), - run_metadata={ - SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000", - SLURM_ISOLATED_JOB_IDS_METADATA_KEY: {str(step_run_id): "1001"}, - }, + run_metadata={SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: "1000"}, ) op.stop_run(run) - assert "scancel 1001" in runner.commands assert "scancel 1000" in runner.commands - assert ( - runner.files[ - f"{op._isolated_run_dir(str(run_id), str(step_run_id))}/cancelled" - ] - == "1\n" - ) assert ( runner.files[f"{op._orchestration_run_dir(str(run_id))}/cancelled"] == "1\n" diff --git a/tests/unit/integrations/slurm/test_slurm_step_operator.py b/tests/unit/integrations/slurm/test_slurm_step_operator.py index ffce853a6c6..60a540057d6 100644 --- a/tests/unit/integrations/slurm/test_slurm_step_operator.py +++ b/tests/unit/integrations/slurm/test_slurm_step_operator.py @@ -46,6 +46,7 @@ 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" @@ -84,9 +85,7 @@ def run(self, command: str) -> CommandResult: 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 "" - ) + 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="") @@ -493,6 +492,11 @@ def op_and_runner(monkeypatch): ".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 @@ -547,6 +551,40 @@ def test_get_status_without_queue_or_sentinel_fails(op_and_runner): ) +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 From 782d2ed137a6d99a63984fe42d21505ba3d58b16 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 17 Jul 2026 11:43:58 +0200 Subject: [PATCH 18/22] Keep dynamic runs unknown while the exit sentinel is unflushed Between a job leaving squeue --states=all and its exit-code sentinel becoming visible on the shared filesystem, the run's outcome is not knowable. Dynamic reconciliation previously passed cleanup_complete=True and declared FAILED on first sight of that window, terminally mis-marking runs whose sentinel merely lagged the queue purge; the static DAG path already treats this case as unknown. Mirror it, and warn on each poll so a job killed before its EXIT trap ran (NODE_FAIL) is visible to operators instead of silently pending. Co-Authored-By: Claude Fable 5 --- .../slurm/orchestrators/slurm_orchestrator.py | 20 ++++++++++++++++- .../slurm/test_slurm_orchestrator.py | 22 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 4890243e1ae..97a89e18841 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -726,8 +726,26 @@ def _fetch_dynamic_status( client=client, state=state, run_dir=self._orchestration_run_dir(str(run.id)), - cleanup_complete=True, + # The orchestration job writes its own exit-code sentinel + # via its EXIT trap. Until that sentinel is visible on the + # shared filesystem, a job that has vanished from the queue + # is unknown, not failed - mirroring the static DAG + # reconciliation. Declaring FAILED on first sight would + # terminally mis-mark a run whose sentinel merely lags the + # queue purge (NFS/Lustre flush). + cleanup_complete=False, ) + 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. If " + "this persists, the job was likely killed before its " + "EXIT trap ran (e.g. NODE_FAIL) - stop the run " + "manually.", + job_id, + run.id, + ) finally: client.runner.close() diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index 2f517ebbf0c..41501eee68b 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -733,6 +733,28 @@ def test_fetch_status_reconciles_dynamic_orchestration_job(monkeypatch): 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 + mis-mark 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_does_not_support_isolated_steps(): """Dynamic steps run inline in the orchestration job's allocation. From e64d1724ef217670346b2aa6fb42ad8f62a71a7d Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 17 Jul 2026 12:09:20 +0200 Subject: [PATCH 19/22] Fix CI: unused import, typos hit, and stale docs base The isolated-step removal orphaned the ORCHESTRATOR_DOCKER_IMAGE_KEY import, and "mis-mark" trips the typos checker. The develop merge also resolves the GitBook redirect check, which compared this branch's docs tree against a develop that had since gained the DigitalOcean pages. Co-Authored-By: Claude Fable 5 --- .../integrations/slurm/orchestrators/slurm_orchestrator.py | 7 ++----- tests/unit/integrations/slurm/test_slurm_orchestrator.py | 2 +- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 97a89e18841..ba2de78040d 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -46,10 +46,7 @@ from zenml.config.base_settings import BaseSettings from zenml.config.resource_settings import ResourceSettings -from zenml.constants import ( - METADATA_ORCHESTRATOR_RUN_ID, - ORCHESTRATOR_DOCKER_IMAGE_KEY, -) +from zenml.constants import METADATA_ORCHESTRATOR_RUN_ID from zenml.entrypoints.step_entrypoint_configuration import ( StepEntrypointConfiguration, ) @@ -731,7 +728,7 @@ def _fetch_dynamic_status( # shared filesystem, a job that has vanished from the queue # is unknown, not failed - mirroring the static DAG # reconciliation. Declaring FAILED on first sight would - # terminally mis-mark a run whose sentinel merely lags the + # terminally mislabel a run whose sentinel merely lags the # queue purge (NFS/Lustre flush). cleanup_complete=False, ) diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index 41501eee68b..a353140f68e 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -739,7 +739,7 @@ def test_fetch_status_keeps_dynamic_run_unknown_without_sentinel(monkeypatch): 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 - mis-mark runs that actually completed. + mislabel runs that actually completed. """ op = _build_orchestrator() runner = FakeRunner() From 6a1a9cf25430fa620cad59d4e169c8255f52b718 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 17 Jul 2026 12:17:15 +0200 Subject: [PATCH 20/22] Give dynamic runs the same credential reaper as static runs The orchestration job's EXIT trap scrubs its staged credentials on clean exits, but a hard kill (NODE_FAIL, OOM, hard timeout) never runs it, stranding the env file and registry auth on the shared filesystem. Submit the same afterany cleanup job static runs get; if the reaper cannot be queued, cancel the orchestration job and scrub immediately rather than run without one. The cleanup job's completion marker also restores conclusive terminal resolution for dynamic reconciliation: a job gone from the queue with no exit sentinel stays unknown only until the cleanup job has run, then becomes FAILED - the same semantics the static DAG path has, replacing the previous never-resolving warning-only behavior. Also drop the typing.Any import orphaned by the isolated-step removal (CI lints with --isolated, which local per-file ignores had masked). Co-Authored-By: Claude Fable 5 --- .../slurm/orchestrators/slurm_orchestrator.py | 65 +++++++++++++++---- .../slurm/test_slurm_orchestrator.py | 28 +++++++- 2 files changed, 81 insertions(+), 12 deletions(-) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index ba2de78040d..4932b06d4ca 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -35,7 +35,6 @@ import shlex from typing import ( TYPE_CHECKING, - Any, Callable, Dict, List, @@ -374,7 +373,7 @@ def submit_dynamic_pipeline( client = build_slurm_client(self.config) try: - job_id, _ = self._submit_container_job( + job_id, sensitive_paths = self._submit_container_job( client=client, run_id=run_id, run_dir=self._orchestration_run_dir(run_id), @@ -389,6 +388,39 @@ def submit_dynamic_pipeline( 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() @@ -401,6 +433,7 @@ def submit_dynamic_pipeline( metadata={ METADATA_ORCHESTRATOR_RUN_ID: run_id, SLURM_ORCHESTRATION_JOB_ID_METADATA_KEY: job_id, + SLURM_CLEANUP_JOB_ID_METADATA_KEY: cleanup_job_id, } ) @@ -716,30 +749,40 @@ def _fetch_dynamic_status( 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. Until that sentinel is visible on the - # shared filesystem, a job that has vanished from the queue - # is unknown, not failed - mirroring the static DAG - # reconciliation. Declaring FAILED on first sight would + # 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=False, + 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. If " - "this persists, the job was likely killed before its " - "EXIT trap ran (e.g. NODE_FAIL) - stop the run " - "manually.", + "sentinel yet; keeping the run status unchanged until " + "the cleanup job resolves it.", job_id, run.id, ) diff --git a/tests/unit/integrations/slurm/test_slurm_orchestrator.py b/tests/unit/integrations/slurm/test_slurm_orchestrator.py index a353140f68e..f295f0a6cf0 100644 --- a/tests/unit/integrations/slurm/test_slurm_orchestrator.py +++ b/tests/unit/integrations/slurm/test_slurm_orchestrator.py @@ -480,15 +480,22 @@ def test_submit_dynamic_pipeline_submits_orchestration_job(monkeypatch): 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) == 1 + 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): @@ -755,6 +762,25 @@ def test_fetch_status_keeps_dynamic_run_unknown_without_sentinel(monkeypatch): 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. From 6f172da3f518e5e2def7b355ada3ba618b00d061 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 17 Jul 2026 12:31:47 +0200 Subject: [PATCH 21/22] Document the cleanup-submission re-raise in the dynamic docstring pydoclint (DOC503) requires the bare re-raise in the credential-reaper error path to appear in the Raises section. Co-Authored-By: Claude Fable 5 --- .../integrations/slurm/orchestrators/slurm_orchestrator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py index 4932b06d4ca..3698f8d375c 100644 --- a/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py +++ b/src/zenml/integrations/slurm/orchestrators/slurm_orchestrator.py @@ -345,6 +345,8 @@ def submit_dynamic_pipeline( 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, From 1590627c1e5ebe448b5d655a63d2b678c98dc979 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 17 Jul 2026 16:26:32 +0200 Subject: [PATCH 22/22] Format slurm_client to the CI ruff line length ruff format --check in CI flags the get_job_states signature; the local ruff version had accepted it. Co-Authored-By: Claude Fable 5 --- src/zenml/integrations/slurm/slurm_client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/zenml/integrations/slurm/slurm_client.py b/src/zenml/integrations/slurm/slurm_client.py index af64db1c81b..9882fe30f1e 100644 --- a/src/zenml/integrations/slurm/slurm_client.py +++ b/src/zenml/integrations/slurm/slurm_client.py @@ -283,7 +283,9 @@ def submit( ) return job_id - def get_job_states(self, job_ids: Sequence[str]) -> Dict[str, Optional[str]]: + def get_job_states( + self, job_ids: Sequence[str] + ) -> Dict[str, Optional[str]]: """Get the queue states of multiple jobs by their IDs. Args: