feat: construct IFRT clients with distributed options

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -81,6 +81,7 @@
 
 // IFRT
 #include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/attribute_map.h"
 #include "xla/python/ifrt/basic_device_list.h"
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt/compiler.h"
@@ -99,7 +100,6 @@
 #include "xla/python/ifrt/topology.h"
 #include "xla/python/ifrt/tuple.h"
 #include "xla/python/ifrt/value.h"
-#include "xla/python/ifrt/attribute_map.h"
 
 // IFRT - PJRT
 #include "xla/python/pjrt_ifrt/pjrt_array.h"
@@ -343,14 +343,12 @@ extern "C" void ProfilerServerStop(tsl::profiler::ProfilerServer *server) {
   delete server;
 }
 
-extern "C" PjRtClient *MakeCPUClient(uint8_t asynchronous, int node_id,
-                                     int num_nodes) {
+extern "C" PjRtClient *MakeCPUClient(uint8_t asynchronous, int node_id) {
   CpuClientOptions options;
-  // options.kv_store = "etcd";
+
   options.process_id = node_id;
-  // options.num_nodes = num_nodes;
-  // options.collectives = num_nodes;
   options.asynchronous = asynchronous != 0;
+
   auto client = MyValueOrThrow(GetTfrtCpuClient(options));
   return client.release();
 }
@@ -1272,28 +1270,6 @@ extern "C" MlirOperation LinkInModule(MlirModule prevModC, MlirModule newModC,
   return wrap(entryFn);
 }
 
-extern "C" HeldPjRtClient *
-pjrt_make_cpu_client_shared(uint8_t asynchronous, int node_id, int num_nodes) {
-  PjRtClient *client = MakeCPUClient(asynchronous, node_id, num_nodes);
-  return reactant::capture(std::shared_ptr<PjRtClient>(client));
-}
-
-extern "C" HeldPjRtClient *pjrt_make_gpu_client_shared(
-    int node_id, int num_nodes, int *allowed_devices, int num_allowed_devices,
-    double memory_fraction, bool preallocate, const char *platform_name,
-    const char **error, void *distributed_runtime_client) {
-  PjRtClient *client = MakeGPUClient(
-      node_id, num_nodes, allowed_devices, num_allowed_devices, memory_fraction,
-      preallocate, platform_name, error, distributed_runtime_client);
-  return reactant::capture(std::shared_ptr<PjRtClient>(client));
-}
-
-extern "C" HeldPjRtClient *pjrt_make_tpu_client_shared(const char *tpu_path,
-                                                       const char **error) {
-  PjRtClient *client = MakeTPUClient(tpu_path, error);
-  return reactant::capture(std::shared_ptr<PjRtClient>(client));
-}
-
 extern "C" void pjrt_client_dtor(HeldPjRtClient *client) { delete client; }
 
 extern "C" int pjrt_client_num_devices(HeldPjRtClient *client) {
@@ -1370,11 +1346,6 @@ extern "C" HeldPjRtClient *pjrt_buffer_get_client(HeldPjRtBuffer *buffer) {
       std::shared_ptr<PjRtClient>(buffer->ptr()->client()));
 }
 
-extern "C" ifrt::Client *ifrt_pjrt_make_client(HeldPjRtClient *pjrt_client) {
-  xla::ifrt::PjRtClient::CreateOptions options = {pjrt_client->obj()};
-  return MyValueOrThrow(xla::ifrt::PjRtClient::Create(options)).release();
-}
-
 extern "C" void ifrt_client_dtor(ifrt::Client *client) { delete client; }
 
 // generic version, but IFRT-PjRt backend only supports SingleDeviceSharding
@@ -1576,59 +1547,65 @@ FreeHloModule(HeldValue<std::shared_ptr<xla::HloModule>> *hlo_module) {
 
 #pragma region IfRtClient
 
-extern "C" ifrt::proxy::GrpcServer *
-ifrt_proxy_grpc_server_create_from_ifrt_client_factory_cpu(
-    const char *c_address, uint8_t asynchronous, int node_id, int num_nodes) {
-  std::string address = c_address;
-
-  return MyValueOrThrow(
-             ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
-                 address,
-                 [asynchronous, node_id, num_nodes]()
-                     -> absl::StatusOr<std::shared_ptr<ifrt::Client>> {
-                   auto pjrt_client = std::shared_ptr<PjRtClient>(
-                       MakeCPUClient(asynchronous, node_id, num_nodes));
-                   return std::shared_ptr<ifrt::Client>(
-                       xla::ifrt::PjRtClient::Create(pjrt_client).release());
-                 }))
-      .release();
-}
-
-extern "C" ifrt::proxy::GrpcServer *
-ifrt_proxy_grpc_server_create_from_ifrt_client_factory_gpu(
-    int node_id, int num_nodes, int *allowed_devices, int num_allowed_devices,
-    double memory_fraction, bool preallocate, const char *platform_name,
-    const char **error) {
-  return MyValueOrThrow(
-             ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
-                 std::string(),
-                 [node_id, num_nodes, allowed_devices, num_allowed_devices,
-                  memory_fraction, preallocate, platform_name,
-                  error]() -> absl::StatusOr<std::shared_ptr<ifrt::Client>> {
-                   auto pjrt_client = std::shared_ptr<PjRtClient>(MakeGPUClient(
-                       node_id, num_nodes, allowed_devices, num_allowed_devices,
-                       memory_fraction, preallocate, platform_name, error));
-                   return std::shared_ptr<ifrt::Client>(
-                       xla::ifrt::PjRtClient::Create(pjrt_client).release());
-                 }))
-      .release();
-}
+// XXX: Bring back with the correct API
+// extern "C" ifrt::proxy::GrpcServer *
+// ifrt_proxy_grpc_server_create_from_ifrt_client_factory_cpu(
+//     const char *c_address, uint8_t asynchronous, int node_id) {
+//   std::string address = c_address;
+
+//   return MyValueOrThrow(
+//              ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
+//                  address,
+//                  [asynchronous,
+//                   node_id]() -> absl::StatusOr<std::shared_ptr<ifrt::Client>>
+//                   {
+//                    auto pjrt_client = std::shared_ptr<PjRtClient>(
+//                        MakeCPUClient(asynchronous, node_id));
+//                    return std::shared_ptr<ifrt::Client>(
+//                        xla::ifrt::PjRtClient::Create(pjrt_client).release());
+//                  }))
+//       .release();
+// }
 
