xmanager-slurm 0.4.15__py3-none-any.whl → 0.4.17__py3-none-any.whl

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",

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,

xm_slurm/executors.py

@@ -37,6 +37,7 @@ class Slurm(xm.Executor):
         switches: Maximum count of leaf switches desired for the job allocation.
         switches_grace_period: Maximum time to wait for that number of switches.
         bind: How to bind tasks to resource (memory, GPU, or generic resource).
+        bind_flag: Generic resource task binding options.
         account: The account to charge the job to.
         partition: The partition to run the job in.
         qos: The quality of service to run the job with.
@@ -59,6 +60,7 @@ class Slurm(xm.Executor):
     requirements: resources.JobRequirements
     time: dt.timedelta
     bind: tp.Mapping[ResourceBindType | str, str | None] | None = None
+    bind_flag: str | None = None
 
     # Placement
     account: str | None = None
@@ -109,6 +111,8 @@ class Slurm(xm.Executor):
                     )
                 if value is not None and not isinstance(value, str):
                     raise TypeError(f"bind value must be None or a string, got {type(value)}")
+        if self.bind_flag is not None and not isinstance(self.bind_flag, str):
+            raise TypeError(f"bind_flag must be a string, got {type(self.bind_flag)}")
 
         if not isinstance(self.timeout_signal, signal.Signals):
             raise TypeError(
@@ -133,28 +137,13 @@ class Slurm(xm.Executor):
     def Spec(cls, tag: str | None = None) -> SlurmSpec:
         return SlurmSpec(tag=tag)
 
-    def to_directives(self) -> list[str]:
+    def batch_directives(self) -> list[str]:
         # Job requirements
-        directives = self.requirements.to_directives()
+        directives = self.requirements.batch_directives()
 
         # Time
         directives.append(f"--time={utils.timestr_from_timedelta(self.time)}")
 
-        # Resource binding
-        if self.bind is not None:
-            for resource, value in self.bind.items():
-                if value is None:
-                    value = "none"
-                match resource:
-                    case resources.ResourceType.MEMORY | resources.ResourceType.RAM:
-                        directives.append(f"--mem-bind={value}")
-                    case resources.ResourceType.GPU:
-                        directives.append(f"--gpu-bind={value}")
-                    case str():
-                        directives.append(f"--tres-bind=gres/{resource}:{value}")
-                    case _:
-                        raise ValueError(f"Unsupported resource type {resource!r} for binding.")
-
         # Job dependency handling
         directives.append(
             f"--kill-on-invalid-dep={'yes' if self.kill_on_invalid_dependencies else 'no'}"
@@ -196,3 +185,26 @@ class Slurm(xm.Executor):
             directives.append("--no-requeue")
 
         return directives
+
+    def step_directives(self) -> list[str]:
+        directives = self.requirements.step_directives()
+
+        # Resource binding
+        if self.bind is not None:
+            for resource, value in self.bind.items():
+                if value is None:
+                    value = "none"
+                match resource:
+                    case resources.ResourceType.MEMORY | resources.ResourceType.RAM:
+                        directives.append(f"--mem-bind={value}")
+                    case resources.ResourceType.GPU:
+                        directives.append(f"--gpu-bind={value}")
+                    case str():
+                        directives.append(f"--tres-bind=gres/{resource}:{value}")
+                    case _:
+                        raise ValueError(f"Unsupported resource type {resource!r} for binding.")
+
+        if self.bind_flag is not None:
+            directives.append(f"--gres-flags={self.bind_flag}")
+
+        return directives

xm_slurm/resources.py

@@ -232,7 +232,7 @@ class JobRequirements:
             raise ValueError(f"Replicas must be a positive integer, got {replicas!r}")
         self.replicas = replicas or 1
 
-    def to_directives(self) -> list[str]:
+    def batch_directives(self) -> list[str]:
         directives = []
 
         for resource, value in self.task_requirements.items():
@@ -302,6 +302,9 @@ class JobRequirements:
 
         return directives
 
+    def step_directives(self) -> list[str]:
+        return []
+
     def replace(
         self,
         replicas: int | None = None,

xm_slurm/templates/slurm/job-array.bash.j2

@@ -9,8 +9,11 @@
 srun \
   --label \
   --unbuffered \
-  --kill-on-bad-exit=0 \
+  --kill-on-bad-exit=1 \
   --export="ALL" \
+{% for directive in job.executor.step_directives() %}
+  {{ directive }} \
+{% endfor %}
   bash <<'SRUN_EOF' &
 set -Eeuxo pipefail
 

xm_slurm/templates/slurm/job-group.bash.j2

@@ -21,7 +21,9 @@
 {% else %}
 #SBATCH --job-name=xm[{{ job_name }}@{{ experiment_id }}]
 {% endif %}
-{{ job.executor.to_directives() | join("\n") }}
+{% for directive in job.executor.batch_directives() %}
+#SBATCH {{ directive }}
+{% endfor %}
 {{ "\n#SBATCH hetjob\n" if not loop.last }}
 {% endfor %}
 {% endblock directives %}
@@ -31,8 +33,11 @@
 srun \
   --label \
   --unbuffered \
-  --kill-on-bad-exit=0 \
+  --kill-on-bad-exit=1 \
   --export="ALL" \
+{% for directive in job.executor.step_directives() %}
+  {{ directive }} \
+{% endfor %}
   --het-group={{ loop.index0 }} \
   bash <<'SRUN_EOF' &
 set -Eeuxo pipefail

xm_slurm/templates/slurm/job.bash.j2

@@ -21,7 +21,7 @@
 {% endif %}
 #SBATCH --job-name=xm[{{ experiment_id }}]
 {% endif %}
-{% for directive in job.executor.to_directives() %}
+{% for directive in job.executor.batch_directives() %}
 #SBATCH {{ directive }}
 {% endfor %}
 {% endblock directives %}
@@ -61,8 +61,11 @@ export {{ key }}="{{ value }}"
 srun \
   --label \
   --unbuffered \
-  --kill-on-bad-exit=0 \
+  --kill-on-bad-exit=1 \
   --export="ALL" \
+{% for directive in job.executor.step_directives() %}
+  {{ directive }} \
+{% endfor %}
   bash <<'SRUN_EOF' &
 set -Eeuxo pipefail
 {{ run(cluster, job) }}

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

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 \

{xmanager_slurm-0.4.15.dist-info → xmanager_slurm-0.4.17.dist-info}/METADATA

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

@@ -6,13 +6,13 @@ xm_slurm/constants.py,sha256=zefVtlFdflgSolie5g_rVxWV-Zpydxapchm3y0a2FDc,999
 xm_slurm/dependencies.py,sha256=G-8vfmvSptZH6c_Ow51SwT84Dr6LI1clRj8F8wOUkiw,6421
 xm_slurm/executables.py,sha256=fGmrFBl-258bMn6ip5adYeM7xxUHAeIbDN9zD2FDGtY,6373
 xm_slurm/execution.py,sha256=mTy5u2oP2StIbGzjaSiGCUAwXuBFOiaJ5ephWoc25hI,31799
-xm_slurm/executors.py,sha256=karM5u2UEG2IWi0z548_vasyBACrXGV675rCllJmwZw,8616
+xm_slurm/executors.py,sha256=27oiMwF84axeTcrcwL0f5seeLL_1j79OjiM_JZjioFs,9112
 xm_slurm/experiment.py,sha256=94r0mhtUPUzw4eaUEz0kpsufC25wEGqlDhV4Fcr1ukY,39883
 xm_slurm/filesystem.py,sha256=4rKtq3t-KDgxJbSGt6JVyRJT_3lCN_vIKTcwKHpTo3I,4389
 xm_slurm/job_blocks.py,sha256=BFOOYgeodoGIQsB5PdC7SsOUou5aZx-1qbQ7lcqqylI,604
 xm_slurm/metadata_context.py,sha256=mksVRbVUuistL1uE7TC-fkW-Y69On52jN_svP1e1kiQ,7841
 xm_slurm/packageables.py,sha256=aEZUQpddfq4FK6h4f6kgGEI4XcOufhm68MjoDFOYR4U,12261
-xm_slurm/resources.py,sha256=aC8MzO_7fB9IAdTACvhwVOaNDjLOlWnCh428-8_IDYA,12322
+xm_slurm/resources.py,sha256=sTfwPc0QHRgfckOFq300FZ4fvtPfE4hq8B27DIvf6m4,12388
 xm_slurm/status.py,sha256=JIBCJPOYsmeJOQbzdACXA2vTWK7g8YWWhzpGP79e7JE,6911
 xm_slurm/types.py,sha256=TsVykDm-LazVkrjeJrTwCMs4Q8APKhy7BTk0yKIhFNg,805
 xm_slurm/utils.py,sha256=9w98HlXF0U9cKKtoB8QtGm0CnB0MnnzBARKlbbVNNpU,6211
@@ -22,8 +22,8 @@ xm_slurm/api/models.py,sha256=_INVh0j-4-rRs0WASyg4fNB6NF1L1nUeGgQ6-XnbwsM,1610
 xm_slurm/api/sqlite/client.py,sha256=jAesCKDuYwnNcAxwJk_1b1TB8cT_QGbSjo1UE3mZjEQ,14037
 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=_iubsmzz5tK2KTaKqSuykS3IDtsdXqJ0MXep1THSJUM,7020
+xm_slurm/contrib/clusters/__init__.py,sha256=oYkuo4jAA6_VLIJs4k1Rjcp6l3qN0qmJExfkPWfmho4,2281
+xm_slurm/contrib/clusters/drac.py,sha256=-Nawzyx3p471f8Fti_tbweHn8h4ItRN4ec-ktHnUE4Q,7087
 xm_slurm/experimental/parameter_controller.py,sha256=IrFzq104LkZrhzuirit5GUZDXDvv2bBSYNMh3orsiPY,8518
 xm_slurm/packaging/__init__.py,sha256=dh307yLpUT9KN7rJ1e9fYC6hegGKfZcGboUq9nGpDVQ,233
 xm_slurm/packaging/docker.py,sha256=-DWcB9qqbeHmIEqyfF0-v6xOT25ae90u2x-QZ7kluOw,13579
@@ -37,16 +37,16 @@ xm_slurm/templates/docker/mamba.Dockerfile,sha256=Sgxr5IA5T-pT1Shumb5k3JngoG4pgC
 xm_slurm/templates/docker/python.Dockerfile,sha256=U4b4QVkopckQ0o9jJIE7d_M6TvExEYlYDirNwCoZ7W4,865
 xm_slurm/templates/docker/uv.Dockerfile,sha256=L2UJMX2c8waMdrRhiqPytQe3pTBu6u5PpMhJYsKkbEg,1040
 xm_slurm/templates/slurm/entrypoint.bash.j2,sha256=MRdSVwgGrgQdpEhqfkP35IidgsblrtVXB1YWzvE9hkk,666
-xm_slurm/templates/slurm/job-array.bash.j2,sha256=j7jkJjSbe39XvSTJ9rmK2oVnHdntElIhdS5PFpZzpFs,550
-xm_slurm/templates/slurm/job-group.bash.j2,sha256=vH5HwneVsVSHx6dPZwbLa4KT9NedRbrZ7cWNE5pXi-M,1113
-xm_slurm/templates/slurm/job.bash.j2,sha256=JnK0D8_3tVNpnvPwM5yL_rjLcjqhuHiCtolDjUGAwpk,2084
+xm_slurm/templates/slurm/job-array.bash.j2,sha256=7cc0nZvEcHhZoo7jXI3fJWgMcc6z5H5FmopPRaklylI,637
+xm_slurm/templates/slurm/job-group.bash.j2,sha256=9H3zfJy8RZGFf00ZQJGmMEPyWQ9YMZfvGoD4Q8hMx9Y,1244
+xm_slurm/templates/slurm/job.bash.j2,sha256=GBKY3DPCODPTtEBfuvfaZAua_ZEd5cqPrShtPGE_IpY,2174
 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/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.15.dist-info/METADATA,sha256=xgJUFDConlb4R5W0cOK3xb2fuIR7x_tjNpOybMagT_A,1007
-xmanager_slurm-0.4.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xmanager_slurm-0.4.15.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
-xmanager_slurm-0.4.15.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
-xmanager_slurm-0.4.15.dist-info/RECORD,,
+xm_slurm/templates/slurm/runtimes/apptainer.bash.j2,sha256=XxAQWLxZogL7zjn7tuzKn-DkYUJMx_HjaRzpVkz97lM,2414
+xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=8N1ZtwHyXxP-Cjo4HBPsJiZXcTvf7q2GzvW9ao8_aok,1208
+xmanager_slurm-0.4.17.dist-info/METADATA,sha256=YUBZ6woSk-9-0GzFuFlTPy6hhaxFI-WPGVt7APaviT0,1007
+xmanager_slurm-0.4.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xmanager_slurm-0.4.17.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.17.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.17.dist-info/RECORD,,