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

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

@@ -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 → xmanager_slurm-0.4.17}/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.17"
 # readme = "README.md"
 requires-python = ">=3.10"
 license = { text = "MIT" }

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.17}/tests/test_executors.py

@@ -226,7 +226,7 @@ def test_slurm_executor_bind_accepts_custom_strings(
 def test_slurm_executor_to_directives_default(basic_requirements):
     """Test directive generation with minimal configuration."""
     executor = executors.Slurm(requirements=basic_requirements, time=dt.timedelta(hours=1))
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     assert any(d.startswith("--time=") for d in directives)
     assert any(d.startswith("--kill-on-invalid-dep=") for d in directives)
     assert any(d.startswith("--signal=") for d in directives)
@@ -244,7 +244,7 @@ def test_slurm_executor_to_directives_default(basic_requirements):
 def test_slurm_executor_to_directives_time(basic_requirements, time, expected):
     """Test time directive generation."""
     executor = executors.Slurm(requirements=basic_requirements, time=time)
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     time_directive = [d for d in directives if d.startswith("--time=")][0]
     assert time_directive == expected
 
@@ -278,7 +278,7 @@ def test_slurm_executor_to_directives_executor_params(
         time=dt.timedelta(hours=1),
         **executor_kwargs,
     )
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     assert expected_directive in directives
 
 
@@ -294,7 +294,7 @@ def test_slurm_executor_to_directives_reservation(basic_requirements, reservatio
     executor = executors.Slurm(
         requirements=basic_requirements, time=dt.timedelta(hours=1), reservation=reservation
     )
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     reservation_directive = [d for d in directives if d.startswith("--reservation=")][0]
     assert reservation_directive == expected
 
@@ -321,7 +321,7 @@ def test_slurm_executor_to_directives_boolean_flags(
         time=dt.timedelta(hours=1),
         **flag_kwargs,
     )
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     for expected in expected_directives:
         assert any(expected in d for d in directives)
     for unexpected in ["--exclusive", "--oversubscribe", "--overcommit"]:
@@ -348,7 +348,7 @@ def test_slurm_executor_to_directives_timeout_signal(
         timeout_signal=timeout_signal,
         timeout_signal_grace_period=grace_period,
     )
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     signal_directive = [d for d in directives if d.startswith("--signal=")][0]
     assert signal_directive == f"--signal={expected_signal}"
 
@@ -369,7 +369,7 @@ def test_slurm_executor_to_directives_requeue(basic_requirements, requeue, max_a
         requeue=requeue,
         requeue_max_attempts=max_attempts,
     )
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     requeue_directives = [d for d in directives if "requeue" in d]
     assert expected in requeue_directives
 
@@ -381,7 +381,7 @@ def test_slurm_executor_to_directives_bind_gpu(basic_requirements):
         time=dt.timedelta(hours=1),
         bind={resources.ResourceType.GPU: "closest"},
     )
-    directives = executor.to_directives()
+    directives = executor.step_directives()
     gpu_bind_directive = [d for d in directives if d.startswith("--gpu-bind=")][0]
     assert gpu_bind_directive == "--gpu-bind=closest"
 
@@ -404,7 +404,7 @@ def test_slurm_executor_to_directives_bind_resource_types(
         time=dt.timedelta(hours=1),
         bind={resource: value},
     )
-    directives = executor.to_directives()
+    directives = executor.step_directives()
     assert any(expected in d for d in directives)
 
 
@@ -425,14 +425,14 @@ def test_slurm_executor_to_directives_bind_custom_gres(
         time=dt.timedelta(hours=1),
         bind={custom_gres: value},
     )
-    directives = executor.to_directives()
+    directives = executor.step_directives()
     assert any(expected in d for d in directives)
 
 
 def test_slurm_executor_to_directives_includes_requirements_directives(basic_requirements):
     """Test that to_directives includes directives from requirements."""
     executor = executors.Slurm(requirements=basic_requirements, time=dt.timedelta(hours=1))
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     assert any(d.startswith("--cpus-per-task=") for d in directives)
     assert any(d.startswith("--mem-per-cpu=") for d in directives)
 
@@ -476,11 +476,22 @@ def test_slurm_executor_with_multiple_binds(basic_requirements):
             resources.ResourceType.MEMORY: "local",
         },
     )
