Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vineethsaivs wants to merge 2 commits into
base: master
Choose a base branch
from
+82 −22
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9a3bf84
Honor --include/--exclude in the SLURM launcher
vineethsaivs Aug 25, 2026
92a4379
Merge branch 'master' into fix/slurm-include-exclude
sfc-gh-truwase Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 16 additions & 7 deletions deepspeed/launcher/multinode_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())

@ebarkhordar ebarkhordar Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I ran this against 8fa2478 in a clean container. The --include path does what the description says. One thing the description does not name: this line is unconditional, so it also changes -n for --num_gpus and --num_nodes.

runner.py trims active_resources for those two flags as well (--num_nodes at runner.py:521, --num_gpus at runner.py:529), and get_cmd now sizes the job from active_resources instead of self.resource_pool. Building the srun command at this SHA and at the merge base aa3914d, on the same two-node four-slot pool the launcher tests use:

                           base    head
control (no flags)         -n 8    -n 8
--include worker-1:0,2     -n 8    -n 2
--num_gpus 2               -n 8    -n 4
--num_nodes 1              -n 8    -n 4

The bottom two look correct to me rather than wrong: world_info_base64 is encoded from the same trimmed active_resources at runner.py:543, so -n 8 against a 4 rank world info was already inconsistent. My point is only that they are unpinned. test_slurm_runner_resource_filter parametrizes include and exclude, and test_slurm_runner asserts the unfiltered 8, so nothing covers the two flags that now also move.

Two more cases in that parametrize would cover it, though they need main()'s trim replicated the way parse_inclusion_exclusion already is, since it happens after it. Naming the wider scope in the description would work too.

How I got the numbers: I called parse_inclusion_exclusion and then replicated runner.py:521-533 in a probe, then called SlurmRunner.get_cmd. I did not run main(), and I have no slurm cluster here, so the --nodelist and --nodes semantics in your description are the only part I did not check.

srun_cmd = [
'srun',
'-n',
Expand All @@ -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)
Comment on lines +377 to +378

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Require every filtered Slurm host

When a filter leaves multiple active hosts but reduces the task count, such as the new --exclude worker-1:0 case, passing only --nodelist does not force Slurm to launch on all of those hosts; SchedMD documents that a lower node or processor count may use only a subset of the supplied nodelist (https://slurm.schedmd.com/srun.html). In clusters where one listed node can satisfy -n, this can silently run all ranks on a subset of active_resources instead of honoring the host set the user requested, so the Slurm translation needs to also constrain the node count/task placement for the filtered hosts.

Useful? React with 👍 / 👎.

# --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}')
Expand Down
38 changes: 25 additions & 13 deletions deepspeed/launcher/runner.py
100755 → 100644
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down Expand Up @@ -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"
Expand Down
43 changes: 41 additions & 2 deletions tests/unit/launcher/test_multinode_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down
Loading