feat: add slurm.exclude option to skip nodes

src/post_training/config.py

@@ -236,6 +236,7 @@ class SlurmConfig:
     partition: str = "gpu"
     account: str | None = None
     qos: str | None = None
+    exclude: str | None = None
     num_nodes: int = 1
     gpus_per_node: int = 4
     cpus_per_task: int = 32

src/post_training/slurm/job.sh.jinja

@@ -12,6 +12,9 @@
 {% if qos -%}
 #SBATCH --qos={{ qos }}
 {% endif -%}
+{% if exclude -%}
+#SBATCH --exclude={{ exclude }}
+{% endif -%}
 #SBATCH --nodes={{ num_nodes }}
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:{{ gpus_per_node }}

src/post_training/slurm/job_llamafactory.sh.jinja

@@ -12,6 +12,9 @@
 {% if qos -%}
 #SBATCH --qos={{ qos }}
 {% endif -%}
+{% if exclude -%}
+#SBATCH --exclude={{ exclude }}
+{% endif -%}
 #SBATCH --nodes={{ num_nodes }}
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:{{ gpus_per_node }}

src/post_training/slurm/job_trl_container.sh.jinja

@@ -12,6 +12,9 @@
 {% if qos -%}
 #SBATCH --qos={{ qos }}
 {% endif -%}
+{% if exclude -%}
+#SBATCH --exclude={{ exclude }}
+{% endif -%}
 #SBATCH --nodes={{ num_nodes }}
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:{{ gpus_per_node }}

src/post_training/slurm/launcher.py

@@ -61,6 +61,7 @@ def render_trl_slurm_script(
         partition=config.slurm.partition,
         account=config.slurm.account,
         qos=config.slurm.qos,
+        exclude=config.slurm.exclude,
         num_nodes=config.slurm.num_nodes,
         gpus_per_node=config.slurm.gpus_per_node,
         cpus_per_task=config.slurm.cpus_per_task,
@@ -120,6 +121,7 @@ def render_trl_container_slurm_script(
         partition=config.slurm.partition,
         account=config.slurm.account,
         qos=config.slurm.qos,
+        exclude=config.slurm.exclude,
         num_nodes=config.slurm.num_nodes,
         gpus_per_node=config.slurm.gpus_per_node,
         cpus_per_task=config.slurm.cpus_per_task,
@@ -177,6 +179,7 @@ def render_llamafactory_slurm_script(
         partition=config.slurm.partition,
         account=config.slurm.account,
         qos=config.slurm.qos,
+        exclude=config.slurm.exclude,
         num_nodes=config.slurm.num_nodes,
         gpus_per_node=config.slurm.gpus_per_node,
         cpus_per_task=config.slurm.cpus_per_task,

tests/test_slurm_render.py

@@ -421,3 +421,35 @@ def test_llamafactory_container_custom_path(tmp_path, config):
     content = render_llamafactory_slurm_script(config, run_dir).read_text()
 
     assert 'export PATH=\\"/opt/conda/bin:/usr/local/bin:/usr/bin:/bin\\"' in content
+
+
+# ---------------------------------------------------------------------------
+# exclude field — all three templates
+# ---------------------------------------------------------------------------
+
+
+def _render_all(config, run_dir):
+    return [
+        render_trl_slurm_script(config, run_dir, "configs/trl/sft.yaml").read_text(),
+        render_trl_container_slurm_script(config, run_dir, "configs/trl/sft.yaml").read_text(),
+        render_llamafactory_slurm_script(config, run_dir).read_text(),
+    ]
+
+
+def test_exclude_rendered(tmp_path, config):
+    """exclude appears as an #SBATCH directive when set."""
+    config.slurm.exclude = "node[01,03]"
+    run_dir = tmp_path / "outputs" / "my-run"
+    run_dir.mkdir(parents=True)
+
+    for content in _render_all(config, run_dir):
+        assert "#SBATCH --exclude=node[01,03]" in content
+
+
+def test_exclude_absent_when_none(tmp_path, config):
+    """The exclude directive is suppressed when not set."""
+    run_dir = tmp_path / "outputs" / "my-run"
+    run_dir.mkdir(parents=True)
+
+    for content in _render_all(config, run_dir):
+        assert "--exclude" not in content