Filter --include against the real slots, not against itself

[vineethsaivs]

Problem

parse_inclusion_exclusion() builds the host_info dict that parse_resource_filter() validates against, and it seeded that dict from the inclusion string:

node_configs = parse_node_config_list(inclusion)
for hostname, slots in resource_pool.items():
    active_resources[hostname] = node_configs[hostname] if hostname in node_configs else list(range(slots))

So for any host named in --include, the filter checked the request against a copy of itself. Two things follow.

A bare hostname resolves to no slots at all. parse_node_config_list("worker-0") gives {"worker-0": []}, so active_resources["worker-0"] became [], parse_resource_filter's whole-node branch copied that empty list back, and the post-processing dropped the host for being empty. parse_resource_filter's own docstring uses worker-0@worker-1:0,2 as its example of "use all slots on worker-0 and slots [0, 2] on worker-1":

pool = {"worker-0": 4, "worker-1": 4}

--include worker-0                  ->  {}                                     (expected all 4 slots)
--include worker-0@worker-1:0,2     ->  {'worker-1': [0, 2]}                   (expected worker-0 too)

The job then launches on 2 GPUs instead of 6, with no error and no warning, and the node the user asked for first is the one that disappears.

The slot check could never fire for --include. The same input is rejected through --exclude and accepted through --include:

--include worker-0:99   ->  ACCEPTED, {'worker-0': [99]}
--exclude worker-0:99   ->  ValueError: No slot '99' specified on host 'worker-0'
--include worker-0:0,4  ->  ACCEPTED, {'worker-0': [0, 4]}    (off-by-one on a 4 GPU host)
--exclude worker-0:0,4  ->  ValueError: No slot '4' specified on host 'worker-0'

{'worker-0': [99]} goes on to CUDA_VISIBLE_DEVICES, so a typo surfaces as a CUDA error from inside torch rather than as the launcher error that already exists for it. Hostname validation is symmetric and works; only the slot check is affected.

Fix

Hand parse_resource_filter() the slots the machines actually have and let it do the filtering it was written to do. parse_resource_filter is unchanged: it already sets filtered_hosts[hostname] = slots for an explicit slot list and host_info[hostname] for a bare hostname, both of which are now correct.

--include worker-0                  ->  {'worker-0': [0, 1, 2, 3]}
--include worker-0@worker-1:0,2     ->  {'worker-0': [0, 1, 2, 3], 'worker-1': [0, 2]}
--include worker-0:99               ->  ValueError: No slot '99' specified on host 'worker-0'

Test

tests/unit/launcher/test_run.py has good coverage of parse_resource_filter, but every one of those tests hands it a correct host_info dict directly, so nothing exercised the wrapper that builds it. That is why this was invisible: the function under test was fine, and the caller was not.

Two tests added, test_parse_inclusion_exclusion and test_parse_inclusion_exclusion_errors, covering the bare hostname, the docstring's own example, the mixed form, exclusion, and out-of-range slots through both --include and --exclude.

                                     before   after
existing 7 tests                     pass     pass
test_parse_inclusion_exclusion       FAIL     pass     (OrderedDict() != {'worker-0': [0, 1, 2, 3]})
test_parse_inclusion_exclusion_errors FAIL    pass     (DID NOT RAISE ValueError)

python -m pytest tests/unit/launcher/test_run.py gives 2 failed / 7 passed with the source change reverted and 9 passed with it, on CPU. yapf --style .style.yapf and flake8 --config .flake8 are clean on both changed files, and clean on the unmodified tree as a control.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5aaeb8083a

ℹ️ 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]

Add the required Signed-off-by trailer

This non-merge commit's message does not contain a Signed-off-by: trailer, so it violates the repository's DCO/commit requirement and may be rejected by CI or maintainers; please amend the commit with --signoff using the configured author identity.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Preserve auto-masked CUDA device IDs

When no hostfile is present, main() first reads CUDA_VISIBLE_DEVICES, converts it to args.include, and only then builds resource_pool; because get_accelerator() has already queried the CUDA backend while the mask is set, the later device count can still be the masked count. With a common mask such as CUDA_VISIBLE_DEVICES=4,5, this line can build localhost: [0, 1], so the generated include localhost:4,5 is rejected as out of range instead of launching on the user-selected GPUs. Please either preserve the auto-generated visible-device list for this path or validate it against the physical slot universe.

