(host_id, server_ip, n_hosts, server_port=5005)
Make JAX GPU Cluster.
Make JAX GPU Cluster.
Make JAX GPU Cluster .
def _make_jax_gpu_cluster(host_id, server_ip, n_hosts, server_port=5005):
  """Make JAX GPU Cluster."""

  addr = f'{server_ip}:{server_port}'
  if host_id == 0:
    logging.info('starting service on %s', addr)
    service = xc.get_distributed_runtime_service(addr, n_hosts)
    # We add an explicit call to shutdown the service via atexit as Python
    # interpreter may not call the service destructor on process termination.
    atexit.register(service.shutdown)

  logging.info('connecting to service on %s', addr)
  dist_client = xc.get_distributed_runtime_client(addr, host_id)
  dist_client.connect()
  atexit.register(dist_client.shutdown)

  # register dist gpu backend
  factory = functools.partial(jax.lib.xla_client.make_gpu_client,
                              dist_client, host_id)
  jax.lib.xla_bridge.register_backend_factory('gpu', factory, priority=300)
def _make_jax_gpu_cluster ( host_id , server_ip , n_hosts , server_port = 5005 ) : addr = f'{server_ip}:{server_port}' if host_id == 0 : logging . info ( 'starting service on %s' , addr ) service = xc . get_distributed_runtime_service ( addr , n_hosts ) # We add an explicit call to shutdown the service via atexit as Python # interpreter may not call the service destructor on process termination. atexit . register ( service . shutdown ) logging . info ( 'connecting to service on %s' , addr ) dist_client = xc . get_distributed_runtime_client ( addr , host_id ) dist_client . connect ( ) atexit . register ( dist_client . shutdown ) # register dist gpu backend factory = functools . partial ( jax . lib . xla_client . make_gpu_client , dist_client , host_id ) jax . lib . xla_bridge . register_backend_factory ( 'gpu' , factory , priority = 300 )