diff --git a/dvc/checkout.py b/dvc/checkout.py index 606d60755..49f78e06f 100644 --- a/dvc/checkout.py +++ b/dvc/checkout.py @@ -9,7 +9,8 @@ from dvc.exceptions import ( ConfirmRemoveError, DvcException, ) -from dvc.objects import ObjectFormatError, check, load +from dvc.objects import check, load +from dvc.objects.errors import ObjectFormatError from dvc.objects.stage import stage from dvc.remote.slow_link_detection import ( # type: ignore[attr-defined] slow_link_guard, diff --git a/dvc/objects/__init__.py b/dvc/objects/__init__.py index 8b8c9c611..a580030bc 100644 --- a/dvc/objects/__init__.py +++ b/dvc/objects/__init__.py @@ -1,77 +1,13 @@ -import itertools import logging -from concurrent import futures -from dvc.exceptions import DvcException from dvc.progress import Tqdm -from .stage import get_file_hash +from .tree import Tree logger = logging.getLogger(__name__) -class ObjectError(DvcException): - pass - - -class ObjectFormatError(ObjectError): - pass - - -class HashFile: - def __init__(self, path_info, fs, hash_info): - self.path_info = path_info - self.fs = fs - self.hash_info = hash_info - - @property - def size(self): - if not (self.path_info and self.fs): - return None - return self.fs.getsize(self.path_info) - - def __len__(self): - return 1 - - def __str__(self): - return f"object {self.hash_info}" - - def __bool__(self): - return bool(self.hash_info) - - def __eq__(self, other): - if not isinstance(other, HashFile): - return False - return ( - self.path_info == other.path_info - and self.fs == other.fs - and self.hash_info == other.hash_info - ) - - def check(self, odb): - actual = get_file_hash( - self.path_info, self.fs, self.hash_info.name, odb.repo.state - ) - - logger.trace( - "cache '%s' expected '%s' actual '%s'", - self.path_info, - self.hash_info, - actual, - ) - - assert actual.name == self.hash_info.name - if actual.value.split(".")[0] != self.hash_info.value.split(".")[0]: - raise ObjectFormatError(f"{self} is corrupted") - - @classmethod - def load(cls, odb, hash_info): - return odb.get(hash_info) - - def save(odb, obj, **kwargs): - from .tree import Tree - if isinstance(obj, Tree): for _, entry in Tqdm(obj): odb.add(entry.path_info, entry.fs, entry.hash_info, **kwargs) @@ -79,8 +15,6 @@ def save(odb, obj, **kwargs): def check(odb, obj): - from .tree import Tree - odb.check(obj.hash_info) if isinstance(obj, Tree): @@ -89,91 +23,6 @@ def check(odb, obj): def load(odb, hash_info): - from .tree import Tree - if hash_info.isdir: return Tree.load(odb, hash_info) return odb.get(hash_info) - - -def _transfer_file(odb, from_fs, from_info): - from dvc.utils import tmp_fname - from dvc.utils.stream import HashedStreamReader - - tmp_info = odb.fs.path_info / tmp_fname() - with from_fs.open( - from_info, mode="rb", chunk_size=from_fs.CHUNK_SIZE - ) as stream: - stream_reader = HashedStreamReader(stream) - # Since we don't know the hash beforehand, we'll - # upload it to a temporary location and then move - # it. - odb.fs.upload_fobj( - stream_reader, - tmp_info, - total=from_fs.getsize(from_info), - desc=from_info.name, - ) - - hash_info = stream_reader.hash_info - return tmp_info, hash_info - - -def _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar): - rel_path_infos = {} - from_infos = from_fs.walk_files(from_info) - - def create_tasks(executor, amount): - for entry_info in itertools.islice(from_infos, amount): - pbar.total += 1 - task = executor.submit( - pbar.wrap_fn(_transfer_file), odb, from_fs, entry_info - ) - rel_path_infos[task] = entry_info.relative_to(from_info) - yield task - - pbar.total = 0 - with futures.ThreadPoolExecutor(max_workers=jobs) as executor: - tasks = set(create_tasks(executor, jobs * 5)) - - while tasks: - done, tasks = futures.wait( - tasks, return_when=futures.FIRST_COMPLETED - ) - tasks.update(create_tasks(executor, len(done))) - for task in done: - yield rel_path_infos.pop(task), task.result() - - -def _transfer_directory(odb, from_fs, from_info, jobs, no_progress_bar=False): - from .tree import Tree - - tree = Tree(None, None, None) - - with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar: - for ( - entry_info, - (entry_tmp_info, entry_hash), - ) in _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar): - obj = HashFile(entry_tmp_info, odb.fs, entry_hash) - tree.add(entry_info.parts, obj) - - tree.digest() - save(odb, tree) - return tree.hash_info - - -def transfer(odb, from_fs, from_info, jobs=None, no_progress_bar=False): - jobs = jobs or min((from_fs.jobs, odb.fs.jobs)) - - if from_fs.isdir(from_info): - return _transfer_directory( - odb, - from_fs, - from_info, - jobs=jobs, - no_progress_bar=no_progress_bar, - ) - tmp_info, hash_info = _transfer_file(odb, from_fs, from_info) - odb.add(tmp_info, odb.fs, hash_info) - return hash_info diff --git a/dvc/objects/db/base.py b/dvc/objects/db/base.py index cbe5b04ca..4288f0c7c 100644 --- a/dvc/objects/db/base.py +++ b/dvc/objects/db/base.py @@ -4,10 +4,10 @@ from concurrent.futures import ThreadPoolExecutor from copy import copy from typing import Optional +from dvc.objects.errors import ObjectFormatError +from dvc.objects.file import HashFile from dvc.progress import Tqdm -from .. import HashFile, ObjectFormatError - logger = logging.getLogger(__name__) diff --git a/dvc/objects/db/local.py b/dvc/objects/db/local.py index 46a6c296e..5a87084a9 100644 --- a/dvc/objects/db/local.py +++ b/dvc/objects/db/local.py @@ -6,7 +6,7 @@ from funcy import cached_property from shortuuid import uuid from dvc.hash_info import HashInfo -from dvc.objects import ObjectFormatError +from dvc.objects.errors import ObjectFormatError from dvc.path_info import PathInfo from dvc.progress import Tqdm from dvc.utils import relpath diff --git a/dvc/objects/errors.py b/dvc/objects/errors.py new file mode 100644 index 000000000..c7616b65e --- /dev/null +++ b/dvc/objects/errors.py @@ -0,0 +1,9 @@ +from dvc.exceptions import DvcException + + +class ObjectError(DvcException): + pass + + +class ObjectFormatError(ObjectError): + pass diff --git a/dvc/objects/file.py b/dvc/objects/file.py new file mode 100644 index 000000000..933913a83 --- /dev/null +++ b/dvc/objects/file.py @@ -0,0 +1,57 @@ +import logging + +from .errors import ObjectFormatError +from .stage import get_file_hash + +logger = logging.getLogger(__name__) + + +class HashFile: + def __init__(self, path_info, fs, hash_info): + self.path_info = path_info + self.fs = fs + self.hash_info = hash_info + + @property + def size(self): + if not (self.path_info and self.fs): + return None + return self.fs.getsize(self.path_info) + + def __len__(self): + return 1 + + def __str__(self): + return f"object {self.hash_info}" + + def __bool__(self): + return bool(self.hash_info) + + def __eq__(self, other): + if not isinstance(other, HashFile): + return False + return ( + self.path_info == other.path_info + and self.fs == other.fs + and self.hash_info == other.hash_info + ) + + def check(self, odb): + actual = get_file_hash( + self.path_info, self.fs, self.hash_info.name, odb.repo.state + ) + + logger.trace( + "cache '%s' expected '%s' actual '%s'", + self.path_info, + self.hash_info, + actual, + ) + + assert actual.name == self.hash_info.name + if actual.value.split(".")[0] != self.hash_info.value.split(".")[0]: + raise ObjectFormatError(f"{self} is corrupted") + + @classmethod + def load(cls, odb, hash_info): + return odb.get(hash_info) diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py index 71a8a2d3b..ac8313f46 100644 --- a/dvc/objects/stage.py +++ b/dvc/objects/stage.py @@ -77,7 +77,7 @@ def _iter_hashes(path_info, fs, name, state, **kwargs): def _build_tree(path_info, fs, name, state, **kwargs): - from . import HashFile + from .file import HashFile from .tree import Tree tree = Tree(None, None, None) @@ -176,7 +176,7 @@ def stage(odb, path_info, fs, name, **kwargs): if fs.isdir(path_info): obj = _get_tree_obj(path_info, fs, name, odb, state, **kwargs) else: - from . import HashFile + from .file import HashFile hash_info = get_file_hash(path_info, fs, name, state) obj = HashFile(path_info, fs, hash_info) diff --git a/dvc/objects/transfer.py b/dvc/objects/transfer.py new file mode 100644 index 000000000..c6fc0268e --- /dev/null +++ b/dvc/objects/transfer.py @@ -0,0 +1,92 @@ +import itertools +import logging +from concurrent import futures + +from dvc.progress import Tqdm + +from . import save +from .file import HashFile +from .tree import Tree + +logger = logging.getLogger(__name__) + + +def _transfer_file(odb, from_fs, from_info): + from dvc.utils import tmp_fname + from dvc.utils.stream import HashedStreamReader + + tmp_info = odb.fs.path_info / tmp_fname() + with from_fs.open( + from_info, mode="rb", chunk_size=from_fs.CHUNK_SIZE + ) as stream: + stream_reader = HashedStreamReader(stream) + # Since we don't know the hash beforehand, we'll + # upload it to a temporary location and then move + # it. + odb.fs.upload_fobj( + stream_reader, + tmp_info, + total=from_fs.getsize(from_info), + desc=from_info.name, + ) + + hash_info = stream_reader.hash_info + return tmp_info, hash_info + + +def _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar): + rel_path_infos = {} + from_infos = from_fs.walk_files(from_info) + + def create_tasks(executor, amount): + for entry_info in itertools.islice(from_infos, amount): + pbar.total += 1 + task = executor.submit( + pbar.wrap_fn(_transfer_file), odb, from_fs, entry_info + ) + rel_path_infos[task] = entry_info.relative_to(from_info) + yield task + + pbar.total = 0 + with futures.ThreadPoolExecutor(max_workers=jobs) as executor: + tasks = set(create_tasks(executor, jobs * 5)) + + while tasks: + done, tasks = futures.wait( + tasks, return_when=futures.FIRST_COMPLETED + ) + tasks.update(create_tasks(executor, len(done))) + for task in done: + yield rel_path_infos.pop(task), task.result() + + +def _transfer_directory(odb, from_fs, from_info, jobs, no_progress_bar=False): + tree = Tree(None, None, None) + + with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar: + for ( + entry_info, + (entry_tmp_info, entry_hash), + ) in _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar): + obj = HashFile(entry_tmp_info, odb.fs, entry_hash) + tree.add(entry_info.parts, obj) + + tree.digest() + save(odb, tree) + return tree.hash_info + + +def transfer(odb, from_fs, from_info, jobs=None, no_progress_bar=False): + jobs = jobs or min((from_fs.jobs, odb.fs.jobs)) + + if from_fs.isdir(from_info): + return _transfer_directory( + odb, + from_fs, + from_info, + jobs=jobs, + no_progress_bar=no_progress_bar, + ) + tmp_info, hash_info = _transfer_file(odb, from_fs, from_info) + odb.add(tmp_info, odb.fs, hash_info) + return hash_info diff --git a/dvc/objects/tree.py b/dvc/objects/tree.py index b885515ac..c146a9880 100644 --- a/dvc/objects/tree.py +++ b/dvc/objects/tree.py @@ -4,7 +4,8 @@ import posixpath from funcy import cached_property -from . import HashFile, ObjectFormatError +from .errors import ObjectFormatError +from .file import HashFile from .stage import get_file_hash logger = logging.getLogger(__name__) diff --git a/dvc/output/base.py b/dvc/output/base.py index 5011e3e0f..e643ef663 100644 --- a/dvc/output/base.py +++ b/dvc/output/base.py @@ -18,6 +18,7 @@ from dvc.exceptions import ( ) from dvc.hash_info import HashInfo from dvc.objects.db import NamedCache +from dvc.objects.errors import ObjectFormatError from dvc.objects.stage import stage as ostage from ..fs.base import BaseFileSystem @@ -220,7 +221,7 @@ class BaseOutput: try: objects.check(self.odb, obj) return False - except (FileNotFoundError, objects.ObjectFormatError): + except (FileNotFoundError, ObjectFormatError): return True def workspace_status(self): @@ -487,7 +488,7 @@ class BaseOutput: try: objects.check(self.odb, self.odb.get(self.hash_info)) - except (FileNotFoundError, objects.ObjectFormatError): + except (FileNotFoundError, ObjectFormatError): self.repo.cloud.pull( NamedCache.make("local", self.hash_info.value, str(self)), show_checksums=False, @@ -496,7 +497,7 @@ class BaseOutput: try: self.obj = objects.load(self.odb, self.hash_info) - except (objects.ObjectFormatError, FileNotFoundError): + except (FileNotFoundError, ObjectFormatError): self.obj = None return self.obj @@ -532,7 +533,7 @@ class BaseOutput: try: objects.check(self.odb, self.odb.get(self.hash_info)) - except (FileNotFoundError, objects.ObjectFormatError): + except (FileNotFoundError, ObjectFormatError): msg = ( "Missing cache for directory '{}'. " "Cache for files inside will be lost. " diff --git a/dvc/remote/base.py b/dvc/remote/base.py index e1d8f14bd..d92e09093 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -506,7 +506,7 @@ class Remote: return ret def transfer(self, from_fs, from_info, jobs=None, no_progress_bar=False): - from dvc.objects import transfer + from dvc.objects.transfer import transfer return transfer( self.odb, diff --git a/dvc/repo/add.py b/dvc/repo/add.py index 5e63f2bf6..f8176bb7b 100644 --- a/dvc/repo/add.py +++ b/dvc/repo/add.py @@ -173,7 +173,7 @@ def _process_stages( ) else: from dvc.fs import get_cloud_fs - from dvc.objects import transfer + from dvc.objects.transfer import transfer from_fs = get_cloud_fs(repo, url=target) out.hash_info = transfer( diff --git a/tests/func/test_odb.py b/tests/func/test_odb.py index 4b7fa2b56..024000d04 100644 --- a/tests/func/test_odb.py +++ b/tests/func/test_odb.py @@ -7,8 +7,8 @@ import pytest from dvc.hash_info import HashInfo from dvc.main import main -from dvc.objects import ObjectFormatError from dvc.objects.db import ODBManager +from dvc.objects.errors import ObjectFormatError from dvc.utils import relpath from tests.basic_env import TestDir, TestDvc diff --git a/tests/unit/objects/test_tree.py b/tests/unit/objects/test_tree.py index 7e57f0fd0..41c567e00 100644 --- a/tests/unit/objects/test_tree.py +++ b/tests/unit/objects/test_tree.py @@ -3,7 +3,7 @@ from operator import itemgetter import pytest from dvc.hash_info import HashInfo -from dvc.objects import HashFile +from dvc.objects.file import HashFile from dvc.objects.tree import Tree, _merge