live/checkpoints: test continue (#5300)

* live: test resuming

* initial test, not working

* experiments: introduce DVCLIVE_RESUME env variable on resume

dvc/env.py

@@ -5,3 +5,4 @@ DVC_ROOT = "DVC_ROOT"
 DVCLIVE_PATH = "DVCLIVE_PATH"
 DVCLIVE_SUMMARY = "DVCLIVE_SUMMARY"
 DVCLIVE_REPORT = "DVCLIVE_REPORT"
+DVCLIVE_RESUME = "DVCLIVE_RESUME"

dvc/repo/experiments/__init__.py

@@ -10,6 +10,7 @@ from typing import Dict, Iterable, Mapping, Optional
 
 from funcy import cached_property, first
 
+from dvc.env import DVCLIVE_RESUME
 from dvc.exceptions import DvcException
 from dvc.path_info import PathInfo
 from dvc.stage.run import CheckpointKilledError
@@ -210,7 +211,8 @@ class Experiments:
                                 )
 
                     # save additional repro command line arguments
-                    self._pack_args(*args, **kwargs)
+                    run_env = {DVCLIVE_RESUME: "1"} if resume_rev else {}
+                    self._pack_args(*args, run_env=run_env, **kwargs)
 
                     # save experiment as a stash commit
                     msg = self._stash_msg(

dvc/stage/__init__.py

@@ -361,7 +361,7 @@ class Stage(params.StageParams):
     def remove_outs(self, ignore_remove=False, force=False):
         """Used mainly for `dvc remove --outs` and :func:`Stage.reproduce`."""
         for out in self.outs:
-            if (out.persist or out.checkpoint) and not force:
+            if (out.persist or out.checkpoint or out.live) and not force:
                 out.unprotect()
                 continue
 

dvc/stage/run.py

@@ -49,10 +49,12 @@ def _enforce_cmd_list(cmd):
     return cmd if isinstance(cmd, list) else cmd.splitlines()
 
 
-def prepare_kwargs(stage, checkpoint_func=None):
+def prepare_kwargs(stage, checkpoint_func=None, run_env=None):
     kwargs = {"cwd": stage.wdir, "env": fix_env(None), "close_fds": True}
 
     kwargs["env"].update(stage.env(checkpoint_func=checkpoint_func))
+    if run_env:
+        kwargs["env"].update(run_env)
 
     # NOTE: when you specify `shell=True`, `Popen` [1] will default to
     # `/bin/sh` on *nix and will add ["/bin/sh", "-c"] to your command.
@@ -108,10 +110,12 @@ def _run(stage, executable, cmd, checkpoint_func, **kwargs):
         raise StageCmdFailedError(cmd, retcode)
 
 
-def cmd_run(stage, dry=False, checkpoint_func=None):
+def cmd_run(stage, dry=False, checkpoint_func=None, run_env=None):
     logger.info("Running stage '%s':", stage.addressing)
     commands = _enforce_cmd_list(stage.cmd)
-    kwargs = prepare_kwargs(stage, checkpoint_func=checkpoint_func)
+    kwargs = prepare_kwargs(
+        stage, checkpoint_func=checkpoint_func, run_env=run_env
+    )
     executable = get_executable()
 
     if not dry:
@@ -125,7 +129,9 @@ def cmd_run(stage, dry=False, checkpoint_func=None):
         _run(stage, executable, cmd, checkpoint_func=checkpoint_func, **kwargs)
 
 
-def run_stage(stage, dry=False, force=False, checkpoint_func=None, **kwargs):
+def run_stage(
+    stage, dry=False, force=False, checkpoint_func=None, run_env=None, **kwargs
+):
     if not (dry or force or checkpoint_func):
         from .cache import RunCacheNotFoundError
 
@@ -136,7 +142,7 @@ def run_stage(stage, dry=False, force=False, checkpoint_func=None, **kwargs):
             stage.save_deps()
 
     run = cmd_run if dry else unlocked_repo(cmd_run)
-    run(stage, dry=dry, checkpoint_func=checkpoint_func)
+    run(stage, dry=dry, checkpoint_func=checkpoint_func, run_env=run_env)
 
 
 @contextmanager

tests/func/test_live.py

@@ -1,6 +1,8 @@
+from copy import deepcopy
 from textwrap import dedent
 
 import pytest
+from funcy import first
 
 from dvc import stage as stage_module
 from dvc.exceptions import MetricsError
@@ -128,3 +130,110 @@ def test_live_report(tmp_dir, dvc, live_stage, report):
     live_stage(report=report)
 
     assert (tmp_dir / "logs.html").is_file() == report
+
+
+@pytest.fixture
+def live_checkpoint_stage(tmp_dir, scm, dvc):
+
+    pytest.skip("dvclive does not exist yet")
+
+    SCRIPT = dedent(
+        """
+            import os
+            import dvclive
+
+            def read(path):
+                value=0
+                if os.path.exists(path):
+                    with open(path, 'r') as fobj:
+                        try:
+                            value = int(fobj.read())
+                        except ValueError:
+                            pass
+                return value
+
+            def dump(value, path):
+                with open(path, "w") as fobj:
+                    fobj.write(str(value))
+
+            r = 3
+            checkpoint_file = "checkpoint"
+
+            value = read(checkpoint_file)
+            for i in range(1,r):
+                m = i + value
+                dump(m, checkpoint_file)
+
+                dvclive.log("metric1", m)
+                dvclive.log("metric2", m * 2)
+                dvclive.next_step()"""
+    )
+
+    tmp_dir.gen("train.py", SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+    stage = dvc.run(
+        cmd="python train.py",
+        params=["foo"],
+        deps=["train.py"],
+        name="live_stage",
+        live="logs",
+        checkpoints=["checkpoint"],
+        no_exec=True,
+    )
+
+    scm.add(["dvc.yaml", "train.py", "params.yaml", ".gitignore"])
+    scm.commit("initial: live_stage")
+    yield stage
+
+
+def checkpoints_metric(show_results, metric_file, metric_name):
+    tmp = deepcopy(show_results)
+    tmp.pop("workspace")
+    tmp = first(tmp.values())
+    tmp.pop("baseline")
+    return list(
+        map(
+            lambda exp: exp["metrics"][metric_file][metric_name],
+            list(tmp.values()),
+        )
+    )
+
+
+def test_live_checkpoints_resume(tmp_dir, scm, dvc, live_checkpoint_stage):
+    results = dvc.experiments.run(
+        live_checkpoint_stage.addressing, params=["foo=2"], tmp_dir=False
+    )
+
+    checkpoint_resume = first(results)
+
+    dvc.experiments.run(
+        live_checkpoint_stage.addressing,
+        checkpoint_resume=checkpoint_resume,
+        tmp_dir=False,
+    )
+
+    results = dvc.experiments.show()
+    assert checkpoints_metric(results, "logs.json", "step") == [
+        3,
+        3,
+        2,
+        1,
+        1,
+        0,
+    ]
+    assert checkpoints_metric(results, "logs.json", "metric1") == [
+        4,
+        4,
+        3,
+        2,
+        2,
+        1,
+    ]
+    assert checkpoints_metric(results, "logs.json", "metric2") == [
+        8,
+        8,
+        6,
+        4,
+        4,
+        2,
+    ]

tests/func/test_repro.py

@@ -1246,4 +1246,6 @@ def test_repro_when_cmd_changes(tmp_dir, dvc, run_copy, mocker):
         stage.addressing: ["changed checksum"]
     }
     assert dvc.reproduce(stage.addressing)[0] == stage
-    m.assert_called_once_with(stage, checkpoint_func=None, dry=False)
+    m.assert_called_once_with(
+        stage, checkpoint_func=None, dry=False, run_env=None
+    )

tests/func/test_repro_multistage.py

@@ -295,7 +295,9 @@ def test_repro_when_cmd_changes(tmp_dir, dvc, run_copy, mocker):
 
     assert dvc.status([target]) == {target: ["changed command"]}
     assert dvc.reproduce(target)[0] == stage
-    m.assert_called_once_with(stage, checkpoint_func=None, dry=False)
+    m.assert_called_once_with(
+        stage, checkpoint_func=None, dry=False, run_env=None
+    )
 
 
 def test_repro_when_new_deps_is_added_in_dvcfile(tmp_dir, dvc, run_copy):

tests/func/test_stage.py

@@ -309,4 +309,6 @@ def test_stage_run_checkpoint(tmp_dir, dvc, mocker, checkpoint):
 
     with lock_repo(dvc):
         run_stage(stage, checkpoint_func=callback)
-    mock_cmd_run.assert_called_with(stage, checkpoint_func=callback, dry=False)
+    mock_cmd_run.assert_called_with(
+        stage, checkpoint_func=callback, dry=False, run_env=None
+    )