-    directives = executor.to_directives()
+    directives = executor.step_directives()
     assert any("--gpu-bind=closest" in d for d in directives)
     assert any("--mem-bind=local" in d for d in directives)
 
 
+def test_slurm_executor_bind_flag_sets_gres_flags(basic_requirements):
+    """Test that bind_flag produces the correct --gres-flags directive."""
+    executor = executors.Slurm(
+        requirements=basic_requirements,
+        time=dt.timedelta(hours=1),
+        bind_flag="enforce-binding",
+    )
+    directives = executor.step_directives()
+    assert "--gres-flags=enforce-binding" in directives
+
+
 @pytest.mark.parametrize("grace_period_secs", [1, 30, 3600])
 def test_slurm_executor_with_various_grace_periods(basic_requirements, grace_period_secs):
     """Test executor with various timeout grace periods."""
@@ -489,7 +500,7 @@ def test_slurm_executor_with_various_grace_periods(basic_requirements, grace_per
         time=dt.timedelta(hours=24),
         timeout_signal_grace_period=dt.timedelta(seconds=grace_period_secs),
     )
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     signal_directive = [d for d in directives if d.startswith("--signal=")][0]
     assert f"@{grace_period_secs}" in signal_directive
 
@@ -520,7 +531,7 @@ def test_slurm_executor_multiple_signal_types(basic_requirements, sig):
         time=dt.timedelta(hours=1),
         timeout_signal=sig,
     )
-    directives = executor.to_directives()
+    directives = executor.batch_directives()
     signal_directive = [d for d in directives if d.startswith("--signal=")][0]
     expected_sig_name = sig.name.removeprefix("SIG")
     assert expected_sig_name in signal_directive

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.17}/tests/test_resources.py

