xmanager-slurm 0.4.12__py3-none-any.whl → 0.4.14__py3-none-any.whl

xm_slurm/contrib/clusters/drac.py

@@ -4,7 +4,7 @@ from typing import Literal
 from xm_slurm import config
 from xm_slurm.resources import FeatureType, ResourceType
 
-__all__ = ["narval", "beluga", "cedar", "graham"]
+__all__ = ["beluga", "cedar", "fir", "graham", "narval"]
 
 
 def _drac_cluster(
@@ -121,6 +121,36 @@ def beluga(
     )
 
 
+def rorqual(
+    *,
+    user: str | None = None,
+    account: str | None = None,
+    proxy: Literal["submission-host"] | str | None = None,
+    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """DRAC Beluga Cluster (https://docs.alliancecan.ca/wiki/Rorqual/en)."""
+    modules = []
+    if proxy != "submission-host":
+        modules.append("httpproxy")
+
+    return _drac_cluster(
+        name="rorqual",
+        host="robot.rorqual.alliancecan.ca",
+        host_public_key=config.PublicKey(
+            "ssh-ed25519", "AAAAC3NzaC1lZDI1NTE5AAAAINME5e9bifKZbuKKOQSpe3xrvC4g1b0QLMYj+AXBQGJe"
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        proxy=proxy,
+        modules=modules,
+        resources={ResourceType.H100: "h100"},
+        features={
+            FeatureType.NVIDIA_NVLINK: "nvlink",
+        },
+    )
+
+
 def cedar(
     *,
     user: str | None = None,
@@ -146,6 +176,27 @@ def cedar(
     )
 
 
+def fir(
+    *,
+    user: str | None = None,
+    account: str | None = None,
+    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """DRAC Fir Cluster (https://docs.alliancecan.ca/wiki/Fir/en)."""
+    return _drac_cluster(
+        name="fir",
+        host="robot.fir.alliancecan.ca",
+        host_public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAIJtenyJz+inwobvlJntWYFNu+ANcVWNcOHRKcEN6zmDo",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        resources={ResourceType.H100: "h100"},
+    )
+
+
 def graham(
     *,
     user: str | None = None,

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__)
 

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.

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 -}}

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

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() %}

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) }}

{xmanager_slurm-0.4.12.dist-info → xmanager_slurm-0.4.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.12
+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.12.dist-info → xmanager_slurm-0.4.14.dist-info}/RECORD

@@ -5,7 +5,7 @@ xm_slurm/console.py,sha256=UpMqeJ0C8i0pkue1AHnnyyX0bFJ9zZeJ7HBR6yhuA8A,54
 xm_slurm/constants.py,sha256=zefVtlFdflgSolie5g_rVxWV-Zpydxapchm3y0a2FDc,999
 xm_slurm/dependencies.py,sha256=-5gN_tpfs3dOA7H5_MIHO2ratb7F5Pm_yjkR5rZcgI8,6421
 xm_slurm/executables.py,sha256=fGmrFBl-258bMn6ip5adYeM7xxUHAeIbDN9zD2FDGtY,6373
-xm_slurm/execution.py,sha256=c0aV1h2tKQFyAGM6JLd16MWFgpRLKAbcutZz17xPUSw,31400
+xm_slurm/execution.py,sha256=mTy5u2oP2StIbGzjaSiGCUAwXuBFOiaJ5ephWoc25hI,31799
 xm_slurm/executors.py,sha256=bUgKcgtvf-nPGjcuHRzUAqD1r3_vwea_h-Y9MAB-Kqo,4887
 xm_slurm/experiment.py,sha256=94r0mhtUPUzw4eaUEz0kpsufC25wEGqlDhV4Fcr1ukY,39883
 xm_slurm/filesystem.py,sha256=4rKtq3t-KDgxJbSGt6JVyRJT_3lCN_vIKTcwKHpTo3I,4389
@@ -23,7 +23,7 @@ xm_slurm/api/sqlite/client.py,sha256=jAesCKDuYwnNcAxwJk_1b1TB8cT_QGbSjo1UE3mZjEQ
 xm_slurm/api/web/client.py,sha256=uO67Y7fnQ-w__Vm_A5BEuy7Qi8wQcWk3vIsBGEBkyfk,6261
 xm_slurm/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xm_slurm/contrib/clusters/__init__.py,sha256=XFCVnkThiU3_8uA_tUgDByOBanXNHrxDvfmuptmQ2KE,2214
