Clear empty thread when graph destroy

.github/workflows/test.yml

@@ -503,7 +503,7 @@ jobs:
         working-directory: ${{ env.ONEFLOW_SRC }}
         run: |
           docker run -d --rm --privileged --shm-size=8g \
-            --pids-limit -1 \
+            --pids-limit 5000 \
             --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
             --runtime=nvidia \
             -v /dataset:/dataset:ro -v /model_zoo:/model_zoo:ro \

oneflow/core/framework/multi_client_session_context.cpp

@@ -135,13 +135,22 @@ Maybe<void> MultiClientSessionContext::TryClose() {
   if (is_inited_) {
     VLOG(2) << "Try to delete multi client session context." << std::endl;
 
+    for (auto wk_graph_ptr : graphs_) {
+      if (auto sh_graph_ptr = wk_graph_ptr.lock()) {
+        VLOG(2) << "grap name " << sh_graph_ptr->job_name() << " not closed, try to close it.";
+        JUST(sh_graph_ptr->Close());
+      }
+    }
+
+    /*
     // sync before NNGraph release to ensure LaunchLazyJob instruction was completed and released
     JUST(vm::ClusterSync());
     for (const auto& graph : graphs_) {
       VLOG(2) << "Try to close graph: " << graph->job_name() << std::endl;
       JUST(graph->Close());
     }
     graphs_.clear();
+    */
     {
       // NOTE(chengcheng): delete runtime global objects
       Global<boxing::collective::Scheduler>::Delete();

oneflow/core/framework/multi_client_session_context.h

@@ -51,7 +51,7 @@ class MultiClientSessionContext {
   bool is_inited_;
   HashMap<std::string, std::vector<std::pair<std::string, std::shared_ptr<one::Tensor>>>>
       graph_name2free_eager_tensors_;
-  std::vector<std::shared_ptr<NNGraph>> graphs_;
+  std::vector<std::weak_ptr<NNGraph>> graphs_;
 };
 
 }  // namespace oneflow

oneflow/core/job/runtime.cpp

@@ -58,10 +58,11 @@ bool HasNonCtrlConsumedRegstDescId(const TaskProto& task) {
 }  // namespace
 
 Runtime::Runtime(const Plan& plan, const HashMap<std::string, Blob*>& variable_op_name2eager_blob) {
+  DumpThreadIdsFromPlan(plan);
   {
     // NOTE(chengcheng): All runtime Global objects AddPlan
     Global<RegstMgr>::Get()->AddPlan(plan, variable_op_name2eager_blob);
-    Global<ThreadMgr>::Get()->AddPlan(plan);
+    Global<ThreadMgr>::Get()->AddThreads(thread_ids_);
     Global<RuntimeJobDescs>::Get()->AddPlan(plan);
     collective_boxing_scheduler_plan_token_ =
         Global<boxing::collective::Scheduler>::Get()->AddPlan(plan);
@@ -106,7 +107,19 @@ Runtime::~Runtime() {
     Global<RuntimeCtx>::Get()->WaitUntilCntEqualZero(GetRunningActorCountKeyByJobId(pair.first));
   }
   OF_SESSION_BARRIER();
+  Global<ThreadMgr>::Get()->TryDeleteThreads(thread_ids_);
   Global<boxing::collective::Scheduler>::Get()->DeletePlan(collective_boxing_scheduler_plan_token_);
 }
 
+void Runtime::DumpThreadIdsFromPlan(const Plan& plan) {
+  const int64_t this_rank = GlobalProcessCtx::Rank();
+  for (const TaskProto& task : plan.task()) {
+    TaskId task_id = DecodeTaskIdFromInt64(task.task_id());
+    StreamId stream_id = task_id.stream_id();
+    if (stream_id.rank() != this_rank) { continue; }
+    int64_t thrd_id = EncodeStreamIdToInt64(stream_id);
+    thread_ids_.insert(thrd_id);
+  }
+}
+
 }  // namespace oneflow

oneflow/core/job/runtime.h

@@ -33,7 +33,10 @@ class Runtime final {
   Runtime(const Plan& plan, const HashMap<std::string, Blob*>& variable_op_name2eager_blob);
 
  private:
+  void DumpThreadIdsFromPlan(const Plan& plan);
+
   HashMap<int64_t, int64_t> job_id2actor_size_;
+  HashSet<int64_t> thread_ids_;
 
   boxing::collective::SchedulerPlanToken* collective_boxing_scheduler_plan_token_;
 };

oneflow/core/thread/thread.cpp

@@ -51,6 +51,8 @@ void Thread::AddTask(const TaskProto& task) {
   CHECK(id2task_.emplace(task.task_id(), task).second);
 }
 
+bool Thread::Empty() const { return id2actor_ptr_.empty(); }
+
 void Thread::PollMsgChannel() {
   while (true) {
     if (local_msg_queue_.empty()) {

oneflow/core/thread/thread.h

@@ -34,6 +34,7 @@ class Thread {
   virtual ~Thread();
 
   void AddTask(const TaskProto&);
+  bool Empty() const;
 
   Channel<ActorMsg>* GetMsgChannelPtr() { return &msg_channel_; }
 

oneflow/core/thread/thread_manager.cpp

@@ -17,16 +17,13 @@ limitations under the License.
 #include "oneflow/core/job/resource_desc.h"
 #include "oneflow/core/job/global_for.h"
 #include "oneflow/core/control/global_process_ctx.h"
-#include "oneflow/core/job/global_for.h"
 
 namespace oneflow {
 
 ThreadMgr::~ThreadMgr() {
   for (auto& thread_pair : threads_) {
-    ActorMsg msg = ActorMsg::BuildCommandMsg(-1, ActorCmd::kStopThread);
-    thread_pair.second->GetMsgChannelPtr()->Send(msg);
-    thread_pair.second.reset();
-    LOG(INFO) << "actor thread " << thread_pair.first << " finish";
+    LOG(FATAL) << " Runtime Error! thread id " << thread_pair.first
+               << " not delete with graph, and it's empty is " << thread_pair.second->Empty();
   }
 }
 
@@ -36,20 +33,42 @@ Thread* ThreadMgr::GetThrd(int64_t thrd_id) {
   return iter->second.get();
 }
 
-void ThreadMgr::AddPlan(const Plan& plan) {
+void ThreadMgr::AddThreads(const HashSet<int64_t>& thread_ids) {
   const int64_t this_rank = GlobalProcessCtx::Rank();
-  for (const TaskProto& task : plan.task()) {
-    TaskId task_id = DecodeTaskIdFromInt64(task.task_id());
-    StreamId stream_id = task_id.stream_id();
+  for (int64_t thrd_id : thread_ids) {
+    const auto& it = threads_.find(thrd_id);
+    if (it != threads_.end()) {
+      // NOTE(chengcheng): check thread is not null.
+      CHECK(it->second);
+      continue;
+    }
+    StreamId stream_id = DecodeStreamIdFromInt64(thrd_id);
     if (stream_id.rank() != this_rank) { continue; }
-    int64_t thrd_id = EncodeStreamIdToInt64(stream_id);
-    if (threads_.find(thrd_id) != threads_.end()) { continue; }
     Thread* thread = new Thread(stream_id);
     CHECK_NOTNULL(thread);
     threads_[thrd_id].reset(thread);
   }
 }
 
+void ThreadMgr::TryDeleteThreads(const HashSet<int64_t>& thread_ids) {
+  for (int64_t thrd_id : thread_ids) {
+    const auto& it = threads_.find(thrd_id);
+    if (it == threads_.end()) { continue; }
+    auto& thread = it->second;
+    CHECK(thread) << " actor thread " << thrd_id << " non-existent but want to delete";
+    if (thread->Empty()) {
+      // NOTE(chengcheng):  Only delete thread when it is empty.
+      ActorMsg msg = ActorMsg::BuildCommandMsg(-1, ActorCmd::kStopThread);
+      thread->GetMsgChannelPtr()->Send(msg);
+      thread.reset();
+      LOG(INFO) << " actor thread " << thrd_id << " finish.";
+      threads_.erase(it);
+    } else {
+      LOG(INFO) << " actor thread " << thrd_id << " not delete because it is not empty.";
+    }
+  }
+}
+
 void SingleThreadLoop(size_t num, std::function<void(size_t i)> Callback) {
   FOR_RANGE(size_t, i, 0, num) { Callback(i); }
 }

oneflow/core/thread/thread_manager.h

@@ -36,7 +36,8 @@ class ThreadMgr final {
   ThreadMgr() = default;
   ~ThreadMgr();
 
-  void AddPlan(const Plan& plan);
+  void AddThreads(const HashSet<int64_t>& thread_ids);
+  void TryDeleteThreads(const HashSet<int64_t>& thread_ids);
   Thread* GetThrd(int64_t thrd_id);
 
  private: