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

dvc: refactor cache saving (#5314)

Browse Source
This commit is contained in:
Ruslan Kuprieiev
2021-01-22 19:07:36 +02:00
committed by GitHub
parent 7b59108379
commit cbd7472079
14 changed files with 145 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/cache/base.py
Vendored
+16 -9
View File
@@ -165,10 +165,16 @@ class CloudCache:
def link(self, from_info, to_info):
self._link(from_info, to_info, self.cache_types)
def move(self, from_info, to_info):
self.tree.move(from_info, to_info)
def makedirs(self, path_info):
self.tree.makedirs(path_info)
def _link(self, from_info, to_info, link_types):
assert self.tree.isfile(from_info)
self.tree.makedirs(to_info.parent)
self.makedirs(to_info.parent)
self._try_links(from_info, to_info, link_types)
@@ -216,7 +222,9 @@ class CloudCache:
cache_info = self.tree.hash_to_path_info(hash_info.value)
if tree == self.tree:
if self.changed_cache(hash_info):
self.tree.move(path_info, cache_info, mode=self.CACHE_MODE)
self.makedirs(cache_info.parent)
self.move(path_info, cache_info)
self.protect(cache_info)
self.link(cache_info, path_info)
elif self.tree.iscopy(path_info) and self._cache_is_copy(
path_info
@@ -287,7 +295,7 @@ class CloudCache:
)
hash_info = stream_reader.hash_info
self.tree.move(tmp_info, self.tree.hash_to_path_info(hash_info.value))
self.move(tmp_info, self.tree.hash_to_path_info(hash_info.value))
return hash_info
def _transfer_file(self, from_tree, from_info):
@@ -391,9 +399,7 @@ class CloudCache:
from_info = PathInfo(tmp)
to_info = self.tree.path_info / tmp_fname("")
self.tree.upload(
from_info, to_info, no_progress_bar=True, file_mode=self.CACHE_MODE
)
self.tree.upload(from_info, to_info, no_progress_bar=True)
hash_info = self.tree.get_file_hash(to_info)
hash_info.value += self.tree.CHECKSUM_DIR_SUFFIX
@@ -414,8 +420,9 @@ class CloudCache:
hi, tmp_info = self._get_dir_info_hash(dir_info)
new_info = self.tree.hash_to_path_info(hi.value)
if self.changed_cache_file(hi):
self.tree.makedirs(new_info.parent)
self.tree.move(tmp_info, new_info, mode=self.CACHE_MODE)
self.makedirs(new_info.parent)
self.move(tmp_info, new_info)
self.protect(new_info)
self.tree.state.save(new_info, hi)
@@ -616,7 +623,7 @@ class CloudCache:
# even if there are no files in it
if not self.tree.exists(path_info):
added = True
self.tree.makedirs(path_info)
self.makedirs(path_info)
dir_info = self.get_dir_cache(hash_info)
dvc/cache/local.py
Vendored
+24 -2
View File
@@ -9,7 +9,7 @@ from dvc.path_info import PathInfo
from dvc.progress import Tqdm
from ..utils import relpath
from ..utils.fs import copyfile, remove, walk_files
from ..utils.fs import copyfile, remove, umask, walk_files
from .base import CloudCache
logger = logging.getLogger(__name__)
@@ -24,6 +24,14 @@ class LocalCache(CloudCache):
super().__init__(tree)
self.cache_dir = tree.config.get("url")
shared = tree.config.get("shared")
if shared:
self._file_mode = 0o664
self._dir_mode = 0o2775
else:
self._file_mode = 0o666 & ~umask
self._dir_mode = 0o777 & ~umask
@property
def cache_dir(self):
return self.tree.path_info.fspath if self.tree.path_info else None
@@ -36,6 +44,20 @@ class LocalCache(CloudCache):
def cache_path(self):
return os.path.abspath(self.cache_dir)
def move(self, from_info, to_info):
super().move(from_info, to_info)
os.chmod(to_info, self._file_mode)
def makedirs(self, path_info):
from dvc.utils.fs import makedirs
makedirs(path_info, exist_ok=True, mode=self._dir_mode)
def link(self, from_info, to_info):
super().link(from_info, to_info)
if self.cache_types[0] not in ("symlink", "hardlink"):
os.chmod(to_info, self._file_mode)
def hash_to_path(self, hash_):
# NOTE: `self.cache_path` is already normalized so we can simply use
# `os.sep` instead of `os.path.join`. This results in this helper
@@ -112,7 +134,7 @@ class LocalCache(CloudCache):
path,
)
os.chmod(path, self.tree.file_mode)
os.chmod(path, self._file_mode)
def _unprotect_dir(self, path):
for fname in self.tree.walk_files(path):
dvc/exceptions.py
-5
View File
@@ -239,11 +239,6 @@ class FileMissingError(DvcException):
)
class FileOwnershipError(DvcException):
def __init__(self, path):
super().__init__(f"file '{path}' not owned by user! ")
class DvcIgnoreInCollectedDirError(DvcException):
def __init__(self, ignore_dirname):
super().__init__(
dvc/remote/base.py
+18 -17
View File
@@ -4,7 +4,7 @@ import itertools
import logging
from concurrent import futures
from concurrent.futures import ThreadPoolExecutor
from functools import partial, wraps
from functools import wraps
from dvc.exceptions import DownloadError, UploadError
from dvc.hash_info import HashInfo
@@ -309,24 +309,11 @@ class Remote:
)
if download:
# NOTE: don't set CACHE_MODE if we don't trust the remote, so that
# cache could verify these files by-hand later.
if self.tree.verify:
mode = self.tree.file_mode
else:
mode = cache.CACHE_MODE
func = partial(
_log_exceptions(self.tree.download, "download"),
dir_mode=cache.tree.dir_mode,
file_mode=mode,
)
func = _log_exceptions(self.tree.download, "download")
status = STATUS_DELETED
desc = "Downloading"
else:
func = partial(
_log_exceptions(self.tree.upload, "upload"),
file_mode=self.cache.CACHE_MODE,
)
func = _log_exceptions(self.tree.upload, "upload")
status = STATUS_NEW
desc = "Uploading"
@@ -463,7 +450,7 @@ class Remote:
@index_locked
def push(self, cache, named_cache, jobs=None, show_checksums=False):
return self._process(
ret = self._process(
cache,
named_cache,
jobs=jobs,
@@ -471,6 +458,19 @@ class Remote:
download=False,
)
if self.tree.scheme == "local":
with self.tree.state:
for checksum in named_cache.scheme_keys("local"):
cache_file = self.tree.hash_to_path_info(checksum)
if self.tree.exists(cache_file):
hash_info = HashInfo(
self.tree.PARAM_CHECKSUM, checksum
)
self.tree.state.save(cache_file, hash_info)
self.cache.protect(cache_file)
return ret
@index_locked
def pull(self, cache, named_cache, jobs=None, show_checksums=False):
ret = self._process(
@@ -494,6 +494,7 @@ class Remote:
cache.tree.PARAM_CHECKSUM, checksum
)
cache.tree.state.save(cache_file, hash_info)
cache.protect(cache_file)
return ret
dvc/stage/cache.py
+2 -6
View File
@@ -60,10 +60,6 @@ class StageCache:
def cache_dir(self):
return os.path.join(self.repo.cache.local.cache_dir, "runs")
@property
def tree(self):
return self.repo.cache.local.tree
def _get_cache_dir(self, key):
return os.path.join(self.cache_dir, key[:2], key)
@@ -170,12 +166,12 @@ class StageCache:
COMPILED_LOCK_FILE_STAGE_SCHEMA(cache)
path = PathInfo(self._get_cache_path(cache_key, cache_value))
self.tree.makedirs(path.parent)
self.repo.cache.local.makedirs(path.parent)
tmp = tempfile.NamedTemporaryFile(delete=False, dir=path.parent).name
assert os.path.exists(path.parent)
assert os.path.isdir(path.parent)
dump_yaml(tmp, cache)
self.tree.move(PathInfo(tmp), path)
self.repo.cache.local.move(PathInfo(tmp), path)
def restore(self, stage, run_cache=True, pull=False):
from .serialize import to_single_stage_lockfile
dvc/tree/base.py
+9 -53
View File
@@ -61,7 +61,6 @@ class BaseTree:
CAN_TRAVERSE = True
CHUNK_SIZE = 64 * 1024 * 1024 # 64 MiB
SHARED_MODE_MAP = {None: (None, None), "group": (None, None)}
PARAM_CHECKSUM: ClassVar[Optional[str]] = None
state = StateNoop()
@@ -72,9 +71,6 @@ class BaseTree:
self._check_requires()
shared = config.get("shared")
self._file_mode, self._dir_mode = self.SHARED_MODE_MAP[shared]
self.verify = config.get("verify", self.DEFAULT_VERIFY)
self.path_info = None
@@ -157,14 +153,6 @@ class BaseTree:
parsed = urlparse(url)
return parsed.scheme == cls.scheme
@property
def file_mode(self):
return self._file_mode
@property
def dir_mode(self):
return self._dir_mode
@property
def cache(self):
return getattr(self.repo.cache, self.scheme)
@@ -231,8 +219,7 @@ class BaseTree:
directories before copying/linking/moving data
"""
def move(self, from_info, to_info, mode=None):
assert mode is None
def move(self, from_info, to_info):
self.copy(from_info, to_info)
self.remove(from_info)
@@ -355,12 +342,7 @@ class BaseTree:
return hash_info
def upload(
self,
from_info,
to_info,
name=None,
no_progress_bar=False,
file_mode=None,
self, from_info, to_info, name=None, no_progress_bar=False,
):
if not hasattr(self, "_upload"):
raise RemoteActionNotImplemented("upload", self.scheme)
@@ -380,7 +362,6 @@ class BaseTree:
to_info,
name=name,
no_progress_bar=no_progress_bar,
file_mode=file_mode,
)
def upload_fobj(self, fobj, to_info, no_progress_bar=False):
@@ -392,8 +373,6 @@ class BaseTree:
to_info,
name=None,
no_progress_bar=False,
file_mode=None,
dir_mode=None,
jobs=None,
**kwargs,
):
@@ -412,29 +391,12 @@ class BaseTree:
if self.isdir(from_info):
return self._download_dir(
from_info,
to_info,
name,
no_progress_bar,
file_mode,
dir_mode,
jobs,
**kwargs,
from_info, to_info, name, no_progress_bar, jobs, **kwargs,
)
return self._download_file(
from_info, to_info, name, no_progress_bar, file_mode, dir_mode,
)
return self._download_file(from_info, to_info, name, no_progress_bar,)
def _download_dir(
self,
from_info,
to_info,
name,
no_progress_bar,
file_mode,
dir_mode,
jobs,
**kwargs,
self, from_info, to_info, name, no_progress_bar, jobs, **kwargs,
):
from_infos = list(self.walk_files(from_info, **kwargs))
to_infos = (
@@ -448,13 +410,7 @@ class BaseTree:
disable=no_progress_bar,
) as pbar:
download_files = pbar.wrap_fn(
partial(
self._download_file,
name=name,
no_progress_bar=True,
file_mode=file_mode,
dir_mode=dir_mode,
)
partial(self._download_file, name=name, no_progress_bar=True,)
)
max_workers = jobs or self.jobs
with ThreadPoolExecutor(max_workers=max_workers) as executor:
@@ -479,9 +435,9 @@ class BaseTree:
raise exc
def _download_file(
self, from_info, to_info, name, no_progress_bar, file_mode, dir_mode
self, from_info, to_info, name, no_progress_bar,
):
makedirs(to_info.parent, exist_ok=True, mode=dir_mode)
makedirs(to_info.parent, exist_ok=True)
logger.debug("Downloading '%s' to '%s'", from_info, to_info)
name = name or to_info.name
@@ -492,4 +448,4 @@ class BaseTree:
from_info, tmp_file, name=name, no_progress_bar=no_progress_bar
)
move(tmp_file, to_info, mode=file_mode)
move(tmp_file, to_info)
dvc/tree/local.py
+5 -34
View File
@@ -2,7 +2,6 @@ import errno
import logging
import os
import stat
from typing import Any, Dict
from funcy import cached_property
@@ -25,11 +24,6 @@ class LocalTree(BaseTree):
PARAM_PATH = "path"
TRAVERSE_PREFIX_LEN = 2
SHARED_MODE_MAP: Dict[Any, Any] = {
None: (0o644, 0o755),
"group": (0o664, 0o775),
}
def __init__(self, repo, config, use_dvcignore=False, dvcignore_root=None):
super().__init__(repo, config)
url = config.get("url")
@@ -140,7 +134,7 @@ class LocalTree(BaseTree):
remove(path_info)
def makedirs(self, path_info):
makedirs(path_info, exist_ok=True, mode=self.dir_mode)
makedirs(path_info, exist_ok=True)
def set_exec(self, path_info):
mode = self.stat(path_info).st_mode
@@ -156,25 +150,17 @@ class LocalTree(BaseTree):
return os.stat(path)
def move(self, from_info, to_info, mode=None):
def move(self, from_info, to_info):
if from_info.scheme != "local" or to_info.scheme != "local":
raise NotImplementedError
self.makedirs(to_info.parent)
if mode is None:
if self.isfile(from_info):
mode = self.file_mode
else:
mode = self.dir_mode
move(from_info, to_info, mode=mode)
move(from_info, to_info)
def copy(self, from_info, to_info):
tmp_info = to_info.parent / tmp_fname("")
try:
System.copy(from_info, tmp_info)
os.chmod(tmp_info, self.file_mode)
os.rename(tmp_info, to_info)
except Exception:
self.remove(tmp_info)
@@ -185,7 +171,6 @@ class LocalTree(BaseTree):
tmp_info = to_info.parent / tmp_fname("")
try:
copy_fobj_to_file(fobj, tmp_info, chunk_size=chunk_size)
os.chmod(tmp_info, self.file_mode)
os.rename(tmp_info, to_info)
except Exception:
self.remove(tmp_info)
@@ -226,12 +211,7 @@ class LocalTree(BaseTree):
return System.is_hardlink(path_info)
def reflink(self, from_info, to_info):
tmp_info = to_info.parent / tmp_fname("")
System.reflink(from_info, tmp_info)
# NOTE: reflink has its own separate inode, so you can set permissions
# that are different from the source.
os.chmod(tmp_info, self.file_mode)
os.rename(tmp_info, to_info)
System.reflink(from_info, to_info)
def chmod(self, path_info, mode):
try:
@@ -263,13 +243,7 @@ class LocalTree(BaseTree):
return os.path.getsize(path_info)
def _upload(
self,
from_file,
to_info,
name=None,
no_progress_bar=False,
file_mode=None,
**_kwargs,
self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs,
):
makedirs(to_info.parent, exist_ok=True)
@@ -277,9 +251,6 @@ class LocalTree(BaseTree):
copyfile(
from_file, tmp_file, name=name, no_progress_bar=no_progress_bar
)
if file_mode is not None:
self.chmod(tmp_file, file_mode)
os.replace(tmp_file, to_info)
def upload_fobj(self, fobj, to_info, no_progress_bar=False, **pbar_args):
dvc/tree/ssh/__init__.py
+1 -2
View File
@@ -187,8 +187,7 @@ class SSHTree(BaseTree):
with self.ssh(path_info) as ssh:
ssh.makedirs(path_info.path)
def move(self, from_info, to_info, mode=None):
assert mode is None
def move(self, from_info, to_info):
if from_info.scheme != self.scheme or to_info.scheme != self.scheme:
raise NotImplementedError
dvc/tree/webdav.py
+1 -1
View File
@@ -188,7 +188,7 @@ class WebDAVTree(BaseTree): # pylint:disable=abstract-method
self._client.mkdir(path_info.path)
# Moves file/directory at remote
def move(self, from_info, to_info, mode=None):
def move(self, from_info, to_info):
# Webdav client move
self._client.move(from_info.path, to_info.path)
dvc/tree/webhdfs.py
+1 -1
View File
@@ -139,7 +139,7 @@ class WebHDFSTree(BaseTree):
with self.hdfs_client.write(to_info.path) as writer:
shutil.copyfileobj(reader, writer)
def move(self, from_info, to_info, mode=None):
def move(self, from_info, to_info):
self.hdfs_client.makedirs(to_info.parent.path)
self.hdfs_client.rename(from_info.path, to_info.path)
dvc/utils/fs.py
+32 -18
View File
@@ -5,7 +5,7 @@ import shutil
import stat
import sys
from dvc.exceptions import DvcException, FileOwnershipError
from dvc.exceptions import DvcException
from dvc.system import System
from dvc.utils import dict_md5
@@ -13,6 +13,9 @@ logger = logging.getLogger(__name__)
LOCAL_CHUNK_SIZE = 2 ** 20 # 1 MB
umask = os.umask(0)
os.umask(umask)
def fs_copy(src, dst):
if os.path.isdir(src):
@@ -82,7 +85,7 @@ def contains_symlink_up_to(path, base_path):
return contains_symlink_up_to(os.path.dirname(path), base_path)
def move(src, dst, mode=None):
def move(src, dst):
"""Atomically move src to dst and chmod it with mode.
Moving is performed in two stages to make the whole operation atomic in
@@ -94,15 +97,6 @@ def move(src, dst, mode=None):
dst = os.path.abspath(dst)
tmp = f"{dst}.{uuid()}"
try:
if mode is not None:
os.chmod(src, mode)
except OSError as e:
if e.errno not in [errno.EACCES, errno.EPERM]:
raise
else:
raise FileOwnershipError(src)
if os.path.islink(src):
shutil.copy(src, tmp)
_unlink(src, _chmod)
@@ -162,14 +156,31 @@ def makedirs(path, exist_ok=False, mode=None):
os.makedirs(path, exist_ok=exist_ok)
return
# utilize umask to set proper permissions since Python 3.7 the `mode`
# `makedirs` argument no longer affects the file permission bits of
# newly-created intermediate-level directories.
umask = os.umask(0o777 - mode)
# Modified version of os.makedirs() with support for extended mode
# (e.g. S_ISGID)
head, tail = os.path.split(path)
if not tail:
head, tail = os.path.split(head)
if head and tail and not os.path.exists(head):
try:
makedirs(head, exist_ok=exist_ok, mode=mode)
except FileExistsError:
# Defeats race condition when another thread created the path
pass
cdir = os.curdir
if isinstance(tail, bytes):
cdir = bytes(os.curdir, "ASCII")
if tail == cdir: # xxx/newdir/. exists if xxx/newdir exists
return
try:
os.makedirs(path, exist_ok=exist_ok)
finally:
os.umask(umask)
os.mkdir(path, mode)
except OSError:
# Cannot rely on checking for EEXIST, since the operating system
# could give priority to other errors like EACCES or EROFS
if not exist_ok or not os.path.isdir(path):
raise
os.chmod(path, mode)
def copyfile(src, dest, no_progress_bar=False, name=None):
@@ -184,6 +195,9 @@ def copyfile(src, dest, no_progress_bar=False, name=None):
try:
System.reflink(src, dest)
# NOTE: reflink has a new inode, but has the same mode as the src,
# so we need to chmod it to look like a normal copy.
os.chmod(dest, 0o666 & ~umask)
except DvcException:
with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
with Tqdm.wrapattr(
tests/func/test_cache.py
+31 -9
View File
@@ -194,26 +194,48 @@ def test_default_cache_type(dvc):
@pytest.mark.skipif(os.name == "nt", reason="Not supported for Windows.")
@pytest.mark.parametrize(
"group, dir_mode", [(False, 0o755), (True, 0o775)],
"group", [False, True],
)
def test_shared_cache(tmp_dir, dvc, group, dir_mode):
def test_shared_cache(tmp_dir, dvc, group):
from dvc.utils.fs import umask
if group:
with dvc.config.edit() as conf:
conf["cache"].update({"shared": "group"})
dvc.cache = Cache(dvc)
cache_dir = dvc.cache.local.cache_dir
assert not os.path.exists(cache_dir)
tmp_dir.dvc_gen(
{"file": "file content", "dir": {"file2": "file 2 " "content"}}
)
for root, dnames, fnames in os.walk(dvc.cache.local.cache_dir):
for dname in dnames:
path = os.path.join(root, dname)
assert stat.S_IMODE(os.stat(path).st_mode) == dir_mode
actual = {}
for root, dnames, fnames in os.walk(cache_dir):
for name in dnames + fnames:
path = os.path.join(root, name)
actual[path] = oct(stat.S_IMODE(os.stat(path).st_mode))
for fname in fnames:
path = os.path.join(root, fname)
assert stat.S_IMODE(os.stat(path).st_mode) == 0o444
file_mode = oct(0o444)
dir_mode = oct(0o2775 if group else (0o777 & ~umask))
expected = {
os.path.join(cache_dir, "17"): dir_mode,
os.path.join(
cache_dir, "17", "4eaa1dd94050255b7b98a7e1924b31.dir"
): file_mode,
os.path.join(cache_dir, "97"): dir_mode,
os.path.join(
cache_dir, "97", "e17781c198500e2766ea56bd697c03"
): file_mode,
os.path.join(cache_dir, "d1"): dir_mode,
os.path.join(
cache_dir, "d1", "0b4c3ff123b26dc068d43a8bef2d23"
): file_mode,
}
assert expected == actual
def test_cache_dir_local(tmp_dir, dvc, caplog):
tests/unit/stage/test_cache.py
+3 -1
View File
@@ -174,6 +174,8 @@ def test_shared_stage_cache(tmp_dir, dvc, run_copy):
dvc.cache = Cache(dvc)
assert not os.path.exists(dvc.cache.local.cache_dir)
run_copy("foo", "bar", name="copy-foo-bar")
parent_cache_dir = os.path.join(dvc.stage_cache.cache_dir, "88",)
@@ -198,7 +200,7 @@ def test_shared_stage_cache(tmp_dir, dvc, run_copy):
dir_mode = 0o777
file_mode = 0o666
else:
dir_mode = 0o775
dir_mode = 0o2775
file_mode = 0o664
assert _mode(dvc.cache.local.cache_dir) == dir_mode
tests/unit/tree/test_tree.py
+2 -24
View File
@@ -1,11 +1,8 @@
import os
import stat
import pytest
from dvc.config import ConfigError
from dvc.path_info import CloudURLInfo, PathInfo
from dvc.tree import LocalTree, get_cloud_tree
from dvc.path_info import CloudURLInfo
from dvc.tree import get_cloud_tree
def test_get_cloud_tree(tmp_dir, dvc):
@@ -49,22 +46,3 @@ def test_get_cloud_tree_validate(tmp_dir, dvc):
with pytest.raises(ConfigError):
get_cloud_tree(dvc, name="second")
@pytest.mark.skipif(os.name == "nt", reason="Not supported for Windows.")
@pytest.mark.parametrize(
"mode, expected", [(None, None), (0o777, 0o777)],
)
def test_upload_file_mode(tmp_dir, mode, expected):
tmp_dir.gen("src", "foo")
src = PathInfo(tmp_dir / "src")
if not expected:
expected = stat.S_IMODE((tmp_dir / "src").stat().st_mode)
dest = PathInfo(tmp_dir / "dest")
tree = LocalTree(None, {"url": os.fspath(tmp_dir)})
tree.upload(src, dest, file_mode=mode)
assert (tmp_dir / "dest").exists()
assert (tmp_dir / "dest").read_text() == "foo"
assert stat.S_IMODE(os.stat(dest).st_mode) == expected