Useful? React with 👍 / 👎.

[ebarkhordar]

The self-filtering you describe is real, and a bare hostname resolving to no slots is clearly wrong. One thing to check before this lands, because the line being removed was put there on purpose.

active_resources[hostname] = node_configs[hostname] if hostname in node_configs else list(range(slots)) came in with 5cbbff4, "Fix device selection using CUDA_VISIBLE_DEVICES (#6530)", whose commit message is "Instead of contiguous numbers based on the device count, this PR uses device indices in --include". That PR closed #5818, where the reported symptom was ValueError: No slot '2' specified on host 'localhost' under CUDA_VISIBLE_DEVICES=0,2.

That is the input whose behaviour changes here. Same container, both SHAs, pip install -e .:

from deepspeed.launcher.runner import parse_inclusion_exclusion
parse_inclusion_exclusion({"localhost": 2}, "localhost:0,2", "")

• master da066407: {'localhost': [0, 2]}
• this PR 5aaeb808: ValueError: No slot '2' specified on host 'localhost'

which is #5818's error string verbatim. test_parse_inclusion_exclusion_errors pins the new behaviour deliberately, so it is a decision rather than an oversight, but I did not see #6530 addressed anywhere in the PR.

Whether the launcher actually reaches that state depends on one value I could not measure: main() builds the pair at runner.py:460 and :474, so the question is whether get_accelerator().device_count() at :471 reports the CVD-reduced count or the full one, given del os.environ[visible_devices_env] runs first at :462. I have no GPU on the machine I tested on, so I only ran the function boundary above. Reading torch, torch.cuda.device_count() now re-reads the variable unless CUDA is already initialized and caches only after init, which would make :471 return the full count and keep the CVD path working. It was decorated with lru_cache in the torch of #6530's era, which is plausibly why #5818 existed at all.

So this may be fine on current torch and broken on older, and neither the existing test_parser_* cases nor the new ones cover the CUDA_VISIBLE_DEVICES path. Would a case that pins device_count() and asserts localhost:0,2 still resolves be worth adding, so #5818 cannot come back unnoticed?

[vineethsaivs]

Two review points, one wrong and one worth answering properly.

The Signed-off-by P1 is a false positive. 5aaeb8083 ends with Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com> and the DCO check on this head is green.

The CUDA_VISIBLE_DEVICES P2 is a real thing to check, and I think the ordering already protects it. The concern is that resource_pool['localhost'] could come back as the masked count while the auto-generated include names physical ids, so --include localhost:4,5 would now be rejected against [0, 1].

The mask is deleted before the count is read, though. In main():

args.include = f"localhost:{visible_devices}"
...
del os.environ[visible_devices_env]        # mask gone here
...
device_count = get_accelerator().device_count()   # read here
resource_pool['localhost'] = device_count

That only helps if the count is not cached from an earlier query, and torch is explicit that it is not:

# NB: Do not cache the device count prior to CUDA initialization, because
# the number of devices can change due to changes to CUDA_VISIBLE_DEVICES
# setting prior to CUDA initialization.
if _initialized:
    _cached_device_count = r

CUDA_Accelerator.device_count() is a straight call to torch.cuda.device_count(), and _device_count_nvml() re-reads CUDA_VISIBLE_DEVICES on each call. The launcher only spawns subprocesses and never creates a CUDA context, so _initialized stays false and nothing is cached. After the del, the count is the physical one and range(count) contains 4 and 5.

Being straight about the limits of that: I established it from torch's source rather than by running it, because I do not have a CUDA box. If someone can run CUDA_VISIBLE_DEVICES=4,5 deepspeed train.py on a multi-GPU machine against this branch, that settles it in one command, and I would rather have that than my reading.

If you would prefer not to depend on the ordering at all, the narrow alternative is to have the VISIBLE_DEVICES branch record the ids it generated and let those through unchecked, keeping the range check for user-supplied --include. Happy to push that instead. I did not do it by default because it re-opens the hole for the one path that generates the string itself, and the ordering above already places the del before the read.