|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Torch-free tests for the nccl-ep single-handle contract. |
| 3 | +
|
| 4 | +The invariant under test: ONE handle per group, rebound per problem shape. Two handles built |
| 5 | +from the same group config resolve to the same LL parity signal slots while advancing their |
| 6 | +parity independently, which corrupts signalling (NVIDIA/nccl#2303) -- so a regression that |
| 7 | +reintroduces per-shape handles must fail loudly here rather than on a cluster. |
| 8 | +
|
| 9 | +torch and nccl are stubbed so this runs without the benchmark image. |
| 10 | +""" |
| 11 | +from __future__ import annotations |
| 12 | + |
| 13 | +import sys |
| 14 | +import types |
| 15 | +import unittest |
| 16 | +from pathlib import Path |
| 17 | +from unittest import mock |
| 18 | + |
| 19 | +ROOT = Path(__file__).resolve().parents[1] |
| 20 | + |
| 21 | + |
| 22 | +def _stub_modules(): |
| 23 | + """Fake torch / nccl modules so `import ep_nccl` succeeds without the benchmark image.""" |
| 24 | + torch = types.ModuleType("torch") |
| 25 | + torch.bfloat16 = "bfloat16" |
| 26 | + torch.int32 = "int32" |
| 27 | + torch.empty = lambda *a, **k: types.SimpleNamespace(shape=a[0] if a else ()) |
| 28 | + torch.zeros = lambda *a, **k: types.SimpleNamespace(item=lambda: 7) |
| 29 | + torch.cuda = types.SimpleNamespace(synchronize=lambda: None) |
| 30 | + dist = types.ModuleType("torch.distributed") |
| 31 | + torch.distributed = dist |
| 32 | + |
| 33 | + ep = types.ModuleType("nccl.ep") |
| 34 | + for name in ( |
| 35 | + "Algorithm", "CombineConfig", "CombineInputs", "CombineOutputs", "DispatchConfig", |
| 36 | + "DispatchInputs", "DispatchOutputs", "GroupConfig", "HandleConfig", "Layout", |
| 37 | + "LayoutInfo", "Tensor", |
| 38 | + ): |
| 39 | + setattr(ep, name, type(name, (), {"__init__": lambda self, *a, **k: None})) |
| 40 | + ep.Algorithm = types.SimpleNamespace(LOW_LATENCY="LL", HIGH_THROUGHPUT="HT") |
| 41 | + ep.Layout = types.SimpleNamespace(EXPERT_MAJOR="EM", FLAT="FLAT") |
| 42 | + core = types.ModuleType("nccl.core") |
| 43 | + pkg = types.ModuleType("nccl") |
| 44 | + pkg.ep, pkg.core = ep, core |
| 45 | + return { |
| 46 | + "torch": torch, "torch.distributed": dist, |
| 47 | + "nccl": pkg, "nccl.ep": ep, "nccl.core": core, |
| 48 | + } |
| 49 | + |
| 50 | + |
| 51 | +sys.path[:0] = [str(ROOT), str(ROOT / "bench")] |
| 52 | + |
| 53 | +# Import ep_nccl against the stubs, then withdraw them: leaving a fake torch in sys.modules |
| 54 | +# makes the genuinely torch-dependent modules in this process (test_runtime, test_ll_oracle) |
| 55 | +# error instead of skipping. Dropping ep_nccl too keeps the stub-built module private to us. |
| 56 | +with mock.patch.dict(sys.modules, _stub_modules()): |
| 57 | + import ep_nccl # noqa: E402 |
| 58 | + |
| 59 | + sys.modules.pop("ep_nccl", None) |
| 60 | + |
| 61 | + |
| 62 | +class FakeHandle: |
| 63 | + """Records every rebind so the tests can assert on the collective call pattern.""" |
| 64 | + |
| 65 | + def __init__(self): |
| 66 | + self.updates = [] |
| 67 | + self.destroyed = False |
| 68 | + |
| 69 | + def update(self, topk_idx, *, layout_info=None, stream=None): |
| 70 | + self.updates.append((topk_idx, layout_info)) |
| 71 | + |
| 72 | + def destroy(self): |
| 73 | + self.destroyed = True |
| 74 | + |
| 75 | + |
| 76 | +class FakeGroup: |
| 77 | + def __init__(self): |
| 78 | + self.created = 0 |
| 79 | + self.handle = FakeHandle() |
| 80 | + |
| 81 | + def create_handle(self, layout, topk_idx, *, layout_info=None, config=None, stream=None): |
| 82 | + self.created += 1 |
| 83 | + return self.handle |
| 84 | + |
| 85 | + |
| 86 | +def backend(ll=True): |
| 87 | + """An NCCLEPBackend with just the fields _ensure_handle touches (no __init__, no GPU).""" |
| 88 | + b = object.__new__(ep_nccl.NCCLEPBackend) |
| 89 | + b._ll = ll |
| 90 | + b._layout = "EM" if ll else "FLAT" |
| 91 | + b._handle = None |
| 92 | + b._bound = None |
| 93 | + b._ep_group = FakeGroup() |
| 94 | + b.device = "cuda:0" |
| 95 | + b.num_local_experts = 4 |
| 96 | + b.args = types.SimpleNamespace(hidden=16) |
| 97 | + b._t = lambda x: x |
| 98 | + b._stream = lambda: 0 |
| 99 | + return b |
| 100 | + |
| 101 | + |
| 102 | +def problem(T): |
| 103 | + return types.SimpleNamespace( |
| 104 | + T=T, dispatch_x=f"x{T}", topk_idx=f"idx{T}", topk_weights=f"w{T}" |
| 105 | + ) |
| 106 | + |
| 107 | + |
| 108 | +class TestSingleHandle(unittest.TestCase): |
| 109 | + def test_one_handle_across_many_shapes(self): |
| 110 | + """Nine ladder rungs must still produce exactly one create_handle.""" |
| 111 | + b = backend() |
| 112 | + for T in (1, 2, 4, 8, 16, 32, 64, 128, 256): |
| 113 | + b._ensure_handle(problem(T)) |
| 114 | + self.assertEqual(b._ep_group.created, 1) |
| 115 | + |
| 116 | + def test_shape_change_rebinds_and_repeat_does_not(self): |
| 117 | + """update() on a shape switch; no collective when the bound shape is re-entered.""" |
| 118 | + b = backend() |
| 119 | + pa, pb = problem(1), problem(2) |
| 120 | + b._ensure_handle(pa) |
| 121 | + self.assertEqual(len(b._ep_group.handle.updates), 0) # first bind is the create |
| 122 | + |
| 123 | + b._ensure_handle(pb) |
| 124 | + self.assertEqual(len(b._ep_group.handle.updates), 1) |
| 125 | + |
| 126 | + # Re-entering the bound problem repeatedly -- the timed loop's steady state -- must not |
| 127 | + # enter a collective, otherwise every iteration gains a rank-synchronising step. |
| 128 | + for _ in range(8): |
| 129 | + b._ensure_handle(pb) |
| 130 | + self.assertEqual(len(b._ep_group.handle.updates), 1) |
| 131 | + |
| 132 | + # Returning to an earlier shape rebinds again (its cached namespace is reused). |
| 133 | + b._ensure_handle(pa) |
| 134 | + self.assertEqual(len(b._ep_group.handle.updates), 2) |
| 135 | + |
| 136 | + def test_every_problem_shares_the_one_handle(self): |
| 137 | + b = backend() |
| 138 | + handles = {id(b._ensure_handle(problem(T)).handle) for T in (1, 2, 4)} |
| 139 | + self.assertEqual(len(handles), 1) |
| 140 | + self.assertIs(b._ensure_handle(problem(1)).handle, b._handle) |
| 141 | + |
| 142 | + def test_ll_never_passes_layout_info_on_rebind(self): |
| 143 | + """The API forbids layout_info on create/update in LL mode.""" |
| 144 | + b = backend(ll=True) |
| 145 | + b._ensure_handle(problem(1)) |
| 146 | + b._ensure_handle(problem(2)) |
| 147 | + self.assertEqual([info for _, info in b._ep_group.handle.updates], [None]) |
| 148 | + |
| 149 | + def test_ht_rebind_carries_that_problems_counters(self): |
| 150 | + """HT re-runs the metadata exchange into the rebound problem's own counter tensors.""" |
| 151 | + b = backend(ll=False) |
| 152 | + ha = b._ensure_handle(problem(1)) |
| 153 | + hb = b._ensure_handle(problem(2)) |
| 154 | + self.assertEqual(len(b._ep_group.handle.updates), 1) |
| 155 | + self.assertIs(b._ep_group.handle.updates[0][1], hb.layout_info) |
| 156 | + self.assertIsNot(ha.layout_info, hb.layout_info) |
| 157 | + self.assertEqual(hb.count, 7) # re-read after the exchange |
| 158 | + |
| 159 | + def test_destroy_releases_the_handle_once(self): |
| 160 | + b = backend() |
| 161 | + b._ensure_handle(problem(1)) |
| 162 | + handle = b._handle |
| 163 | + b._destroy_handles() |
| 164 | + self.assertTrue(handle.destroyed) |
| 165 | + self.assertIsNone(b._handle) |
| 166 | + self.assertIsNone(b._bound) |
| 167 | + b._destroy_handles() # idempotent |
| 168 | + |
| 169 | + |
| 170 | +if __name__ == "__main__": |
| 171 | + unittest.main() |
0 commit comments