diff --git a/src/post_training/config.py b/src/post_training/config.py index 6a20f1a..6b82c4a 100644 --- a/src/post_training/config.py +++ b/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 diff --git a/src/post_training/slurm/job.sh.jinja b/src/post_training/slurm/job.sh.jinja index 8cf99d7..2de2c2b 100644 --- a/src/post_training/slurm/job.sh.jinja +++ b/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 }} diff --git a/src/post_training/slurm/job_llamafactory.sh.jinja b/src/post_training/slurm/job_llamafactory.sh.jinja index edd99cc..c1206f2 100644 --- a/src/post_training/slurm/job_llamafactory.sh.jinja +++ b/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 }} diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja index 989a947..d2d5913 100644 --- a/src/post_training/slurm/job_trl_container.sh.jinja +++ b/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 }} diff --git a/src/post_training/slurm/launcher.py b/src/post_training/slurm/launcher.py index fd8338f..079ff34 100644 --- a/src/post_training/slurm/launcher.py +++ b/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, diff --git a/tests/test_slurm_render.py b/tests/test_slurm_render.py index a0da377..5c40cfc 100644 --- a/tests/test_slurm_render.py +++ b/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