nit: do you still need to make_optional here, or can you directly assign? e.g. allowed_devices = std::set{local_process_rank}

Synced offline - this is great for single-host single-process development, but in cases where there is a single process per host, this would break in a multihost environment. Outside of SPMD, I'm not aware of a use case for a multihost environment with a single process per host (cc @JackCaoG)

Since we don't officially support SPMD on GPU at the moment, this looks fine to me for now. Once we decide on the right entrypoint for SPMD, we'll need to revisit this.

Not for this PR, but IMO we should just delete these tests. Do we support DataParallel anymore @JackCaoG?

Conditioning on SPMD mode here could cause issues using xr.use_spmd() after the runtime has been initialized.

Is it correct to say that allowed_devices is only needed in the MP case? If so, can we invert the logic to check for MP using one of the env vars instead of checking for SPMD mode?

Is there ever a reason to call xr.use_spmd() after the runtime is initialized? In any case, I think we can also assume that if LOCAL_WORLD_SIZE=1, then we can use all of the devices (which should be compatible with SPMD)

AFAIK there's not a strong use case at the moment, but for example our unit tests will check xr.global_runtime_device_count() before calling xr.use_spmd(). Keeping the runtime independent of SPMD mode was something we wanted to maintain, cc @yeounoh

That's a good point!

Conditioning on SPMD mode here could cause issues using xr.use_spmd() after the runtime has been initialized.

This seems to be a downside of using xr.use_spmd() as opposed to a env flag XLA_USE_SPMD=1. With the latter, it's less flexible but less error-prone. It guarantee we'll use spmd mode at the beginning. With the former, it may also impact other SPMD special cases: user does something pytorch ops, then call xr.use_spmd(), then continue to do something else.

can we invert the logic to check for MP using one of the env vars instead of checking for SPMD mode?
I think we can also assume that if LOCAL_WORLD_SIZE=1, then we can use all of the devices

I'm thinking about the case where the user has 2 GPU machines and she wants to use 1 GPU device on each machine and to do multi-host training. In that case (multi-host-single-process), each process has access to all devices and I guess the user can still do multi-host training

I think we can also assume that if LOCAL_WORLD_SIZE=1, then we can use all of the devices

Perhaps we also need to check GLOBAL_WORLD_SIZE:

if LOCAL_WORLD_SIZE==1:
  if GLOBAL_WORLD_SIZE>1: # multi-host-single-process
    initialize coordinator service
  else: # single-host-single-process
    do nothing
else: multi-process for single-host and multi-host
  initialize coordinator service
  allowd_devices={current_device}

Why do we need the coordinator as input here?

We need it get the DistributedRuntimeClient and later create the kv_store below

It looks like it's being recreated on L60 - should we just make this function return the new value?

done

Thanks for generalizing these tests!

Could we change these comments to e.g. # if self.n_devices==4, mesh=(2, 2)? Other device counts will have different meshes.

Yes! Done.

I think we could simplify the logic here some. We want to restrict allowed_devices if local_world_size > 1 and create the coordinator if global_world_size > 1. I'm assuming local_world_size > 1 => global_world_size > 1, would this be equivalent?

if (local_world_size > 1) {
    allowed_devices = std::set{local_process_rank};
} 
if (global_world_size > 1) {
    // We can keep the old initialization block here and remove `SetGpuClientKVCallBack`
}

DebugString should include the type?

Actually, it doesn't. The debugString is empty in that case.