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.

Sizing from active_resources also moves -n for --num_gpus and --num_nodes, which
runner.py trims the same dict for. Those counts were already inconsistent with
world_info_base64, which is encoded from the trimmed dict, so this lines them up;
the two flags now have test coverage. main()'s trim moves to
apply_num_nodes_and_gpus() unchanged so the tests can build the same resource
dict main() hands the backends.

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

deepspeed/launcher/runner.py

@@ -398,6 +398,30 @@ def parse_inclusion_exclusion(resource_pool, inclusion, exclusion):
     return parse_resource_filter(active_resources, include_str=inclusion, exclude_str=exclusion)
 
 
+def apply_num_nodes_and_gpus(active_resources, num_nodes, num_gpus):
+    """Trim resolved resources to the top num_nodes hosts and num_gpus slots per host.
+
+    Split out of main() so the launcher backends can be exercised against the same
+    resource dict main() hands them. Both flags are mutually exclusive with
+    --include/--exclude, which main() enforces before calling this.
+    """
+    if num_nodes > 0:
+        updated_active_resources = collections.OrderedDict()
+        for count, hostname in enumerate(active_resources.keys()):
+            if num_nodes == count:
+                break
+            updated_active_resources[hostname] = active_resources[hostname]
+        active_resources = updated_active_resources
+
+    if num_gpus > 0:
+        updated_active_resources = collections.OrderedDict()
+        for hostname in active_resources.keys():
+            updated_active_resources[hostname] = list(range(num_gpus))
+        active_resources = updated_active_resources
+
+    return active_resources
+
+
 def encode_world_info(world_info):
     world_info_json = json.dumps(world_info).encode('utf-8')
     world_info_base64 = base64.urlsafe_b64encode(world_info_json).decode('utf-8')
@@ -518,19 +542,7 @@ def main(args=None):
         run_autotuning(args, active_resources)
         return
 
-    if args.num_nodes > 0:
-        updated_active_resources = collections.OrderedDict()
-        for count, hostname in enumerate(active_resources.keys()):
-            if args.num_nodes == count:
-                break
-            updated_active_resources[hostname] = active_resources[hostname]
-        active_resources = updated_active_resources
-
-    if args.num_gpus > 0:
-        updated_active_resources = collections.OrderedDict()
-        for hostname in active_resources.keys():
-            updated_active_resources[hostname] = list(range(args.num_gpus))
-        active_resources = updated_active_resources
+    active_resources = apply_num_nodes_and_gpus(active_resources, args.num_nodes, args.num_gpus)
 
     if args.elastic_training:
         assert not args.no_local_rank, "--no_local_rank argument is not supported in Elastic training"

tests/unit/launcher/test_multinode_runner.py

@@ -5,7 +5,8 @@
 
 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,
+                                       apply_num_nodes_and_gpus)
 import os
 import pytest
 
@@ -62,9 +63,47 @@ 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
+
+
+@pytest.mark.parametrize('resource_flag, expected_srun_flag, expected_process_count',
+                         [(['--num_gpus', '2'], ('--gpus', '2'), '4'), (['--num_nodes', '1'], ('--nodes', '1'), '4')])
+def test_slurm_runner_num_nodes_and_gpus(runner_info, resource_flag, expected_srun_flag, expected_process_count):
+    # main() trims active_resources for these two flags as well, so sizing the job from it
+    # moves their task count too. They are mutually exclusive with --include/--exclude, so the
+    # resource-filter branch must stay silent and cannot append a second --nodes.
+    env, resource_pool, world_info, _ = runner_info
+    args = parse_args(resource_flag + ['test_launcher.py'])
+    active_resources = parse_inclusion_exclusion(resource_pool, args.include, args.exclude)
+    active_resources = apply_num_nodes_and_gpus(active_resources, args.num_nodes, args.num_gpus)
+    runner = mnrunner.SlurmRunner(args, world_info, resource_pool)
+    cmd = runner.get_cmd(env, active_resources)
+    assert '--nodelist' not in cmd
+    assert cmd.count(expected_srun_flag[0]) == 1
+    assert cmd[cmd.index(expected_srun_flag[0]) + 1] == expected_srun_flag[1]
+    assert cmd[cmd.index('-n') + 1] == expected_process_count
 
 
 def test_mvapich_runner(runner_info):