Fix bad error message when PjRtComputationClient throws exception (#…

…5946)
pytorch · Dec 12, 2023 · 280ca1d · 280ca1d
1 parent fc713b7
commit 280ca1d
Showing 1 changed file with 15 additions and 21 deletions.
diff --git a/torch_xla/csrc/runtime/runtime.cc b/torch_xla/csrc/runtime/runtime.cc
@@ -8,35 +8,29 @@
 
 namespace torch_xla {
 namespace runtime {
-namespace {
 
 std::atomic<bool> g_computation_client_initialized(false);
 
-ComputationClient* CreateClient() {
-  bool was_initialized = g_computation_client_initialized.exchange(true);
-  XLA_CHECK(!was_initialized) << "ComputationClient already initialized";
-  if (sys_util::GetEnvBool("XLA_DUMP_FATAL_STACK", false)) {
-    tsl::testing::InstallStacktraceHandler();
-  }
-
-  ComputationClient* client;
+ComputationClient* GetComputationClient() {
+  static std::unique_ptr<ComputationClient> client = []() {
+    if (sys_util::GetEnvBool("XLA_DUMP_FATAL_STACK", false)) {
+      tsl::testing::InstallStacktraceHandler();
+    }
 
-  if (sys_util::GetEnvString(env::kEnvPjRtDevice, "") != "") {
-    client = new PjRtComputationClient();
-  } else {
-    g_computation_client_initialized = false;
-    XLA_ERROR() << "$PJRT_DEVICE is not set." << std::endl;
-  }
+    std::unique_ptr<ComputationClient> client;
 
-  XLA_CHECK(client != nullptr);
+    if (sys_util::GetEnvString(env::kEnvPjRtDevice, "") != "") {
+      client = std::make_unique<PjRtComputationClient>();
+    } else {
+      XLA_ERROR() << "$PJRT_DEVICE is not set." << std::endl;
+    }
 
-  return client;
-}
+    XLA_CHECK(client);
 
-}  // namespace
+    g_computation_client_initialized = true;
+    return client;
+  }();
 
-ComputationClient* GetComputationClient() {
-  static auto client = std::unique_ptr<ComputationClient>(CreateClient());
   return client.get();
 }