From 8f433febea4d2e16ab15aaf949c355dd85055de5 Mon Sep 17 00:00:00 2001
From: Andrew Reusch <areusch@octoml.ai>
Date: Wed, 22 Apr 2020 18:30:11 -0700
Subject: [PATCH] Don't remove() TempDirectory in __del__ after atexit hook
 runs. (#5414)

* Use atexit to remove TempDirectory before interpreter shutdown.
 * Can't rely on complex functions from __del__ anyway.
 * Fixes warning message on my box:
       Exception ignored in: <function TempDirectory.__del__ at 0x12be10680>
       Traceback (most recent call last):
        File ".../tvm/python/tvm/contrib/util.py", line 55, in __del__
        File ".../tvm/python/tvm/contrib/util.py", line 51, in remove
        File "/usr/local/opt/python/Frameworks/Python.framework/Versions/3.7/lib/python3.7/shutil.py", line 509, in rmtree
        AttributeError: 'NoneType' object has no attribute 'path'
---
 python/tvm/contrib/util.py | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/python/tvm/contrib/util.py b/python/tvm/contrib/util.py
index 2ebe175e8160..e980e5520802 100644
--- a/python/tvm/contrib/util.py
+++ b/python/tvm/contrib/util.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Common system utilities"""
+import atexit
 import os
 import tempfile
 import shutil
@@ -23,27 +24,46 @@
 except ImportError:
     fcntl = None
 
-
 class TempDirectory(object):
     """Helper object to manage temp directory during testing.
 
     Automatically removes the directory when it went out of scope.
     """
+
+    TEMPDIRS = set()
+    @classmethod
+    def remove_tempdirs(cls):
+        temp_dirs = getattr(cls, 'TEMPDIRS', None)
+        if temp_dirs is None:
+            return
+
+        for path in temp_dirs:
+            shutil.rmtree(path, ignore_errors=True)
+
+        cls.TEMPDIRS = None
+
     def __init__(self, custom_path=None):
         if custom_path:
             os.mkdir(custom_path)
             self.temp_dir = custom_path
         else:
             self.temp_dir = tempfile.mkdtemp()
-        self._rmtree = shutil.rmtree
+
+        self.TEMPDIRS.add(self.temp_dir)
 
     def remove(self):
         """Remote the tmp dir"""
         if self.temp_dir:
-            self._rmtree(self.temp_dir, ignore_errors=True)
+            shutil.rmtree(self.temp_dir, ignore_errors=True)
+            self.TEMPDIRS.remove(self.temp_dir)
             self.temp_dir = None
 
     def __del__(self):
+        temp_dirs = getattr(self, 'TEMPDIRS', None)
+        if temp_dirs is None:
+            # Do nothing if the atexit hook has already run.
+            return
+
         self.remove()
 
     def relpath(self, name):
@@ -72,6 +92,9 @@ def listdir(self):
         return os.listdir(self.temp_dir)
 
 
+atexit.register(TempDirectory.remove_tempdirs)
+
+
 def tempdir(custom_path=None):
     """Create temp dir which deletes the contents when exit.