From dd187df6674688ad82f0e933b589a8953c465e1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Thu, 28 Apr 2022 22:29:34 +0545 Subject: [PATCH] fs.utils: add upload that can be conditionally atomic The PR also makes odb.add operations atomic, but the checkout is still not atomic. --- dvc/data/checkout.py | 9 ++- dvc/data/stage.py | 9 ++- dvc/fs/base.py | 37 +++------ dvc/fs/local.py | 4 +- dvc/fs/ssh.py | 22 ------ dvc/fs/utils.py | 111 ++++++++++++++++++++------- dvc/objects/db.py | 1 + dvc/repo/experiments/executor/ssh.py | 2 +- dvc/stage/cache.py | 67 ++++++++-------- tests/func/objects/db/test_index.py | 2 +- tests/func/test_fs.py | 21 +---- tests/func/test_remote.py | 18 ----- 12 files changed, 146 insertions(+), 157 deletions(-) diff --git a/dvc/data/checkout.py b/dvc/data/checkout.py index 3097ffd77..a683e8925 100644 --- a/dvc/data/checkout.py +++ b/dvc/data/checkout.py @@ -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: diff --git a/dvc/data/stage.py b/dvc/data/stage.py index 04e2d1ad1..662db51f8 100644 --- a/dvc/data/stage.py +++ b/dvc/data/stage.py @@ -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) diff --git a/dvc/fs/base.py b/dvc/fs/base.py index fb923bd80..e9e65ccbc 100644 --- a/dvc/fs/base.py +++ b/dvc/fs/base.py @@ -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) diff --git a/dvc/fs/local.py b/dvc/fs/local.py index a28e2cd8b..ec87ee592 100644 --- a/dvc/fs/local.py +++ b/dvc/fs/local.py @@ -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) diff --git a/dvc/fs/ssh.py b/dvc/fs/ssh.py index 97ac12271..332e58f78 100644 --- a/dvc/fs/ssh.py +++ b/dvc/fs/ssh.py @@ -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 - ) diff --git a/dvc/fs/utils.py b/dvc/fs/utils.py index 3c96791e4..cd75fa3f3 100644 --- a/dvc/fs/utils.py +++ b/dvc/fs/utils.py @@ -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()) diff --git a/dvc/objects/db.py b/dvc/objects/db.py index 1c03a3e6f..0e8e0b9fc 100644 --- a/dvc/objects/db.py +++ b/dvc/objects/db.py @@ -101,6 +101,7 @@ class ObjectDB: self.fs, to_info, hardlink=hardlink, + atomic=True, ) def add( diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py index 29fc79d83..adc9cc7f1 100644 --- a/dvc/repo/experiments/executor/ssh.py +++ b/dvc/repo/experiments/executor/ssh.py @@ -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 diff --git a/dvc/stage/cache.py b/dvc/stage/cache.py index 04e225f93..6e1e6c0b3 100644 --- a/dvc/stage/cache.py +++ b/dvc/stage/cache.py @@ -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.""" diff --git a/tests/func/objects/db/test_index.py b/tests/func/objects/db/test_index.py index 6df598fe5..4aa6d3ca3 100644 --- a/tests/func/objects/db/test_index.py +++ b/tests/func/objects/db/test_index.py @@ -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: diff --git a/tests/func/test_fs.py b/tests/func/test_fs.py index 8988c2237..affbdc4f1 100644 --- a/tests/func/test_fs.py +++ b/tests/func/test_fs.py @@ -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() diff --git a/tests/func/test_remote.py b/tests/func/test_remote.py index ff844ed7e..dbfe080f9 100644 --- a/tests/func/test_remote.py +++ b/tests/func/test_remote.py @@ -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"