diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index 7d50355f991..30dbc049943 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -417,6 +417,15 @@ def select_segment_nodes( domain_nodes: dict[str, list[tuple[str, int]]] = {} for nid, (domain, topo_rank) in topology.items(): + # Skip nodes with no NVLink-domain info. They all collapse into a single + # NVLINK_DOMAIN_UNKNOWN pseudo-domain with TOPO_RANK_UNKNOWN (-1), so they + # would sort first and be selected — but the resulting placement-group + # constraint {NVLINK_DOMAIN_UNKNOWN: 0.001} names a Ray resource that + # ray.sub never registers, so the bundle can never schedule. Excluding + # them here means we only ever pin to real, registered NVLink domains + # (and these nodes fall through to remaining_node_ids). + if domain == NVLINK_DOMAIN_UNKNOWN: + continue domain_nodes.setdefault(domain, []).append((nid, topo_rank)) for domain in domain_nodes: domain_nodes[domain].sort(key=lambda x: x[1]) diff --git a/tests/unit/distributed/test_topology_placement.py b/tests/unit/distributed/test_topology_placement.py index 49efa21f2ea..112e8be5b56 100644 --- a/tests/unit/distributed/test_topology_placement.py +++ b/tests/unit/distributed/test_topology_placement.py @@ -323,6 +323,35 @@ def test_selected_plus_remaining_is_full_topology(self): assert set(selected) | set(remaining) == set(topo.keys()) assert len(set(selected) & set(remaining)) == 0 # no overlap + def test_unknown_domain_nodes_excluded_from_selection(self): + # Regression: nodes with no NVLink-domain info (UNKNOWN, topo_rank -1) + # must never be selected. They sort first (rank -1), so a naive selection + # would pick them and emit a {"unknown": 0.001} placement-group constraint + # that ray.sub never registers as a resource -> unschedulable bundle. + # This happens on heterogeneous / partial-probe clusters (e.g. a GPU-less + # or unprobed head node). + topo = _make_topology( + { + NVLINK_DOMAIN_UNKNOWN: [TOPO_RANK_UNKNOWN, TOPO_RANK_UNKNOWN], + "domain_A": [0, 1], + "domain_B": [2, 3], + } + ) + selected, remaining = select_segment_nodes(topo, segment_size=2, num_nodes=4) + selected_domains = {topo[n][0] for n in selected} + assert NVLINK_DOMAIN_UNKNOWN not in selected_domains + assert selected_domains == {"domain_A", "domain_B"} + # The unknown nodes (e.g. a GPU-less head) fall through to remaining. + assert {topo[n][0] for n in remaining} == {NVLINK_DOMAIN_UNKNOWN} + assert set(selected) | set(remaining) == set(topo.keys()) + + def test_unknown_domain_only_raises_instead_of_unschedulable_pg(self): + # If every node is UNKNOWN there is no real domain to pin to, so we must + # raise rather than emit an unschedulable {"unknown": ...} constraint. + topo = _make_topology({NVLINK_DOMAIN_UNKNOWN: [TOPO_RANK_UNKNOWN] * 4}) + with pytest.raises(ResourceInsufficientError, match="Cannot form"): + select_segment_nodes(topo, segment_size=2, num_nodes=4) + # --------------------------------------------------------------------------- # 5. The 40-node / 5-domain scenario: training + inference placement