@@ -277,7 +277,7 @@ def test_job_requirements_to_directives_cpu(
 ) -> None:
     """Test CPU directive generation."""
     req = resources.JobRequirements(cpu=4, cluster=dummy_cluster_config)
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--cpus-per-task=4" in directives
 
 
@@ -290,7 +290,7 @@ def test_job_requirements_to_directives_memory(
         memory=8 * 1024**3,  # 8GB
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     # 8GB / 2 CPUs / 2^20 bytes per MB = 4096 MB
     assert "--mem-per-cpu=4096M" in directives
 
@@ -303,7 +303,7 @@ def test_job_requirements_to_directives_memory_with_single_cpu(
         memory=2 * 1024**3,  # 2GB
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     # 2GB / 1 CPU / 2^20 = 2048 MB
     assert "--mem-per-cpu=2048M" in directives
 
@@ -317,7 +317,7 @@ def test_job_requirements_to_directives_ram_alias(
         ram=8 * 1024**3,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--mem-per-cpu=4096M" in directives
 
 
@@ -329,7 +329,7 @@ def test_job_requirements_to_directives_disk(
         disk=1024**3,  # 1GB
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     # 1GB / 2^20 MB = 1024 MB
     assert "--tmp=1024M" in directives
 
@@ -342,7 +342,7 @@ def test_job_requirements_to_directives_ephemeral_storage_alias(
         ephemeral_storage=2 * 1024**3,  # 2GB
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     # 2GB / 2^20 MB = 2048 MB
     assert "--tmp=2048M" in directives
 
@@ -355,7 +355,7 @@ def test_job_requirements_to_directives_generic_gpu(
         gpu=4,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--gpus=4" in directives
 
 
@@ -367,7 +367,7 @@ def test_job_requirements_to_directives_specific_gpu(
         a100=2,
         cluster=cluster_with_gpu_mapping,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--gpus=a100:2" in directives
 
 
@@ -394,7 +394,7 @@ def test_job_requirements_to_directives_various_gpu_types(
     )
     kwargs = {accelerator_type.name.lower(): count}
     req = resources.JobRequirements(**kwargs, cluster=cluster)  # type: ignore
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert f"--gpus={gpu_name}:{count}" in directives
 
 
@@ -407,7 +407,7 @@ def test_job_requirements_to_directives_unmapped_gpu_raises_error(
         cluster=dummy_cluster_config,
     )
     with pytest.raises(ValueError, match="does not map resource type"):
-        req.to_directives()
+        req.batch_directives()
 
 
 def test_job_requirements_to_directives_custom_gres(
@@ -418,7 +418,7 @@ def test_job_requirements_to_directives_custom_gres(
         custom_resource="2",  # type: ignore
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--gres=custom_resource:2" in directives
 
 
@@ -431,7 +431,7 @@ def test_job_requirements_to_directives_replicas(
         replicas=8,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--ntasks=8" in directives
 
 
@@ -444,7 +444,7 @@ def test_job_requirements_to_directives_location(
         location="node[001-005]",
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--nodelist=node[001-005]" in directives
 
 
@@ -461,7 +461,7 @@ def test_job_requirements_to_directives_combined_all_resources(
         location="gpu-nodes[001-004]",
         cluster=cluster_with_gpu_mapping,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--cpus-per-task=8" in directives
     assert "--mem-per-cpu=8192M" in directives  # 64GB / 8 CPUs / 2^20
     assert "--tmp=512000M" in directives  # 500GB / 2^20
@@ -655,7 +655,7 @@ def test_job_requirements_invalid_memory_type(
     req = resources.JobRequirements(cluster=dummy_cluster_config)
     req.task_requirements[resources.ResourceType.MEMORY] = "invalid"  # type: ignore
     with pytest.raises(AssertionError, match="Memory must be an integer or float"):
-        req.to_directives()
+        req.batch_directives()
 
 
 def test_job_requirements_invalid_cpu_type(
@@ -665,7 +665,7 @@ def test_job_requirements_invalid_cpu_type(
     req = resources.JobRequirements(cluster=dummy_cluster_config)
     req.task_requirements[resources.ResourceType.CPU] = 4.5  # type: ignore
     with pytest.raises(AssertionError, match="CPU must be an integer"):
-        req.to_directives()
+        req.batch_directives()
 
 
 def test_job_requirements_invalid_gpu_type(
@@ -675,7 +675,7 @@ def test_job_requirements_invalid_gpu_type(
     req = resources.JobRequirements(cluster=dummy_cluster_config)
     req.task_requirements[resources.ResourceType.GPU] = 1.5  # type: ignore
     with pytest.raises(AssertionError, match="GPU must be an integer"):
-        req.to_directives()
+        req.batch_directives()
 
 
 def test_job_requirements_invalid_disk_type(
@@ -685,7 +685,7 @@ def test_job_requirements_invalid_disk_type(
     req = resources.JobRequirements(cluster=dummy_cluster_config)
     req.task_requirements[resources.ResourceType.DISK] = 1.5  # type: ignore
     with pytest.raises(AssertionError, match="Disk space must be an integer"):
-        req.to_directives()
+        req.batch_directives()
 
 
 def test_job_requirements_invalid_replicas(
@@ -695,7 +695,7 @@ def test_job_requirements_invalid_replicas(
     req = resources.JobRequirements(cluster=dummy_cluster_config)
     req.replicas = -1
     with pytest.raises(AssertionError, match="Replicas must be a positive integer"):
-        req.to_directives()
+        req.batch_directives()
 
 
 def test_job_requirements_invalid_replicas_zero(
@@ -705,7 +705,7 @@ def test_job_requirements_invalid_replicas_zero(
     req = resources.JobRequirements(cluster=dummy_cluster_config)
     req.replicas = 0
     with pytest.raises(AssertionError, match="Replicas must be a positive integer"):
-        req.to_directives()
+        req.batch_directives()
 
 
 def test_job_requirements_invalid_location_type(
@@ -715,7 +715,7 @@ def test_job_requirements_invalid_location_type(
     req = resources.JobRequirements(cluster=dummy_cluster_config)
     req.location = 123  # type: ignore
     with pytest.raises(AssertionError, match="Location must be a string"):
-        req.to_directives()
+        req.batch_directives()
 
 
 @pytest.mark.parametrize(
@@ -740,7 +740,7 @@ def test_job_requirements_memory_directive_calculation(
         memory=memory_bytes,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert f"--mem-per-cpu={expected_mem_per_cpu_mb}M" in directives
 
 
@@ -763,7 +763,7 @@ def test_job_requirements_disk_directive_calculation(
         disk=disk_bytes,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert f"--tmp={expected_tmp_mb}M" in directives
 
 
@@ -776,7 +776,7 @@ def test_job_requirements_fractional_memory_rounds_correctly(
         memory=10 * 1024**3,  # 10GB / 3 = 3.33GB per CPU
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     # math.ceil(10 * 1024^3 / 3 / 2^20) = math.ceil(3413.33...) = 3414
     assert "--mem-per-cpu=3414M" in directives
 
@@ -789,7 +789,7 @@ def test_job_requirements_fractional_disk_rounds_correctly(
         disk=5 * 1024**3 + 512 * 1024**2,  # 5.5GB
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     # math.ceil((5*1024^3 + 512*1024^2) / 2^20) = math.ceil(5632) = 5632
     assert "--tmp=5632M" in directives
 
@@ -799,7 +799,7 @@ def test_job_requirements_empty_task_requirements(
 ) -> None:
     """Test job requirements with no specific resources."""
     req = resources.JobRequirements(cluster=dummy_cluster_config)
-    directives = req.to_directives()
+    directives = req.batch_directives()
     # Should still have ntasks directive
     assert "--ntasks=1" in directives
     assert len(directives) == 1
@@ -835,8 +835,8 @@ def test_job_requirements_multiple_directives_order_independence(
         memory=4 * 1024**3,
         cluster=cluster_with_gpu_mapping,
     )
-    directives1 = set(req1.to_directives())
-    directives2 = set(req2.to_directives())
+    directives1 = set(req1.batch_directives())
+    directives2 = set(req2.batch_directives())
     assert directives1 == directives2
 
 
@@ -1150,7 +1150,7 @@ def test_topology_directive_gpus_per_task(dummy_cluster_config: config.SlurmClus
         gpu=topo,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--gpus-per-task=4" in directives
 
 
@@ -1163,7 +1163,7 @@ def test_topology_directive_gpus_per_task_with_specific_accelerator(
         a100=topo,
         cluster=cluster_with_gpu_mapping,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--gpus-per-task=a100:2" in directives
 
 
@@ -1176,7 +1176,7 @@ def test_topology_directive_ntasks_from_topology(
         gpu=topo,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--ntasks=3" in directives
 
 
@@ -1187,7 +1187,7 @@ def test_topology_directive_switches(dummy_cluster_config: config.SlurmClusterCo
         gpu=topo,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert any("--switches=" in d for d in directives)
 
 
@@ -1201,7 +1201,7 @@ def test_topology_directive_switches_with_timeout(
         gpu=topo,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     # Should have format like "--switches=2@5:00" (5 minutes in SLURM format)
     switches_directives = [d for d in directives if "--switches=" in d]
     assert len(switches_directives) > 0
@@ -1218,7 +1218,7 @@ def test_topology_1d_topology_no_ntasks_override(
         replicas=4,
         cluster=dummy_cluster_config,
     )
-    directives = req.to_directives()
+    directives = req.batch_directives()
     assert "--ntasks=4" in directives
 
 

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

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

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.17}/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.17}/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.17}/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

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.17}/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,

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.17}/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
 

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.17}/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

{xmanager_slurm-0.4.15 → xmanager_slurm-0.4.17}/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.15 → xmanager_slurm-0.4.17}/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.17}/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 \