Implement straight-to-remote functionality (#5198)

* Implement export-to-remote functionality

* Redistrubition of functionality from tree to remote

* Pass chunk size into the open_url

* Move functionality away from remote to cache

* Basic progress bars for both transfer types

* Proper job calculation

* Redesign .open() usage in order to work with different trees

* Align the signatures of other open() functions to accept **kwargs

* Implement chunked uploading to the ssh

* Use local cache to compute directory hash (md5)

* S3: multipart upload test

* Wrap local cache's hash_to_path with PathInfo

* Use .isdir() instead of a try/except statement for calculating the from_infos

* New CLI layout (add/import-url)

* Adjust CLI tests to accept new arguments

* tests: import_url/add func/unit tests

* Introduce cache.transfer()

* Invalid CLI combination tests

* Store dir_info key in the transfer_directory

* LocalTree.upload_fobj

* Tests for utils

* Lazy imports on some places

* --to-remote migration (and repo._transfer)

* Use proper temporary location, more efficient pbar calc

* Directory transfer logic rewrite

* Use shutil.copyfileobj()

* Rebase...

* Better tests

* Inline transfer operation to .add() / .transfer()

* Open in the 'rb' mode for the tests

* More tests, better invalid option handling at repo level

* .read() certain amount of bytes to validate the content

* Better invalid argument handling

* Better progress bars by using .getsize() API

* For fallback upload, use the from_info.name instead of temporary naem

* introduce --jobs for add/import-url

Author: isidentical

dvc/cache/base.py

@@ -3,6 +3,7 @@
 import json
 import logging
 import os
+from concurrent import futures
 from concurrent.futures import ThreadPoolExecutor
 from copy import copy
 from typing import Optional
@@ -247,6 +248,116 @@ def _save_file(self, path_info, tree, hash_info, save_link=True, **kwargs):
 
         self.tree.state.save(cache_info, hash_info)
 
+    def _transfer_file_as_whole(self, from_tree, from_info):
+        from dvc.utils import tmp_fname
+
+        # When we can't use the chunked upload, we have to first download
+        # and then calculate the hash as if it were a local file and then
+        # upload it.
+        local_tree = self.repo.cache.local.tree
+        local_info = local_tree.path_info / tmp_fname()
+
+        from_tree.download(from_info, local_info)
+        hash_info = local_tree.get_file_hash(local_info)
+
+        self.tree.upload(
+            local_info,
+            self.tree.hash_to_path_info(hash_info.value),
+            name=from_info.name,
+        )
+        return hash_info
+
+    def _transfer_file_as_chunked(self, from_tree, from_info):
+        from dvc.utils import tmp_fname
+        from dvc.utils.stream import HashedStreamReader
+
+        tmp_info = self.tree.path_info / tmp_fname()
+        with from_tree.open(
+            from_info, mode="rb", chunk_size=from_tree.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.
+            self.tree.upload_fobj(
+                stream_reader,
+                tmp_info,
+                total=from_tree.getsize(from_info),
+                desc=from_info.name,
+            )
+
+        hash_info = stream_reader.hash_info
+        self.tree.move(tmp_info, self.tree.hash_to_path_info(hash_info.value))
+        return hash_info
+
+    def _transfer_file(self, from_tree, from_info):
+        try:
+            hash_info = self._transfer_file_as_chunked(from_tree, from_info)
+        except RemoteActionNotImplemented:
+            hash_info = self._transfer_file_as_whole(from_tree, from_info)
+
+        return hash_info
+
+    def _transfer_directory_contents(self, from_tree, from_info, jobs, pbar):
+        rel_path_infos = {}
+        from_infos = from_tree.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(self._transfer_file), from_tree, entry_info
+                )
+                rel_path_infos[task] = entry_info.relative_to(from_info)
+                yield task
+
+        pbar.total = 0
+        with 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(
+        self, from_tree, from_info, jobs, no_progress_bar=False
+    ):
+        dir_info = DirInfo()
+
+        with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar:
+            for entry_info, entry_hash in self._transfer_directory_contents(
+                from_tree, from_info, jobs, pbar
+            ):
+                dir_info.trie[entry_info.parts] = entry_hash
+
+        local_cache = self.repo.cache.local
+        (
+            hash_info,
+            to_info,
+        ) = local_cache._get_dir_info_hash(  # pylint: disable=protected-access
+            dir_info
+        )
+
+        self.tree.upload(to_info, self.tree.hash_to_path_info(hash_info.value))
+        return hash_info
+
+    def transfer(self, from_tree, from_info, jobs=None, no_progress_bar=False):
+        jobs = jobs or min((from_tree.jobs, self.tree.jobs))
+
+        if from_tree.isdir(from_info):
+            return self._transfer_directory(
+                from_tree,
+                from_info,
+                jobs=jobs,
+                no_progress_bar=no_progress_bar,
+            )
+        else:
+            return self._transfer_file(from_tree, from_info)
+
     def _cache_is_copy(self, path_info):
         """Checks whether cache uses copies."""
         if self.cache_type_confirmed:

dvc/command/add.py

@@ -26,6 +26,10 @@ def run(self):
                 external=self.args.external,
                 glob=self.args.glob,
                 desc=self.args.desc,
+                out=self.args.out,
+                remote=self.args.remote,
+                to_remote=self.args.to_remote,
+                jobs=self.args.jobs,
             )
 
         except DvcException:
