From f1c48f85bfb8f32194dc44b298a3f3539ad2afce Mon Sep 17 00:00:00 2001
From: Thunderbrook <52529258+Thunderbrook@users.noreply.github.com>
Date: Thu, 2 Dec 2021 21:15:57 +0800
Subject: [PATCH] release dataset (#37790)

---
 paddle/fluid/framework/data_set.cc | 5 +++++
 paddle/fluid/framework/data_set.h  | 9 ++++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/framework/data_set.cc b/paddle/fluid/framework/data_set.cc
index c511526c315..1b5db838051 100644
--- a/paddle/fluid/framework/data_set.cc
+++ b/paddle/fluid/framework/data_set.cc
@@ -463,6 +463,11 @@ void DatasetImpl<T>::WaitPreLoadDone() {
 // release memory data
 template <typename T>
 void DatasetImpl<T>::ReleaseMemory() {
+  release_thread_ = new std::thread(&DatasetImpl<T>::ReleaseMemoryFun, this);
+}
+
+template <typename T>
+void DatasetImpl<T>::ReleaseMemoryFun() {
   VLOG(3) << "DatasetImpl<T>::ReleaseMemory() begin";
   if (input_channel_) {
     input_channel_->Clear();
diff --git a/paddle/fluid/framework/data_set.h b/paddle/fluid/framework/data_set.h
index b41f701548f..58223a2f28b 100644
--- a/paddle/fluid/framework/data_set.h
+++ b/paddle/fluid/framework/data_set.h
@@ -63,6 +63,7 @@ class Dataset {
   virtual void SetTrainerNum(int trainer_num) = 0;
   // set fleet send batch size
   virtual void SetFleetSendBatchSize(int64_t size) = 0;
+  virtual void ReleaseMemoryFun() = 0;
   // set fs name and ugi
   virtual void SetHdfsConfig(const std::string& fs_name,
                              const std::string& fs_ugi) = 0;
@@ -168,8 +169,13 @@ template <typename T>
 class DatasetImpl : public Dataset {
  public:
   DatasetImpl();
-  virtual ~DatasetImpl() {}
+  virtual ~DatasetImpl() {
+    if (release_thread_ != nullptr) {
+      release_thread_->join();
+    }
+  }
   virtual void SetFileList(const std::vector<std::string>& filelist);
+  virtual void ReleaseMemoryFun();
   virtual void SetThreadNum(int thread_num);
   virtual void SetTrainerNum(int trainer_num);
   virtual void SetFleetSendBatchSize(int64_t size);
@@ -295,6 +301,7 @@ class DatasetImpl : public Dataset {
   int64_t fleet_send_batch_size_;
   int64_t fleet_send_sleep_seconds_;
   std::vector<std::thread> preload_threads_;
+  std::thread* release_thread_ = nullptr;
   bool merge_by_insid_;
   bool parse_ins_id_;
   bool parse_content_;
-- 
GitLab