Destroy the ComputationClient when the program exits (#5750)

* Destroy the ComputationClient when the program exits * Fix extra error when PJRT_DEVICE is not set
pytorch · Jan 12, 2024 · c46da91 · c46da91
1 parent 0d83c1b
commit c46da91
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 11 deletions.
diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -95,7 +95,6 @@ void PrepareToExit() {
       runtime::GetComputationClientIfInitialized();
   if (client != nullptr) {
     XLAGraphExecutor::Get()->WaitDeviceOps({});
-    client->PrepareToExit();
   }
 }
 

diff --git a/torch_xla/csrc/runtime/computation_client.h b/torch_xla/csrc/runtime/computation_client.h
@@ -344,8 +344,6 @@ class ComputationClient {
 
   virtual MemoryInfo GetMemoryInfo(const std::string& device) = 0;
 
-  virtual void PrepareToExit() = 0;
-
   // Block until pass in devices' async operation are finished. If empty, all
   // the local devices will be waited for.
   virtual void WaitDeviceOps(const std::vector<std::string>& devices) = 0;

diff --git a/torch_xla/csrc/runtime/pjrt_computation_client.h b/torch_xla/csrc/runtime/pjrt_computation_client.h
@@ -85,8 +85,6 @@ class PjRtComputationClient : public ComputationClient {
 
   std::shared_ptr<std::vector<std::string>> GetReplicationDevices() override;
 
-  void PrepareToExit() override { return; };
-
   void WaitDeviceOps(const std::vector<std::string>& devices) override;
 
   std::map<std::string, Metric> GetMetrics() const override;

diff --git a/torch_xla/csrc/runtime/runtime.cc b/torch_xla/csrc/runtime/runtime.cc
@@ -10,10 +10,11 @@ namespace torch_xla {
 namespace runtime {
 namespace {
 
-std::atomic<ComputationClient*> g_computation_client(nullptr);
-std::once_flag g_computation_client_once;
+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();
   }
@@ -23,6 +24,7 @@ ComputationClient* CreateClient() {
   if (sys_util::GetEnvString(env::kEnvPjRtDevice, "") != "") {
     client = new PjRtComputationClient();
   } else {
+    g_computation_client_initialized = false;
     XLA_ERROR() << "$PJRT_DEVICE is not set." << std::endl;
   }
 
@@ -34,13 +36,12 @@ ComputationClient* CreateClient() {
 }  // namespace
 
 ComputationClient* GetComputationClient() {
-  std::call_once(g_computation_client_once,
-                 [&]() { g_computation_client = std::move(CreateClient()); });
-  return g_computation_client.load();
+  static auto client = std::unique_ptr<ComputationClient>(CreateClient());
+  return client.get();
 }
 
 ComputationClient* GetComputationClientIfInitialized() {
-  return g_computation_client.load();
+  return g_computation_client_initialized ? GetComputationClient() : nullptr;
 }
 
 }  // namespace runtime