The coordination handshake

The chapter above establishes that stock JAX has no coordinator handing out work; every host runs the same script. What it takes to make that safe is one call: jax.distributed.initialize(coordinator_address, num_processes, process_id). One process's address is nominated as the meeting point, every process states how many peers exist, and each names its own index. On Slurm, Kubernetes, or a Cloud TPU deployment the three arguments populate themselves and the call takes none.

The timing rule is strict and the docs state it as a must: initialize runs before jax.devices(), before jax.local_devices(), before any computation touches a device. The handshake is how a process learns the cluster exists; ask about devices earlier and you get a single-process answer that poisons everything after it.

During initialization the processes discover each other and exchange device information, and afterward each one can answer a question it could not answer alone: what does the whole cluster look like? The mechanism underneath is one you have already met. The PJRT boundary lesson read PJRT_Client_Create_Args field by field and found key-value store callbacks, kv_get and kv_put, sitting in the middle of the struct. This handshake is the far side of those callbacks: the coordination service is the store, and a plugin on host three learns what host zero decided by reading keys host zero wrote.

After the handshake, two device lists coexist and confusing them is the classic multi-host bug. jax.local_devices() is the hardware attached to this process; jax.devices() is every device across every process, and jax.process_index() says which slice of the world is yours. Global operations are written against the global list even though each process only ever touches its local slice.

The rule that keeps a thousand processes from deadlocking is behavioral, not enforced: all processes run the same computation in the same order. A collective is a meeting; if one process compiles a different program or reaches the meeting in a different order, the others wait forever for a participant that is never coming. The chapter's framing holds: isolated except when a collective says otherwise, and the handshake is what makes the collectives able to say it.