Add debug mode to tempdir()

python/tvm/contrib/util.py

@@ -16,8 +16,10 @@
 # under the License.
 """Common system utilities"""
 import atexit
+import datetime
 import os
 import tempfile
+import threading
 import shutil
 try:
     import fcntl
@@ -30,6 +32,32 @@ class TempDirectory(object):
     Automatically removes the directory when it went out of scope.
     """
 
+    # When True, all TempDirectory are *NOT* deleted and instead live inside a predicable directory
+    # tree.
+    DEBUG_MODE = False
+
+    # In debug mode, each tempdir is named after the sequence
+    _NUM_TEMPDIR_CREATED = 0
+    _NUM_TEMPDIR_CREATED_LOCK = threading.Lock()
+    @classmethod
+    def _increment_num_tempdir_created(cls):
+        with cls._NUM_TEMPDIR_CREATED_LOCK:
+            to_return = cls._NUM_TEMPDIR_CREATED
+            cls._NUM_TEMPDIR_CREATED += 1
+
+        return to_return
+
+    _DEBUG_PARENT_DIR = None
+    @classmethod
+    def _get_debug_parent_dir(cls):
+        if cls._DEBUG_PARENT_DIR is None:
+            all_parents = f'{tempfile.gettempdir()}/tvm-debug-mode-tempdirs'
+            if not os.path.isdir(all_parents):
+                os.makedirs(all_parents)
+            cls._DEBUG_PARENT_DIR = tempfile.mkdtemp(
+                prefix=datetime.datetime.now().strftime('%Y-%m-%dT%H-%M-%S___'), dir=all_parents)
+        return cls._DEBUG_PARENT_DIR
+
     TEMPDIRS = set()
     @classmethod
     def remove_tempdirs(cls):
@@ -43,19 +71,27 @@ def remove_tempdirs(cls):
         cls.TEMPDIRS = None
 
     def __init__(self, custom_path=None):
+        self._created_with_debug_mode = self.DEBUG_MODE
         if custom_path:
             os.mkdir(custom_path)
             self.temp_dir = custom_path
         else:
-            self.temp_dir = tempfile.mkdtemp()
+            if self._created_with_debug_mode:
+                parent_dir = self._get_debug_parent_dir()
+                self.temp_dir = f'{parent_dir}/{self._increment_num_tempdir_created():05d}'
+                os.mkdir(self.temp_dir)
+            else:
+                self.temp_dir = tempfile.mkdtemp()
 
-        self.TEMPDIRS.add(self.temp_dir)
+        if not self._created_with_debug_mode:
+            self.TEMPDIRS.add(self.temp_dir)
 
     def remove(self):
         """Remote the tmp dir"""
         if self.temp_dir:
-            shutil.rmtree(self.temp_dir, ignore_errors=True)
-            self.TEMPDIRS.remove(self.temp_dir)
+            if not self._created_with_debug_mode:
+                shutil.rmtree(self.temp_dir, ignore_errors=True)
+                self.TEMPDIRS.remove(self.temp_dir)
             self.temp_dir = None
 
     def __del__(self):

tests/python/contrib/test_util.py

@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for functions in tvm/python/tvm/contrib/util.py."""
+
+import datetime
+import os
+import shutil
+from tvm.contrib import util
+
+
+def test_tempdir():
+  assert util.TempDirectory.DEBUG_MODE == False, "don't submit with DEBUG_MODE == True"
+
+  temp_dir = util.tempdir()
+  assert os.path.exists(temp_dir.temp_dir)
+
+  old_debug_mode = util.TempDirectory.DEBUG_MODE
+  try:
+    util.TempDirectory.DEBUG_MODE = True
+
+    for temp_dir_number in range(0, 3):
+      debug_temp_dir = util.tempdir()
+      try:
+        dirname, basename = os.path.split(debug_temp_dir.temp_dir)
+        assert basename == ('0000' + str(temp_dir_number)), 'unexpected basename: %s' % (basename,)
+
+        parent_dir = os.path.basename(dirname)
+        create_time = datetime.datetime.strptime(parent_dir.split('___', 1)[0], '%Y-%m-%dT%H-%M-%S')
+        assert abs(datetime.datetime.now() - create_time) < datetime.timedelta(seconds=60)
+
+      finally:
+        shutil.rmtree(debug_temp_dir.temp_dir)
+
+  finally:
+    util.TempDirectory.DEBUG_MODE = old_debug_mode
+
+
+if __name__ == '__main__':
+  test_tempdir()