stepup-queue 1.0.6__tar.gz → 1.0.7__tar.gz

{stepup_queue-1.0.6/stepup_queue.egg-info → stepup_queue-1.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stepup-queue
-Version: 1.0.6
+Version: 1.0.7
 Summary: StepUp Queue integrates queued jobs into a StepUp workflow.
 Author-email: Toon Verstraelen <toon.verstraelen@ugent.be>
 License-Expression: GPL-3.0-or-later
@@ -24,7 +24,7 @@ Classifier: Topic :: Software Development :: Build Tools
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: stepup<4.0.0,>=3.1.3
+Requires-Dist: stepup<4.0.0,>=3.1.4
 Provides-Extra: dev
 Requires-Dist: psutil; extra == "dev"
 Requires-Dist: pytest; extra == "dev"

{stepup_queue-1.0.6 → stepup_queue-1.0.7}/pyproject.toml

@@ -28,7 +28,7 @@ classifiers = [
 ]
 dependencies = [
     # Ensure changes to these dependencies are reflected in .github/requirements-old.txt
-    "stepup>=3.1.3,<4.0.0",
+    "stepup>=3.1.4,<4.0.0",
 ]
 dynamic = ["version"]
 

{stepup_queue-1.0.6 → stepup_queue-1.0.7}/stepup/queue/actions.py

@@ -49,6 +49,9 @@ def sbatch(argstr: str, work_thread: WorkThread) -> int:
         # Cancel running job (if any), clean log and resubmit
         path_log = Path("slurmjob.log")
         job_id, cluster = read_jobid_cluster(path_log)
-        work_thread.runsh(f"scancel -M {cluster} {job_id}")
+        if cluster is None:
+            work_thread.runsh(f"scancel {job_id}")
+        else:
+            work_thread.runsh(f"scancel -M {cluster} {job_id}")
         path_log.remove_p()
     return submit_once_and_wait(work_thread, args.ext, args.rc, args.onchange != "ignore")

{stepup_queue-1.0.6 → stepup_queue-1.0.7}/stepup/queue/canceljobs.py

@@ -20,16 +20,17 @@
 """Tool to cancel jobs."""
 
 import argparse
-import os
+import subprocess
 
 from path import Path
 
-from .sbatch import FIRST_LINE
+from .sbatch import FIRST_LINE, parse_sbatch
 
 
 def canceljobs_tool(args: argparse.Namespace) -> int:
     if len(args.paths) == 0:
         args.paths = [Path(".")]
+
     # Iterate over all slurmjob.log files in the specified directories, and kill them.
     job_ids = {}
     for path in args.paths:
@@ -39,18 +40,42 @@ def canceljobs_tool(args: argparse.Namespace) -> int:
         if not path.is_dir():
             print(f"Path {path} is not a directory.")
             continue
-        for job_log in path.glob("**/slurmjob.log"):
-            job_id, cluster = read_jobid_cluster(job_log)
-            print(f"Found job {job_id} on cluster {cluster} in {job_log}")
-            job_ids.setdefault(cluster, []).append(job_id)
-    # Cancel 100 at a time to avoid exceeding the command line length limit.
+        print(f"Searching recursively in {path}")
+        paths_log = list(path.glob("**/slurmjob.log"))
+        if (path / "slurmjob.log").is_file():
+            paths_log.append(path / "slurmjob.log")
+        for job_log in paths_log:
+            try:
+                job_id, cluster = read_jobid_cluster(job_log)
+                msg = f"Found job {job_id} in {job_log}"
+                if cluster is not None:
+                    msg += f" on cluster {cluster}"
+                print(msg)
+                job_ids.setdefault(cluster, []).append(job_id)
+            except ValueError as e:
+                print(f"Warning: Could not read job ID from {job_log}: {e}")
+                continue
+
+    returncode = 0
+    # Cancel at most 100 at a time to avoid exceeding the command line length limit,
+    # and to play nice with SLURM.
     for cluster, cluster_job_ids in job_ids.items():
         while len(cluster_job_ids) > 0:
-            command = f"scancel -M {cluster} " + " ".join(cluster_job_ids[:100])
-            print(command)
-            os.system(command)
+            cancel_ids = cluster_job_ids[:100]
             cluster_job_ids[:] = cluster_job_ids[100:]
 
+            command_args = ["scancel"]
+            if cluster is not None:
+                command_args.extend(["-M", cluster])
+            command_args.extend(str(job_id) for job_id in cancel_ids)
+
+            # Using subprocess.run for better control and error handling
+            print(f"Executing: {' '.join(command_args)}")
+            result = subprocess.run(command_args, check=False)
+            if result.returncode != 0:
+                returncode = 1
+    return returncode
+
 
 def read_jobid_cluster(job_log: Path) -> tuple[str, str]:
     """Read the job ID and cluster from the job log file."""
@@ -58,8 +83,7 @@ def read_jobid_cluster(job_log: Path) -> tuple[str, str]:
         lines = f.readlines()
         if len(lines) < 3 or lines[0][:-1] != FIRST_LINE:
             raise ValueError(f"Invalid first line in {job_log}.")
-        job_id, cluster = lines[2].split()[-1].split(";")
-    return job_id, cluster
+        return parse_sbatch(lines[2].split()[-1])
 
 
 def canceljobs_subcommand(subparser: argparse.ArgumentParser) -> callable:

{stepup_queue-1.0.6 → stepup_queue-1.0.7}/stepup/queue/sbatch.py

@@ -28,15 +28,17 @@ from datetime import datetime
 
 from path import Path
 
-from stepup.core.utils import string_to_bool
 from stepup.core.worker import WorkThread
 
 FIRST_LINE = "StepUp Queue sbatch wait log format version 2"
-SCONTROL_FAILED = "The command `scontrol show job` failed!\n"
-DEBUG = string_to_bool(os.getenv("STEPUP_SBATCH_DEBUG", "0"))
+SBATCH_RETRY_NUM = int(os.getenv("STEPUP_SBATCH_RETRY_NUM", "5"))
+SBATCH_RETRY_DELAY_MIN = int(os.getenv("STEPUP_SBATCH_RETRY_DELAY_MIN", "60"))
+SBATCH_RETRY_DELAY_MAX = int(os.getenv("STEPUP_SBATCH_RETRY_DELAY_MAX", "120"))
 CACHE_TIMEOUT = int(os.getenv("STEPUP_SBATCH_CACHE_TIMEOUT", "30"))
-POLLING_INTERVAL = int(os.getenv("STEPUP_SBATCH_POLLING_INTERVAL", "10"))
-TIME_MARGIN = int(os.getenv("STEPUP_SBATCH_TIME_MARGIN", "15"))
+POLLING_MIN = int(os.getenv("STEPUP_SBATCH_POLLING_MIN", "10"))
+POLLING_MAX = max(int(os.getenv("STEPUP_SBATCH_POLLING_MAX", "20")), POLLING_MIN)
+SACCT_START = os.getenv("STEPUP_SACCT_START_TIME", "now-7days")
+UNLISTED_TIMEOUT = int(os.getenv("STEPUP_SBATCH_UNLISTED_TIMEOUT", "600"))
 
 
 def submit_once_and_wait(
@@ -68,11 +70,7 @@ def submit_once_and_wait(
     """
     # Read previously logged steps
     path_log = Path("slurmjob.log")
-    if path_log.is_file():
-        previous_lines = read_log(path_log, validate_inp_digest)
-    else:
-        previous_lines = []
-        _init_log(path_log)
+    previous_lines = read_log(path_log, validate_inp_digest) if path_log.is_file() else []
 
     # Go through or skip steps.
     submit_time, status = read_step(previous_lines)
@@ -80,6 +78,8 @@ def submit_once_and_wait(
         # A new job must be submitted.
         submit_time = time.time()
         sbatch_stdout = submit_job(work_thread, job_ext, sbatch_rc)
+        # Create a new log file after submitting the job.
+        _init_log(path_log)
         log_step(path_log, f"Submitted {sbatch_stdout}")
         rndsleep()
     else:
@@ -103,12 +103,17 @@ def submit_once_and_wait(
             work_thread, submit_time, jobid, cluster, previous_lines, path_log, status
         )
 
-    # Get the return code from the job
-    with open("slurmjob.ret") as fh:
-        returncode = fh.read().strip()
-    if returncode == "":
-        raise ValueError("The job did not return a return code, e.g. because it was cancelled.")
-    return int(returncode)
+    if status == "COMPLETED":
+        # Get the return code from the job
+        with open("slurmjob.ret") as fh:
+            returncode = fh.read().strip()
+        try:
+            return int(returncode)
+        except ValueError as exc:
+            raise ValueError(
+                f"Could not parse return code from slurmjob.ret. Got '{returncode}'"
+            ) from exc
+    raise RuntimeError(f"Job ended with status '{status}'.")
 
 
 def read_log(path_log: str, do_inp_digest: bool = True) -> list[str]:
@@ -141,6 +146,50 @@ def _init_log(path_log: str):
         print(inp_digest, file=fh)
 
 
+# From: https://slurm.schedmd.com/job_state_codes.html
+KNOWN_JOB_STATES = [
+    # -- Job states
+    # done
+    "BOOT_FAIL",
+    "CANCELLED",
+    "COMPLETED",
+    "DEADLINE",
+    "FAILED",
+    "NODE_FAIL",
+    "OUT_OF_MEMORY",
+    "PREEMPTED",
+    "TIMEOUT",
+    # waiting or running
+    "PENDING",
+    "RUNNING",
+    "SUSPENDED",
+    # -- Job flags
+    # done
+    "LAUNCH_FAILED",
+    "RECONFIG_FAIL",
+    "REVOKED",
+    "STOPPED",
+    # waiting or running
+    "COMPLETING",
+    "CONFIGURING",
+    "EXPEDITING",
+    "POWER_UP_NODE",
+    "REQUEUED",
+    "REQUEUE_FED",
+    "REQUEUE_HOLD",
+    "RESIZING",
+    "RESV_DEL_HOLD",
+    "SIGNALING",
+    "SPECIAL_EXIT",
+    "STAGE_OUT",
+    "UPDATE_DB",
+    # -- Specific to this script
+    # to be ignored (same as waiting or running), must not be logged
+    "invalid",
+    "unlisted",
+]
+
+
 def _read_or_poll_status(
     work_thread: WorkThread,
     submit_time: float,
@@ -155,7 +204,7 @@ def _read_or_poll_status(
     Parameters
     ----------
     work_thread
-        The work thread to use for launching the scontrol command.
+        The work thread to use for launching the sacct command.
     submit_time
         The timestamp when the job was submitted.
     jobid
@@ -165,7 +214,6 @@ def _read_or_poll_status(
     previous_lines
         Lines from an existing log file to be processed first.
         (It will be gradually emptied.)
-    path_log
         The log file to write new polling results to.
     last_status
         The status from the previous iteration.
@@ -179,19 +227,40 @@ def _read_or_poll_status(
         True when the waiting is over.
     """
     # First try to replay previously logged steps
-    status_time, status = read_step(previous_lines)
+    _, status = read_step(previous_lines)
     if status is None:
         # All previously logged steps are processed.
-        # Call scontrol and parse its response.
+        # Call sacct and parse its response.
         rndsleep()
-        status_time, status = get_status(work_thread, jobid, cluster)
+        _, status = get_status(work_thread, jobid, cluster)
         # Log only if the status changed, and is not invalid or unlisted.
         # These two statuses are (potentially) transient and should not be logged.
         if status != last_status and status not in ["invalid", "unlisted"]:
             log_step(path_log, status)
-    done = (status_time > submit_time + TIME_MARGIN) and (
-        status not in ["PENDING", "CONFIGURING", "RUNNING", "invalid"]
-    )
+    if status not in KNOWN_JOB_STATES:
+        raise ValueError(f"Unknown job status '{status}' obtained from scheduler.")
+
+    # Determine if the job is done
+    done = status in [
+        "BOOT_FAIL",
+        "CANCELLED",
+        "COMPLETED",
+        "DEADLINE",
+        "FAILED",
+        "NODE_FAIL",
+        "OUT_OF_MEMORY",
+        "PREEMPTED",
+        "TIMEOUT",
+        "LAUNCH_FAILED",
+        "RECONFIG_FAIL",
+        "REVOKED",
+        "STOPPED",
+    ]
+    if status == "unlisted" and time.time() > submit_time + UNLISTED_TIMEOUT:
+        # If the job remains unlisted for too long, we declare it failed.
+        # This prevents an infinite loop if the job ID was wrong or purged.
+        done = True
+
     return status, done
 
 
@@ -232,7 +301,7 @@ def read_step(lines: list[str]) -> str | None:
 
 def rndsleep():
     """Randomized sleep to distribute I/O load evenly."""
-    sleep_seconds = 1 if DEBUG else random.randint(POLLING_INTERVAL, POLLING_INTERVAL + TIME_MARGIN)
+    sleep_seconds = random.randint(POLLING_MIN, POLLING_MAX)
     time.sleep(sleep_seconds)
 
 
@@ -241,36 +310,56 @@ JOB_SCRIPT_WRAPPER = """\
 {sbatch_header}
 
 touch slurmjob.ret
-chmod +x '{job_script}'
 ./'{job_script}'
 RETURN_CODE=$?
 echo $RETURN_CODE > slurmjob.ret
 exit $RETURN_CODE
 """
 
+RE_SBATCH_STDOUT = re.compile(r"#\s*SBATCH\b.*(--output|-o)")
+RE_SBATCH_STDERR = re.compile(r"#\s*SBATCH\b.*(--error|-e)")
+RE_SBATCH_ARRAY = re.compile(r"#\s*SBATCH\b.*(--array|-a)")
+RE_SBATCH = re.compile(r"#\s*SBATCH\b")
+
 
 def submit_job(work_thread: WorkThread, job_ext: str, sbatch_rc: str | None = None) -> str:
     """Submit a job with sbatch."""
-    # Copy the #SBATCH lines from the job script.
+    # Verify that the job script is executable.
     path_job = f"slurmjob{job_ext}"
+    if not os.access(path_job, os.X_OK):
+        raise ValueError("The job script must be executable.")
+
+    # Copy the #SBATCH lines from the job script and perform some checks.
     with open(path_job) as f:
-        sbatch_header = "\n".join(line for line in f if line.startswith("#SBATCH"))
+        sbatch_header = []
+        first_line = next(f)
+        if not first_line.startswith("#!"):
+            raise ValueError("The job script must start with a shebang line.")
+        for line in f:
+            if RE_SBATCH_STDOUT.match(line):
+                raise ValueError("The job script must not contain a #SBATCH --output/-o line.")
+            if RE_SBATCH_STDERR.match(line):
+                raise ValueError("The job script must not contain a #SBATCH --error/-e line.")
+            if RE_SBATCH_ARRAY.match(line):
+                raise ValueError("StepUp Queue does not support array jobs. (Found -a or --array)")
+            if RE_SBATCH.match(line):
+                sbatch_header.append(line.strip())
+        sbatch_header = "\n".join(sbatch_header)
 
     command = "sbatch --parsable -o slurmjob.out -e slurmjob.err"
     if sbatch_rc is not None:
         command = f"{sbatch_rc} < /dev/null && {command}"
-    returncode, stdout, stderr = work_thread.runsh(
-        command,
-        stdin=JOB_SCRIPT_WRAPPER.format(
-            sbatch_header=sbatch_header,
-            job_script=path_job,
-        ),
-    )
-    if returncode != 0:
+    stdin = JOB_SCRIPT_WRAPPER.format(sbatch_header=sbatch_header, job_script=path_job)
+    for _ in range(SBATCH_RETRY_NUM):
+        returncode, stdout, stderr = work_thread.runsh(command, stdin=stdin)
+        if returncode == 0:
+            return stdout.strip()
         if not (stderr is None or stderr == ""):
             print(stderr)
-        raise RuntimeError(f"sbatch failed with return code {returncode}.")
-    return stdout.strip()
+        delay = random.randint(SBATCH_RETRY_DELAY_MIN, SBATCH_RETRY_DELAY_MAX)
+        print(f"sbatch failed with return code {returncode}. Retrying in {delay} seconds.")
+        time.sleep(delay)
+    raise RuntimeError(f"sbatch failed {SBATCH_RETRY_NUM} times. Giving up.")
 
 
 def log_step(path_log: Path, step: str):
@@ -292,12 +381,12 @@ def parse_sbatch(stdout: str) -> tuple[int, str | None]:
 
 
 def get_status(work_thread: WorkThread, jobid: int, cluster: str | None) -> str:
-    """Load cached scontrol output or run scontrol if outdated.
+    """Load cached sacct output or run sacct if outdated.
 
     Parameters
     ----------
     work_thread
-        The work thread to use for launching the scontrol command.
+        The work thread to use for launching the sacct command.
     jobid
         The job to wait for.
     cluster
@@ -306,24 +395,22 @@ def get_status(work_thread: WorkThread, jobid: int, cluster: str | None) -> str:
     Returns
     -------
     status
-        A status reported by scontrol,
-        or `invalid` if scontrol failed (retry scontrol later),
+        A status reported by sacct,
+        or `invalid` if sacct failed (retry sacct later),
         or `unlisted` if the job is not found (probably ended long ago).
     """
     # Load cached output or run again
-    command = "scontrol show job"
-    path_out = Path(os.getenv("HOME")) / ".cache/stepup-queue"
+    command = f"sacct -o 'jobid,state' -PXn -S {SACCT_START}"
+    path_out = Path(os.getenv("ROOT")) / ".stepup/queue"
     if cluster is None:
-        path_out /= "sbatch_wait.out"
+        path_out /= "sbatch_wait_sacct.out"
     else:
         command += f" --cluster={cluster}"
-        path_out /= f"sbatch_wait.{cluster}.out"
-    status_time, scontrol_out, returncode = cached_run(
-        work_thread, command, path_out, CACHE_TIMEOUT
-    )
+        path_out /= f"sbatch_wait_sacct.{cluster}.out"
+    status_time, sacct_out, returncode = cached_run(work_thread, command, path_out, CACHE_TIMEOUT)
     if returncode != 0:
         return status_time, "invalid"
-    return status_time, parse_scontrol_out(scontrol_out, jobid)
+    return status_time, parse_sacct_out(sacct_out, jobid)
 
 
 def cached_run(
@@ -405,13 +492,13 @@ def parse_cache_header(header: str) -> tuple[float, int]:
 CACHE_HEADER_LENGTH = len(make_cache_header(time.time(), 0))
 
 
-def parse_scontrol_out(scontrol_out: str, jobid: int) -> str:
-    """Get the job state for a specific from from the output of ``scontrol show job``.
+def parse_sacct_out(sacct_out: str, jobid: int) -> str:
+    """Get the job state for a specific from from the output of ``sacct -o 'jobid,state' -PXn``.
 
     Parameters
     ----------
-    scontrol_out
-        A string with the output of ``scontrol show job``.
+    sacct_out
+        A string with the output of ``sacct -o 'jobid,state' -PXn``.
     jobid
         The jobid of interest.
 
@@ -423,12 +510,13 @@ def parse_scontrol_out(scontrol_out: str, jobid: int) -> str:
         - Any of the SLURM job states.
         - `unlisted` if the job cannot be found,
           which practically means it has ended long ago.
+        - `invalid` if the sacct output cannot be parsed.
     """
-    match = re.search(
-        f"JobId={jobid}.*?JobState=(?P<state>[A-Z]+)",
-        scontrol_out,
-        flags=re.MULTILINE | re.DOTALL,
-    )
-    if match is not None:
-        return match.group("state")
+    try:
+        for line in sacct_out.splitlines():
+            columns = line.strip().split("|")
+            if int(columns[0]) == jobid:
+                return columns[1].strip().split()[0]
+    except (ValueError, IndexError):
+        return "invalid"
     return "unlisted"

{stepup_queue-1.0.6 → stepup_queue-1.0.7/stepup_queue.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stepup-queue
-Version: 1.0.6
+Version: 1.0.7
 Summary: StepUp Queue integrates queued jobs into a StepUp workflow.
 Author-email: Toon Verstraelen <toon.verstraelen@ugent.be>
 License-Expression: GPL-3.0-or-later
@@ -24,7 +24,7 @@ Classifier: Topic :: Software Development :: Build Tools
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: stepup<4.0.0,>=3.1.3
+Requires-Dist: stepup<4.0.0,>=3.1.4
 Provides-Extra: dev
 Requires-Dist: psutil; extra == "dev"
 Requires-Dist: pytest; extra == "dev"

{stepup_queue-1.0.6 → stepup_queue-1.0.7}/stepup_queue.egg-info/requires.txt

@@ -1,4 +1,4 @@
-stepup<4.0.0,>=3.1.3
+stepup<4.0.0,>=3.1.4
 
 [dev]
 psutil