diff --git a/tests/integration/ops/test_rng_dispatch.py b/tests/integration/ops/test_rng_dispatch.py index 95317140..dd7d27e6 100644 --- a/tests/integration/ops/test_rng_dispatch.py +++ b/tests/integration/ops/test_rng_dispatch.py @@ -219,6 +219,23 @@ def test_default_generators_populated(self): "default generator state is not philox-shaped (2 x int64)" ) + @pytest.mark.anyplatform + @pytest.mark.main_ops + @pytest.mark.parametrize("attr", ["cuda", "flagos"]) + def test_default_generators_iterable(self, attr): + # Upstream types default_generators as a *tuple*, so callers iterate it, + # slice it and list() it. Our shims are list-like proxies with no + # __iter__, which sends Python to the legacy protocol: __getitem__(0, 1, + # 2, ...) until IndexError. Unbounded __getitem__ therefore made + # `for g in default_generators` an infinite loop that allocated a fresh + # generator per step -- a hang, not an error, so nothing surfaced it. + gens = getattr(torch if attr == "cuda" else torch_fl, attr).default_generators + n = len(gens) + assert len(list(gens)) == n, "iteration does not stop at device_count" + assert len(gens[:2]) == min(2, n), "slicing is not supported" + with pytest.raises(IndexError): + gens[n] + @pytest.mark.anyplatform @pytest.mark.main_ops def test_manual_seed_reaches_flagos_module(self): diff --git a/torch_fl/accelerator/cuda/_cuda_compat.py b/torch_fl/accelerator/cuda/_cuda_compat.py index 81124d77..e196d592 100644 --- a/torch_fl/accelerator/cuda/_cuda_compat.py +++ b/torch_fl/accelerator/cuda/_cuda_compat.py @@ -85,10 +85,28 @@ class _CudaDefaultGenerators: Indexing yields a real (lazily created) per-device CUDA generator; ``len`` reports the device count so flag_gems' ``len(default_generators) == 0`` guard is False and it uses the generator instead of erroring. + + Upstream declares ``default_generators`` as a *tuple*, so callers are + entitled to iterate it, slice it or wrap it in ``list()``. Bounds-checking + ``__getitem__`` is what makes that safe: with no ``__iter__``, Python falls + back to the legacy protocol of calling ``__getitem__(0, 1, 2, ...)`` until + IndexError, so an unchecked index turned ``for g in default_generators`` + into an infinite loop that allocated a fresh CUDA generator per step. """ + def __iter__(self): + return (self[i] for i in range(len(self))) + def __getitem__(self, idx): - return _get_cuda_generator(int(idx)) + n = len(self) + if isinstance(idx, slice): + return tuple(self[i] for i in range(*idx.indices(n))) + idx = int(idx) + if idx < 0: # negative indices wrap, as on the tuple this replaces + idx += n + if not 0 <= idx < n: + raise IndexError(f"device index {idx} out of range for {n} device(s)") + return _get_cuda_generator(idx) def __len__(self): try: diff --git a/torch_fl/accelerator/metax/_metax_compat.py b/torch_fl/accelerator/metax/_metax_compat.py index eb7e0369..b1e3e45f 100644 --- a/torch_fl/accelerator/metax/_metax_compat.py +++ b/torch_fl/accelerator/metax/_metax_compat.py @@ -182,10 +182,31 @@ def _get_cuda_generator(idx): class _CudaDefaultGenerators: """list-like stand-in for ``torch.cuda.default_generators`` (see _get_cuda_generator). Indexing yields a per-device CUDA generator; ``len`` - reports the device count so flag_gems' empty-tuple guard is False.""" + reports the device count so flag_gems' empty-tuple guard is False. + + Upstream declares ``default_generators`` as a *tuple*, so callers are + entitled to iterate it, slice it or wrap it in ``list()``. Bounds-checking + ``__getitem__`` is what makes that safe: with no ``__iter__``, Python falls + back to the legacy protocol of calling ``__getitem__(0, 1, 2, ...)`` until + IndexError, so an unchecked index turned ``for g in default_generators`` + into an infinite loop that allocated a fresh CUDA generator per step. + """ + + def __iter__(self): + return (self[i] for i in range(len(self))) def __getitem__(self, idx): - return _get_cuda_generator(int(idx)) + n = len(self) + if isinstance(idx, slice): + return tuple(self[i] for i in range(*idx.indices(n))) + idx = int(idx) + if idx < 0: # negative indices wrap, as on the tuple this replaces + idx += n + if not 0 <= idx < n: + raise IndexError( + f"device index {idx} out of range for {n} flagos device(s)" + ) + return _get_cuda_generator(idx) def __len__(self): try: diff --git a/torch_fl/flagos/__init__.py b/torch_fl/flagos/__init__.py index 74810d4b..77a68896 100644 --- a/torch_fl/flagos/__init__.py +++ b/torch_fl/flagos/__init__.py @@ -146,9 +146,28 @@ def _patched_getitem(self, indices): # default_generators: list of one Generator per device, required by FlagGems class _DefaultGenerators: - """Lazy list-like accessor for per-device default generators.""" + """Lazy list-like accessor for per-device default generators. + + Bounds-checked so iteration terminates: with no ``__iter__``, Python's + legacy protocol calls ``__getitem__(0, 1, 2, ...)`` until IndexError, and + an out-of-range index otherwise surfaces as a RuntimeError from C++ (which + aborts iteration instead of ending it). + """ + + def __iter__(self): + return (self[i] for i in range(len(self))) def __getitem__(self, device): + n = len(self) + if isinstance(device, slice): + return tuple(self[i] for i in range(*device.indices(n))) + device = int(device) + if device < 0: # negative indices wrap, as on a list + device += n + if not 0 <= device < n: + raise IndexError( + f"device index {device} out of range for {n} flagos device(s)" + ) return _C._get_default_generator(device) def __len__(self):