run-cache: support uncached files (#3778)

Currently we only restore `--outs`, but `--outs-no-cache` are also
useful, especially when they are metrics files
(`-M|--metrics-no-cache`). So this PR makes run-cache save uncached
files to our cache and also tries to checkout them back during the
`restore` phase.

dvc/stage/cache.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from contextlib import contextmanager
 
 import yaml
 from funcy import first
@@ -76,6 +77,45 @@ class StageCache:
 
         return None
 
+    def _create_stage(self, cache):
+        from dvc.stage import create_stage, PipelineStage
+
+        stage = create_stage(
+            PipelineStage,
+            repo=self.repo,
+            path="dvc.yaml",
+            cmd=cache["cmd"],
+            deps=[dep["path"] for dep in cache["deps"]],
+            outs=[out["path"] for out in cache["outs"]],
+        )
+        StageLoader.fill_from_lock(stage, cache)
+        return stage
+
+    @contextmanager
+    def _cache_type_copy(self):
+        cache_types = self.repo.cache.local.cache_types
+        self.repo.cache.local.cache_types = ["copy"]
+        try:
+            yield
+        finally:
+            self.repo.cache.local.cache_types = cache_types
+
+    def _uncached_outs(self, stage, cache):
+        # NOTE: using temporary stage to avoid accidentally modifying original
+        # stage and to workaround `commit/checkout` not working for uncached
+        # outputs.
+        cached_stage = self._create_stage(cache)
+
+        outs_no_cache = [
+            out.def_path for out in stage.outs if not out.use_cache
+        ]
+
+        # NOTE: using copy link to make it look like a git-tracked file
+        with self._cache_type_copy():
+            for out in cached_stage.outs:
+                if out.def_path in outs_no_cache:
+                    yield out
+
     def save(self, stage):
         cache_key = _get_stage_hash(stage)
         if not cache_key:
@@ -84,7 +124,13 @@ class StageCache:
         cache = to_single_stage_lockfile(stage)
         cache_value = _get_cache_hash(cache)
 
-        if self._load_cache(cache_key, cache_value):
+        existing_cache = self._load_cache(cache_key, cache_value)
+        cache = existing_cache or cache
+
+        for out in self._uncached_outs(stage, cache):
+            out.commit()
+
+        if existing_cache:
             return
 
         # sanity check
@@ -104,6 +150,9 @@ class StageCache:
             return
         StageLoader.fill_from_lock(stage, cache)
 
+        for out in self._uncached_outs(stage, cache):
+            out.checkout()
+
     @staticmethod
     def _transfer(func, from_remote, to_remote):
         ret = []
@@ -134,7 +183,6 @@ class StageCache:
 
     def get_used_cache(self, used_run_cache, *args, **kwargs):
         from dvc.cache import NamedCache
-        from dvc.stage import create_stage, PipelineStage
 
         cache = NamedCache()
 
@@ -142,14 +190,6 @@ class StageCache:
             entry = self._load_cache(key, value)
             if not entry:
                 continue
-            stage = create_stage(
-                PipelineStage,
-                repo=self.repo,
-                path="dvc.yaml",
-                cmd=entry["cmd"],
-                deps=[dep["path"] for dep in entry["deps"]],
-                outs=[out["path"] for out in entry["outs"]],
-            )
-            StageLoader.fill_from_lock(stage, entry)
+            stage = self._create_stage(entry)
             cache.update(stage.get_used_cache(*args, **kwargs))
         return cache

tests/unit/test_stage.py

@@ -106,22 +106,36 @@ def test_always_changed(dvc):
 
 def test_stage_cache(tmp_dir, dvc, run_copy, mocker):
     tmp_dir.gen("dep", "dep")
-    stage = run_copy("dep", "out", single_stage=True)
+    tmp_dir.gen(
+        "script.py",
+        (
+            'open("out", "w+").write("out"); '
+            'open("out_no_cache", "w+").write("out_no_cache")'
+        ),
+    )
+    stage = dvc.run(
+        cmd="python script.py",
+        deps=["script.py", "dep"],
+        outs=["out"],
+        outs_no_cache=["out_no_cache"],
+        single_stage=True,
+    )
 
     with dvc.lock, dvc.state:
         stage.remove(remove_outs=True, force=True)
 
     assert not (tmp_dir / "out").exists()
+    assert not (tmp_dir / "out_no_cache").exists()
     assert not (tmp_dir / "out.dvc").exists()
 
     cache_dir = os.path.join(
         dvc.stage_cache.cache_dir,
-        "75",
-        "75f8a9097d76293ff4b3684d52e4ad0e83686d31196f27eb0b2ea9fd5085565e",
+        "10",
+        "10b45372fdf4ec14d3f779c5b256378d7a12780e4c7f549a44138e492f098bfe",
     )
     cache_file = os.path.join(
         cache_dir,
-        "c1747e52065bc7801262fdaed4d63f5775e5da304008bd35e2fea4e6b1ccb272",
+        "bb32e04c6da96a7192513390acedbe4cd6123f8fe5b0ba5fffe39716fe87f6f4",
     )
 
     assert os.path.isdir(cache_dir)
@@ -137,4 +151,6 @@ def test_stage_cache(tmp_dir, dvc, run_copy, mocker):
     assert checkout_spy.call_count == 1
 
     assert (tmp_dir / "out").exists()
-    assert (tmp_dir / "out").read_text() == "dep"
+    assert (tmp_dir / "out_no_cache").exists()
+    assert (tmp_dir / "out").read_text() == "out"
+    assert (tmp_dir / "out_no_cache").read_text() == "out_no_cache"