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

[vineethsaivs]

Problem

SlurmRunner.get_cmd() forwards DeepSpeed's own --include / --exclude strings to srun and sizes the job from the unfiltered hostfile. With a two-node, four-slot hostfile:

deepspeed --launcher slurm --include worker-1:0,2 train.py

builds

srun -n 8 --include worker-1:0,2 --export=ALL python -u train.py

Two things are wrong with that line:

1. srun has no --include option, so it exits with unrecognized option '--include' and the job never starts.
2. -n 8 counts every slot in the pool. The user asked for 2.

--exclude is broken the same way. srun does accept that flag, but DeepSpeed's NAME[:SLOT[,SLOT ...]] syntax is not a slurm hostlist (--exclude worker-1:0 names no node slurm knows), and -n still ignores the filter, so --exclude worker-1 asks for 8 tasks on the 4 slots that are left.

Fix

runner.py already resolves both flags into active_resources before calling get_cmd(), and it passes that dict in as the second argument. The other runners either use it (PDSHRunner) or explicitly reject the flags in validate_args(); SlurmRunner was the only one accepting them and then ignoring the resolved result.

So take the process count and the host list from active_resources, and give srun --nodelist, which is the flag it actually has:

srun -n 2 --nodelist worker-1 --nodes 1 --export=ALL python -u train.py

--nodelist alone is only an upper bound. SchedMD documents that "a lower node or processor count may only require a subset of the supplied node list", so -n 7 --nodelist worker-0,worker-1 could put all seven ranks on worker-0 and silently drop a host the filter kept. --nodes pins the count (a single number is used as both minimum and maximum). It cannot collide with the --nodes the branch below emits, because runner.py raises Cannot specify num_nodes/gpus with include/exclude, so args.num_nodes is always -1 on this path.

Slot-level selection still cannot be expressed to srun, but the task count and the node set are now what the user asked for, instead of an invalid command line.

Wider scope, stated rather than left implicit

Sizing from active_resources is unconditional, so it also moves -n for --num_gpus and --num_nodes, which runner.py trims the same dict for (L521 and L529):

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

The bottom two are corrections rather than regressions: world_info_base64 is encoded from that same trimmed dict at L543, so -n 8 against a four-rank world info was already inconsistent. Both are now covered by tests.

To pin them without a hand-rolled copy of main()'s trim in the test, those lines move verbatim into apply_num_nodes_and_gpus() next to parse_inclusion_exclusion, and main() calls it. Pure extraction, no behaviour change.

Test

Added test_slurm_runner_resource_filter to tests/unit/launcher/test_multinode_runner.py, covering --include with slots, --exclude with slots, and --exclude of a whole host, asserting the host list, the node count and the task count for each. test_slurm_runner also now asserts the unfiltered count and passes get_cmd() the slot lists that runner.py really hands it, rather than the raw hostfile counts.

test_slurm_runner_num_nodes_and_gpus covers the two flags above, asserting -n 4 for each plus that the resource-filter branch stays silent, so --nodelist is absent and --nodes appears exactly once.

# before the fix, with the new tests
5 failed, 7 passed

# after
12 passed

The rest of tests/unit/launcher/ is unchanged at 9 failed, 23 passed on a clean tree and 9 failed, 28 passed here; the 9 are pre-existing test_user_args.py failures in this environment because the deepspeed console script is not installed, and they reproduce identically on a clean tree.

yapf --style .style.yapf -d and flake8 --config .flake8 are clean on both changed files.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 78dbdc1416

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

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 👍 / 👎.

[vineethsaivs]

Good catch, and it is right. --nodelist is only an upper bound: srun's documentation says "a lower node or processor count may only require a subset of the supplied node list", so --exclude worker-1:0 producing srun -n 7 --nodelist worker-0,worker-1 could put all seven ranks on worker-0 and silently drop a host the filter kept, which is the same class of bug this PR is fixing.

Fixed in 8fa2478 by also passing --nodes with the count of filtered hosts. A single number there is used as both the minimum and the maximum, so the host count is pinned rather than merely permitted:

--include worker-1:0,2   ->  srun -n 2 --nodelist worker-1 --nodes 1 ...
--exclude worker-1:0     ->  srun -n 7 --nodelist worker-0,worker-1 --nodes 2 ...
--exclude worker-1       ->  srun -n 4 --nodelist worker-0 --nodes 1 ...

It cannot collide with the --nodes the args.num_nodes branch below emits: runner.py raises Cannot specify num_nodes/gpus with include/exclude, so num_nodes is always -1 whenever this branch runs.

The three parametrized cases now assert the node count alongside the host list and the task count: 3 failed / 7 passed against the unpatched source, 10 passed after. The unfiltered path is byte-for-byte unchanged and still emits no --nodelist or --nodes at all.

[ebarkhordar]

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.

[vineethsaivs]

Thanks, this is right and I have covered it in 9a3bf84.

You are correct that the line is unconditional and that --num_gpus / --num_nodes move with it, and correct that the new values are the consistent ones: world_info_base64 is encoded from the same trimmed dict at runner.py:543, so -n 8 against a 4-rank world info was the inconsistency. I reproduced your table exactly on the two-node four-slot pool:

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

Did both things you offered rather than picking one.

Pinned. New test_slurm_runner_num_nodes_and_gpus parametrizes --num_gpus 2 and --num_nodes 1, asserting -n 4 for each plus that the resource-filter branch stays silent, so --nodelist is absent and --nodes appears exactly once (the --num_nodes case is the one that could have grown a second one). Against the unpatched multinode_runner.py both fail with assert '8' == '4', which is the gap you found.

Not replicated in the test. You are right that replicating main()'s trim in a test would drift. main() lines 521-533 move verbatim into apply_num_nodes_and_gpus() next to parse_inclusion_exclusion, and main() calls it; no behaviour change, and the test now builds the same dict main() hands the backends instead of a hand-rolled copy.

Named in the description. The PR body and commit message now say that sizing from active_resources moves -n for those two flags too, and why that is the consistent direction.

Whole file: 5 failed / 7 passed before, 12 passed after. Whole tests/unit/launcher/: 9 failed / 23 passed on a clean tree, 9 failed / 28 passed here, the 9 being the pre-existing test_user_args.py failures where the deepspeed console script is not installed.

On the part you flagged as unchecked: I have no slurm cluster either, so the --nodelist / --nodes semantics come from the srun manual rather than a run. The specific sentence the --nodes addition rests on is "a lower node or processor count may only require a subset of the supplied node list", which is why --nodelist alone is not enough, and "if only one number is specified, this is used as both the minimum and maximum node count" for pinning with a bare --nodes N.

[ebarkhordar]

That covers it. Moving the trim into apply_num_nodes_and_gpus() and calling it from main() is the better of the two options, since the test now exercises the same dict the backends get rather than a copy that can drift, and test_slurm_runner_num_nodes_and_gpus pins the two cases that were unpinned. Nothing further from me.