snakemake-executor-plugin-slurm 0.4.2__tar.gz → 0.12.1__tar.gz

{snakemake_executor_plugin_slurm-0.4.2 → snakemake_executor_plugin_slurm-0.12.1}/PKG-INFO

@@ -1,8 +1,7 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: snakemake-executor-plugin-slurm
-Version: 0.4.2
+Version: 0.12.1
 Summary: A Snakemake executor plugin for submitting jobs to a SLURM cluster.
-Home-page: https://github.com/snakemake/snakemake-executor-plugin-slurm
 License: MIT
 Keywords: snakemake,plugin,executor,cluster,slurm
 Author: Christian Meesters
@@ -12,9 +11,10 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: snakemake-executor-plugin-slurm-jobstep (>=0.1.10,<0.2.0)
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: snakemake-executor-plugin-slurm-jobstep (>=0.2.0,<0.3.0)
 Requires-Dist: snakemake-interface-common (>=1.13.0,<2.0.0)
-Requires-Dist: snakemake-interface-executor-plugins (>=9.0.0,<10.0.0)
+Requires-Dist: snakemake-interface-executor-plugins (>=9.1.1,<10.0.0)
 Requires-Dist: throttler (>=1.2.2,<2.0.0)
 Project-URL: Documentation, https://snakemake.github.io/snakemake-plugin-catalog/plugins/executor/slurm.html
 Project-URL: Repository, https://github.com/snakemake/snakemake-executor-plugin-slurm

