feat(placement): bond parallel RDMA links in JACCL device matrix

src/exo/master/placement_utils.py

@@ -296,15 +296,27 @@ def get_shard_assignments(
 def get_mlx_jaccl_devices_matrix(
     selected_cycle: list[NodeId],
     cycle_digraph: Topology,
-) -> list[list[str | None]]:
+) -> list[list[list[str] | None]]:
     """Build connectivity matrix mapping device i to device j via RDMA interface names.
 
-    The matrix element [i][j] contains the interface name on device i that connects
-    to device j, or None if no connection exists or no interface name is found.
-    Diagonal elements are always None.
+    The matrix element [i][j] contains the list of interface names on device i that
+    connect to device j, or None if no connection exists. Diagonal elements are always
+    None.
+
+    Multiple parallel RDMA links between the same pair are returned as a list so the
+    JACCL backend can bond them. The backend's ring path stripes across up to
+    RING_MAX_CONNS (=4) wires per direction; extra links beyond that are dropped here
+    to stay within the backend's limit. Note the mesh path only consumes the first
+    interface per pair, so bonding is effective for ring/tensor-parallel workloads.
+
+    The JACCL ring validator requires every ring edge to carry the same number of
+    connections, so the counts are normalised down to the minimum across all pairs to
+    avoid the backend rejecting an asymmetric matrix.
     """
     num_nodes = len(selected_cycle)
