diff --git a/dvc/data/checkout.py b/dvc/data/checkout.py index a683e8925..3097ffd77 100644 --- a/dvc/data/checkout.py +++ b/dvc/data/checkout.py @@ -138,14 +138,7 @@ class Link: cache.makedirs(cache.fs.path.parent(to_path)) try: - transfer( - cache.fs, - from_path, - to_fs, - to_path, - links=self._links, - atomic=False, - ) + transfer(cache.fs, from_path, to_fs, to_path, links=self._links) 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 662db51f8..04e2d1ad1 100644 --- a/dvc/data/stage.py +++ b/dvc/data/stage.py @@ -27,22 +27,21 @@ _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 fobj: - stream = HashedStreamReader(fobj) - size = fileobj_size(fs, fobj, from_fs_path) + with fs.open(from_fs_path, mode="rb") as stream: + stream = HashedStreamReader(stream) + size = fs.size(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.put_file(stream, tmp_info, size=size, callback=cb) + upload_odb.fs.upload(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 e9e65ccbc..fb923bd80 100644 --- a/dvc/fs/base.py +++ b/dvc/fs/base.py @@ -1,7 +1,6 @@ import logging import os import shutil -from contextlib import nullcontext from multiprocessing import cpu_count from typing import ( IO, @@ -361,6 +360,31 @@ 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, @@ -402,8 +426,7 @@ class FileSystem: ) as executor: list( executor.imap_unordered( - lambda args: download_files(*args, atomic=True), - pairs.items(), + lambda args: download_files(*args), pairs.items() ) ) @@ -412,7 +435,6 @@ class FileSystem: from_info: AnyFSPath, to_info: AnyFSPath, callback: FsspecCallback = None, - atomic: bool = False, ): from .local import localfs @@ -422,12 +444,7 @@ class FileSystem: desc=self.path.name(from_info), bytes=True, ) as cb: - context = ( - as_atomic(localfs, to_info, create_parents=True) - if atomic - else nullcontext(to_info) - ) - with context as tmp_file: + with as_atomic(localfs, to_info, create_parents=True) 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 ec87ee592..a28e2cd8b 100644 --- a/dvc/fs/local.py +++ b/dvc/fs/local.py @@ -87,7 +87,9 @@ class FsspecLocalFileSystem(AbstractFileSystem): def put_file(self, lpath, rpath, callback=None, **kwargs): parent = self._parent(rpath) makedirs(parent, exist_ok=True) - copyfile(lpath, rpath, callback=callback) + tmp_file = os.path.join(parent, tmp_fname()) + copyfile(lpath, tmp_file, callback=callback) + os.replace(tmp_file, rpath) 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 332e58f78..97ac12271 100644 --- a/dvc/fs/ssh.py +++ b/dvc/fs/ssh.py @@ -6,7 +6,9 @@ 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 @@ -116,3 +118,23 @@ 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 cd75fa3f3..3c96791e4 100644 --- a/dvc/fs/utils.py +++ b/dvc/fs/utils.py @@ -1,40 +1,20 @@ import errno import logging import os -from contextlib import nullcontext -from typing import IO, TYPE_CHECKING, List, Optional, Union, cast +from typing import TYPE_CHECKING, List, Optional -from ..utils.fs import as_atomic -from ._callback import DEFAULT_CALLBACK, FsspecCallback from .base import RemoteActionNotImplemented from .local import LocalFileSystem if TYPE_CHECKING: - from dvc.types import AnyPath - - from .base import AnyFSPath, FileSystem + from .base import FileSystem + from .types import AnyPath 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: "str", + link: "AnyPath", from_fs: "FileSystem", from_path: "AnyPath", to_fs: "FileSystem", @@ -54,74 +34,36 @@ def _link( ) from exc -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( +def _copy( from_fs: "FileSystem", - from_path: "AnyFSPath", + from_path: "AnyPath", to_fs: "FileSystem", - to_path: "AnyFSPath", - callback: "FsspecCallback" = DEFAULT_CALLBACK, - atomic: bool = False, + to_path: "AnyPath", ) -> None: - logger.debug("Transferring '%s' to '%s'", from_path, to_path) - if isinstance(from_fs, LocalFileSystem): - return upload( - to_fs, from_path, to_path, callback=callback, atomic=atomic - ) + return to_fs.upload(from_path, to_path) if isinstance(to_fs, LocalFileSystem): - return from_fs.download_file( - from_path, to_path, callback=callback, atomic=atomic - ) + return from_fs.download_file(from_path, to_path) with from_fs.open(from_path, mode="rb") as fobj: - size = fileobj_size(from_fs, fobj, from_path) - return upload( - to_fs, fobj, to_path, callback=callback, size=size, atomic=atomic - ) + size = from_fs.size(from_path) + return to_fs.upload(fobj, to_path, size=size) def _try_links( - links: List["str"], + links: List["AnyPath"], from_fs: "FileSystem", - from_path: "AnyFSPath", + from_path: "AnyPath", to_fs: "FileSystem", - to_path: "AnyFSPath", - callback: "FsspecCallback" = None, - atomic: bool = False, + to_path: "AnyPath", ) -> None: error = None while links: link = links[0] if link == "copy": - 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, - ) + return _copy(from_fs, from_path, to_fs, to_path) try: return _link(link, from_fs, from_path, to_fs, to_path) @@ -139,12 +81,11 @@ def _try_links( def transfer( from_fs: "FileSystem", - from_path: "AnyFSPath", + from_path: "AnyPath", to_fs: "FileSystem", - to_path: "AnyFSPath", + to_path: "AnyPath", hardlink: bool = False, - links: List["str"] = None, - atomic: bool = False, + links: Optional[List["AnyPath"]] = None, ) -> None: try: assert not (hardlink and links) @@ -152,7 +93,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, atomic=atomic) + _try_links(links, from_fs, from_path, to_fs, to_path) except OSError as exc: # If the target file already exists, we are going to simply # ignore the exception (#4992). @@ -173,11 +114,11 @@ def transfer( def _test_link( - link: "str", + link: "AnyPath", from_fs: "FileSystem", - from_file: "AnyFSPath", + from_file: "AnyPath", to_fs: "FileSystem", - to_file: "AnyFSPath", + to_file: "AnyPath", ) -> bool: try: _try_links([link], from_fs, from_file, to_fs, to_file) @@ -195,12 +136,12 @@ def _test_link( def test_links( - links: List["str"], + links: List["AnyPath"], from_fs: "FileSystem", - from_path: "AnyFSPath", + from_path: "AnyPath", to_fs: "FileSystem", - to_path: "AnyFSPath", -) -> List["AnyFSPath"]: + to_path: "AnyPath", +) -> List["AnyPath"]: 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 0e8e0b9fc..1c03a3e6f 100644 --- a/dvc/objects/db.py +++ b/dvc/objects/db.py @@ -101,7 +101,6 @@ 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 adc9cc7f1..29fc79d83 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.put_file(self._setup_script, script_path) + fs.upload(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 6e1e6c0b3..04e225f93 100644 --- a/dvc/stage/cache.py +++ b/dvc/stage/cache.py @@ -6,13 +6,11 @@ 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 import ObjectDB + from dvc.objects.db.base import ObjectDB logger = logging.getLogger(__name__) @@ -212,46 +210,51 @@ class StageCache: ) cached_stage.checkout() - 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") - + @staticmethod + def _transfer(func, from_remote, to_remote): ret = [] - if not from_fs.exists(runs): - return ret - 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) + 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) # 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_fs.exists(key) and first(to_fs.find(key)): + if to_remote.fs.exists(key) and first(to_remote.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): - dest_odb = odb or self.repo.cloud.get_remote_odb(remote) - return self.transfer(self.repo.odb.local, dest_odb) + from dvc.data.transfer import _log_exceptions - 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) + 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 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 4aa6d3ca3..6df598fe5 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.put_file + original = LocalFileSystem.upload 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 affbdc4f1..8988c2237 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.put_file(stream, (cloud / "dir" / "foo").fs_path) + fs.upload(stream, (cloud / "dir" / "foo").fs_path) assert fs.isdir((cloud / "dir").fs_path) assert fs.exists((cloud / "dir" / "foo").fs_path) @@ -242,12 +242,29 @@ 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.put_file((tmp_dir / "to_upload").fs_path, (cloud / "foo").fs_path) + fs.upload((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 dbfe080f9..ff844ed7e 100644 --- a/tests/func/test_remote.py +++ b/tests/func/test_remote.py @@ -430,3 +430,21 @@ 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"