@@ -74,6 +78,35 @@ def add_parser(subparsers, parent_parser):
         help="Specify name of the DVC-file this command will generate.",
         metavar="<filename>",
     )
+    parser.add_argument(
+        "-o",
+        "--out",
+        help="Destination path to put files to.",
+        metavar="<path>",
+    )
+    parser.add_argument(
+        "--to-remote",
+        action="store_true",
+        default=False,
+        help="Download it directly to the remote",
+    )
+    parser.add_argument(
+        "-r",
+        "--remote",
+        help="Remote storage to download to",
+        metavar="<name>",
+    )
+    parser.add_argument(
+        "-j",
+        "--jobs",
+        type=int,
+        help=(
+            "Number of jobs to run simultaneously. "
+            "The default value is 4 * cpu_count(). "
+            "For SSH remotes, the default is 4. "
+        ),
+        metavar="<number>",
+    )
     parser.add_argument(
         "--desc",
         type=str,

dvc/command/imp_url.py

@@ -16,7 +16,10 @@ def run(self):
                 out=self.args.out,
                 fname=self.args.file,
                 no_exec=self.args.no_exec,
+                remote=self.args.remote,
+                to_remote=self.args.to_remote,
                 desc=self.args.desc,
+                jobs=self.args.jobs,
             )
         except DvcException:
             logger.exception(
@@ -68,6 +71,29 @@ def add_parser(subparsers, parent_parser):
         default=False,
         help="Only create DVC-file without actually downloading it.",
     )
+    import_parser.add_argument(
+        "--to-remote",
+        action="store_true",
+        default=False,
+        help="Download it directly to the remote",
+    )
+    import_parser.add_argument(
+        "-r",
+        "--remote",
+        help="Remote storage to download to",
+        metavar="<name>",
+    )
+    import_parser.add_argument(
+        "-j",
+        "--jobs",
+        type=int,
+        help=(
+            "Number of jobs to run simultaneously. "
+            "The default value is 4 * cpu_count(). "
+            "For SSH remotes, the default is 4. "
+        ),
+        metavar="<number>",
+    )
     import_parser.add_argument(
         "--desc",
         type=str,

dvc/data_cloud.py

@@ -92,6 +92,24 @@ def pull(
             show_checksums=show_checksums,
         )
 
+    def transfer(self, source, jobs=None, remote=None, command=None):
+        """Transfer data items in a cloud-agnostic way.
+
+        Args:
+            source (str): url for the source location.
+            jobs (int): number of jobs that can be running simultaneously.
+            remote (dvc.remote.base.BaseRemote): optional remote to compare
+                cache to. By default remote from core.remote config option
+                is used.
+            command (bool): the command which is benefitting from this function
+                (to be used for reporting better error messages).
+        """
+        from dvc.tree import get_cloud_tree
+
+        from_tree = get_cloud_tree(self.repo, url=source)
+        remote = self.get_remote(remote, command)
+        return remote.transfer(from_tree, from_tree.path_info, jobs=jobs)
+
     def status(
         self,
         cache,

dvc/istextfile.py

@@ -7,19 +7,7 @@
 TEXT_CHARS = bytes(range(32, 127)) + b"\n\r\t\f\b"
 
 
-def istextfile(fname, blocksize=512, tree=None):
-    """ Uses heuristics to guess whether the given file is text or binary,
-        by reading a single block of bytes from the file.
-        If more than 30% of the chars in the block are non-text, or there
-        are NUL ('\x00') bytes in the block, assume this is a binary file.
-    """
-    if tree:
-        open_func = tree.open
-    else:
-        open_func = open
-    with open_func(fname, "rb") as fobj:
-        block = fobj.read(blocksize)
-
+def istextblock(block):
     if not block:
         # An empty file is considered a valid text file
         return True
@@ -32,3 +20,18 @@ def istextfile(fname, blocksize=512, tree=None):
     # occurrences of TEXT_CHARS from the block
     nontext = block.translate(None, TEXT_CHARS)
     return float(len(nontext)) / len(block) <= 0.30
+
+
+def istextfile(fname, blocksize=512, tree=None):
+    """ Uses heuristics to guess whether the given file is text or binary,
+        by reading a single block of bytes from the file.
+        If more than 30% of the chars in the block are non-text, or there
+        are NUL ('\x00') bytes in the block, assume this is a binary file.
+    """
+    if tree:
+        open_func = tree.open
+    else:
+        open_func = open
+    with open_func(fname, "rb") as fobj:
+        block = fobj.read(blocksize)
+    return istextblock(block)

dvc/remote/base.py

@@ -472,6 +472,11 @@ def pull(self, cache, named_cache, jobs=None, show_checksums=False):
 
         return ret
 
+    def transfer(self, from_tree, from_info, jobs=None, no_progress_bar=False):
+        return self.cache.transfer(
+            from_tree, from_info, jobs=jobs, no_progress_bar=no_progress_bar
+        )
+
     @staticmethod
     def _log_missing_caches(hash_info_dict):
         missing_caches = [