{snakemake_executor_plugin_slurm-0.4.2 → snakemake_executor_plugin_slurm-0.12.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "snakemake-executor-plugin-slurm"
-version = "0.4.2"
+version = "0.12.1"
 description = "A Snakemake executor plugin for submitting jobs to a SLURM cluster."
 authors = [
     "Christian Meesters <meesters@uni-mainz.de>",
@@ -16,8 +16,8 @@ keywords = ["snakemake", "plugin", "executor", "cluster", "slurm"]
 [tool.poetry.dependencies]
 python = "^3.11"
 snakemake-interface-common = "^1.13.0"
-snakemake-interface-executor-plugins = "^9.0.0"
-snakemake-executor-plugin-slurm-jobstep = "^0.1.10"
+snakemake-interface-executor-plugins = "^9.1.1"
+snakemake-executor-plugin-slurm-jobstep = "^0.2.0"
 throttler = "^1.2.2"
 
 [tool.poetry.group.dev.dependencies]
@@ -25,7 +25,7 @@ black = "^23.7.0"
 flake8 = "^6.1.0"
 coverage = "^7.3.1"
 pytest = "^7.4.2"
-snakemake = {git = "https://github.com/snakemake/snakemake.git"}
+snakemake = "^8.20.0"
 
 [tool.coverage.run]
 omit = [".*", "*/site-packages/*", "Snakefile"]

{snakemake_executor_plugin_slurm-0.4.2 → snakemake_executor_plugin_slurm-0.12.1}/snakemake_executor_plugin_slurm/__init__.py

@@ -3,21 +3,88 @@ __copyright__ = "Copyright 2023, David Lähnemann, Johannes Köster, Christian M
 __email__ = "johannes.koester@uni-due.de"
 __license__ = "MIT"
 
+import atexit
 import csv
 from io import StringIO
 import os
+from pathlib import Path
+import re
+import shlex
 import subprocess
 import time
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta
-from typing import List, Generator
+from typing import List, Generator, Optional
 import uuid
 from snakemake_interface_executor_plugins.executors.base import SubmittedJobInfo
 from snakemake_interface_executor_plugins.executors.remote import RemoteExecutor
-from snakemake_interface_executor_plugins.settings import CommonSettings
+from snakemake_interface_executor_plugins.settings import (
+    ExecutorSettingsBase,
+    CommonSettings,
+)
 from snakemake_interface_executor_plugins.jobs import (
     JobExecutorInterface,
 )
 from snakemake_interface_common.exceptions import WorkflowError
+from snakemake_executor_plugin_slurm_jobstep import get_cpus_per_task
+
+from .utils import delete_slurm_environment, delete_empty_dirs
+
+
+@dataclass
+class ExecutorSettings(ExecutorSettingsBase):
+    logdir: Optional[Path] = field(
+        default=None,
+        metadata={
+            "help": "Per default the SLURM log directory is relative to "
+            "the working directory."
+            "This flag allows to set an alternative directory.",
+            "env_var": False,
+            "required": False,
+        },
+    )
+    keep_successful_logs: bool = field(
+        default=False,
+        metadata={
+            "help": "Per default SLURM log files will be deleted upon sucessful "
+            "completion of a job. Whenever a SLURM job fails, its log "
+            "file will be preserved. "
+            "This flag allows to keep all SLURM log files, even those "
+            "of successful jobs.",
+            "env_var": False,
+            "required": False,
+        },
+    )
+    delete_logfiles_older_than: Optional[int] = field(
+        default=10,
+        metadata={
+            "help": "Per default SLURM log files in the SLURM log directory "
+            "of a workflow will be deleted after 10 days. For this, "
+            "best leave the default log directory unaltered. "
+            "Setting this flag allows to change this behaviour. "
+            "If set to <=0, no old files will be deleted. ",
+        },
+    )
+    init_seconds_before_status_checks: Optional[int] = field(
+        default=40,
+        metadata={
+            "help": "Defines the time in seconds before the first status "
+            "check is performed after job submission.",
+            "env_var": False,
+            "required": False,
+        },
+    )
+    requeue: bool = field(
+        default=False,
+        metadata={
+            "help": "Allow requeuing preempted of failed jobs, "
+            "if no cluster default. Results in "
+            "`sbatch ... --requeue ...` "
+            "This flag has no effect, if not set.",
+            "env_var": False,
+            "required": False,
+        },
+    )
 
 
 # Required:
@@ -48,10 +115,51 @@ common_settings = CommonSettings(
 # Implementation of your executor
 class Executor(RemoteExecutor):
     def __post_init__(self):
+        # run check whether we are running in a SLURM job context
+        self.warn_on_jobcontext()
         self.run_uuid = str(uuid.uuid4())
         self.logger.info(f"SLURM run ID: {self.run_uuid}")
         self._fallback_account_arg = None
         self._fallback_partition = None
+        self._preemption_warning = False  # no preemption warning has been issued
+        self.slurm_logdir = None
+        atexit.register(self.clean_old_logs)
+
+    def clean_old_logs(self) -> None:
+        """Delete files older than specified age from the SLURM log directory."""
+        # shorthands:
+        age_cutoff = self.workflow.executor_settings.delete_logfiles_older_than
+        keep_all = self.workflow.executor_settings.keep_successful_logs
+        if age_cutoff <= 0 or keep_all:
+            return
+        cutoff_secs = age_cutoff * 86400
+        current_time = time.time()
+        self.logger.info(f"Cleaning up log files older than {age_cutoff} day(s)")
+        for path in self.slurm_logdir.rglob("*.log"):
+            if path.is_file():
+                try:
+                    file_age = current_time - path.stat().st_mtime
+                    if file_age > cutoff_secs:
+                        path.unlink()
+                except (OSError, FileNotFoundError) as e:
+                    self.logger.warning(f"Could not delete logfile {path}: {e}")
+        # we need a 2nd iteration to remove putatively empty directories
+        try:
+            delete_empty_dirs(self.slurm_logdir)
+        except (OSError, FileNotFoundError) as e:
+            self.logger.warning(f"Could not delete empty directory {path}: {e}")
+
+    def warn_on_jobcontext(self, done=None):
+        if not done:
+            if "SLURM_JOB_ID" in os.environ:
+                self.logger.warning(
+                    "You are running snakemake in a SLURM job context. "
+                    "This is not recommended, as it may lead to unexpected behavior. "
+                    "Please run Snakemake directly on the login node."
+                )
+                time.sleep(5)
+                delete_slurm_environment()
+        done = True
 
     def additional_general_args(self):
         return "--executor slurm-jobstep --jobs 1"
@@ -65,29 +173,55 @@ class Executor(RemoteExecutor):
         # with job_info being of type
         # snakemake_interface_executor_plugins.executors.base.SubmittedJobInfo.
 
-        log_folder = f"group_{job.name}" if job.is_group() else f"rule_{job.name}"
+        group_or_rule = f"group_{job.name}" if job.is_group() else f"rule_{job.name}"
 
         try:
-            wildcard_str = f"_{'_'.join(job.wildcards)}" if job.wildcards else ""
+            wildcard_str = "_".join(job.wildcards) if job.wildcards else ""
         except AttributeError:
             wildcard_str = ""
 
-        slurm_logfile = os.path.abspath(
-            f".snakemake/slurm_logs/{log_folder}/%j{wildcard_str}.log"
+        self.slurm_logdir = (
+            Path(self.workflow.executor_settings.logdir)
+            if self.workflow.executor_settings.logdir
+            else Path(".snakemake/slurm_logs").resolve()
+        )
+
+        self.slurm_logdir.mkdir(parents=True, exist_ok=True)
+        slurm_logfile = self.slurm_logdir / group_or_rule / wildcard_str / "%j.log"
+        slurm_logfile.parent.mkdir(parents=True, exist_ok=True)
+        # this behavior has been fixed in slurm 23.02, but there might be plenty of
+        # older versions around, hence we should rather be conservative here.
+        assert "%j" not in str(self.slurm_logdir), (
+            "bug: jobid placeholder in parent dir of logfile. This does not work as "
+            "we have to create that dir before submission in order to make sbatch "
+            "happy. Otherwise we get silent fails without logfiles being created."
         )
-        os.makedirs(os.path.dirname(slurm_logfile), exist_ok=True)
 
         # generic part of a submission string:
         # we use a run_uuid as the job-name, to allow `--name`-based
         # filtering in the job status checks (`sacct --name` and `squeue --name`)
+        if wildcard_str == "":
+            comment_str = f"rule_{job.name}"
+        else:
+            comment_str = f"rule_{job.name}_wildcards_{wildcard_str}"
         call = (
-            f"sbatch --job-name {self.run_uuid} --output {slurm_logfile} --export=ALL "
-            f"--comment {job.name}"
+            f"sbatch "
+            f"--parsable "
+            f"--job-name {self.run_uuid} "
+            f"--output '{slurm_logfile}' "
+            f"--export=ALL "
+            f"--comment {comment_str}"
         )
 
         call += self.get_account_arg(job)
         call += self.get_partition_arg(job)
 
+        if self.workflow.executor_settings.requeue:
+            call += " --requeue"
+
+        if job.resources.get("clusters"):
+            call += f" --clusters {job.resources.clusters}"
+
         if job.resources.get("runtime"):
             call += f" -t {job.resources.runtime}"
         else:
@@ -99,7 +233,7 @@ class Executor(RemoteExecutor):
             )
 
         if job.resources.get("constraint"):
-            call += f" -C {job.resources.constraint}"
+            call += f" -C '{job.resources.constraint}'"
         if job.resources.get("mem_mb_per_cpu"):
             call += f" --mem-per-cpu {job.resources.mem_mb_per_cpu}"
         elif job.resources.get("mem_mb"):
@@ -110,28 +244,27 @@ class Executor(RemoteExecutor):
                 "- submitting without. This might or might not work on your cluster."
             )
 
-        # MPI job
-        if job.resources.get("mpi", False):
-            if job.resources.get("nodes", False):
-                call += f" --nodes={job.resources.get('nodes', 1)}"
+        if job.resources.get("nodes", False):
+            call += f" --nodes={job.resources.get('nodes', 1)}"
 
-        # fixes #40 - set ntasks regarlless of mpi, because
+        # fixes #40 - set ntasks regardless of mpi, because
         # SLURM v22.05 will require it for all jobs
         call += f" --ntasks={job.resources.get('tasks', 1)}"
-
-        cpus_per_task = job.threads
-        if job.resources.get("cpus_per_task"):
-            if not isinstance(cpus_per_task, int):
-                raise WorkflowError(
-                    f"cpus_per_task must be an integer, but is {cpus_per_task}"
+        # MPI job
+        if job.resources.get("mpi", False):
+            if not job.resources.get("tasks_per_node") and not job.resources.get(
+                "nodes"
+            ):
+                self.logger.warning(
+                    "MPI job detected, but no 'tasks_per_node' or 'nodes' "
+                    "specified. Assuming 'tasks_per_node=1'."
+                    "Probably not what you want."
                 )
-            cpus_per_task = job.resources.cpus_per_task
-        # ensure that at least 1 cpu is requested
-        # because 0 is not allowed by slurm
-        cpus_per_task = max(1, cpus_per_task)
-        call += f" --cpus-per-task={cpus_per_task}"
+
+        call += f" --cpus-per-task={get_cpus_per_task(job)}"
 
         if job.resources.get("slurm_extra"):
+            self.check_slurm_extra(job)
             call += f" {job.resources.slurm_extra}"
 
         exec_job = self.format_job_exec(job)
@@ -145,16 +278,39 @@ class Executor(RemoteExecutor):
 
         self.logger.debug(f"sbatch call: {call}")
         try:
-            out = subprocess.check_output(
-                call, shell=True, text=True, stderr=subprocess.STDOUT
-            ).strip()
+            process = subprocess.Popen(
+                call,
+                shell=True,
+                text=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            out, err = process.communicate()
+            if process.returncode != 0:
+                raise subprocess.CalledProcessError(
+                    process.returncode, call, output=err
+                )
         except subprocess.CalledProcessError as e:
             raise WorkflowError(
-                f"SLURM job submission failed. The error message was {e.output}"
+                f"SLURM sbatch failed. The error message was {e.output}"
+            )
+        # any other error message indicating failure?
+        if "submission failed" in err:
+            raise WorkflowError(
+                f"SLURM job submission failed. The error message was {err}"
             )
 
-        slurm_jobid = out.split(" ")[-1]
-        slurm_logfile = slurm_logfile.replace("%j", slurm_jobid)
+        # multicluster submissions yield submission infos like
+        # "Submitted batch job <id> on cluster <name>" by default, but with the
+        # --parsable option it simply yields "<id>;<name>".
+        # To extract the job id we split by semicolon and take the first element
+        # (this also works if no cluster name was provided)
+        slurm_jobid = out.strip().split(";")[0]
+        if not slurm_jobid:
+            raise WorkflowError("Failed to retrieve SLURM job ID from sbatch output.")
+        slurm_logfile = slurm_logfile.with_name(
+            slurm_logfile.name.replace("%j", slurm_jobid)
+        )
         self.logger.info(
             f"Job {job.jobid} has been submitted with SLURM jobid {slurm_jobid} "
             f"(log: {slurm_logfile})."
@@ -169,7 +325,6 @@ class Executor(RemoteExecutor):
         self, active_jobs: List[SubmittedJobInfo]
     ) -> Generator[SubmittedJobInfo, None, None]:
         # Check the status of active jobs.
-
         # You have to iterate over the given list active_jobs.
         # For jobs that have finished successfully, you have to call
         # self.report_job_success(job).
@@ -189,7 +344,6 @@ class Executor(RemoteExecutor):
             "FAILED",
             "NODE_FAIL",
             "OUT_OF_MEMORY",
-            "PREEMPTED",
             "TIMEOUT",
             "ERROR",
         )
@@ -212,21 +366,28 @@ class Executor(RemoteExecutor):
 
         # We use this sacct syntax for argument 'starttime' to keep it compatible
         # with slurm < 20.11
-        sacct_starttime = f"{datetime.now() - timedelta(days=2):%Y-%m-%dT%H:00}"
+        sacct_starttime = f"{datetime.now() - timedelta(days = 2):%Y-%m-%dT%H:00}"
         # previously we had
         # f"--starttime now-2days --endtime now --name {self.run_uuid}"
         # in line 218 - once v20.11 is definitively not in use any more,
         # the more readable version ought to be re-adapted
 
+        # -X: only show main job, no substeps
+        sacct_command = f"""sacct -X --parsable2 \
+                        --clusters all \
+                        --noheader --format=JobIdRaw,State \
+                        --starttime {sacct_starttime} \
+                        --endtime now --name {self.run_uuid}"""
+
+        # for better redability in verbose output
+        sacct_command = " ".join(shlex.split(sacct_command))
+
         # this code is inspired by the snakemake profile:
         # https://github.com/Snakemake-Profiles/slurm
         for i in range(status_attempts):
             async with self.status_rate_limiter:
                 (status_of_jobs, sacct_query_duration) = await self.job_stati(
-                    # -X: only show main job, no substeps
-                    f"sacct -X --parsable2 --noheader --format=JobIdRaw,State "
-                    f"--starttime {sacct_starttime} "
-                    f"--endtime now --name {self.run_uuid}"
+                    sacct_command
                 )
                 if status_of_jobs is None and sacct_query_duration is None:
                     self.logger.debug(f"could not check status of job {self.run_uuid}")
@@ -282,6 +443,30 @@ class Executor(RemoteExecutor):
                     self.report_job_success(j)
                     any_finished = True
                     active_jobs_seen_by_sacct.remove(j.external_jobid)
+                    if not self.workflow.executor_settings.keep_successful_logs:
+                        self.logger.debug(
+                            "removing log for successful job "
+                            f"with SLURM ID '{j.external_jobid}'"
+                        )
+                        try:
+                            if j.aux["slurm_logfile"].exists():
+                                j.aux["slurm_logfile"].unlink()
+                        except (OSError, FileNotFoundError) as e:
+                            self.logger.warning(
+                                "Could not remove log file"
+                                f" {j.aux['slurm_logfile']}: {e}"
+                            )
+                elif status == "PREEMPTED" and not self._preemption_warning:
+                    self._preemption_warning = True
+                    self.logger.warning(
+                        """
+===== A Job preemption  occured! =====
+Leave Snakemake running, if possible. Otherwise Snakemake
+needs to restart this job upon a Snakemake restart.
+
+We leave it to SLURM to resume your job(s)"""
+                    )
+                    yield j
                 elif status == "UNKNOWN":
                     # the job probably does not exist anymore, but 'sacct' did not work
                     # so we assume it is finished
@@ -291,9 +476,13 @@ class Executor(RemoteExecutor):
                 elif status in fail_stati:
                     msg = (
                         f"SLURM-job '{j.external_jobid}' failed, SLURM status is: "
-                        f"'{status}'"
+                        # message ends with '. ', because it is proceeded
+                        # with a new sentence
+                        f"'{status}'. "
+                    )
+                    self.report_job_error(
+                        j, msg=msg, aux_logs=[j.aux["slurm_logfile"]._str]
                     )
-                    self.report_job_error(j, msg=msg, aux_logs=[j.aux["slurm_logfile"]])
                     active_jobs_seen_by_sacct.remove(j.external_jobid)
                 else:  # still running?
                     yield j
@@ -316,8 +505,10 @@ class Executor(RemoteExecutor):
                 # about 30 sec, but can be longer in extreme cases.
                 # Under 'normal' circumstances, 'scancel' is executed in
                 # virtually no time.
+                scancel_command = f"scancel {jobids} --clusters=all"
+
                 subprocess.check_output(
-                    f"scancel {jobids}",
+                    scancel_command,
                     text=True,
                     shell=True,
                     timeout=60,
@@ -325,6 +516,14 @@ class Executor(RemoteExecutor):
                 )
             except subprocess.TimeoutExpired:
                 self.logger.warning("Unable to cancel jobs within a minute.")
+            except subprocess.CalledProcessError as e:
+                msg = e.stderr.strip()
+                if msg:
+                    msg = f": {msg}"
+                raise WorkflowError(
+                    "Unable to cancel jobs with scancel "
+                    f"(exit code {e.returncode}){msg}"
+                ) from e
 
     async def job_stati(self, command):
         """Obtain SLURM job status of all submitted jobs with sacct
@@ -371,13 +570,14 @@ class Executor(RemoteExecutor):
             # here, we check whether the given or guessed account is valid
             # if not, a WorkflowError is raised
             self.test_account(job.resources.slurm_account)
-            return f" -A {job.resources.slurm_account}"
+            return f" -A '{job.resources.slurm_account}'"
         else:
             if self._fallback_account_arg is None:
                 self.logger.warning("No SLURM account given, trying to guess.")
                 account = self.get_account()
                 if account:
                     self.logger.warning(f"Guessed SLURM account: {account}")
+                    self.test_account(f"{account}")
                     self._fallback_account_arg = f" -A {account}"
                 else:
                     self.logger.warning(
@@ -415,7 +615,7 @@ class Executor(RemoteExecutor):
             sacct_out = subprocess.check_output(
                 cmd, shell=True, text=True, stderr=subprocess.PIPE
             )
-            return sacct_out.strip()
+            return sacct_out.replace("(null)", "").strip()
         except subprocess.CalledProcessError as e:
             self.logger.warning(
                 f"No account was given, not able to get a SLURM account via sacct: "
@@ -433,12 +633,28 @@ class Executor(RemoteExecutor):
                 cmd, shell=True, text=True, stderr=subprocess.PIPE
             )
         except subprocess.CalledProcessError as e:
-            raise WorkflowError(
-                f"Unable to test the validity of the given or guessed SLURM account "
-                f"'{account}' with sacctmgr: {e.stderr}"
+            sacctmgr_report = (
+                "Unable to test the validity of the given or guessed "
+                f"SLURM account '{account}' with sacctmgr: {e.stderr}."
             )
+            try:
+                cmd = "sshare -U --format Account --noheader"
+                accounts = subprocess.check_output(
+                    cmd, shell=True, text=True, stderr=subprocess.PIPE
+                )
+            except subprocess.CalledProcessError as e2:
+                sshare_report = (
+                    "Unable to test the validity of the given or guessed"
+                    f" SLURM account '{account}' with sshare: {e2.stderr}."
+                )
+                raise WorkflowError(
+                    f"The 'sacctmgr' reported: '{sacctmgr_report}' "
+                    f"and likewise 'sshare' reported: '{sshare_report}'."
+                )
 
-        accounts = accounts.split()
+        # The set() has been introduced during review to eliminate
+        # duplicates. They are not harmful, but disturbing to read.
+        accounts = set(_.strip() for _ in accounts.split("\n") if _)
 
         if account not in accounts:
             raise WorkflowError(
@@ -473,3 +689,15 @@ class Executor(RemoteExecutor):
             "'slurm_partition=<your default partition>'."
         )
         return ""
+
+    def check_slurm_extra(self, job):
+        jobname = re.compile(r"--job-name[=?|\s+]|-J\s?")
+        if re.search(jobname, job.resources.slurm_extra):
+            raise WorkflowError(
+                "The --job-name option is not allowed in the 'slurm_extra' "
+                "parameter. The job name is set by snakemake and must not be "
+                "overwritten. It is internally used to check the stati of the "
+                "all submitted jobs by this workflow."
+                "Please consult the documentation if you are unsure how to "
+                "query the status of your jobs."
+            )

snakemake_executor_plugin_slurm-0.12.1/snakemake_executor_plugin_slurm/utils.py

@@ -0,0 +1,42 @@
+# utility functions for the SLURM executor plugin
+
+import os
+from pathlib import Path
+
+
+def delete_slurm_environment():
+    """
+    Function to delete all environment variables
+    starting with 'SLURM_'. The parent shell will
+    still have this environment. This is needed to
+    submit within a SLURM job context to avoid
+    conflicting environments.
+    """
+    for var in os.environ:
+        if var.startswith("SLURM_"):
+            del os.environ[var]
+
+
+def delete_empty_dirs(path: Path) -> None:
+    """
+    Function to delete all empty directories in a given path.
+    This is needed to clean up the working directory after
+    a job has sucessfully finished. This function is needed because
+    the shutil.rmtree() function does not delete empty
+    directories.
+    """
+    if not path.is_dir():
+        return
+
+    # Process subdirectories first (bottom-up)
+    for child in path.iterdir():
+        if child.is_dir():
+            delete_empty_dirs(child)
+
+    try:
+        # Check if directory is now empty after processing children
+        if not any(path.iterdir()):
+            path.rmdir()
+    except (OSError, FileNotFoundError) as e:
+        # Provide more context in the error message
+        raise OSError(f"Failed to remove empty directory {path}: {e}") from e