Unverified Commit 8f433feb by Andrew Reusch Committed by GitHub

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'
parent 3ab37512
...@@ -15,6 +15,7 @@ ...@@ -15,6 +15,7 @@
# specific language governing permissions and limitations # specific language governing permissions and limitations
# under the License. # under the License.
"""Common system utilities""" """Common system utilities"""
import atexit
import os import os
import tempfile import tempfile
import shutil import shutil
...@@ -23,27 +24,46 @@ try: ...@@ -23,27 +24,46 @@ try:
except ImportError: except ImportError:
fcntl = None fcntl = None
class TempDirectory(object): class TempDirectory(object):
"""Helper object to manage temp directory during testing. """Helper object to manage temp directory during testing.
Automatically removes the directory when it went out of scope. 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): def __init__(self, custom_path=None):
if custom_path: if custom_path:
os.mkdir(custom_path) os.mkdir(custom_path)
self.temp_dir = custom_path self.temp_dir = custom_path
else: else:
self.temp_dir = tempfile.mkdtemp() self.temp_dir = tempfile.mkdtemp()
self._rmtree = shutil.rmtree
self.TEMPDIRS.add(self.temp_dir)
def remove(self): def remove(self):
"""Remote the tmp dir""" """Remote the tmp dir"""
if self.temp_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 self.temp_dir = None
def __del__(self): 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() self.remove()
def relpath(self, name): def relpath(self, name):
...@@ -72,6 +92,9 @@ class TempDirectory(object): ...@@ -72,6 +92,9 @@ class TempDirectory(object):
return os.listdir(self.temp_dir) return os.listdir(self.temp_dir)
atexit.register(TempDirectory.remove_tempdirs)
def tempdir(custom_path=None): def tempdir(custom_path=None):
"""Create temp dir which deletes the contents when exit. """Create temp dir which deletes the contents when exit.
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment