diff --git a/dvc/stage/cache.py b/dvc/stage/cache.py index c383d8725..18c2cb2f5 100644 --- a/dvc/stage/cache.py +++ b/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 diff --git a/tests/unit/test_stage.py b/tests/unit/test_stage.py index dcfb771c4..09121da2b 100644 --- a/tests/unit/test_stage.py +++ b/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"