-xm_slurm/contrib/clusters/drac.py,sha256=ViLYerYBMSuZXnWVbz9RDIPPV7JA8BgBpgTfj1wPP28,5881
+xm_slurm/contrib/clusters/drac.py,sha256=vY3dxrNUk12H9Gq-tuCcqo2YcdTGq-4LJnQF6DzD4_k,7431
 xm_slurm/experimental/parameter_controller.py,sha256=b5LfglHV307F6QcPrHeZX5GJBtyOK9aQydke_SZ3Wto,8457
 xm_slurm/packaging/__init__.py,sha256=dh307yLpUT9KN7rJ1e9fYC6hegGKfZcGboUq9nGpDVQ,233
 xm_slurm/packaging/docker.py,sha256=-DWcB9qqbeHmIEqyfF0-v6xOT25ae90u2x-QZ7kluOw,13579
@@ -39,13 +39,14 @@ xm_slurm/templates/docker/uv.Dockerfile,sha256=L2UJMX2c8waMdrRhiqPytQe3pTBu6u5Pp
 xm_slurm/templates/slurm/entrypoint.bash.j2,sha256=MRdSVwgGrgQdpEhqfkP35IidgsblrtVXB1YWzvE9hkk,666
 xm_slurm/templates/slurm/job-array.bash.j2,sha256=smxmSSzBEUHm6MJF-nYPVVjK6CLKrb1fRxF_tfrzAX8,552
 xm_slurm/templates/slurm/job-group.bash.j2,sha256=Cp8YhNOxYqaOkl4MFjQlcaLMGZwdDh97m8OGT5RWbAo,1101
-xm_slurm/templates/slurm/job.bash.j2,sha256=pNKir1tkmRTGDiGxlQ3DkUaW9Zos_gdkkXJC_xX5Cxo,1985
-xm_slurm/templates/slurm/fragments/monitor.bash.j2,sha256=BJ1brSjhESOe9VX_OYaPyy9-qE3uiFlzxp8ZkFcTw8Y,2504
+xm_slurm/templates/slurm/job.bash.j2,sha256=DrDipliaEfiHbq9vDfOdfD8zBVFLy1jjlvCV-9-6k9s,2086
+xm_slurm/templates/slurm/fragments/monitor.bash.j2,sha256=ri5FgoKs6_bQVf5DO8SL4rJf4UsLxV34aOV-OD8VWDU,2526
 xm_slurm/templates/slurm/fragments/proxy.bash.j2,sha256=VJLglZo-Nvx9R-qe3rHTxr07CylTQ6Z9NwBzvIpAZrA,814
-xm_slurm/templates/slurm/runtimes/apptainer.bash.j2,sha256=lE2EWVCK2O-n08RL4_MJYIikVTvODjcYKuv7Eh73Q2w,1932
-xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=3j7K5eyXt_WhXK0EoMlxnhlmFVJ2JyxRKbsMRaDqzSs,1148
-xmanager_slurm-0.4.12.dist-info/METADATA,sha256=ttPSMz6bQs8BEQVCurQ9lFWV7MuzLOHWU7lISOOoufA,1007
-xmanager_slurm-0.4.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xmanager_slurm-0.4.12.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
-xmanager_slurm-0.4.12.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
-xmanager_slurm-0.4.12.dist-info/RECORD,,
+xm_slurm/templates/slurm/library/retry.bash,sha256=bLe59qvfWEk17rE1wZ4EHiHba3RvR2WWZPq-kSe8RAA,2164
+xm_slurm/templates/slurm/runtimes/apptainer.bash.j2,sha256=v0LwHM-kBW8sJqVcVA2jYr1n44imDSZrJqmqlr5uTGc,1980
+xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=zWLsFEuVzOMSETOmv4A5ZCV4oQHwCipiR6wi79XVzNI,1188
+xmanager_slurm-0.4.14.dist-info/METADATA,sha256=T7xNy0jmrKhQemaDhCg9E-J64gWkabjJGpxDYgdsBx8,1007
+xmanager_slurm-0.4.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xmanager_slurm-0.4.14.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.14.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.14.dist-info/RECORD,,