xmanager-slurm 0.4.13__tar.gz → 0.4.14__tar.gz

{xmanager_slurm-0.4.13 → xmanager_slurm-0.4.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.13
+Version: 0.4.14
 Summary: Slurm backend for XManager.
 Project-URL: GitHub, https://github.com/jessefarebro/xm-slurm
 Author-email: Jesse Farebrother <jfarebro@cs.mcgill.ca>

{xmanager_slurm-0.4.13 → xmanager_slurm-0.4.14}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "xmanager-slurm"
 description = "Slurm backend for XManager."
-version = "0.4.13"
+version = "0.4.14"
 # readme = "README.md"
 requires-python = ">=3.10"
 license = { text = "MIT" }

{xmanager_slurm-0.4.13 → xmanager_slurm-0.4.14}/uv.lock

@@ -2251,7 +2251,7 @@ wheels = [
 
 [[package]]
 name = "xmanager-slurm"
-version = "0.4.13"
+version = "0.4.14"
 source = { editable = "." }
 dependencies = [
     { name = "aiofile" },

{xmanager_slurm-0.4.13 → xmanager_slurm-0.4.14}/xm_slurm/execution.py

@@ -3,6 +3,8 @@ import collections.abc
 import dataclasses
 import functools
 import hashlib
+import importlib
+import importlib.resources
 import logging
 import operator
 import os
@@ -311,6 +313,14 @@ def get_template_env(runtime: ContainerRuntime) -> j2.Environment:
     template_env.globals["raise"] = _raise_template_exception
     template_env.globals["operator"] = operator
 
+    # Iterate over stdlib files and insert them into the template environment
+    stdlib = []
+    for file in importlib.resources.files("xm_slurm.templates.slurm.library").iterdir():
+        if not file.is_file() or not file.name.endswith(".bash"):
+            continue
+        stdlib.append(file.read_text())
+    template_env.globals["stdlib"] = stdlib
+
     entrypoint_template = template_env.get_template("entrypoint.bash.j2")
     template_env.globals.update(entrypoint_template.module.__dict__)
 

{xmanager_slurm-0.4.13 → xmanager_slurm-0.4.14}/xm_slurm/templates/slurm/fragments/monitor.bash.j2

@@ -25,6 +25,7 @@ __xm_slurm_wait_for_children() {
   while [ ${#children[@]} -gt 0 ]; do
 {% endraw %}
     echo "INFO: Waiting for child processes to finish..."
+    set +e
 {% if requeue_on_timeout %}
     # Wait on either one of the child processes or the timeout process.
     wait -n -p child_pid "${children[@]}" "${timeout_pid}"
@@ -32,6 +33,7 @@ __xm_slurm_wait_for_children() {
     wait -n -p child_pid "${children[@]}"
 {% endif %}
     local child_exit_code=$?
+    set -e
 
 {% if requeue_on_timeout %}
     # If the finished process is the watchdog, trigger the timeout handling.

{xmanager_slurm-0.4.13 → xmanager_slurm-0.4.14}/xm_slurm/templates/slurm/job.bash.j2

@@ -27,6 +27,13 @@
 {% endblock directives %}
 set -Eeuxo pipefail
 
+{% if stdlib %}
+# --- Helper functions ---
+{% for fn in stdlib %}
+{{ fn }}
+{% endfor %}
+{% endif %}
+
 {% block prolog %}
 {% if cluster.prolog %}
 {{- cluster.prolog -}}

xmanager_slurm-0.4.14/xm_slurm/templates/slurm/library/retry.bash

@@ -0,0 +1,62 @@
+# retry: rerun a command if it exits with certain codes
+# Options:
+#   -c CODE   Retry on this exit code (repeatable).
+#   -n N      Max attempts (incl. first). Default: unlimited
+#   -d SECS   Initial delay before first retry. Default: 1
+#   -b FACTOR Integer backoff multiplier per retry. Default: 1 (no backoff)
+#   -q        Quiet (no logs)
+# Usage:
+#   retry [-c CODE ...] [-n N] [-d SECS] [-b FACTOR] [-q] -- cmd arg1 arg2 ...
+retry() {
+  local -a codes=()
+  local -i max=-1 delay=1 backoff=1 quiet=0 status
+  local opt OPTIND=1
+
+  while getopts ":c:n:d:b:q" opt; do
+    case "$opt" in
+      c) codes+=("$OPTARG") ;;
+      n) max=$OPTARG ;;
+      d) delay=$OPTARG ;;
+      b) backoff=$OPTARG ;;
+      q) quiet=1 ;;
+      :) printf 'retry: option -%s requires an argument\n' "$OPTARG" >&2; return 2 ;;
+      \?) printf 'retry: invalid option -- %s\n' "$OPTARG" >&2; return 2 ;;
+    esac
+  done
+  shift $((OPTIND-1))
+  (( $# )) || { printf 'retry: missing command\n' >&2; return 2; }
+
+  ((${#codes[@]})) || { printf 'retry: no return codes specified\n' >&2; return 2; }
+
+  for ((attempt=1; ; attempt++)); do
+    if "$@"; then                    # safe with set -e (exception context)
+      return 0
+    else
+      status=$?                       # capture failing status immediately
+    fi
+
+    # retryable?
+    local retryable=0 c
+    for c in "${codes[@]}"; do
+      (( status == c )) && { retryable=1; break; }
+    done
+
+    # stop if not retryable OR we've just hit the max attempt
+    if (( !retryable )) || (( max >= 0 && attempt >= max )); then
+      (( quiet )) || {
+        if (( attempt > 1 )); then
+          printf 'retry: giving up after %d attempts; last exit=%d\n' "$attempt" "$status" >&2
+        else
+          printf 'retry: command failed; exit=%d\n' "$status" >&2
+        fi
+      }
+      return "$status"               # propagate exact code; errexit will catch
+    fi
+
+    (( quiet )) || printf 'retry: attempt %d failed with %d; retrying in %ds...\n' \
+                          "$attempt" "$status" "$delay" >&2
+    sleep "$delay" || :              # never trip set -e if sleep errors
+    (( delay *= backoff ))
+  done
+}
+export -f retry

{xmanager_slurm-0.4.13 → xmanager_slurm-0.4.14}/xm_slurm/templates/slurm/runtimes/apptainer.bash.j2

@@ -5,16 +5,17 @@
 # container-workdir will be our container's scratch directory
 mkdir -p "$SLURM_TMPDIR"/{container,container-workdir,container-overlay}
 
+retry -c 255 -n 10 -d 1 -b 2 -- \
 {% if job.executable.credentials %}
-env {{ runtime | upper }}_DOCKER_USERNAME="{{ job.executable.credentials.username }}" {{ runtime | upper }}_DOCKER_PASSWORD="{{ job.executable.credentials.password }}" time {{ runtime }} build \
+  env {{ runtime | upper }}_DOCKER_USERNAME="{{ job.executable.credentials.username }}" {{ runtime | upper }}_DOCKER_PASSWORD="{{ job.executable.credentials.password }}" time {{ runtime }} build \
 {% else %}
-time {{ runtime }} build \
+  time {{ runtime }} build \
 {% endif %}
-  --force \
-  --sandbox \
-  --fix-perms \
-  "$SLURM_TMPDIR"/container \
-  docker://{{ job.executable.image }}
+    --force \
+    --sandbox \
+    --fix-perms \
+    "$SLURM_TMPDIR"/container \
+    docker://{{ job.executable.image }}
 
 {% if runtime == "singularity" and cluster.mounts %}
 {% for source, dest in cluster.mounts.items() %}

{xmanager_slurm-0.4.13 → xmanager_slurm-0.4.14}/xm_slurm/templates/slurm/runtimes/podman.bash.j2

@@ -1,9 +1,10 @@
 {% macro run(cluster, job) -%}
-time podman pull \
+retry -c 255 -n 10 -d 1 -b 2 -- \
+  time podman pull \
 {% if job.executable.credentials %}
-  --creds {{ job.executable.credentials.username }}:{{ job.executable.credentials.password }} \
+    --creds {{ job.executable.credentials.username }}:{{ job.executable.credentials.password }} \
 {% endif %}
-  {{ job.executable.image }}
+    {{ job.executable.image }}
 
 cat << 'ENTRYPOINT_EOF' > "$SLURM_TMPDIR"/xm-slurm-entrypoint.sh
 {{ entrypoint(cluster, job) }}