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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.16
+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.16.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
@@ -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=XxAQWLxZogL7zjn7tuzKn-DkYUJMx_HjaRzpVkz97lM,2414
 xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=8N1ZtwHyXxP-Cjo4HBPsJiZXcTvf7q2GzvW9ao8_aok,1208
-xmanager_slurm-0.4.16.dist-info/METADATA,sha256=j0282EV56cTAG80Q4R419IE-Z74OJVvAkP--IPMCwuo,1007
-xmanager_slurm-0.4.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xmanager_slurm-0.4.16.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
-xmanager_slurm-0.4.16.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
-xmanager_slurm-0.4.16.dist-info/RECORD,,
+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,,