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

fs.utils: add upload that can be conditionally atomic

Browse Source 
The PR also makes odb.add operations atomic, but the checkout is still
not atomic.
This commit is contained in:
Saugat Pachhai (सौगात)
2022-04-29 23:02:15 +03:00
committed by Ruslan Kuprieiev
parent b71ea349d2
commit dd187df667
12 changed files with 146 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/data/checkout.py
+8 -1
View File
@@ -138,7 +138,14 @@ class Link:
cache.makedirs(cache.fs.path.parent(to_path))
try:
transfer(cache.fs, from_path, to_fs, to_path, links=self._links)
transfer(
cache.fs,
from_path,
to_fs,
to_path,
links=self._links,
atomic=False,
)
except FileNotFoundError as exc:
raise CheckoutError([to_path]) from exc
except OSError as exc:
dvc/data/stage.py
+5 -4
View File
@@ -27,21 +27,22 @@ _STAGING_MEMFS_PATH = "dvc-staging"
def _upload_file(from_fs_path, fs, odb, upload_odb, callback=None):
from dvc.fs._callback import FsspecCallback
from dvc.fs.utils import fileobj_size
from dvc.utils import tmp_fname
from dvc.utils.stream import HashedStreamReader
fs_path = upload_odb.fs.path
tmp_info = fs_path.join(upload_odb.fs_path, tmp_fname())
with fs.open(from_fs_path, mode="rb") as stream:
stream = HashedStreamReader(stream)
size = fs.size(from_fs_path)
with fs.open(from_fs_path, mode="rb") as fobj:
stream = HashedStreamReader(fobj)
size = fileobj_size(fs, fobj, from_fs_path)
with FsspecCallback.as_tqdm_callback(
callback,
desc=fs_path.name(from_fs_path),
bytes=True,
total=size,
) as cb:
upload_odb.fs.upload(stream, tmp_info, size=size, callback=cb)
upload_odb.fs.put_file(stream, tmp_info, size=size, callback=cb)
odb.add(tmp_info, upload_odb.fs, stream.hash_info)
meta = Meta(size=size)
dvc/fs/base.py
+10 -27
View File
@@ -1,6 +1,7 @@
import logging
import os
import shutil
from contextlib import nullcontext
from multiprocessing import cpu_count
from typing import (
IO,
@@ -360,31 +361,6 @@ class FileSystem:
def size(self, path: AnyFSPath) -> Optional[int]:
return self.fs.size(path)
# pylint: enable=unused-argument
def upload(
self,
from_info: Union[AnyFSPath, IO],
to_info: AnyFSPath,
size: int = None,
callback: FsspecCallback = None,
):
from .local import localfs
if not hasattr(from_info, "read"):
logger.debug("Uploading '%s' to '%s'", from_info, to_info)
desc = localfs.path.name(from_info)
else:
desc = self.path.name(to_info)
with FsspecCallback.as_tqdm_callback(
callback,
desc=desc,
bytes=True,
total=size or -1,
) as cb:
return self.put_file(from_info, to_info, callback=cb, size=size)
def download(
self,
from_info: AnyFSPath,
@@ -426,7 +402,8 @@ class FileSystem:
) as executor:
list(
executor.imap_unordered(
lambda args: download_files(*args), pairs.items()
lambda args: download_files(*args, atomic=True),
pairs.items(),
)
)
@@ -435,6 +412,7 @@ class FileSystem:
from_info: AnyFSPath,
to_info: AnyFSPath,
callback: FsspecCallback = None,
atomic: bool = False,
):
from .local import localfs
@@ -444,7 +422,12 @@ class FileSystem:
desc=self.path.name(from_info),
bytes=True,
) as cb:
with as_atomic(localfs, to_info, create_parents=True) as tmp_file:
context = (
as_atomic(localfs, to_info, create_parents=True)
if atomic
else nullcontext(to_info)
)
with context as tmp_file:
self.get_file(from_info, tmp_file, callback=cb)
dvc/fs/local.py
+1 -3
View File
@@ -87,9 +87,7 @@ class FsspecLocalFileSystem(AbstractFileSystem):
def put_file(self, lpath, rpath, callback=None, **kwargs):
parent = self._parent(rpath)
makedirs(parent, exist_ok=True)
tmp_file = os.path.join(parent, tmp_fname())
copyfile(lpath, tmp_file, callback=callback)
os.replace(tmp_file, rpath)
copyfile(lpath, rpath, callback=callback)
def get_file(self, rpath, lpath, callback=None, **kwargs):
copyfile(rpath, lpath, callback=callback)
dvc/fs/ssh.py
-22
View File
@@ -6,9 +6,7 @@ from funcy import cached_property, memoize, silent, wrap_prop, wrap_with
from dvc import prompt
from dvc.scheme import Schemes
from dvc.utils.fs import as_atomic
from ._callback import DEFAULT_CALLBACK
from .base import FileSystem
DEFAULT_PORT = 22
@@ -118,23 +116,3 @@ class SSHFileSystem(FileSystem):
from sshfs import SSHFileSystem as _SSHFileSystem
return _SSHFileSystem(**self.fs_args)
# Ensure that if an interrupt happens during the transfer, we don't
# pollute the cache.
def upload_fobj(self, fobj, to_info, **kwargs):
with as_atomic(self, to_info) as tmp_file:
super().upload_fobj(fobj, tmp_file, **kwargs)
def put_file(
self,
from_file,
to_info,
callback=DEFAULT_CALLBACK,
size=None,
**kwargs,
):
with as_atomic(self, to_info) as tmp_file:
super().put_file(
from_file, tmp_file, callback=callback, size=size, **kwargs
)
dvc/fs/utils.py
+85 -26
View File
@@ -1,20 +1,40 @@
import errno
import logging
import os
from typing import TYPE_CHECKING, List, Optional
from contextlib import nullcontext
from typing import IO, TYPE_CHECKING, List, Optional, Union, cast
from ..utils.fs import as_atomic
from ._callback import DEFAULT_CALLBACK, FsspecCallback
from .base import RemoteActionNotImplemented
from .local import LocalFileSystem
if TYPE_CHECKING:
from .base import FileSystem
from .types import AnyPath
from dvc.types import AnyPath
from .base import AnyFSPath, FileSystem
logger = logging.getLogger(__name__)
def upload(
fs: "FileSystem",
lpath: Union["AnyFSPath", IO],
rpath: "AnyFSPath",
callback: "FsspecCallback" = DEFAULT_CALLBACK,
size: int = None,
atomic: bool = False,
):
"""Atomic version of `fs.put_file`."""
# no need for `create_parents` here
# it should depend on `fs.put_file` to handle that.
context = as_atomic(fs, rpath) if atomic else nullcontext(rpath)
with context as tmp_file:
return fs.put_file(lpath, tmp_file, callback=callback, size=size)
def _link(
link: "AnyPath",
link: "str",
from_fs: "FileSystem",
from_path: "AnyPath",
to_fs: "FileSystem",
@@ -34,36 +54,74 @@ def _link(
) from exc
def _copy(
def fileobj_size(
fs: "FileSystem", fileobj: "IO", path: "AnyFSPath"
) -> Optional[int]:
"""Try to deduce size of the file object.
This does not use `.seek`, instead it tries to deduce that from
AbstractBufferredFile.size if it exists, otherwise it fetches from
the filesystem.
"""
size = cast(Optional[int], getattr(fileobj, "size", None))
return size or fs.size(path)
def copy(
from_fs: "FileSystem",
from_path: "AnyPath",
from_path: "AnyFSPath",
to_fs: "FileSystem",
to_path: "AnyPath",
to_path: "AnyFSPath",
callback: "FsspecCallback" = DEFAULT_CALLBACK,
atomic: bool = False,
) -> None:
logger.debug("Transferring '%s' to '%s'", from_path, to_path)
if isinstance(from_fs, LocalFileSystem):
return to_fs.upload(from_path, to_path)
return upload(
to_fs, from_path, to_path, callback=callback, atomic=atomic
)
if isinstance(to_fs, LocalFileSystem):
return from_fs.download_file(from_path, to_path)
return from_fs.download_file(
from_path, to_path, callback=callback, atomic=atomic
)
with from_fs.open(from_path, mode="rb") as fobj:
size = from_fs.size(from_path)
return to_fs.upload(fobj, to_path, size=size)
size = fileobj_size(from_fs, fobj, from_path)
return upload(
to_fs, fobj, to_path, callback=callback, size=size, atomic=atomic
)
def _try_links(
links: List["AnyPath"],
links: List["str"],
from_fs: "FileSystem",
from_path: "AnyPath",
from_path: "AnyFSPath",
to_fs: "FileSystem",
to_path: "AnyPath",
to_path: "AnyFSPath",
callback: "FsspecCallback" = None,
atomic: bool = False,
) -> None:
error = None
while links:
link = links[0]
if link == "copy":
return _copy(from_fs, from_path, to_fs, to_path)
with FsspecCallback.as_tqdm_callback(
callback,
desc=from_fs.path.name(from_path),
bytes=True,
total=-1,
) as cb:
return copy(
from_fs,
from_path,
to_fs,
to_path,
callback=cb,
atomic=atomic,
)
try:
return _link(link, from_fs, from_path, to_fs, to_path)
@@ -81,11 +139,12 @@ def _try_links(
def transfer(
from_fs: "FileSystem",
from_path: "AnyPath",
from_path: "AnyFSPath",
to_fs: "FileSystem",
to_path: "AnyPath",
to_path: "AnyFSPath",
hardlink: bool = False,
links: Optional[List["AnyPath"]] = None,
links: List["str"] = None,
atomic: bool = False,
) -> None:
try:
assert not (hardlink and links)
@@ -93,7 +152,7 @@ def transfer(
links = links or ["reflink", "hardlink", "copy"]
else:
links = links or ["reflink", "copy"]
_try_links(links, from_fs, from_path, to_fs, to_path)
_try_links(links, from_fs, from_path, to_fs, to_path, atomic=atomic)
except OSError as exc:
# If the target file already exists, we are going to simply
# ignore the exception (#4992).
@@ -114,11 +173,11 @@ def transfer(
def _test_link(
link: "AnyPath",
link: "str",
from_fs: "FileSystem",
from_file: "AnyPath",
from_file: "AnyFSPath",
to_fs: "FileSystem",
to_file: "AnyPath",
to_file: "AnyFSPath",
) -> bool:
try:
_try_links([link], from_fs, from_file, to_fs, to_file)
@@ -136,12 +195,12 @@ def _test_link(
def test_links(
links: List["AnyPath"],
links: List["str"],
from_fs: "FileSystem",
from_path: "AnyPath",
from_path: "AnyFSPath",
to_fs: "FileSystem",
to_path: "AnyPath",
) -> List["AnyPath"]:
to_path: "AnyFSPath",
) -> List["AnyFSPath"]:
from dvc.utils import tmp_fname
from_file = from_fs.path.join(from_path, tmp_fname())
dvc/objects/db.py
+1
View File
@@ -101,6 +101,7 @@ class ObjectDB:
self.fs,
to_info,
hardlink=hardlink,
atomic=True,
)
def add(
dvc/repo/experiments/executor/ssh.py
+1 -1
View File
@@ -175,7 +175,7 @@ class SSHExecutor(BaseExecutor):
posixpath.join(self._repo_abspath, self.dvc_dir)
)
assert self._setup_script
fs.upload(self._setup_script, script_path)
fs.put_file(self._setup_script, script_path)
def _ssh_cmd(self, sshfs, cmd, chdir=None, **kwargs):
working_dir = chdir or self.root_dir
dvc/stage/cache.py
+32 -35
View File
@@ -6,11 +6,13 @@ from typing import TYPE_CHECKING, Optional
from funcy import cached_property, first
from dvc import fs
from dvc.data.transfer import _log_exceptions
from dvc.exceptions import DvcException
from dvc.utils import dict_sha256, relpath
if TYPE_CHECKING:
from dvc.objects.db.base import ObjectDB
from dvc.objects.db import ObjectDB
logger = logging.getLogger(__name__)
@@ -210,51 +212,46 @@ class StageCache:
)
cached_stage.checkout()
@staticmethod
def _transfer(func, from_remote, to_remote):
def transfer(self, from_odb, to_odb):
from dvc.fs._callback import FsspecCallback
from_fs = from_odb.fs
to_fs = from_odb.fs
func = _log_exceptions(fs.utils.copy)
runs = from_fs.path.join(from_odb.fs_path, "runs")
ret = []
if not from_fs.exists(runs):
return ret
runs = from_remote.fs.path.join(from_remote.fs_path, "runs")
if not from_remote.fs.exists(runs):
return []
from_path = from_remote.fs.path
for src in from_remote.fs.find(runs):
rel = from_path.relpath(src, from_remote.fs_path)
dst = to_remote.fs.path.join(to_remote.fs_path, rel)
key = to_remote.fs.path.parent(dst)
for src in from_fs.find(runs):
rel = from_fs.path.relpath(src, from_odb.fs_path)
dst = to_fs.path.join(to_odb.fs_path, rel)
key = to_fs.path.parent(dst)
# check if any build cache already exists for this key
# TODO: check if MaxKeys=1 or something like that applies
# or otherwise this will take a lot of time!
if to_remote.fs.exists(key) and first(to_remote.fs.find(key)):
if to_fs.exists(key) and first(to_fs.find(key)):
continue
func(src, dst)
ret.append(
(from_path.name(from_path.parent(src)), from_path.name(src))
)
src_name = from_fs.path.name(src)
parent_name = from_fs.path.name(from_fs.path.parent(src))
with FsspecCallback.as_tqdm_callback(
desc=src_name,
bytes=True,
total=-1,
) as cb:
func(from_fs, src, to_fs, dst, callback=cb, atomic=True)
ret.append((parent_name, src_name))
return ret
def push(self, remote: Optional[str], odb: Optional["ObjectDB"] = None):
from dvc.data.transfer import _log_exceptions
dest_odb = odb or self.repo.cloud.get_remote_odb(remote)
return self.transfer(self.repo.odb.local, dest_odb)
if odb is None:
odb = self.repo.cloud.get_remote_odb(remote)
return self._transfer(
_log_exceptions(odb.fs.upload),
self.repo.odb.local,
odb,
)
def pull(self, remote: Optional[str]):
from dvc.data.transfer import _log_exceptions
odb = self.repo.cloud.get_remote_odb(remote)
return self._transfer(
_log_exceptions(odb.fs.download),
odb,
self.repo.odb.local,
)
def pull(self, remote: Optional[str], odb: Optional["ObjectDB"] = None):
odb = odb or self.repo.cloud.get_remote_odb(remote)
return self.transfer(odb, self.repo.odb.local)
def get_used_objs(self, used_run_cache, *args, **kwargs):
"""Return used cache for the specified run-cached stages."""
tests/func/objects/db/test_index.py
+1 -1
View File
@@ -76,7 +76,7 @@ def test_partial_upload(tmp_dir, dvc, index, mocker):
tmp_dir.dvc_gen({"foo": "foo content"})
tmp_dir.dvc_gen({"bar": {"baz": "baz content"}})
original = LocalFileSystem.upload
original = LocalFileSystem.put_file
def unreliable_upload(self, from_file, to_info, name=None, **kwargs):
if "baz" in name:
tests/func/test_fs.py
+2 -19
View File
@@ -231,7 +231,7 @@ def test_fs_makedirs_on_upload_and_copy(dvc, cloud):
fs = cls(**config)
with io.BytesIO(b"foo") as stream:
fs.upload(stream, (cloud / "dir" / "foo").fs_path)
fs.put_file(stream, (cloud / "dir" / "foo").fs_path)
assert fs.isdir((cloud / "dir").fs_path)
assert fs.exists((cloud / "dir" / "foo").fs_path)
@@ -242,29 +242,12 @@ def test_fs_makedirs_on_upload_and_copy(dvc, cloud):
assert fs.exists((cloud / "dir2" / "foo").fs_path)
def test_upload_callback(tmp_dir, dvc, cloud):
tmp_dir.gen("foo", "foo")
cls, config, _ = get_cloud_fs(dvc, **cloud.config)
fs = cls(**config)
expected_size = os.path.getsize(tmp_dir / "foo")
callback = FsspecCallback()
fs.upload(
(tmp_dir / "foo").fs_path,
(cloud / "foo").fs_path,
callback=callback,
)
assert callback.size == expected_size
assert callback.value == expected_size
def test_download_callback(tmp_dir, dvc, cloud, local_cloud):
cls, config, _ = get_cloud_fs(dvc, **cloud.config)
fs = cls(**config)
(tmp_dir / "to_upload").write_text("foo")
fs.upload((tmp_dir / "to_upload").fs_path, (cloud / "foo").fs_path)
fs.put_file((tmp_dir / "to_upload").fs_path, (cloud / "foo").fs_path)
expected_size = fs.size((cloud / "foo").fs_path)
callback = FsspecCallback()
tests/func/test_remote.py
-18
View File
@@ -430,21 +430,3 @@ def test_push_incomplete_dir(tmp_dir, dvc, mocker, local_remote):
assert not remote_odb.exists(out.hash_info)
assert not remote_odb.exists(file_objs[0])
assert remote_odb.exists(file_objs[1])
def test_upload_exists(tmp_dir, dvc, local_remote):
tmp_dir.gen("foo", "foo")
odb = dvc.cloud.get_remote_odb("upstream")
# allow uploaded files to be writable for this test,
# normally they are set to read-only for DVC remotes
odb.fs.CACHE_MODE = 0o644
from_info = (tmp_dir / "foo").fs_path
to_info = odb.fs.path.join(odb.fs_path, "foo")
odb.fs.upload(from_info, to_info)
assert odb.fs.exists(to_info)
tmp_dir.gen("foo", "bar")
odb.fs.upload(from_info, to_info)
with odb.fs.open(to_info) as fobj:
assert fobj.read() == "bar"