Skip to content

Commit

Permalink
Don't remove() TempDirectory in __del__ after atexit hook runs. (apac…
Browse files Browse the repository at this point in the history
…he#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'
  • Loading branch information
areusch authored and trevor-m committed Jun 18, 2020
1 parent 9251efd commit 0a57ae6
Showing 1 changed file with 26 additions and 3 deletions.
29 changes: 26 additions & 3 deletions python/tvm/contrib/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Common system utilities"""
import atexit
import os
import tempfile
import shutil
Expand All @@ -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):
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 0a57ae6

Please sign in to comment.