Improve distributed dataloader handling

@@ -1482,8 +1482,9 @@ The :class:`~nemo.collections.common.data.lhotse.broadcasting.BroadcastingDataLo fixes this at the data layer: construct the real Lhotse loader on a single DP-source rank (``cp_rank == 0`` and ``tp_rank == 0``) and let the wrapper broadcast each batch to the other ranks in the ``(cp, tp)``-sub-mesh over NCCL. Iteration ends in lockstep via a continue/stop-broadcast — no length needs to be known up-front.+sub-mesh over NCCL. Each step broadcasts one framed message containing a+batch, clean-exhaustion signal, or source-side error, so all ranks finish+or fail in lockstep without requiring the loader length up-front.  .. code-block:: python 

@@ -40,13 +40,15 @@ ``torchdata.StatefulDataLoader``, or any other source object that implements those methods. -Iteration termination is handled with two broadcasts per step: a-continue/stop boolean followed by the batch. This works regardless of+Each iteration broadcasts one framed message containing a data batch, a stop+signal, or details about a source-side error. This works regardless of whether the source loader exposes ``__len__`` (Lhotse training loaders-typically don't).+typically don't) and prevents receivers from waiting indefinitely when the+source loader fails. """ from __future__ import annotations +import pickle from typing import Any, Iterable, Iterator, Sequence  import torch@@ -102,9 +104,8 @@ def broadcast_batch(     if resolved is None:         return batch     group, src = resolved-    obj_list = [batch]-    dist.broadcast_object_list(obj_list, src=src, group=group, device=_broadcast_device(group))-    return obj_list[0]+    packet = (_PACKET_DATA, batch) if dist.get_rank() == src else None+    return _packet_payload(_broadcast_packet(packet, group, src))   class BroadcastingDataLoader:@@ -114,10 +115,10 @@ class BroadcastingDataLoader:      Pass ``source=real_loader`` on the DP source rank (``cp_rank == 0`` and     ``tp_rank == 0``); pass ``source=None`` on every other rank. Iteration-    issues two broadcasts per step on every rank: a continue/stop boolean-    followed by the batch. After the source loader is exhausted, the-    continue broadcast is False and iteration ends in lockstep on all-    ranks regardless of whether the source exposes ``__len__``.+    issues one framed broadcast per step on every rank. The frame contains a+    data batch, clean-exhaustion signal, or source-side error. This lets all+    ranks finish or fail in lockstep regardless of whether the source exposes+    ``__len__``.      ``state_dict`` / ``load_state_dict`` are delegated to the source on the     source rank (no-ops on non-source ranks), so checkpoint/resume keeps@@ -138,6 +139,7 @@ def __init__(         self._source = source         self._mesh = device_mesh         self._axes = axes+        self._group_and_source = None         if not _is_noop(device_mesh, axes):             self._is_source = is_dp_source_rank(device_mesh, axes)             if self._is_source and source is None:@@ -149,19 +151,45 @@ def __iter__(self) -> Iterator[Any]:                 return             yield from self._source             return+        if not (dist.is_available() and dist.is_initialized()):+            if self._source is None:+                return+            yield from self._source+            return++        if self._group_and_source is None:+            self._group_and_source = _resolve_group_and_source(self._mesh, self._axes)+        if self._group_and_source is None:+            if self._source is None:+                return+            yield from self._source+            return++        group, src = self._group_and_source         if self._is_source:-            for batch in self._source:-                broadcast_batch(True, self._mesh, self._axes)-                broadcast_batch(batch, self._mesh, self._axes)+            try:+                source_iterator = iter(self._source)+            except Exception as error:+                _broadcast_packet(_error_packet(error, "source iterator creation failed"), group, src)+                raise+            while True:+                try:+                    batch = next(source_iterator)+                except StopIteration:+                    _broadcast_packet((_PACKET_STOP, None), group, src)+                    return+                except Exception as error:+                    _broadcast_packet(_error_packet(error, "source iterator failed"), group, src)+                    raise++                _broadcast_packet((_PACKET_DATA, batch), group, src)                 yield batch-            broadcast_batch(False, self._mesh, self._axes)         else:             while True:-                keep_iterating = broadcast_batch(None, self._mesh, self._axes)-                if not keep_iterating:+                packet = _broadcast_packet(None, group, src)+                if packet[0] == _PACKET_STOP:                     return-                batch = broadcast_batch(None, self._mesh, self._axes)-                yield batch+                yield _packet_payload(packet)      def __len__(self) -> int:         # Pass-through when the source defines __len__; raise TypeError@@ -186,10 +214,9 @@ def load_state_dict(self, state_dict) -> None: # ---------------------------------------------------------------------------  -# Cache: (id(device_mesh), tuple_of_axes) -> (process_group, source_global_rank).-# Sub-mesh creation calls ``_flatten`` which materializes a process group;-# we don't want to repeat that per training step.-_GROUP_CACHE: dict[tuple[int, tuple[str, ...]], tuple[Any, int]] = {}+_PACKET_DATA = "data"+_PACKET_STOP = "stop"+_PACKET_ERROR = "error"   def _present_axes(device_mesh, axes: Sequence[str]) -> tuple[str, ...]:@@ -219,10 +246,6 @@ def _resolve_group_and_source(device_mesh, axes: Sequence[str]):     if _is_noop(device_mesh, axes):         return None     present = _present_axes(device_mesh, axes)-    cache_key = (id(device_mesh), present)-    cached = _GROUP_CACHE.get(cache_key)-    if cached is not None:-        return cached      if len(present) == 1:         sub = device_mesh[present[0]]@@ -231,5 +254,67 @@ def _resolve_group_and_source(device_mesh, axes: Sequence[str]):      group = sub.get_group()     source_global_rank = int(sub.mesh.flatten()[0].item())-    _GROUP_CACHE[cache_key] = (group, source_global_rank)     return group, source_global_rank+++def _broadcast_packet(packet, group, src: int):+    """Broadcast one pre-serialized protocol packet over ``group``.++    Serializing before the first tensor collective lets the source replace an+    unpicklable data packet with a small error packet while receivers are still+    waiting for the packet size. Failures after a collective begins must be+    handled by the process-group timeout and distributed job teardown.+    """+    is_source = dist.get_rank() == src+    source_error = None+    if is_source:+        try:+            serialized = pickle.dumps(packet, protocol=pickle.HIGHEST_PROTOCOL)+        except Exception as error:+            source_error = error+            packet = _error_packet(error, "batch serialization failed")+            serialized = pickle.dumps(packet, protocol=pickle.HIGHEST_PROTOCOL)+    else:+        serialized = b""++    device = _broadcast_device(group)+    serialized_size = torch.tensor([len(serialized)], dtype=torch.long, device=device)+    dist.broadcast(serialized_size, src=src, group=group)++    if is_source:+        serialized_buffer = bytearray(serialized)+        serialized_tensor = torch.frombuffer(serialized_buffer, dtype=torch.uint8).to(device)+    else:+        serialized_tensor = torch.empty(int(serialized_size.item()), dtype=torch.uint8, device=device)+    dist.broadcast(serialized_tensor, src=src, group=group)++    if source_error is not None:+        raise source_error+    if not is_source:+        packet = pickle.loads(serialized_tensor.cpu().numpy().tobytes())+    _validate_packet(packet)+    return packet+++def _packet_payload(packet):+    kind, payload = packet+    if kind == _PACKET_ERROR:+        raise RuntimeError(f"BroadcastingDataLoader source error: {payload}")+    if kind != _PACKET_DATA:+        raise RuntimeError(f"Expected a data packet, received {kind!r}")+    return payload+++def _error_packet(error: Exception, context: str):+    try:+        error_text = f"{type(error).__module__}.{type(error).__qualname__}: {error}"+    except Exception:+        error_text = f"{type(error).__module__}.{type(error).__qualname__}"+    return _PACKET_ERROR, f"{context}: {error_text}"+++def _validate_packet(packet) -> None:+    if not isinstance(packet, tuple) or len(packet) != 2:+        raise RuntimeError(f"Received an invalid broadcast packet: {type(packet).__name__}")+    if packet[0] not in (_PACKET_DATA, _PACKET_STOP, _PACKET_ERROR):+        raise RuntimeError(f"Received an unknown broadcast packet kind: {packet[0]!r}")

@@ -30,6 +30,7 @@ import torch.distributed as dist import torch.multiprocessing as mp +from nemo.collections.common.data.lhotse import broadcasting from nemo.collections.common.data.lhotse.broadcasting import BroadcastingDataLoader, broadcast_batch, is_dp_source_rank  # ---------------------------------------------------------------------------@@ -38,30 +39,42 @@   class _FakeAxis:-    def __init__(self, size: int, local_rank: int):+    def __init__(self, size: int, local_rank: int, group=None, ranks=None):         self._size = size         self._local_rank = local_rank+        self._group = group+        self.mesh = torch.tensor(range(size) if ranks is None else ranks)      def size(self) -> int:         return self._size      def get_local_rank(self) -> int:         return self._local_rank +    def get_group(self):+        return self._group+  class _FakeMesh:     """Minimal DeviceMesh stand-in covering ``mesh_dim_names`` + ``__getitem__``.""" -    def __init__(self, sizes: dict[str, int], coords: dict[str, int]):+    def __init__(self, sizes: dict[str, int], coords: dict[str, int], groups=None, ranks=None):         assert sizes.keys() == coords.keys()         self.mesh_dim_names = tuple(sizes.keys())         self._sizes = sizes         self._coords = coords+        self._groups = groups or {}+        self._ranks = ranks or {}      def __getitem__(self, name):         if isinstance(name, tuple):             raise NotImplementedError("multi-axis slicing not needed for fake-mesh tests")-        return _FakeAxis(self._sizes[name], self._coords[name])+        return _FakeAxis(+            self._sizes[name],+            self._coords[name],+            group=self._groups.get(name),+            ranks=self._ranks.get(name),+        )   def test_is_dp_source_rank_none_mesh():@@ -109,6 +122,27 @@ def test_broadcast_batch_noop_when_axes_size_one():     assert broadcast_batch(payload, mesh) is payload  +def test_group_resolution_does_not_cache_by_object_id(monkeypatch):+    group_a = object()+    group_b = object()+    mesh_a = _FakeMesh(+        {"cp": 2},+        {"cp": 0},+        groups={"cp": group_a},+        ranks={"cp": [2, 3]},+    )+    mesh_b = _FakeMesh(+        {"cp": 2},+        {"cp": 0},+        groups={"cp": group_b},+        ranks={"cp": [6, 7]},+    )+    monkeypatch.setattr(broadcasting, "id", lambda _: 1, raising=False)++    assert broadcasting._resolve_group_and_source(mesh_a, ("cp",)) == (group_a, 2)+    assert broadcasting._resolve_group_and_source(mesh_b, ("cp",)) == (group_b, 6)++ def test_broadcasting_dataloader_noop_iterates_source():     real = [{"i": i} for i in range(4)]     loader = BroadcastingDataLoader(source=real, device_mesh=None)@@ -221,21 +255,91 @@ def _broadcasting_loader_worker(rank: int, world_size: int, port: int, queue: mp             dist.destroy_process_group()  -def _spawn_workers(target, world_size: int) -> list:+class _RaisesAfterOneBatch:+    def __iter__(self):+        yield {"i": 0}+        raise ValueError("source iterator exploded")+++class _RaisesOnIteratorCreation:+    def __iter__(self):+        raise ValueError("iterator creation failed")+++class _UnpicklableBatch:+    def __reduce__(self):+        raise TypeError("batch is intentionally unpicklable")+++def _source_iterator_error_worker(+    rank: int, world_size: int, port: int, queue: mp.Queue, source_type=_RaisesAfterOneBatch+) -> None:+    try:+        _init_gloo(rank, world_size, port)+        mesh = _build_cp_mesh(world_size)+        source = source_type() if is_dp_source_rank(mesh) else None+        loader = BroadcastingDataLoader(source=source, device_mesh=mesh)+        received = []+        try:+            received.extend(batch["i"] for batch in loader)+        except Exception as error:+            queue.put(("err", rank, received, type(error).__name__, str(error)))+        else:+            queue.put(("unexpected-success", rank, received, None, None))+    finally:+        if dist.is_initialized():+            dist.destroy_process_group()+++def _serialization_error_worker(rank: int, world_size: int, port: int, queue: mp.Queue) -> None:+    try:+        _init_gloo(rank, world_size, port)+        mesh = _build_cp_mesh(world_size)+        source = [_UnpicklableBatch()] if is_dp_source_rank(mesh) else None+        loader = BroadcastingDataLoader(source=source, device_mesh=mesh)+        try:+            list(loader)+        except Exception as error:+            queue.put(("err", rank, type(error).__name__, str(error)))+        else:+            queue.put(("unexpected-success", rank, None, None))+    finally:+        if dist.is_initialized():+            dist.destroy_process_group()+++def _early_break_worker(rank: int, world_size: int, port: int, queue: mp.Queue) -> None:+    try:+        _init_gloo(rank, world_size, port)+        mesh = _build_cp_mesh(world_size)+        source = [{"i": i} for i in range(3)] if is_dp_source_rank(mesh) else None+        loader = BroadcastingDataLoader(source=source, device_mesh=mesh)+        received = []+        for batch in loader:+            received.append(batch["i"])+            break+        queue.put(("ok", rank, received))+    finally:+        if dist.is_initialized():+            dist.destroy_process_group()+++def _spawn_workers(target, world_size: int, timeout: int = 120) -> list:     ctx = mp.get_context("spawn")     queue = ctx.Queue()     port = _get_free_port()     procs = [ctx.Process(target=target, args=(rank, world_size, port, queue)) for rank in range(world_size)]     for p in procs:         p.start()     for p in procs:-        p.join(timeout=120)+        p.join(timeout=timeout)     results = []     while not queue.empty():         results.append(queue.get())     for p in procs:-        if p.exitcode != 0 and p.is_alive():+        if p.is_alive():             p.terminate()+            p.join(timeout=5)     return results  @@ -254,3 +358,63 @@ def test_broadcasting_dataloader_iterates_in_lockstep_across_ranks():     for status, received in results:         assert status == "ok", results         assert received == [0, 1, 2]+++def test_broadcasting_dataloader_propagates_source_iterator_error():+    results = _spawn_workers(_source_iterator_error_worker, world_size=2, timeout=15)+    assert len(results) == 2, results+    results_by_rank = {result[1]: result for result in results}+    assert results_by_rank[0][:4] == ("err", 0, [0], "ValueError")+    assert "source iterator exploded" in results_by_rank[0][4]+    assert results_by_rank[1][:4] == ("err", 1, [0], "RuntimeError")+    assert "source iterator failed" in results_by_rank[1][4]+    assert "source iterator exploded" in results_by_rank[1][4]+++def test_broadcasting_dataloader_propagates_batch_serialization_error():+    results = _spawn_workers(_serialization_error_worker, world_size=2, timeout=15)+    assert len(results) == 2, results+    results_by_rank = {result[1]: result for result in results}+    assert results_by_rank[0][:3] == ("err", 0, "TypeError")+    assert "batch is intentionally unpicklable" in results_by_rank[0][3]+    assert results_by_rank[1][:3] == ("err", 1, "RuntimeError")+    assert "batch serialization failed" in results_by_rank[1][3]+    assert "batch is intentionally unpicklable" in results_by_rank[1][3]+++def test_broadcasting_dataloader_allows_synchronized_early_break():+    results = _spawn_workers(_early_break_worker, world_size=2, timeout=15)+    assert len(results) == 2, results+    for status, _, received in results:+        assert status == "ok", results+        assert received == [0]+++def test_broadcasting_dataloader_reports_iterator_creation_failure(monkeypatch):+    class FailingSource:+        def __iter__(self):+            raise ValueError('iterator creation failed')++    mesh = _FakeMesh({'cp': 2}, {'cp': 0}, groups={'cp': object()}, ranks={'cp': [0, 1]})+    monkeypatch.setattr(dist, 'is_initialized', lambda: True)+    packets = []+    monkeypatch.setattr(broadcasting, '_broadcast_packet', lambda packet, group, src: packets.append(packet))+    with pytest.raises(ValueError, match='iterator creation failed'):+        list(BroadcastingDataLoader(FailingSource(), mesh))+    assert len(packets) == 1+    assert packets[0][0] == 'error'+    assert 'iterator creation failed' in packets[0][1]+++def _source_iterator_creation_error_worker(rank, world_size, port, queue):+    _source_iterator_error_worker(rank, world_size, port, queue, source_type=_RaisesOnIteratorCreation)+++def test_broadcasting_dataloader_propagates_iterator_creation_error():+    results = _spawn_workers(_source_iterator_creation_error_worker, world_size=2, timeout=15)+    assert len(results) == 2, results+    results_by_rank = {result[1]: result for result in results}+    assert results_by_rank[0][:4] == ("err", 0, [], "ValueError")+    assert "iterator creation failed" in results_by_rank[0][4]+    assert results_by_rank[1][:4] == ("err", 1, [], "RuntimeError")+    assert "source iterator creation failed" in results_by_rank[1][4]