Error Handling: refactor the existing ComputationClient implementations to use status QOL functions. (#9420)

Author: ysiraichi

torch_xla/csrc/runtime/BUILD

@@ -24,6 +24,7 @@ cc_library(
         ":env_vars",
         ":ifrt_computation_client",
         ":pjrt_computation_client",
+        "//torch_xla/csrc:status",
         "@com_google_absl//absl/log:absl_check",
         "@tsl//tsl/platform:stacktrace",
     ],
@@ -497,7 +498,7 @@ ptxla_cc_test(
         ":pjrt_computation_client",
         ":operation_manager",
         ":tensor_source",
-        "@com_google_absl//absl/status",
+        "//torch_xla/csrc:status",
         "@tsl//tsl/platform:test_main",
         "@xla//xla:literal",
         "@xla//xla:literal_util",

torch_xla/csrc/runtime/ifrt_computation_client.cpp

@@ -120,11 +120,13 @@ std::vector<std::string> IfrtComputationClient::IfrtDevicesToString(
   return strs;
 }
 
-IfrtComputationClient::IfrtComputationClient() {
+IfrtComputationClient::IfrtComputationClient(PrivateUse) {}
+
+absl::Status IfrtComputationClient::Initialize() {
   std::string device_type = sys_util::GetEnvString(env::kEnvPjRtDevice, "");
   std::unique_ptr<xla::PjRtClient> pjrt_client;
-  std::tie(pjrt_client, coordinator_) =
-      GetValueOrThrow(InitializePjRt(device_type));
+  XLA_ASSIGN_OR_RETURN(std::tie(pjrt_client, coordinator_),
+                       InitializePjRt(device_type));
 
   client_ = xla::ifrt::PjRtClient::Create(std::move(pjrt_client));
 
@@ -146,6 +148,15 @@ IfrtComputationClient::IfrtComputationClient() {
   auto tracked_devices = GetLocalDevices();
   tracked_devices.emplace_back(spmd_device_str);
   operation_manager_ = std::move(OperationManager(std::move(tracked_devices)));
+
+  return absl::OkStatus();
+}
+
+absl::StatusOr<absl_nonnull std::unique_ptr<IfrtComputationClient>>
+IfrtComputationClient::Create() {
+  auto ifrt_client = std::make_unique<IfrtComputationClient>(PrivateUse());
+  XLA_RETURN_IF_ERROR(ifrt_client->Initialize());
+  return std::move(ifrt_client);
 }
 
 IfrtComputationClient::~IfrtComputationClient() {
@@ -652,7 +663,7 @@ int IfrtComputationClient::GetNumProcesses() const {
   }
 
   return max_process_index + 1;
-};
+}
 
 std::string IfrtComputationClient::GetDeviceKind(const std::string& device) {
   return std::string(StringToIfrtDevice(device)->Kind());

torch_xla/csrc/runtime/ifrt_computation_client.h

@@ -27,8 +27,17 @@ namespace torch_xla {
 namespace runtime {
 
 class IfrtComputationClient : public ComputationClient {
+ private:
+  // Private struct for making the constructor private, but still callable
+  // as: `std::make_unique<IfrtComputationClient>(PrivateUse())`.
+  struct PrivateUse {
+    // Constructor needs to be explicit for disallowing implicit construction
+    // from `{}`.
+    explicit PrivateUse() = default;
+  };
+
  public:
-  IfrtComputationClient();
+  IfrtComputationClient(PrivateUse);
   ~IfrtComputationClient();
 
   DataPtr CreateDataPlaceholder(
@@ -165,7 +174,15 @@ class IfrtComputationClient : public ComputationClient {
     XLA_ERROR() << __FUNCTION__ << " not implemented";
   }
 
+  // Creates a new instance of IfrtComputationClient and initializes it.
+  static absl::StatusOr<absl_nonnull std::unique_ptr<IfrtComputationClient>>
+  Create();
+
  private:
+  // Convenience function called by `Create()` that initializes the current
+  // IfrtComputationClient.
+  absl::Status Initialize();
+
   std::shared_ptr<xla::ifrt::PjRtClient> client_;
   std::unique_ptr<XlaCoordinator> coordinator_;
   // global_ordinals_ tracks a map from PjRtDeviceId to the device's

torch_xla/csrc/runtime/ifrt_computation_client_test.cpp

@@ -6,9 +6,9 @@
 #include <string>
 #include <vector>
 
-#include "absl/status/status.h"
 #include "torch_xla/csrc/runtime/computation_client.h"
 #include "torch_xla/csrc/runtime/tensor_source.h"
+#include "torch_xla/csrc/status.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/statusor.h"
@@ -36,7 +36,7 @@ absl::StatusOr<xla::XlaComputation> MakeComputation() {
 TEST(PjRtComputationClientTest, Init) {
   // Get a CPU client.
   tsl::setenv("PJRT_DEVICE", "CPU", true);
-  auto client = std::make_unique<IfrtComputationClient>();
+  auto client = GetValueOrThrow(IfrtComputationClient::Create());
   std::string device = client->GetDefaultDevice();
 
   // Compose a computation.

torch_xla/csrc/runtime/pjrt_computation_client.cpp

@@ -115,10 +115,12 @@ std::vector<std::string> PjRtComputationClient::PjRtDevicesToString(
   return strs;
 }
 
-PjRtComputationClient::PjRtComputationClient() {
+PjRtComputationClient::PjRtComputationClient(PrivateUse) {}
+
+absl::Status PjRtComputationClient::Initialize() {
   std::string device_type = sys_util::GetEnvString(env::kEnvPjRtDevice, "");
-  std::tie(client_, coordinator_) =
-      GetValueOrThrow(InitializePjRt(device_type));
+  XLA_ASSIGN_OR_RETURN(std::tie(client_, coordinator_),
+                       InitializePjRt(device_type));
 
   // PjRtDevice IDs are not guaranteed to be dense, so we need to track
   // a device's global ordinal separately from its device ID. Order the
@@ -137,6 +139,15 @@ PjRtComputationClient::PjRtComputationClient() {
   auto tracked_devices = GetLocalDevices();
   tracked_devices.emplace_back(spmd_device_str);
   operation_manager_ = std::move(OperationManager(std::move(tracked_devices)));
+
+  return absl::OkStatus();
+}
+
+absl::StatusOr<absl_nonnull std::unique_ptr<PjRtComputationClient>>
+PjRtComputationClient::Create() {
+  auto pjrt_client = std::make_unique<PjRtComputationClient>(PrivateUse());
+  XLA_RETURN_IF_ERROR(pjrt_client->Initialize());
+  return std::move(pjrt_client);
 }
 
 PjRtComputationClient::~PjRtComputationClient() {
@@ -837,7 +848,7 @@ PjRtComputationClient::ExecuteReplicated(
               argument_handles[d][i] = shard->buffer.get();
             }
             counter.DecrementCount();
-          };
+          }
         });
     counter.Wait();
   }
@@ -962,7 +973,7 @@ int PjRtComputationClient::GetNumProcesses() const {
   }
 
   return max_process_index + 1;
-};
+}
 
 std::string PjRtComputationClient::GetDeviceKind(const std::string& device) {
   return std::string(StringToPjRtDevice(device)->device_kind());

torch_xla/csrc/runtime/pjrt_computation_client.h

@@ -25,8 +25,17 @@ namespace torch_xla {
 namespace runtime {
 
 class PjRtComputationClient : public ComputationClient {
+ private:
+  // Private struct for making the constructor private, but still callable
+  // as: `std::make_unique<PjRtComputationClient>(PrivateUse())`.
+  struct PrivateUse {
+    // Constructor needs to be explicit for disallowing implicit construction
+    // from `{}`.
+    explicit PrivateUse() = default;
+  };
+
  public:
-  PjRtComputationClient();
+  PjRtComputationClient(PrivateUse);
   ~PjRtComputationClient() override;
 
   DataPtr CreateDataPlaceholder(
@@ -163,6 +172,10 @@ class PjRtComputationClient : public ComputationClient {
   void OnReadyCallback(DataPtr data,
                        const std::function<void()>& callback) override;
 
+  // Creates a new instance of PjRtComputationClient and initializes it.
+  static absl::StatusOr<absl_nonnull std::unique_ptr<PjRtComputationClient>>
+  Create();
+
  private:
   friend class PjRtComputationClientTest;
 
@@ -172,6 +185,10 @@ class PjRtComputationClient : public ComputationClient {
     fake_xla_compile_ = std::move(function);
   }
 
+  // Convenience function called by `Create()` that initializes the current
+  // PjRtComputationClient.
+  absl::Status Initialize();
+
   std::unique_ptr<xla::PjRtClient> client_;
   std::unique_ptr<XlaCoordinator> coordinator_;
   // global_ordinals_ tracks a map from PjRtDeviceId to the device's

torch_xla/csrc/runtime/pjrt_computation_client_test.cpp

@@ -8,10 +8,9 @@
 #include <string>
 #include <vector>
 
-#include "absl/status/status.h"
 #include "torch_xla/csrc/runtime/computation_client.h"
-#include "torch_xla/csrc/runtime/pjrt_computation_client.h"
 #include "torch_xla/csrc/runtime/tensor_source.h"
+#include "torch_xla/csrc/status.h"
 #include "xla/hlo/builder/xla_builder.h"
 #include "xla/hlo/builder/xla_computation.h"
 #include "xla/literal.h"
@@ -26,7 +25,7 @@ class PjRtComputationClientTest : public ::testing::Test {
   PjRtComputationClientTest() {
     // Get a CPU client.
     tsl::setenv("PJRT_DEVICE", "CPU", true);
-    client_ = std::make_unique<PjRtComputationClient>();
+    client_ = GetValueOrThrow(PjRtComputationClient::Create());
     device_ = client_->GetDefaultDevice();
   }
 

torch_xla/csrc/runtime/runtime.cpp

@@ -7,6 +7,7 @@
 #include "torch_xla/csrc/runtime/env_vars.h"
 #include "torch_xla/csrc/runtime/ifrt_computation_client.h"
 #include "torch_xla/csrc/runtime/pjrt_computation_client.h"
+#include "torch_xla/csrc/status.h"
 #include "tsl/platform/stacktrace_handler.h"
 
 namespace torch_xla::runtime {
@@ -32,13 +33,18 @@ InitializeComputationClient() {
   // static bool use_ifrt = sys_util::GetEnvBool("XLA_USE_IFRT", false);
   const bool use_ifrt = false;
   if (sys_util::GetEnvString(env::kEnvPjRtDevice, "") == "") {
-    return absl::FailedPreconditionError("$PJRT_DEVICE is not set.");
+    return XLA_ERROR_WITH_LOCATION(
+        absl::FailedPreconditionError("$PJRT_DEVICE is not set."));
   }
 
+  std::unique_ptr<ComputationClient> client;
   if (use_ifrt) {
-    return new IfrtComputationClient();
+    XLA_ASSIGN_OR_RETURN(client, IfrtComputationClient::Create());
+  } else {
+    XLA_ASSIGN_OR_RETURN(client, PjRtComputationClient::Create());
   }
-  return new PjRtComputationClient();
+
+  return client.release();
 }
 
 const absl::StatusOr<ComputationClient * absl_nonnull>& GetComputationClient() {