[Backport] Fix bad error message when PjRtComputationClient throws …

…exception (#6144)

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();
 }