refactor: Distributed setup is not PJRT specific

Author: avik-pal

src/Distributed.jl

@@ -24,11 +24,11 @@ function initialize(;
 
     @debug "Detected Reactant distributed params" coordinator_address num_processes process_id local_device_ids
 
-    Reactant.XLA.PJRT.update_global_state!(;
+    Reactant.XLA.update_global_state!(;
         coordinator_address, num_processes, process_id, local_device_ids, kwargs...
     )
 
-    @debug "New Global State" Reactant.XLA.PJRT.global_state
+    @debug "New Global State" Reactant.XLA.global_state
 
     initialized[] = true
     return nothing

src/xla/Client.jl

@@ -28,8 +28,8 @@ function GPUClient(
     node_id=0,
     num_nodes=1,
     platform="gpu";
-    allowed_devices=nothing,
-    distributed_runtime_client=nothing,
+    allowed_devices::Union{Nothing,Vector{Int}}=nothing,
+    distributed_runtime_client::Union{Nothing,DistributedRuntimeClient}=nothing,
 )
     f = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, string(cfunc))
     refstr = Ref{Cstring}()

src/xla/PJRT/Distributed.jl src/xla/Distributed.jl

@@ -175,8 +175,8 @@ function update!(
     if length(proxy_vars) > 0
         vars = join(proxy_vars, ", ")
         @warn "Reactant detected proxy variable(s) in the environment as distributed \
-               setup: $(vars). On some systems, this may cause a hang of \
-               `XLA.PJRT.update!` and you may need to unset the proxy variables."
+               setup: $(vars). On some systems, this may cause a hang of `XLA.update!` and \
+               you may need to unset the proxy variables."
     end
 
     @assert state.client === nothing "`Reactant.Distributed.initialize` should only be \

src/xla/PJRT/PJRT.jl

@@ -6,29 +6,11 @@ using Reactant_jll: Reactant_jll
 
 using Libdl: Libdl
 
-include("Distributed.jl")
 include("Client.jl")
 include("Device.jl")
 include("Future.jl")
 include("Buffer.jl")
 include("AsyncBuffer.jl")
 include("LoadedExecutable.jl")
 
-const global_state = State()
-
-function __init__()
-    if haskey(ENV, "REACTANT_VISIBLE_GPU_DEVICES")
-        global_state.local_device_ids =
-            parse.(Int, split(ENV["REACTANT_VISIBLE_GPU_DEVICES"], ","))
-    end
-    return nothing
-end
-
-function update_global_state!(args...; kwargs...)
-    update!(global_state, args...; kwargs...)
-    # We need to update the clients based on the new state
-    XLA.initialize_default_clients()
-    return nothing
-end
-
 end

src/xla/XLA.jl

@@ -21,6 +21,7 @@ function LLVMclopts(opts...)
     )::Cvoid
 end
 
+include("Distributed.jl")
 include("Client.jl")
 include("Device.jl")
 include("Sharding.jl")
@@ -53,6 +54,7 @@ function Base.setproperty!(bs::BackendState, sym::Symbol, val)
 end
 
 const global_backend_state = BackendState()
+const global_state = State()
 
 client(backend::String) = global_backend_state.clients[backend]
 default_backend() = global_backend_state.default_client
@@ -68,6 +70,13 @@ function set_default_backend(backend::String)
     return nothing
 end
 
+function update_global_state!(args...; kwargs...)
+    update!(global_state, args...; kwargs...)
+    # We need to update the clients based on the new state
+    initialize_default_clients!(global_backend_state)
+    return nothing
+end
+
 function __init__()
     # This must be the very first thing initialized (otherwise we can't throw errors)
     errptr = cglobal((:ReactantThrowError, MLIR.API.mlir_c), Ptr{Ptr{Cvoid}})
@@ -90,16 +99,17 @@ function __init__()
         @debug "XLA_REACTANT_GPU_PREALLOCATE: " XLA_REACTANT_GPU_PREALLOCATE[]
     end
 
+    if haskey(ENV, "REACTANT_VISIBLE_GPU_DEVICES")
+        global_state.local_device_ids =
+            parse.(Int, split(ENV["REACTANT_VISIBLE_GPU_DEVICES"], ","))
+        @debug "REACTANT_VISIBLE_GPU_DEVICES: " global_state.local_device_ids
+    end
+
     @ccall MLIR.API.mlir_c.RegisterEnzymeXLACPUHandler()::Cvoid
     @ccall MLIR.API.mlir_c.RegisterEnzymeXLAGPUHandler()::Cvoid
     return nothing
 end
 
-function initialize_default_clients()
-    initialize_default_clients!(global_backend_state)
-    return nothing
-end
-
 function initialize_default_clients!(state::BackendState)
     was_initialized = state.initialized
     state.initialized = true
@@ -109,7 +119,7 @@ function initialize_default_clients!(state::BackendState)
         XLA.free_client(state.clients["cpu"])
         XLA.PJRT.cpu_client_count[] -= 1
     end
-    cpu = PJRT.CPUClient(PJRT.global_state.process_id, PJRT.global_state.num_processes)
+    cpu = PJRT.CPUClient(global_state.process_id, global_state.num_processes)
     state.clients["cpu"] = cpu
     state.default_client = cpu
 
@@ -142,9 +152,9 @@ function initialize_default_clients!(state::BackendState)
         else
             if !Reactant.precompiling()
                 try
-                    distributed_runtime_client = if PJRT.global_state.num_processes > 1
-                        @assert PJRT.global_state.client !== nothing
-                        PJRT.global_state.client
+                    distributed_runtime_client = if global_state.num_processes > 1
+                        @assert global_state.client !== nothing
+                        global_state.client
                     else
                         nothing
                     end
@@ -154,9 +164,9 @@ function initialize_default_clients!(state::BackendState)
                         XLA.PJRT.gpu_client_count[] -= 1
                     end
                     gpu = PJRT.GPUClient(
-                        PJRT.global_state.process_id,
-                        PJRT.global_state.num_processes;
-                        allowed_devices=PJRT.global_state.local_device_ids,
+                        global_state.process_id,
+                        global_state.num_processes;
+                        allowed_devices=global_state.local_device_ids,
                         distributed_runtime_client,
                     )
                     state.clients["gpu"] = gpu