Honor --include/--exclude in the SLURM launcher

The SLURM runner passed DeepSpeed's own resource-filter strings straight to srun
and sized the job from the unfiltered hostfile, so both filters were broken:

    deepspeed --launcher slurm --include worker-1:0,2 train.py
    -> srun -n 8 --include worker-1:0,2 ... train.py

srun has no --include, so it exits with "unrecognized option" and the job never
starts. --exclude fares no better: srun does have that flag, but DeepSpeed's
NAME[:SLOT,...] syntax is not a slurm hostlist, and -n still counts every slot in
the pool rather than the ones left after filtering.

runner.py already resolves both flags into active_resources before calling
get_cmd(), so take the process count and the node list from there and pass srun
--nodelist, which is the flag it actually has. --nodelist alone is only an upper
bound, since srun documents that a lower task count "may only require a subset of
the supplied node list", so pass --nodes as well to pin the host count; runner.py
forbids --num_nodes alongside a resource filter, so nothing else sets it.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>

Author: vineethsaivs

deepspeed/launcher/multinode_runner.py

@@ -358,7 +358,9 @@ def name(self):
     def get_cmd(self, environment, active_resources):
         assert not getattr(self.args, 'detect_nvlink_pairs',
                            False), "slurm backend does not support remapping visible devices"
-        total_process_count = sum(self.resource_pool.values())
+        # --include/--exclude are already resolved into active_resources, so counting the
+        # whole pool here would ask srun for slots the user filtered out.
+        total_process_count = sum(len(slots) for slots in active_resources.values())
         srun_cmd = [
             'srun',
             '-n',
@@ -368,12 +370,19 @@ def get_cmd(self, environment, active_resources):
         if getattr(self.args, 'slurm_comment', ''):
             srun_cmd += ['--comment', self.args.slurm_comment]
 
-        if self.args.include != "":
-            srun_cmd.append('--include')
-            srun_cmd.append(f'{self.args.include}')
-        if self.args.exclude != "":
-            srun_cmd.append('--exclude')
-            srun_cmd.append(f'{self.args.exclude}')
+        if self.args.include != "" or self.args.exclude != "":
+            # srun has no --include, and the NAME[:SLOT,...] syntax DeepSpeed accepts is not
+            # a slurm hostlist, so name the hosts that survived the filter instead.
+            active_hosts = ",".join(active_resources.keys())
+            srun_cmd.append('--nodelist')
+            srun_cmd.append(active_hosts)
+            # --nodelist alone is only an upper bound: srun documents that a lower task count
+            # "may only require a subset of the supplied node list", so it could pack every
+            # rank onto one host and silently drop a host the filter kept. --nodes pins the
+            # count, and runner.py forbids --num_nodes alongside a resource filter, so the
+            # branch below cannot also set it.
+            srun_cmd.append('--nodes')
+            srun_cmd.append(f'{len(active_resources)}')
         if self.args.num_nodes > 0:
             srun_cmd.append('--nodes')
             srun_cmd.append(f'{self.args.num_nodes}')

tests/unit/launcher/test_multinode_runner.py

@@ -5,7 +5,7 @@
 
 from copy import deepcopy
 from deepspeed.launcher import multinode_runner as mnrunner
-from deepspeed.launcher.runner import encode_world_info, parse_args
+from deepspeed.launcher.runner import encode_world_info, parse_args, parse_inclusion_exclusion
 import os
 import pytest
 
@@ -62,9 +62,29 @@ def test_mpich_runner(runner_info):
 
 def test_slurm_runner(runner_info):
     env, resource_pool, world_info, args = runner_info
+    active_resources = parse_inclusion_exclusion(resource_pool, args.include, args.exclude)
     runner = mnrunner.SlurmRunner(args, world_info, resource_pool)
-    cmd = runner.get_cmd(env, resource_pool)
+    cmd = runner.get_cmd(env, active_resources)
     assert cmd[0] == 'srun'
+    assert cmd[cmd.index('-n') + 1] == '8'
+
+
+@pytest.mark.parametrize('resource_filter, expected_hosts, expected_node_count, expected_process_count',
+                         [(['--include', 'worker-1:0,2'], 'worker-1', '1', '2'),
+                          (['--exclude', 'worker-1:0'], 'worker-0,worker-1', '2', '7'),
+                          (['--exclude', 'worker-1'], 'worker-0', '1', '4')])
+def test_slurm_runner_resource_filter(runner_info, resource_filter, expected_hosts, expected_node_count,
+                                      expected_process_count):
+    env, resource_pool, world_info, _ = runner_info
+    args = parse_args(resource_filter + ['test_launcher.py'])
+    active_resources = parse_inclusion_exclusion(resource_pool, args.include, args.exclude)
+    runner = mnrunner.SlurmRunner(args, world_info, resource_pool)
+    cmd = runner.get_cmd(env, active_resources)
+    assert '--include' not in cmd
+    assert cmd[cmd.index('--nodelist') + 1] == expected_hosts
+    # Without --nodes, srun may satisfy -n from a subset of --nodelist and drop a kept host.
+    assert cmd[cmd.index('--nodes') + 1] == expected_node_count
+    assert cmd[cmd.index('-n') + 1] == expected_process_count
 
 
 def test_mvapich_runner(runner_info):