From cc9ded2081ddfcc0a6e05dc3386139c1fb93fa82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Thu, 11 Mar 2021 14:54:36 +0900 Subject: [PATCH] exp: subclass tempfile.TemporaryDirectory to handle cleanup (#5598) * exp: subclass tempfile.TemporaryDirectory to handle cleanup * override cleanup() since cls._rmtree does not exist in <3.8 --- dvc/repo/experiments/executor/local.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index f584799a0..f2b6ef358 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -3,13 +3,29 @@ import os from tempfile import TemporaryDirectory from typing import Optional -from dvc.utils.fs import remove - from .base import BaseExecutor logger = logging.getLogger(__name__) +class ExpTemporaryDirectory(TemporaryDirectory): + # Python's TemporaryDirectory cleanup shutil.rmtree wrapper does not handle + # git read-only dirs cleanly in Windows on Python <3.8, so we use our own + # remove(). See: + # https://github.com/iterative/dvc/pull/5425 + # https://bugs.python.org/issue26660 + + @classmethod + def _rmtree(cls, name): + from dvc.utils.fs import remove + + remove(name) + + def cleanup(self): + if self._finalizer.detach(): + self._rmtree(self.name) + + class BaseLocalExecutor(BaseExecutor): """Base local machine executor.""" @@ -37,7 +53,7 @@ class TempDirExecutor(BaseLocalExecutor): cache_dir: Optional[str] = None, **kwargs, ): - self._tmp_dir = TemporaryDirectory(dir=tmp_dir) + self._tmp_dir = ExpTemporaryDirectory(dir=tmp_dir) kwargs["root_dir"] = self._tmp_dir.name super().__init__(*args, **kwargs) if cache_dir: @@ -55,4 +71,4 @@ class TempDirExecutor(BaseLocalExecutor): def cleanup(self): super().cleanup() logger.debug("Removing tmpdir '%s'", self._tmp_dir) - remove(self._tmp_dir.name) + self._tmp_dir.cleanup()