Delete ComputationClient in PrepareToExit

torch_xla/csrc/init_python_bindings.cpp

@@ -97,6 +97,7 @@ void PrepareToExit() {
   if (client != nullptr) {
     XLAGraphExecutor::Get()->WaitDeviceOps({});
     client->PrepareToExit();
+    runtime::DeleteComputationClient();
   }
 }
 

torch_xla/csrc/runtime/pjrt_computation_client.cc

@@ -189,6 +189,8 @@ PjRtComputationClient::PjRtComputationClient() {
 }
 
 PjRtComputationClient::~PjRtComputationClient() {
+  // In the GPU case, the PjRtClient depends on the DistributedRuntimeClient
+  // tracked in XlaCoordinator, so the PjRtClient must be destroyed first.
   client_ = nullptr;
   coordinator_ = nullptr;
 }

torch_xla/csrc/runtime/runtime.cc

@@ -1,7 +1,5 @@
 #include <torch/csrc/lazy/backend/backend_device.h>
 
-#include <cstdlib>
-
 #include "torch_xla/csrc/device.h"
 #include "torch_xla/csrc/runtime/computation_client.h"
 #include "torch_xla/csrc/runtime/env_vars.h"
@@ -35,11 +33,17 @@ ComputationClient* CreateClient() {
 
 }  // namespace
 
+void DeleteComputationClient() {
+  ComputationClient* client = g_computation_client.load();
+  if (client != nullptr) {
+    g_computation_client = nullptr;
+    delete client;
+  }
+}
+
 ComputationClient* GetComputationClient() {
-  std::call_once(g_computation_client_once, [&]() {
-    g_computation_client = std::move(CreateClient());
-    std::atexit([]() { delete g_computation_client.load(); });
-  });
+  std::call_once(g_computation_client_once,
+                 [&]() { g_computation_client = std::move(CreateClient()); });
   return g_computation_client.load();
 }
 

torch_xla/csrc/runtime/runtime.h

@@ -11,6 +11,10 @@ ComputationClient* GetComputationClient();
 
 ComputationClient* GetComputationClientIfInitialized();
 
+// Delete the ComputationClient. This should be called when the program is about
+// to exit.
+void DeleteComputationClient();
+
 // Run the XRT local service, this will block the caller unitl the server
 // being stopped.
 void RunLocalService(uint64_t service_port);