xmanager-slurm 0.4.15__tar.gz → 0.4.16__tar.gz

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.15
+Version: 0.4.16
 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.15 → xmanager_slurm-0.4.16}/pyproject.toml

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

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.16}/uv.lock

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

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.16}/xm_slurm/contrib/clusters/__init__.py

@@ -26,6 +26,7 @@ def mila(
             "/home/mila/${USER:0:1}/$USER/.local/state/xm-slurm": "/xm-slurm-state",
             "/home/mila/${USER:0:1}/$USER/.ssh": "/home/mila/${USER:0:1}/$USER/.ssh",
         }
+    mounts = dict(mounts) | {"/dev/infiniband": "/dev/infiniband"}
 
     return config.SlurmClusterConfig(
         name="mila",

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.16}/xm_slurm/contrib/clusters/drac.py

@@ -29,6 +29,7 @@ def _drac_cluster(
             "/home/$USER/.ssh": "/home/$USER/.ssh",
             "/home/$USER/.local/state/xm-slurm": "/xm-slurm-state",
         }
+    mounts = dict(mounts) | {"/dev/infiniband": "/dev/infiniband"}
 
     return config.SlurmClusterConfig(
         name=name,

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.16}/xm_slurm/templates/slurm/runtimes/apptainer.bash.j2

@@ -3,7 +3,9 @@
 
 # Bundle will be where our built sandbox image is stored
 # container-workdir will be our container's scratch directory
-mkdir -p "$SLURM_TMPDIR"/{container,container-workdir,container-overlay}
+# TODO(jfarebro): We can make this more efficient by doing an srun per node and downloading the container once per node.
+# but this requires apptainer support to have an overlay per procid
+mkdir -p "$SLURM_TMPDIR"/{container-"$SLURM_PROCID",container-workdir-"$SLURM_PROCID",container-overlay-"$SLURM_PROCID"}
 
 retry -c 255 -n 10 -d 1 -b 2 -- \
 {% if job.executable.credentials %}
@@ -14,19 +16,21 @@ retry -c 255 -n 10 -d 1 -b 2 -- \
     --force \
     --sandbox \
     --fix-perms \
-    "$SLURM_TMPDIR"/container \
+    "$SLURM_TMPDIR"/container-"$SLURM_PROCID" \
     docker://{{ job.executable.image }}
 
 {% if runtime == "singularity" and cluster.mounts %}
 {% for source, dest in cluster.mounts.items() %}
-mkdir -p "$SLURM_TMPDIR"/container/{{ dest | trim('/') }}
+mkdir -p "$SLURM_TMPDIR"/container-"$SLURM_PROCID"/{{ dest | trim('/') }}
 {% endfor %}
 {% endif %}
 
-cat << 'ENTRYPOINT_EOF' > "$SLURM_TMPDIR"/container/xm-slurm-entrypoint.sh
+cat << 'ENTRYPOINT_EOF' > "$SLURM_TMPDIR"/container-"$SLURM_PROCID"/xm-slurm-entrypoint.sh
 {{ entrypoint(cluster, job) }}
 ENTRYPOINT_EOF
-chmod +x "$SLURM_TMPDIR"/container/xm-slurm-entrypoint.sh
+chmod +x "$SLURM_TMPDIR"/container-"$SLURM_PROCID"/xm-slurm-entrypoint.sh
+
+for var in "${!SLURM_@}"; do export "{{ runtime | upper }}ENV_${var}=${!var}"; done
 
 exec {{ runtime }} exec \
 {% if job.executor.requirements.accelerator %}
@@ -45,16 +49,16 @@ exec {{ runtime }} exec \
   --bind {{ source }}:{{ dest }} \
 {% endfor %}
 {% endif %}
-  --workdir "$SLURM_TMPDIR"/container-workdir \
+  --workdir "$SLURM_TMPDIR"/container-workdir-"$SLURM_PROCID" \
 {% if (cluster.runtime | string) == "apptainer" %}
-  --overlay "$SLURM_TMPDIR"/container-overlay \
+  --overlay "$SLURM_TMPDIR"/container-overlay-"$SLURM_PROCID" \
 {% else %}
   --writable \
 {% endif %}
 {% if job.executable.workdir %}
   --pwd {{ job.executable.workdir }} \
 {% endif %}
-  "$SLURM_TMPDIR"/container \
+  "$SLURM_TMPDIR"/container-"$SLURM_PROCID" \
   /xm-slurm-entrypoint.sh \
 {% for arg in job.executable.args.to_list() %}
   {{ arg }} \

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.16}/xm_slurm/templates/slurm/runtimes/podman.bash.j2

@@ -16,6 +16,7 @@ exec podman run \
   --entrypoint /xm-slurm-entrypoint.sh \
   --pull never \
   --restart no \
+  --env "SLURM_*" \
   --rm \
 {% if job.executor.requirements.accelerator %}
   --device nvidia.com/gpu=all \