From 23c036d62192bc76c7657ac87ab75ea682760eb5 Mon Sep 17 00:00:00 2001
From: ronnywang <524019753@qq.com>
Date: Wed, 16 Mar 2022 17:33:56 +0800
Subject: [PATCH] clean up DeviceManager in advance manually (#40504)

---
 paddle/fluid/pybind/pybind.cc         | 6 ++++++
 paddle/phi/backends/device_manager.cc | 5 +++++
 paddle/phi/backends/device_manager.h  | 2 ++
 python/paddle/fluid/__init__.py       | 2 ++
 4 files changed, 15 insertions(+)

diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
index 21bbc7f3e36..ed42d0792ea 100644
--- a/paddle/fluid/pybind/pybind.cc
+++ b/paddle/fluid/pybind/pybind.cc
@@ -114,6 +114,7 @@ limitations under the License. */
 #include "paddle/fluid/pybind/metrics_py.h"
 #include "paddle/fluid/pybind/ps_gpu_wrapper_py.h"
 #include "paddle/fluid/pybind/pybind_boost_headers.h"
+#include "paddle/phi/backends/device_manager.h"
 
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
 #include "paddle/fluid/pybind/nccl_wrapper_py.h"
@@ -742,6 +743,11 @@ PYBIND11_MODULE(core_noavx, m) {
   // stored in this static instance to avoid illegal memory access.
   m.def("clear_kernel_factory",
         []() { phi::KernelFactory::Instance().kernels().clear(); });
+  m.def("clear_device_manager", []() {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+    phi::DeviceManager::Clear();
+#endif
+  });
 
   // NOTE(zjl): ctest would load environment variables at the beginning even
   // though we have not `import paddle.fluid as fluid`. So we add this API
diff --git a/paddle/phi/backends/device_manager.cc b/paddle/phi/backends/device_manager.cc
index 1ffe38d8e1f..35339aed0f3 100644
--- a/paddle/phi/backends/device_manager.cc
+++ b/paddle/phi/backends/device_manager.cc
@@ -393,6 +393,11 @@ DeviceManager& DeviceManager::Instance() {
   return platform_manager;
 }
 
+void DeviceManager::Clear() {
+  Instance().device_map_.clear();
+  Instance().device_impl_map_.clear();
+}
+
 std::vector<std::string> ListAllLibraries(const std::string& library_dir) {
   std::vector<std::string> libraries;
   std::regex express(".*\\.so");
diff --git a/paddle/phi/backends/device_manager.h b/paddle/phi/backends/device_manager.h
index c0911a0f8d5..39eef27b4a6 100644
--- a/paddle/phi/backends/device_manager.h
+++ b/paddle/phi/backends/device_manager.h
@@ -158,6 +158,8 @@ class DeviceManager {
 
   static std::vector<size_t> GetDeviceList(const std::string& device_type);
 
+  static void Clear();
+
  private:
   DISABLE_COPY_AND_ASSIGN(DeviceManager);
   DeviceManager() {}
diff --git a/python/paddle/fluid/__init__.py b/python/paddle/fluid/__init__.py
index 7480909a2d8..fb9e8d8ece1 100644
--- a/python/paddle/fluid/__init__.py
+++ b/python/paddle/fluid/__init__.py
@@ -228,3 +228,5 @@ if core.is_compiled_with_npu():
 atexit.register(core.clear_executor_cache)
 # NOTE(Aganlengzi): clean up KernelFactory in advance manually.
 atexit.register(core.clear_kernel_factory)
+# NOTE(wangran16): clean up DeviceManger in advance manually.
+atexit.register(core.clear_device_manager)
-- 
GitLab