Fix deadlock in instruction done

oneflow/core/vm/instruction.cpp

@@ -41,16 +41,21 @@ void Instruction::InitStatus() { instruction_policy_->InitInstructionStatusIf(th
 Maybe<void> Instruction::Prepare() { return instruction_policy_->PrepareIf(this); }
 void Instruction::Compute() { return instruction_policy_->ComputeIf(this); }
 
-void Instruction::DeleteStatusAndClearEdges() {
-  OF_PROFILER_RANGE_GUARD("Instruction::DeleteStatusAndClearEdges");
+void Instruction::DeleteStatusAndCheckEdges() {
+  OF_PROFILER_RANGE_GUARD("Instruction::DeleteStatusAndCheckEdges");
   instruction_policy_->DeleteInstructionStatusIf(this);
-  mut_in_edges()->Clear();
-  mut_out_edges()->Clear();
+  INTRUSIVE_FOR_EACH_PTR(edge, mut_in_edges()) {
+    Instruction* in_instruction = edge->mut_src_instruction();
+    LOG(FATAL) << "unerased edge: " << in_instruction->DebugName() << " -> " << this->DebugName();
+  }
+  INTRUSIVE_FOR_EACH_PTR(edge, mut_out_edges()) {
+    Instruction* out_instruction = edge->mut_dst_instruction();
+    LOG(FATAL) << "unerased edge: " << this->DebugName() << " -> " << out_instruction->DebugName();
+  }
 }
 
 bool Instruction::Done() const {
-  return stream_policy().QueryInstructionStatusDone(stream(), status_buffer())
-         && in_edges().empty();
+  return stream_policy().QueryInstructionStatusDone(stream(), status_buffer());
 }
 
 StreamPolicy* Instruction::mut_stream_policy() { return mut_stream()->mut_stream_policy(); }

oneflow/core/vm/instruction.h

@@ -138,7 +138,7 @@ class Instruction final : public intrusive::Base {
 
   // methods
   void InitStatus();
-  void DeleteStatusAndClearEdges();
+  void DeleteStatusAndCheckEdges();
   bool Done() const;
   StreamPolicy* mut_stream_policy();
   const StreamPolicy& stream_policy() const;

oneflow/core/vm/virtual_machine_engine.cpp

@@ -177,13 +177,14 @@ void VirtualMachineEngine::ReleaseFinishedInstructions(const ScheduleCtx& schedu
   INTRUSIVE_FOR_EACH_PTR(stream, mut_active_stream_list()) {
     while (true) {
       auto* instruction_ptr = stream->mut_running_instruction_list()->Begin();
-      if (instruction_ptr == nullptr || !instruction_ptr->Done()) { break; }
+      if (instruction_ptr == nullptr) { break; }
+      if (!(instruction_ptr->in_edges().empty() && instruction_ptr->Done())) { break; }
       ReleaseInstruction(instruction_ptr);
       // Prevent destructing instruction_ptr.
       intrusive::shared_ptr<Instruction> instruction =
           stream->mut_running_instruction_list()->Erase(instruction_ptr);
       LivelyInstructionListErase(instruction_ptr);
-      instruction_ptr->DeleteStatusAndClearEdges();
+      instruction_ptr->DeleteStatusAndCheckEdges();
     }
     if (stream->running_instruction_list().empty()) { mut_active_stream_list()->Erase(stream); }
   }