From 1cfbcb14fcb4d7da500f13bd8771a17d13575b5b Mon Sep 17 00:00:00 2001
From: danleifeng <52735331+danleifeng@users.noreply.github.com>
Date: Tue, 22 Jun 2021 20:26:00 +0800
Subject: [PATCH] fix exit bug for heterps;test=develop (#33724)

---
 paddle/fluid/framework/fleet/ps_gpu_wrapper.h | 9 +++++++++
 paddle/fluid/pybind/ps_gpu_wrapper_py.cc      | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/paddle/fluid/framework/fleet/ps_gpu_wrapper.h b/paddle/fluid/framework/fleet/ps_gpu_wrapper.h
index cfb23d1be2a..81b2b0a12b2 100644
--- a/paddle/fluid/framework/fleet/ps_gpu_wrapper.h
+++ b/paddle/fluid/framework/fleet/ps_gpu_wrapper.h
@@ -209,6 +209,15 @@ class PSGPUWrapper {
   void EndPass() { HeterPs_->end_pass(); }
   void ShowOneTable(int index) { HeterPs_->show_one_table(index); }
 
+  void Finalize() {
+    VLOG(3) << "PSGPUWrapper Begin Finalize.";
+    if (s_instance_ == nullptr) {
+      return;
+    }
+    s_instance_ = nullptr;
+    VLOG(3) << "PSGPUWrapper Finalize Finished.";
+  }
+
  private:
   static std::shared_ptr<PSGPUWrapper> s_instance_;
   Dataset* dataset_;
diff --git a/paddle/fluid/pybind/ps_gpu_wrapper_py.cc b/paddle/fluid/pybind/ps_gpu_wrapper_py.cc
index 0c239f8157e..bdd7abe1d83 100644
--- a/paddle/fluid/pybind/ps_gpu_wrapper_py.cc
+++ b/paddle/fluid/pybind/ps_gpu_wrapper_py.cc
@@ -48,6 +48,8 @@ void BindPSGPUWrapper(py::module* m) {
       .def("end_pass", &framework::PSGPUWrapper::EndPass,
            py::call_guard<py::gil_scoped_release>())
       .def("build_gpu_ps", &framework::PSGPUWrapper::BuildGPUPS,
+           py::call_guard<py::gil_scoped_release>())
+      .def("finalize", &framework::PSGPUWrapper::Finalize,
            py::call_guard<py::gil_scoped_release>());
 }  // end PSGPUWrapper
 #endif
-- 
GitLab