-    matrix: list[list[str | None]] = [
+    # First pass: collect every RDMA interface per pair (order is stable: insertion
+    # order from get_all_connections_between).
+    interfaces_per_pair: list[list[list[str] | None]] = [
         [None for _ in range(num_nodes)] for _ in range(num_nodes)
     ]
 
@@ -313,14 +325,35 @@ def get_mlx_jaccl_devices_matrix(
             if i == j:
                 continue
 
+            names: list[str] = []
             for conn in cycle_digraph.get_all_connections_between(node_i, node_j):
                 if isinstance(conn, RDMAConnection):
-                    matrix[i][j] = conn.source_rdma_iface
-                    break
-            else:
+                    names.append(conn.source_rdma_iface)
+
+            if not names:
                 raise ValueError(
                     "Current jaccl backend requires all-to-all RDMA connections"
                 )
+            interfaces_per_pair[i][j] = names
+
+    # Second pass: normalise to a uniform connection count across all pairs. The
+    # backend's ring validator rejects matrices where ring edges have differing
+    # numbers of connections, and bonding only helps when it's consistent, so trim
+    # every pair down to the smallest count observed anywhere in the matrix.
+    non_null_counts = [
+        len(names) for row in interfaces_per_pair for names in row if names is not None
+    ]
+    target_count = min(non_null_counts) if non_null_counts else 0
+
+    matrix: list[list[list[str] | None]] = [
+        [None for _ in range(num_nodes)] for _ in range(num_nodes)
+    ]
+    for i in range(num_nodes):
+        for j in range(num_nodes):
+            if i == j:
+                continue
+            assert interfaces_per_pair[i][j] is not None
+            matrix[i][j] = interfaces_per_pair[i][j][:target_count]  # type: ignore[slice]
 
     return matrix
 

src/exo/master/tests/test_placement.py

@@ -505,9 +505,9 @@ def test_tensor_rdma_backend_connectivity_matrix(
     idx_b = node_to_idx[node_b]
     idx_c = node_to_idx[node_c]
 
-    assert matrix[idx_a][idx_b] == "rdma_en3"
-    assert matrix[idx_b][idx_c] == "rdma_en4"
-    assert matrix[idx_c][idx_a] == "rdma_en5"
+    assert matrix[idx_a][idx_b] == ["rdma_en3"]
+    assert matrix[idx_b][idx_c] == ["rdma_en4"]
+    assert matrix[idx_c][idx_a] == ["rdma_en5"]
 
     # Verify coordinators are set for all nodes
     assert len(instance.jaccl_coordinators) == 3
@@ -1056,3 +1056,185 @@ def test_mlx_jaccl_rejects_cuda_only_cycle(model_card: ModelCard):
             node_backends,
             node_rdma_ctl=node_rdma_ctl,
         )
+
+
+def _build_two_node_bonded_rdma_topology(
+    links_per_pair: int = 2,
+) -> tuple[Topology, NodeId, NodeId, dict[NodeId, NodeNetworkInfo]]:
+    """Two-node topology with multiple parallel RDMA links between the pair.
+
+    Each direction gets ``links_per_pair`` RDMA edges using distinct interfaces
+    (rdma_en3, rdma_en4, ...). This mirrors two Macs connected by more than one
+    Thunderbolt cable.
+    """
+    topology = Topology()
+    node_a = NodeId()
+    node_b = NodeId()
+
+    ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1")
+    ethernet_conn = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
+    )
+    node_network = {
+        node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
+    }
+
+    for n in (node_a, node_b):
+        topology.add_node(n)
+
+    # Parallel RDMA links in both directions, each on its own interface.
+    for offset in range(links_per_pair):
+        iface = 3 + offset  # rdma_en3, rdma_en4, ...
+        topology.add_connection(
+            Connection(source=node_a, sink=node_b, edge=create_rdma_connection(iface))
+        )
+        topology.add_connection(
+            Connection(source=node_b, sink=node_a, edge=create_rdma_connection(iface))
+        )
+
+    topology.add_connection(Connection(source=node_a, sink=node_b, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_b, sink=node_a, edge=ethernet_conn))
+
+    return topology, node_a, node_b, node_network
+
+
+def test_jaccl_devices_matrix_bonds_parallel_rdma_links(model_card: ModelCard):
+    """Multiple Thunderbolt links between a pair should be emitted as a list so the
+    JACCL ring backend can stripe across them, rather than picking one and discarding
+    the rest."""
+    # arrange
+    model_card = model_card.model_copy(
+        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
+    )
+    topology, node_a, node_b, node_network = _build_two_node_bonded_rdma_topology(
+        links_per_pair=2
+    )
+    node_memory = {
+        node_a: create_node_memory(1000),
+        node_b: create_node_memory(1000),
+    }
+    node_rdma_ctl = {
+        node_a: NodeRdmaCtlStatus(enabled=True),
+        node_b: NodeRdmaCtlStatus(enabled=True),
+    }
+
+    cic = PlaceInstance(
+        sharding=Sharding.Tensor,
+        instance_meta=InstanceMeta.MlxJaccl,
+        command_id=CommandId(),
+        model_card=model_card,
+        min_nodes=1,
+    )
+
+    placements = place_instance(
+        cic,
+        topology,
+        {},
+        node_memory,
+        node_network,
+        _metal_only(node_memory),
+        node_rdma_ctl=node_rdma_ctl,
+    )
+
+    assert len(placements) == 1
+    instance = next(iter(placements.values()))
+    assert isinstance(instance, MlxJacclInstance)
+
+    matrix = instance.jaccl_devices
+    assert len(matrix) == 2
+    assert matrix[0][0] is None
+    assert matrix[1][1] is None
+    # Both parallel interfaces survive in the cell. Order is not guaranteed (the
+    # topology graph returns edges in reverse insertion order), so compare as sets.
+    assert set(matrix[0][1]) == {"rdma_en3", "rdma_en4"} and len(matrix[0][1]) == 2
+    assert set(matrix[1][0]) == {"rdma_en3", "rdma_en4"} and len(matrix[1][0]) == 2
+
+
+def test_jaccl_devices_matrix_normalises_asymmetric_link_counts(model_card: ModelCard):
+    """When one pair has more parallel links than another, the matrix is normalised
+    down to the minimum count so the JACCL ring validator (which requires every ring
+    edge to carry the same number of connections) accepts it."""
+    # arrange: three nodes where the A<->B pair has 2 links but B<->C and A<->C have 1
+    model_card = model_card.model_copy(
+        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
+    )
+    topology = Topology()
+    node_a, node_b, node_c = NodeId(), NodeId(), NodeId()
+
+    ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1")
+    ethernet_conn = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
+    )
+    node_network = {
+        n: NodeNetworkInfo(interfaces=[ethernet_interface])
+        for n in (node_a, node_b, node_c)
+    }
+    for n in (node_a, node_b, node_c):
+        topology.add_node(n)
+
+    # A<->B gets two parallel links; the other two pairs get one each.
+    for iface in (3, 4):
+        topology.add_connection(
+            Connection(source=node_a, sink=node_b, edge=create_rdma_connection(iface))
+        )
+        topology.add_connection(
+            Connection(source=node_b, sink=node_a, edge=create_rdma_connection(iface))
+        )
+    for src, sink, iface in [
+        (node_b, node_c, 5),
+        (node_c, node_b, 5),
+        (node_a, node_c, 6),
+        (node_c, node_a, 6),
+    ]:
+        topology.add_connection(
+            Connection(source=src, sink=sink, edge=create_rdma_connection(iface))
+        )
+    for src, sink in [
+        (node_a, node_b),
+        (node_b, node_a),
+        (node_b, node_c),
+        (node_c, node_b),
+        (node_a, node_c),
+        (node_c, node_a),
+    ]:
+        topology.add_connection(Connection(source=src, sink=sink, edge=ethernet_conn))
+
+    node_memory = {n: create_node_memory(1000) for n in (node_a, node_b, node_c)}
+    node_rdma_ctl = {
+        n: NodeRdmaCtlStatus(enabled=True) for n in (node_a, node_b, node_c)
+    }
+
+    cic = PlaceInstance(
+        sharding=Sharding.Tensor,
+        instance_meta=InstanceMeta.MlxJaccl,
+        command_id=CommandId(),
+        model_card=model_card,
+        min_nodes=1,
+    )
+
+    placements = place_instance(
+        cic,
+        topology,
+        {},
+        node_memory,
+        node_network,
+        _metal_only(node_memory),
+        node_rdma_ctl=node_rdma_ctl,
+    )
+
+    instance = next(iter(placements.values()))
+    assert isinstance(instance, MlxJacclInstance)
+    matrix = instance.jaccl_devices
+
+    # place_instance selects the smallest sufficient cycle, so the matrix may be
+    # 2x2 (one pair) rather than the full 3x3. Either way, every off-diagonal cell
+    # must be normalised to a uniform length (the ring validator requires this).
+    n = len(matrix)
+    off_diag_lengths = {len(matrix[i][j]) for i in range(n) for j in range(n) if i != j}
+    assert len(off_diag_lengths) == 1, (
+        f"ring edges must carry a uniform connection count, got {off_diag_lengths}"
+    )
+    # The only single-link pairs are count 1; the A<->B pair (2 physical links) must
+    # have been trimmed down to that minimum, not left at 2.
+    assert off_diag_lengths == {1}

src/exo/shared/types/worker/instances.py

@@ -31,7 +31,10 @@ class MlxRingInstance(BaseInstance):
 
 
 class MlxJacclInstance(BaseInstance):
-    jaccl_devices: list[list[str | None]]
+    # Each cell is the list of RDMA interface names from this device to another
+    # (multiple when parallel Thunderbolt links are bonded), or None on the
+    # diagonal. The JACCL backend accepts either a single name or a list per cell.
+    jaccl_devices: list[list[list[str] | None]]
     jaccl_coordinators: dict[NodeId, str]