From e0ee6c25bbef01331c91b2cb3523f96a4c5f15e1 Mon Sep 17 00:00:00 2001
From: barry-jin <barryjin1995@gmail.com>
Date: Tue, 27 Oct 2020 10:00:13 -0700
Subject: [PATCH] Revert "Remove cleanup on side threads (#19378)"

This reverts commit 43750c8bfed6ca91fc47fd1fa6d620197e26c84c.
---
 src/engine/naive_engine.cc              | 3 +++
 src/engine/threaded_engine_perdevice.cc | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/src/engine/naive_engine.cc b/src/engine/naive_engine.cc
index 95528a934c6c..9d737143994a 100644
--- a/src/engine/naive_engine.cc
+++ b/src/engine/naive_engine.cc
@@ -74,11 +74,14 @@ class NaiveEngine final : public Engine {
     LOG(INFO) << "Engine shutdown";
     for (size_t i = 0; i < streams_.size(); ++i) {
       if (streams_[i] != nullptr) {
+        // Catch exception for CUDA driver shutdown
+        MSHADOW_CATCH_ERROR(mshadow::DeleteStream(streams_[i]));
         streams_[i] = nullptr;
       }
     }
     for (size_t i = 0; i < aux_streams_.size(); ++i) {
       if (aux_streams_[i] != nullptr) {
+        delete aux_streams_[i];
         aux_streams_[i] = nullptr;
       }
     }
diff --git a/src/engine/threaded_engine_perdevice.cc b/src/engine/threaded_engine_perdevice.cc
index 4c5d1befb8b3..81494ec66096 100644
--- a/src/engine/threaded_engine_perdevice.cc
+++ b/src/engine/threaded_engine_perdevice.cc
@@ -273,6 +273,10 @@ class ThreadedEnginePerDevice : public ThreadedEngine {
     while (task_queue->Pop(&opr_block)) {
       this->ExecuteOprBlock(run_ctx, opr_block);
     }
+    // Catch exception for CUDA driver shutdown
+    MSHADOW_CATCH_ERROR(mshadow::DeleteStream<gpu>(stream));
+    if (aux_stream != nullptr)
+      delete aux_stream;
 #else
     ready_event->signal();
 #endif