Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-10 15:28:34 +01:00
Code Issues Packages Projects Releases Wiki Activity

dvc.objects: extend diskcache to handle pickle errors (#7734)

Browse Source 
dvc.objects: extend diskcache to handle pickle errors
This commit is contained in:
skshetry
2022-05-12 21:12:33 +05:45
committed by GitHub
parent ee29888e18
commit b3fb0dfdad
8 changed files with 151 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/cli/__init__.py
+39 -21
View File
@@ -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:
dvc/data/db/index.py
+6 -17
View File
@@ -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())
dvc/objects/cache.py
+61
View File
@@ -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
)
dvc/objects/state.py
+2 -11
View File
@@ -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
dvc/utils/decorators.py
-30
View File
@@ -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
tests/unit/cli/__init__.py
View File
tests/unit/cli/test_main.py
+26
View File
@@ -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 <https://error.dvc.org/pickle> for more information."
) in caplog.text
tests/unit/utils/test_decorators.py → tests/unit/objects/test_cache.py
+17 -15
View File
@@ -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: