Making JACCL coordinator optional

[angeloskath]

This allows the user to define their own side channel used to exchange RDMA metadata to initialize the connection.

This can be used for apps that have a different way to establish a slow connection between the nodes that can be used to exchange metadata. It avoids the need for opening a port and listening for connections etc. It could use a shared filesystem or iCloud or any centralized service etc

Here is an example I asked Claude to write that uses ZeroMQ to handle connection retries etc.

ZeroMQ all gather example

# Copyright © 2025 Apple Inc.

"""A minimal ZeroMQ based all-gather side channel for the jaccl backend.

The jaccl backend bootstraps its RDMA connections by performing a couple of
byte-level all-gathers to exchange connection metadata. This module provides a
simple :func:`zmq_all_gather_factory` that implements that all-gather over TCP
using ZeroMQ. Rank 0 acts as a coordinator: every other rank connects to it,
sends its chunk and receives the fully assembled buffer back.

Usage:

    import mlx.core as mx
    from mlx._distributed_utils.zmq_all_gather import zmq_all_gather_factory

    group = mx.distributed.init(
        backend="jaccl",
        all_gather_factory=zmq_all_gather_factory(),
    )
"""

import os
import struct

import zmq

# The environment variable set by ``mlx.launch`` that holds the ``ip:port`` of
# rank 0 which acts as the all-gather coordinator.
_COORDINATOR_ENV = ("JACCL_COORDINATOR", "MLX_JACCL_COORDINATOR")


def _coordinator_address(coordinator):
    if coordinator is not None:
        return coordinator
    for name in _COORDINATOR_ENV:
        value = os.environ.get(name)
        if value:
            return value
    raise RuntimeError(
        "No coordinator address provided and neither of "
        f"{_COORDINATOR_ENV} is set in the environment."
    )


def zmq_all_gather_factory(coordinator=None):
    """Create an all-gather factory for the jaccl backend backed by ZeroMQ.

    Args:
        coordinator (str, optional): The ``ip:port`` of rank 0. If ``None`` it
            is read from the ``MLX_JACCL_COORDINATOR`` environment variable.

    Returns:
        Callable[[int, int], Callable[[bytes, int], bytes]]: A factory that
        ``mx.distributed.init`` calls once per rank with ``(rank, size)`` and
        returns the actual all-gather callable.
    """
    ip, port = _coordinator_address(coordinator).rsplit(":", 1)
    endpoint = f"tcp://{ip}:{port}"

    def factory(rank, size):
        context = zmq.Context.instance()

        if rank == 0:
            # The coordinator uses a ROUTER socket so it can gather every
            # rank's chunk before replying with the assembled buffer.
            socket = context.socket(zmq.ROUTER)
            socket.bind(f"tcp://*:{port}")

            def all_gather(src, n_bytes):
                chunks = [None] * size
                chunks[0] = src
                envelopes = []
                for _ in range(size - 1):
                    frames = socket.recv_multipart()
                    # REQ sockets add an empty delimiter frame; everything
                    # before the payload is the routing envelope.
                    envelope, msg = frames[:-1], frames[-1]
                    peer = struct.unpack("<i", msg[:4])[0]
                    chunks[peer] = msg[4:]
                    envelopes.append(envelope)
                assembled = b"".join(chunks)
                for envelope in envelopes:
                    socket.send_multipart(envelope + [assembled])
                return assembled

            return all_gather

        # Non-coordinator ranks send their chunk and receive the full buffer.
        socket = context.socket(zmq.REQ)
        socket.connect(endpoint)

        def all_gather(src, n_bytes):
            socket.send(struct.pack("<i", rank) + src)
            return socket.recv()

        return all_gather

    return factory