From 86811aac209180dd1e516fef3087d904b9e7f4eb Mon Sep 17 00:00:00 2001
From: Alexander Schepanovski <suor.web@gmail.com>
Date: Wed, 23 Oct 2019 22:02:50 +0700
Subject: [PATCH] imp: dedup and parallelize fetching imported deps (#2654)

* imp: dedup and parallelize external fetches

This prevents same fetches for external artifacts, groups them by source
and fetches in parallel.

* imp: test that fetch works on non workspace deps
---
 dvc/cache.py              |  8 ++++--
 dvc/dependency/repo.py    |  5 ++++
 dvc/repo/__init__.py      |  3 ++-
 dvc/repo/fetch.py         | 53 +++++++++++++++++++++++++++++----------
 tests/func/test_import.py | 17 +++++++++++++
 5 files changed, 70 insertions(+), 16 deletions(-)

diff --git a/dvc/cache.py b/dvc/cache.py
index 0b4c90404..4ca3b4f62 100644
--- a/dvc/cache.py
+++ b/dvc/cache.py
@@ -109,7 +109,7 @@ class Cache(object):
 class NamedCache(object):
     def __init__(self):
         self._items = defaultdict(lambda: defaultdict(set))
-        self.repo = []
+        self.external = defaultdict(set)
 
     @classmethod
     def make(cls, scheme, checksum, name):
@@ -123,6 +123,9 @@ class NamedCache(object):
     def add(self, scheme, checksum, name):
         self._items[scheme][checksum].add(name)
 
+    def add_external(self, url, rev, path):
+        self.external[url, rev].add(path)
+
     def update(self, cache, suffix=""):
         for scheme, src in cache._items.items():
             dst = self._items[scheme]
@@ -132,4 +135,5 @@ class NamedCache(object):
                 else:
                     dst[checksum].update(names)
 
-        self.repo.extend(cache.repo)
+        for repo_pair, files in cache.external.items():
+            self.external[repo_pair].update(files)
diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py
index b56bc8008..376a38a6c 100644
--- a/dvc/dependency/repo.py
+++ b/dvc/dependency/repo.py
@@ -35,6 +35,11 @@ class DependencyREPO(DependencyLOCAL):
     def is_in_repo(self):
         return False
 
+    @property
+    def repo_pair(self):
+        d = self.def_repo
+        return d[self.PARAM_URL], d[self.PARAM_REV_LOCK] or d[self.PARAM_REV]
+
     def __str__(self):
         return "{} ({})".format(self.def_path, self.def_repo[self.PARAM_URL])
 
diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py
index 97dfa3103..017d7b4d4 100644
--- a/dvc/repo/__init__.py
+++ b/dvc/repo/__init__.py
@@ -253,7 +253,8 @@ class Repo(object):
 
             for stage in stages:
                 if stage.is_repo_import:
-                    cache.repo += stage.deps
+                    dep, = stage.deps
+                    cache.external[dep.repo_pair].add(dep.def_path)
                     continue
 
                 for out in stage.outs:
diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py
index 3f0bce1eb..702946163 100644
--- a/dvc/repo/fetch.py
+++ b/dvc/repo/fetch.py
@@ -5,6 +5,8 @@ import logging
 from dvc.config import NoRemoteError
 from dvc.exceptions import DownloadError, OutputNotFoundError
 from dvc.scm.base import CloneError
+from dvc.external_repo import external_repo
+from dvc.cache import NamedCache
 
 
 logger = logging.getLogger(__name__)
@@ -52,25 +54,50 @@ def _fetch(
             used, jobs, remote=remote, show_checksums=show_checksums
         )
     except NoRemoteError:
-        if not used.repo and used["local"]:
+        if not used.external and used["local"]:
             raise
-
     except DownloadError as exc:
         failed += exc.amount
 
-    for dep in used.repo:
-        try:
-            out = dep.fetch()
-            downloaded += out.get_files_number()
-        except DownloadError as exc:
-            failed += exc.amount
-        except (CloneError, OutputNotFoundError):
-            failed += 1
-            logger.exception(
-                "failed to fetch data for '{}'".format(dep.stage.outs[0])
-            )
+    for (repo_url, repo_rev), files in used.external.items():
+        d, f = _fetch_external(self, repo_url, repo_rev, files)
+        downloaded += d
+        failed += f
 
     if failed:
         raise DownloadError(failed)
 
     return downloaded
+
+
+def _fetch_external(self, repo_url, repo_rev, files):
+    failed = 0
+
+    cache_dir = self.cache.local.cache_dir
+    try:
+        with external_repo(repo_url, repo_rev, cache_dir=cache_dir) as repo:
+            cache = NamedCache()
+            for name in files:
+                try:
+                    out = repo.find_out_by_relpath(name)
+                except OutputNotFoundError:
+                    failed += 1
+                    logger.exception(
+                        "failed to fetch data for '{}'".format(name)
+                    )
+                    continue
+                else:
+                    cache.update(out.get_used_cache())
+
+            with repo.state:
+                try:
+                    return repo.cloud.pull(cache), failed
+                except DownloadError as exc:
+                    failed += exc.amount
+    except CloneError:
+        failed += 1
+        logger.exception(
+            "failed to fetch data for '{}'".format(", ".join(files))
+        )
+
+    return 0, failed
diff --git a/tests/func/test_import.py b/tests/func/test_import.py
index 0d1f17e99..fc00ee7b1 100644
--- a/tests/func/test_import.py
+++ b/tests/func/test_import.py
@@ -98,3 +98,20 @@ def test_import_to_dir(dname, repo_dir, dvc_repo, erepo):
     assert stage.outs[0].fspath == os.path.abspath(dst)
     assert os.path.isdir(dname)
     assert filecmp.cmp(repo_dir.FOO, dst, shallow=False)
+
+
+def test_pull_non_workspace(git, dvc_repo, erepo):
+    src = "version"
+    dst = src
+
+    stage = dvc_repo.imp(erepo.root_dir, src, dst, rev="branch")
+    dvc_repo.scm.add([stage.relpath])
+    dvc_repo.scm.commit("imported branch")
+    dvc_repo.scm.tag("ref-to-branch")
+
+    # Ovewrite via import
+    dvc_repo.imp(erepo.root_dir, src, dst, rev="master")
+
+    os.remove(stage.outs[0].cache_path)
+    dvc_repo.fetch(all_tags=True)
+    assert os.path.exists(stage.outs[0].cache_path)