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

[vineethsaivs]

Ran this down rather than taking it either way, because whether it happens depends on how torch counts devices.

main() deletes the mask before asking for the count:

        del os.environ[visible_devices_env]
    ...
        device_count = get_accelerator().device_count()

and CUDA_Accelerator.device_count() is torch.cuda.device_count(), which is explicit about not caching across a mask change:

    if _cached_device_count is not None:
        return _cached_device_count
    nvml_count = _device_count_amdsmi() if torch.version.hip else _device_count_nvml()
    r = torch._C._cuda_getDeviceCount() if nvml_count < 0 else nvml_count
    # 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

The launcher never initializes CUDA, so _initialized is False and _cached_device_count stays None. _device_count_nvml() re-reads CUDA_VISIBLE_DEVICES from the environment on every call, so once the mask is gone it reports the physical count. That is the normal path, and test_visible_devices_resolve_to_the_requested_slots pins it: mask 0,2, pool {'localhost': 4}, resolved to {'localhost': [0, 2]}.

The narrow case where you are right is NVML discovery failing (nvml_count < 0), where the count falls back to torch._C._cuda_getDeviceCount(), which the driver does cache from the first call, and the accelerator detection in real_accelerator.py makes that first call while the mask is still set. There the pool is {'localhost': 2} while --include names slot 2.

Both behaviours are wrong there, and they are wrong differently: before this PR the include was filtered against itself, so localhost:4,5 passed a check that never looked at the real slots; after it, the same request is rejected by name. Preferring the loud one is the trade this PR makes, and test_visible_devices_are_rejected_when_the_device_count_is_masked says so rather than leaving it implicit.

What would actually fix that corner is sizing the pool from the mask when the mask is what produced --include, which is a change in main() and not in this filter. Happy to send it separately; it needs care because CUDA_VISIBLE_DEVICES entries can be GPU UUIDs and not just indices, so max(index) + 1 is not always available.

(The sibling P1 on this file asks for a Signed-off-by trailer that is already present: the DCO check is green on every commit here.)

[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.

[tohtana]

Hi @vineethsaivs, thank you for submitting this fix!
Thank you @ebarkhordar for the insightful discussion.

To clarify the behavior, I ran the check on a 4x H100 node (torch 2.9.1): with CUDA_VISIBLE_DEVICES=0,2. device_count after the del returns 4, and the include resolves to {'localhost': [0, 2]}. So the ordering holds on current torch.

If a user uses an old version or a backend caches the count again, this path regresses to #5818. But the error is clear, and this PR fixes real bugs, so I'm fine taking that trade.
Just to detect the regression, would it be possible for you to add a test for the CUDA_VISIBLE_DEVICES=0,2 case before merge? @vineethsaivs

[vineethsaivs]

Thank you @tohtana, and thank you for actually running it on a 4x H100 node rather than taking the reasoning on trust. Test added in 83daee0.

Two cases, both CPU-only, in tests/unit/launcher/test_run.py. Neither launches anything: a helper drives main() as far as parse_inclusion_exclusion() and captures what it was handed, with a fake accelerator supplying visible_devices_envs() and device_count().

1. The case you measured, CUDA_VISIBLE_DEVICES=0,2 with the physical count reported after the del:

main() hands over: resource_pool={'localhost': 4}  include='localhost:0,2'
which resolves to: {'localhost': [0, 2]}

This passes on master today as well, which is the point of it: it is the regression detector for the ordering, so if a backend ever caches the count again this goes red instead of the failure resurfacing as #5818.

2. The trade you accepted, pinned so it is visible. Same mask, but with the accelerator reporting the masked count of 2:

                master   ->  {'localhost': [0, 2]}          # slot 2 is not in the pool
                this PR  ->  ValueError: No slot '2' specified on host 'localhost'

So the regression, if it ever comes back, is now a named error rather than a launch on the wrong GPUs, and the test says which of the two behaviours is intended. Run against unpatched master this second test is the only one that fails (DID NOT RAISE <class 'ValueError'>), which is what makes it a real assertion rather than a restatement.

Whole file: 11 passed on this branch, 1 failed, 8 passed with deepspeed/launcher/runner.py reverted to master's version. ruff reports the identical 6 pre-existing findings on the file before and after, and yapf is clean on both changed files.

On the codex bot's P2 above, for completeness: it describes this same masked-count path and asks to either preserve the auto-generated list or validate it against the physical slot universe. This PR does the second, and case 2 is now the test for it.

[tohtana]

Sorry for the delay, and thank you for adding the tests!
I still see the test has a small gap. The fake always returns the injected device count regardless of whether CUDA_VISIBLE_DEVICES has already been cleared, so the test would still pass if the ordering regressed. Could the test record the environment state when device_count() is called and verify that the variable is absent at that point? Otherwise, this looks good to me.

[vineethsaivs]

Thank you, and you are right: the fake returned a constant, so the ordering was never actually pinned.

device_count() now records os.environ.get('CUDA_VISIBLE_DEVICES') each time it is called, and both cases assert it was already gone:

    def device_count(self):
        self.masks_at_device_count.append(os.environ.get('CUDA_VISIBLE_DEVICES'))
        return self._device_count
...
    # the count has to be asked for after the mask is gone, or it is the masked one
    assert captured['masks_at_device_count'] == [None]

Checked that the new assertion is load-bearing rather than taking it on faith. Moving the count above the del in runner.py, which is the regression you described:

            print(f"{detected_str}: setting --include={args.include}")
+       _regression_masked_count = get_accelerator().device_count()
        del os.environ[visible_devices_env]

unit/launcher/test_run.py::test_visible_devices_resolve_to_the_requested_slots
  AssertionError: assert ['0,2', None] == [None]
unit/launcher/test_run.py::test_visible_devices_are_rejected_when_the_device_count_is_masked
  AssertionError: assert ['0,2', None] == [None]
2 failed

Three states of tests/unit/launcher/test_run.py, all CPU:

runner.py	result
this PR	11 passed
this PR, count moved above the del	2 failed, as above
upstream master	3 failed (test_parse_inclusion_exclusion, test_parse_inclusion_exclusion_errors, test_visible_devices_are_rejected_when_the_device_count_is_masked)

[tohtana]

Thank you for the update! Looks good to me.