Record every cabled pair in a jaccl-ring hostfile

python/mlx/_distributed_utils/config.py

@@ -483,26 +483,47 @@ def configure_jaccl(args, hosts, ips, sshinfo):
     save_hostfile(args, hostfile)
 
 
+def jaccl_ring_devices(ring, count, ips):
+    """The device matrix for a jaccl-ring hostfile, in ring order.
+
+    Every cabled pair is recorded, not only the ring neighbours. The ring
+    neighbours are what the data plane uses, and `MLX_JACCL_RING` is what
+    selects the ring, so the extra entries change nothing about how the group
+    runs. They are the links a subgroup would need: a child can only contain
+    members that are directly connected, and a pair dropped here cannot be
+    recovered later because the hostfile is all the runtime ever sees.
+
+    `ips` is keyed by every physically connected pair, so this is a matter of
+    writing down what has already been discovered.
+    """
+    n = len(ring)
+    matrix = []
+    for i, node in enumerate(ring):
+        peers = {ring[i - 1], ring[(i + 1) % n]}
+        rdmas = []
+        for other in ring:
+            devices = [] if node == other else ips.get((node, other), [])
+            if not devices:
+                rdmas.append(None)
+                continue
+            # A ring neighbour keeps the width the ring was built with. Any
+            # other pair reports the cables it actually has.
+            width = count if other in peers else len(devices)
+            rdma = [f"rdma_{devices[c][0]}" for c in range(min(width, len(devices)))]
+            rdmas.append(rdma[0] if len(rdma) == 1 else rdma)
+        matrix.append(rdmas)
+    return matrix
+
+
 def configure_jaccl_ring(args, hosts, ips, ring, sshinfo):
     log(args.verbose, "Prepare a jaccl-ring hostfile")
     add_ips(hosts, args.verbose)
 
     jaccl_hosts = []
-    num_nodes = len(hosts)
     ring, count = ring
-    for i, node in enumerate(ring):
+    matrix = jaccl_ring_devices(ring, count, ips.ips)
+    for i, (node, rdmas) in enumerate(zip(ring, matrix)):
         h = hosts[node]
-        peer_left = ring[i - 1]
-        peer_right = ring[(i + 1) % num_nodes]
-        rdmas = []
-        for other in ring:
-            if other not in (peer_left, peer_right):
-                rdmas.append(None)
-            else:
-                rdma = []
-                for c in range(count):
-                    rdma.append(f"rdma_{ips.ips[node, other][c][0]}")
-                rdmas.append(rdma[0] if count == 1 else rdma)
         jaccl_hosts.append(Host(i, h.ssh_hostname, h.ips, rdmas))
     hostfile = Hostfile(jaccl_hosts, "jaccl-ring", args.env)
 

python/tests/test_distributed_config.py

@@ -0,0 +1,72 @@
+# Copyright © 2026 Apple Inc.
+
+import unittest
+
+from mlx._distributed_utils.config import jaccl_ring_devices
+
+
+def links(*pairs):
+    """A device map keyed the way IPConfigurator keys it, one cable per pair."""
+    ips = {}
+    for a, b in pairs:
+        ips.setdefault((a, b), []).append((f"en{a}{b}", "192.168.0.1"))
+        ips.setdefault((b, a), []).append((f"en{b}{a}", "192.168.0.2"))
+    return ips
+
+
+class TestJacclRingDevices(unittest.TestCase):
+    def test_ring_only(self):
+        # Four nodes cabled as a ring and nothing else. Non-adjacent pairs have
+        # no cable, so they stay null and the hostfile is what it always was.
+        m = jaccl_ring_devices([0, 1, 2, 3], 1, links((0, 1), (1, 2), (2, 3), (3, 0)))
+        self.assertEqual(
+            m,
+            [
+                [None, "rdma_en01", None, "rdma_en03"],
+                ["rdma_en10", None, "rdma_en12", None],
+                [None, "rdma_en21", None, "rdma_en23"],
+                ["rdma_en30", None, "rdma_en32", None],
+            ],
+        )
+
+    def test_extra_cable_is_recorded(self):
+        # The same ring with one chord. The chord is not part of the data plane
+        # and is still written down: a subgroup can only hold members that are
+        # directly connected, and a link dropped here cannot be recovered,
+        # because the hostfile is everything the runtime is given.
+        m = jaccl_ring_devices(
+            [0, 1, 2, 3], 1, links((0, 1), (1, 2), (2, 3), (3, 0), (0, 2))
+        )
+        self.assertEqual(m[0][2], "rdma_en02")
+        self.assertEqual(m[2][0], "rdma_en20")
+
+    def test_shape_contract(self):
+        # launch_jaccl rejects a hostfile that is not square with a null
+        # diagonal, so both hold whatever the cabling looks like.
+        for extra in ([], [(0, 2)], [(0, 2), (1, 3)]):
+            m = jaccl_ring_devices(
+                [0, 1, 2, 3], 1, links((0, 1), (1, 2), (2, 3), (3, 0), *extra)
+            )
+            self.assertTrue(all(len(row) == 4 for row in m))
+            self.assertTrue(all(m[i][i] is None for i in range(4)))
+
+    def test_ring_order_is_kept(self):
+        # Rows come out in ring order, not host order, and rank i is ring[i].
+        m = jaccl_ring_devices([2, 0, 3, 1], 1, links((2, 0), (0, 3), (3, 1), (1, 2)))
+        self.assertEqual(m[0][1], "rdma_en20")
+        self.assertEqual(m[1][2], "rdma_en03")
+        self.assertEqual(m[3][0], "rdma_en12")
+
+    def test_multiple_cables_between_neighbours(self):
+        # Two cables between every ring neighbour: neighbours keep the width the
+        # ring was built with.
+        ips = links((0, 1), (1, 2), (2, 3), (3, 0))
+        for a, b in [(0, 1), (1, 2), (2, 3), (3, 0)]:
+            ips[(a, b)].append((f"en{a}{b}x", "192.168.0.3"))
+            ips[(b, a)].append((f"en{b}{a}x", "192.168.0.4"))
+        m = jaccl_ring_devices([0, 1, 2, 3], 2, ips)
+        self.assertEqual(m[0][1], ["rdma_en01", "rdma_en01x"])
+
+
+if __name__ == "__main__":
+    unittest.main()