-extern "C" ifrt::proxy::GrpcServer *
-ifrt_proxy_grpc_server_create_from_ifrt_client_factory_tpu(
-    const char *c_address, const char *tpu_path, const char **error) {
-  std::string address = c_address;
+// extern "C" ifrt::proxy::GrpcServer *
+// ifrt_proxy_grpc_server_create_from_ifrt_client_factory_gpu(
+//     int node_id, int num_nodes, int *allowed_devices, int
+//     num_allowed_devices, double memory_fraction, bool preallocate, const char
+//     *platform_name, const char **error) {
+//   return MyValueOrThrow(
+//              ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
+//                  std::string(),
+//                  [node_id, num_nodes, allowed_devices, num_allowed_devices,
+//                   memory_fraction, preallocate, platform_name,
+//                   error]() -> absl::StatusOr<std::shared_ptr<ifrt::Client>> {
+//                    auto pjrt_client =
+//                    std::shared_ptr<PjRtClient>(MakeGPUClient(
+//                        node_id, num_nodes, allowed_devices,
+//                        num_allowed_devices, memory_fraction, preallocate,
+//                        platform_name, error));
+//                    return std::shared_ptr<ifrt::Client>(
+//                        xla::ifrt::PjRtClient::Create(pjrt_client).release());
+//                  }))
+//       .release();
+// }
 
-  return MyValueOrThrow(
-             xla::ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
-                 address,
-                 [](xla::ifrt::AttributeMap initialization_data) -> absl::StatusOr<std::shared_ptr<xla::ifrt::Client>> {
-                   auto pjrt_client =
-                       std::shared_ptr<xla::PjRtClient>(GetCApiClient("TPU"));
-                   return xla::ifrt::PjRtClient::Create(std::move(pjrt_client));
-                 }))
-      .release();
-}
+// extern "C" ifrt::proxy::GrpcServer *
+// ifrt_proxy_grpc_server_create_from_ifrt_client_factory_tpu(
+//     const char *c_address, const char *tpu_path, const char **error) {
+//   std::string address = c_address;
+//
+//   return MyValueOrThrow(
+//              xla::ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
+//                  address,
+//                  [](xla::ifrt::AttributeMap initialization_data) ->
+//                  absl::StatusOr<std::shared_ptr<xla::ifrt::Client>> {
+//                    auto pjrt_client =
+//                        std::shared_ptr<xla::PjRtClient>(GetCApiClient("TPU"));
+//                    return
+//                    xla::ifrt::PjRtClient::Create(std::move(pjrt_client));
+//                  }))
+//       .release();
+// }
 
 extern "C" void ifrt_proxy_grpc_server_dtor(ifrt::proxy::GrpcServer *server) {
   delete server;
@@ -1661,24 +1638,88 @@ ifrt_proxy_create_client(const char *c_proxy_server_address,
       .release();
 }
 
-extern "C" ifrt::Client *ifrt_make_pjrt_cpu_client(uint8_t asynchronous,
-                                                   int node_id, int num_nodes) {
-  return ifrt_pjrt_make_client(
-      pjrt_make_cpu_client_shared(asynchronous, node_id, num_nodes));
+extern "C" ifrt::Client *ifrt_pjrt_make_client(HeldPjRtClient *pjrt_client,
+                                               int node_id, int num_nodes,
+                                               void *distributed_runtime_client,
+                                               const char **error,
+                                               std::string key_prefix) {
+  ifrt::PjRtClient::CreateOptions options;
+  options.pjrt_client = pjrt_client->obj();
+
+  if (num_nodes > 1) {
+    if (distributed_runtime_client == nullptr) {
+      *error =
+          "`distributed_runtime_client` must be non-null if `num_nodes` > 1";
+      return nullptr;
+    }
+    auto typed_distributed_runtime_client = static_cast<
+        HeldValue<std::shared_ptr<xla::DistributedRuntimeClient>> *>(
+        distributed_runtime_client);
+    options.kv_store = GetDistributedKeyValueStore(
+        typed_distributed_runtime_client->obj(), key_prefix);
+  }
+
+  options.process_id = node_id;
+  options.num_processes = num_nodes;
+
+  return MyValueOrThrow(xla::ifrt::PjRtClient::Create(options)).release();
+}
+
+extern "C" HeldPjRtClient *pjrt_make_cpu_client_shared(uint8_t asynchronous,
+                                                       int node_id) {
+  PjRtClient *client = MakeCPUClient(asynchronous, node_id);
+  return reactant::capture(std::shared_ptr<PjRtClient>(client));
+}
+
+extern "C" ifrt::Client *
+ifrt_make_pjrt_cpu_client(uint8_t asynchronous, int node_id, int num_nodes,
+                          void *distributed_runtime_client,
+                          const char **error) {
+  HeldPjRtClient *pjrt_client =
+      pjrt_make_cpu_client_shared(asynchronous, node_id);
+  if (pjrt_client == nullptr)
+    return nullptr;
+  return ifrt_pjrt_make_client(pjrt_client, node_id, num_nodes,
+                               distributed_runtime_client, error, "cpu");
+}
+
+extern "C" HeldPjRtClient *pjrt_make_gpu_client_shared(
+    int node_id, int num_nodes, int *allowed_devices, int num_allowed_devices,
+    double memory_fraction, bool preallocate, const char *platform_name,
+    const char **error, void *distributed_runtime_client) {
+  PjRtClient *client = MakeGPUClient(
+      node_id, num_nodes, allowed_devices, num_allowed_devices, memory_fraction,
+      preallocate, platform_name, error, distributed_runtime_client);
+  return reactant::capture(std::shared_ptr<PjRtClient>(client));
 }
 
 extern "C" ifrt::Client *ifrt_make_pjrt_gpu_client(
     int node_id, int num_nodes, int *allowed_devices, int num_allowed_devices,
     double memory_fraction, bool preallocate, const char *platform_name,
     const char **error, void *distributed_runtime_client) {
-  return ifrt_pjrt_make_client(pjrt_make_gpu_client_shared(
+  HeldPjRtClient *pjrt_client = pjrt_make_gpu_client_shared(
       node_id, num_nodes, allowed_devices, num_allowed_devices, memory_fraction,
-      preallocate, platform_name, error, distributed_runtime_client));
+      preallocate, platform_name, error, distributed_runtime_client);
+  if (pjrt_client == nullptr)
+    return nullptr;
+  return ifrt_pjrt_make_client(pjrt_client, node_id, num_nodes,
+                               distributed_runtime_client, error, "gpu");
 }
 
-extern "C" ifrt::Client *ifrt_make_pjrt_tpu_client(const char *tpu_path,
-                                                   const char **error) {
-  return ifrt_pjrt_make_client(pjrt_make_tpu_client_shared(tpu_path, error));
+extern "C" HeldPjRtClient *pjrt_make_tpu_client_shared(const char *tpu_path,
+                                                       const char **error) {
+  PjRtClient *client = MakeTPUClient(tpu_path, error);
+  return reactant::capture(std::shared_ptr<PjRtClient>(client));
+}
+
+extern "C" ifrt::Client *
+ifrt_make_pjrt_tpu_client(const char *tpu_path, const char **error, int node_id,
+                          int num_nodes, void *distributed_runtime_client) {
+  HeldPjRtClient *pjrt_client = pjrt_make_tpu_client_shared(tpu_path, error);
+  if (pjrt_client == nullptr)
+    return nullptr;
+  return ifrt_pjrt_make_client(pjrt_client, node_id, num_nodes,
+                               distributed_runtime_client, error, "tpu");
 }
 
 extern "C" void ifrt_FreeClient(ifrt::Client *client) { delete client; }
@@ -1942,7 +1983,7 @@ extern "C" void ifrt_array_copy_to_host_buffer(HeldIfrtArray *array,
 
 #pragma endregion
 
-#pragma region PjRtDistributed
+#pragma region xla::Distributed
 
 extern "C" HeldValue<std::shared_ptr<xla::DistributedRuntimeClient>> *
 GetDistributedRuntimeClient(char *c_address, int32_t node_id,

src/xla/Client.jl

@@ -14,55 +14,3 @@ function get_addressable_device end
 function platform_name end
 
 default_device(client::AbstractClient) = first(addressable_devices(client))
-
-# Clients for Different Backends
-function CPUClient(cfunc, node_id=0, num_nodes=1; asynchronous=true)
-    f = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, string(cfunc))
-    client = ccall(f, Ptr{Cvoid}, (UInt, Cint, Cint), asynchronous, node_id, num_nodes)
-    LLVMclopts("-nvptx-fma-level=1")
-    return client
-end
-
-function GPUClient(
-    cfunc,
-    node_id=0,
-    num_nodes=1,
-    platform="gpu";
-    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}()
-
-    num_allowed_devices = allowed_devices === nothing ? 0 : length(allowed_devices)
-    allowed_devices = allowed_devices === nothing ? C_NULL : allowed_devices
-    distributed_runtime_client =
-        distributed_runtime_client === nothing ? C_NULL : distributed_runtime_client.client
-
-    client = ccall(
-        f,
-        Ptr{Cvoid},
-        (Cint, Cint, Ptr{Cvoid}, Cint, Cdouble, Bool, Cstring, Ptr{Cstring}, Ptr{Cvoid}),
-        node_id,
-        num_nodes,
-        allowed_devices,
-        num_allowed_devices,
-        XLA_REACTANT_GPU_MEM_FRACTION[],
-        false,
-        platform,
-        refstr,
-        distributed_runtime_client,
-    )
-    client == C_NULL && throw(AssertionError(unsafe_string(refstr[])))
-    LLVMclopts("-nvptx-fma-level=1")
-    return client
-end
-
-function TPUClient(cfunc, tpu_path::String)
-    f = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, string(cfunc))
-    refstr = Ref{Cstring}()
-    client = ccall(f, Ptr{Cvoid}, (Cstring, Ptr{Cstring}), tpu_path, refstr)
-    client == C_NULL && throw(AssertionError(unsafe_string(refstr[])))
-    LLVMclopts("-nvptx-fma-level=1")
-    return client
-end