From b3fb0dfdad82a347534e5c7d14a13aa1e2b90e46 Mon Sep 17 00:00:00 2001 From: skshetry Date: Thu, 12 May 2022 21:12:33 +0545 Subject: [PATCH] dvc.objects: extend diskcache to handle pickle errors (#7734) dvc.objects: extend diskcache to handle pickle errors --- dvc/cli/__init__.py | 60 +++++++++++------- dvc/data/db/index.py | 23 ++----- dvc/objects/cache.py | 61 +++++++++++++++++++ dvc/objects/state.py | 13 +--- dvc/utils/decorators.py | 30 --------- tests/unit/cli/__init__.py | 0 tests/unit/cli/test_main.py | 26 ++++++++ .../test_cache.py} | 32 +++++----- 8 files changed, 151 insertions(+), 94 deletions(-) create mode 100644 dvc/objects/cache.py delete mode 100644 dvc/utils/decorators.py create mode 100644 tests/unit/cli/__init__.py create mode 100644 tests/unit/cli/test_main.py rename tests/unit/{utils/test_decorators.py => objects/test_cache.py} (52%) diff --git a/dvc/cli/__init__.py b/dvc/cli/__init__.py index bcdb1e5e6..2a5cb9ce7 100644 --- a/dvc/cli/__init__.py +++ b/dvc/cli/__init__.py @@ -36,6 +36,43 @@ def parse_args(argv=None): return args +def _log_exceptions(exc: Exception): + """Try to log some known exceptions, that are not DVCExceptions.""" + from dvc.objects.cache import DiskError + from dvc.utils import error_link + + if isinstance(exc, DiskError): + from dvc.utils import relpath + + directory = relpath(exc.directory) + logger.exception( + f"Could not open pickled '{exc.type}' cache.\n" + f"Remove the '{directory}' directory and then retry this command." + f"\nSee {error_link('pickle')} for more information.", + extra={"tb_only": True}, + ) + return + + if isinstance(exc, OSError): + import errno + + if exc.errno == errno.EMFILE: + logger.exception( + "too many open files, please visit " + "{} to see how to handle this " + "problem".format(error_link("many-files")), + extra={"tb_only": True}, + ) + return + + from dvc.info import get_dvc_info + from dvc.logger import FOOTER + + logger.exception("unexpected error") + logger.debug("Version info for developers:\n%s", get_dvc_info()) + logger.info(FOOTER) + + def main(argv=None): # noqa: C901 """Main entry point for dvc CLI. @@ -48,7 +85,7 @@ def main(argv=None): # noqa: C901 from dvc._debug import debugtools from dvc.config import ConfigError from dvc.exceptions import DvcException, NotDvcRepoError - from dvc.logger import FOOTER, disable_other_loggers + from dvc.logger import disable_other_loggers # NOTE: stderr/stdout may be closed if we are running from dvc.daemon. # On Linux we directly call cli.main after double forking and closing @@ -104,26 +141,7 @@ def main(argv=None): # noqa: C901 ret = 254 except Exception as exc: # noqa, pylint: disable=broad-except # pylint: disable=no-member - import errno - - if isinstance(exc, OSError) and exc.errno == errno.EMFILE: - from dvc.utils import error_link - - logger.exception( - "too many open files, please visit " - "{} to see how to handle this " - "problem".format(error_link("many-files")), - extra={"tb_only": True}, - ) - else: - from dvc.info import get_dvc_info - - logger.exception("unexpected error") - - dvc_info = get_dvc_info() - logger.debug("Version info for developers:\n%s", dvc_info) - - logger.info(FOOTER) + _log_exceptions(exc) ret = 255 try: diff --git a/dvc/data/db/index.py b/dvc/data/db/index.py index 25f42285c..60c095c8c 100644 --- a/dvc/data/db/index.py +++ b/dvc/data/db/index.py @@ -4,7 +4,6 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Iterable, Set from dvc.objects.errors import ObjectDBError -from dvc.utils.decorators import with_diskcache if TYPE_CHECKING: from dvc.types import StrPath @@ -89,49 +88,40 @@ class ObjectDBIndex(ObjectDBIndexBase): tmp_dir: "StrPath", name: str, ): # pylint: disable=super-init-not-called - from diskcache import Cache, Index - from dvc.fs.local import LocalFileSystem + from dvc.objects.cache import Cache, Index from dvc.utils.fs import makedirs self.index_dir = os.path.join(tmp_dir, self.INDEX_DIR, name) makedirs(self.index_dir, exist_ok=True) self.fs = LocalFileSystem() - self.index = Index.fromcache( - Cache( - self.index_dir, - disk_pickle_protocol=4, - eviction_policy="none", - ) + cache = Cache( + self.index_dir, eviction_policy="none", disk_type="index" ) + self.index = Index.fromcache(cache) - @with_diskcache(name="index") def __iter__(self): return iter(self.index) - @with_diskcache(name="index") def __contains__(self, hash_): return hash_ in self.index - @with_diskcache(name="index") def dir_hashes(self): """Iterate over .dir hashes stored in the index.""" yield from (hash_ for hash_, is_dir in self.index.items() if is_dir) - @with_diskcache(name="index") def clear(self): """Clear this index (to force re-indexing later).""" - from diskcache import Timeout + from dvc.objects.cache import Timeout try: self.index.clear() except Timeout as exc: raise ObjectDBError("Failed to clear ODB index") from exc - @with_diskcache(name="index") def update(self, dir_hashes: Iterable[str], file_hashes: Iterable[str]): """Update this index, adding the specified hashes.""" - from diskcache import Timeout + from dvc.objects.cache import Timeout try: with self.index.transact(): @@ -143,7 +133,6 @@ class ObjectDBIndex(ObjectDBIndexBase): except Timeout as exc: raise ObjectDBError("Failed to update ODB index") from exc - @with_diskcache(name="index") def intersection(self, hashes: Set[str]): """Iterate over values from `hashes` which exist in the index.""" yield from hashes.intersection(self.index.keys()) diff --git a/dvc/objects/cache.py b/dvc/objects/cache.py new file mode 100644 index 000000000..175fb91c2 --- /dev/null +++ b/dvc/objects/cache.py @@ -0,0 +1,61 @@ +import os +import pickle +from functools import wraps +from typing import Any + +import diskcache +from diskcache import Disk as disk +from diskcache import Index # noqa: F401, pylint: disable=unused-import +from diskcache import Timeout # noqa: F401, pylint: disable=unused-import + +# pylint: disable=redefined-builtin + + +class DiskError(Exception): + def __init__(self, directory: str, type: str) -> None: + self.directory = directory + self.type = type + super().__init__(f"Could not open disk '{type}' in {directory}") + + +def translate_pickle_error(fn): + @wraps(fn) + def wrapped(self, *args, **kwargs): + try: + return fn(self, *args, **kwargs) + except (pickle.PickleError, ValueError) as e: + if isinstance(e, ValueError) and "pickle protocol" not in str(e): + raise + # pylint: disable=protected-access + raise DiskError(self._directory, type=self._type) from e + + return wrapped + + +class Disk(disk): + """Reraise pickle-related errors as DiskError.""" + + def __init__(self, *args, type=None, **kwargs): + super().__init__(*args, **kwargs) + self._type = type or os.path.basename(self._directory) + + put = translate_pickle_error(disk.put) + get = translate_pickle_error(disk.get) + store = translate_pickle_error(disk.store) + fetch = translate_pickle_error(disk.fetch) + + +class Cache(diskcache.Cache): + """Extended to handle pickle errors and use a constant pickle protocol.""" + + def __init__( + self, + directory: str = None, + timeout: int = 60, + disk: disk = Disk, # pylint: disable=redefined-outer-name + **settings: Any, + ) -> None: + settings.setdefault("disk_pickle_protocol", 4) + super().__init__( + directory=directory, timeout=timeout, disk=Disk, **settings + ) diff --git a/dvc/objects/state.py b/dvc/objects/state.py index d09c4a163..fb16acdee 100644 --- a/dvc/objects/state.py +++ b/dvc/objects/state.py @@ -7,7 +7,6 @@ from abc import ABC, abstractmethod from dvc.fs.local import LocalFileSystem from dvc.fs.system import inode as get_inode from dvc.utils import relpath -from dvc.utils.decorators import with_diskcache from .hash_info import HashInfo from .utils import get_mtime_and_size @@ -49,7 +48,7 @@ class StateNoop(StateBase): class State(StateBase): # pylint: disable=too-many-instance-attributes def __init__(self, root_dir=None, tmp_dir=None, dvcignore=None): - from diskcache import Cache + from .cache import Cache super().__init__() @@ -60,10 +59,7 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes if not tmp_dir: return - config = { - "eviction_policy": "least-recently-used", - "disk_pickle_protocol": 4, - } + config = {"eviction_policy": "least-recently-used"} self.links = Cache(directory=os.path.join(tmp_dir, "links"), **config) self.md5s = Cache(directory=os.path.join(tmp_dir, "md5s"), **config) @@ -71,7 +67,6 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes self.md5s.close() self.links.close() - @with_diskcache(name="md5s") def save(self, path, fs, hash_info): """Save hash for the specified path info. @@ -96,7 +91,6 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes self.md5s[inode] = (mtime, str(size), hash_info.value) - @with_diskcache(name="md5s") def get(self, path, fs): """Gets the hash for the specified path info. Hash will be retrieved from the state database if available. @@ -127,7 +121,6 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes return Meta(size=size), HashInfo("md5", value[2]) - @with_diskcache(name="links") def save_link(self, path, fs): """Adds the specified path to the list of links created by dvc. This list is later used on `dvc checkout` to cleanup old links. @@ -149,7 +142,6 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes with self.links as ref: ref[relative_path] = (inode, mtime) - @with_diskcache(name="links") def get_unused_links(self, used, fs): """Removes all saved links except the ones that are used. @@ -177,7 +169,6 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes return unused - @with_diskcache(name="links") def remove_links(self, unused, fs): if not isinstance(fs, LocalFileSystem): return diff --git a/dvc/utils/decorators.py b/dvc/utils/decorators.py deleted file mode 100644 index b9ea532be..000000000 --- a/dvc/utils/decorators.py +++ /dev/null @@ -1,30 +0,0 @@ -import pickle -from typing import Callable, TypeVar - -from funcy import decorator - -from dvc.exceptions import DvcException - -from . import format_link - -_R = TypeVar("_R") - - -@decorator -def with_diskcache(call: Callable[..., _R], name: str) -> _R: - try: - return call() - except (pickle.PickleError, ValueError) as exc: - if isinstance(exc, ValueError) and not str(exc).startswith( - "pickle protocol" - ): - raise - link = format_link( - "https://dvc.org/doc/user-guide/troubleshooting#pickle" - ) - msg = ( - f"Could not open pickled '{name}' cache. Remove the " - f"'.dvc/tmp/{name}' directory and then retry this command. " - f"See {link} for more information." - ) - raise DvcException(msg) from exc diff --git a/tests/unit/cli/__init__.py b/tests/unit/cli/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/cli/test_main.py b/tests/unit/cli/test_main.py new file mode 100644 index 000000000..05fe48d65 --- /dev/null +++ b/tests/unit/cli/test_main.py @@ -0,0 +1,26 @@ +from argparse import Namespace + +from funcy import raiser + +from dvc.cli import main +from dvc.objects.cache import DiskError + + +def test_state_pickle_errors_are_correctly_raised(tmp_dir, caplog, mocker): + path = tmp_dir / "dir" / "test" + mocker.patch( + "dvc.cli.parse_args", + return_value=Namespace( + func=raiser(DiskError(path, "md5s")), + quiet=False, + verbose=True, + ), + ) + + assert main() == 255 + assert ( + f"Could not open pickled 'md5s' cache.\n" + f"Remove the '{path.relative_to(tmp_dir)}' directory " + "and then retry this command.\n" + "See for more information." + ) in caplog.text diff --git a/tests/unit/utils/test_decorators.py b/tests/unit/objects/test_cache.py similarity index 52% rename from tests/unit/utils/test_decorators.py rename to tests/unit/objects/test_cache.py index 5998b7739..39d59d65c 100644 --- a/tests/unit/utils/test_decorators.py +++ b/tests/unit/objects/test_cache.py @@ -1,27 +1,29 @@ import pickle from typing import Any -import diskcache import pytest -from dvc.exceptions import DvcException -from dvc.utils.decorators import with_diskcache +from dvc.objects.cache import Cache, DiskError -@with_diskcache(name="test") -def set_value(cache: diskcache.Cache, key: str, value: Any) -> Any: +def set_value(cache: Cache, key: str, value: Any) -> Any: cache[key] = value return cache[key] -def test_pickle_protocol_error(tmp_dir): - with pytest.raises(DvcException) as exc: - with diskcache.Cache( - directory=(tmp_dir / "test"), - disk_pickle_protocol=pickle.HIGHEST_PROTOCOL + 1, - ) as cache: - set_value(cache, "key", ("value1", "value2")) - assert "troubleshooting#pickle" in str(exc) +@pytest.mark.parametrize("disk_type", [None, "test"]) +def test_pickle_protocol_error(tmp_dir, disk_type): + directory = tmp_dir / "test" + cache = Cache( + directory, + disk_pickle_protocol=pickle.HIGHEST_PROTOCOL + 1, + disk_type=disk_type, + ) + with pytest.raises(DiskError) as exc, cache as cache: + set_value(cache, "key", ("value1", "value2")) + assert exc.value.directory == str(directory) + assert exc.value.type == "test" + assert f"Could not open disk 'test' in {directory}" == str(exc.value) @pytest.mark.parametrize( @@ -32,12 +34,12 @@ def test_pickle_protocol_error(tmp_dir): ], ) def test_pickle_backwards_compat(tmp_dir, proto_a, proto_b): - with diskcache.Cache( + with Cache( directory=(tmp_dir / "test"), disk_pickle_protocol=proto_a, ) as cache: set_value(cache, "key", ("value1", "value2")) - with diskcache.Cache( + with Cache( directory=(tmp_dir / "test"), disk_pickle_protocol=proto_b, ) as cache: