MCPcopy Create free account
hub / github.com/ml-explore/mlx / configure_jaccl_ring

Function configure_jaccl_ring

python/mlx/_distributed_utils/config.py:486–512  ·  view source on GitHub ↗
(args, hosts, ips, ring, sshinfo)

Source from the content-addressed store, hash-verified

484
485
486def configure_jaccl_ring(args, hosts, ips, ring, sshinfo):
487 log(args.verbose, "Prepare a jaccl-ring hostfile")
488 add_ips(hosts, args.verbose)
489
490 jaccl_hosts = []
491 num_nodes = len(hosts)
492 ring, count = ring
493 for i, node in enumerate(ring):
494 h = hosts[node]
495 peer_left = ring[i - 1]
496 peer_right = ring[(i + 1) % num_nodes]
497 rdmas = []
498 for other in ring:
499 if other not in (peer_left, peer_right):
500 rdmas.append(None)
501 else:
502 rdma = []
503 for c in range(count):
504 rdma.append(f"rdma_{ips.ips[node, other][c][0]}")
505 rdmas.append(rdma[0] if count == 1 else rdma)
506 jaccl_hosts.append(Host(i, h.ssh_hostname, h.ips, rdmas))
507 hostfile = Hostfile(jaccl_hosts, "jaccl-ring", args.env)
508
509 has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup)
510 ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo)
511
512 save_hostfile(args, hostfile)
513
514
515def prepare_tb_hostfile(args, hosts, sshinfo):

Callers 1

prepare_tb_hostfileFunction · 0.85

Calls 8

add_ipsFunction · 0.85
HostClass · 0.85
HostfileClass · 0.85
can_auto_setupFunction · 0.85
save_hostfileFunction · 0.85
appendMethod · 0.80
setupMethod · 0.80
logFunction · 0.70

Tested by

no test coverage detected