experiments: checkpoints proof of concept (#4591)

* stage: make runs checkpointable

* api: add make_checkpoint API call

* tests: add test for checkpointed stage

* experiments: implement checkpoints callback chain

* initial working state

* use .dvc/tmp for checkpoint signal file location

* force repro for checkpoint experiments

* checkout: add flag to allow missing persistent outputs

(i.e. not yet created)

* use allow_persist_missing on checkpoint experiment repro

* executor: force checkout before checkpoints experiment repro

* move checkpoint commands to `dvc exp run`

* support resuming checkpoint runs with --continue

* fix git bugs

* experiments: include checkpoint commits in `exp show`

* use 1sec sleep timers

* fix apply_workspace test conflict

* tests: add test for `dvc exp run`

* cleanup signal file behavior

* fix bool conversion

* tests: add tests for checkpoint and checkpoint_continue

* fix styling

* fix checkpoint stage test

* fix checkpoint monitor thread condition/notification

tests/func/experiments/test_experiments.py

@@ -1,8 +1,47 @@
+from textwrap import dedent
+
 import pytest
+from funcy import first
 
 from dvc.utils.serialize import PythonFileCorruptedError
 from tests.func.test_repro_multistage import COPY_SCRIPT
 
+CHECKPOINT_SCRIPT_FORMAT = dedent(
+    """\
+    import os
+    import sys
+    import shutil
+    from time import sleep
+
+    from dvc.api import make_checkpoint
+
+    checkpoint_file = {}
+    checkpoint_iterations = int({})
+    if os.path.exists(checkpoint_file):
+        with open(checkpoint_file) as fobj:
+            try:
+                value = int(fobj.read())
+            except ValueError:
+                value = 0
+    else:
+        with open(checkpoint_file, "w"):
+            pass
+        value = 0
+
+    shutil.copyfile({}, {})
+
+    if os.getenv("DVC_CHECKPOINT"):
+        for _ in range(checkpoint_iterations):
+            value += 1
+            with open(checkpoint_file, "w") as fobj:
+                fobj.write(str(value))
+            make_checkpoint()
+"""
+)
+CHECKPOINT_SCRIPT = CHECKPOINT_SCRIPT_FORMAT.format(
+    "sys.argv[1]", "sys.argv[2]", "sys.argv[3]", "sys.argv[4]"
+)
+
 
 def test_new_simple(tmp_dir, scm, dvc, mocker):
     tmp_dir.gen("copy.py", COPY_SCRIPT)
@@ -236,3 +275,77 @@ def test_detached_parent(tmp_dir, scm, dvc, mocker):
     assert dvc.experiments.get_baseline(exp_rev) == detached_rev
     assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 3"
     assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 3"
+
+
+def test_new_checkpoint(tmp_dir, scm, dvc, mocker):
+    tmp_dir.gen("checkpoint.py", CHECKPOINT_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+    stage = dvc.run(
+        cmd="python checkpoint.py foo 5 params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        outs_persist=["foo"],
+        always_changed=True,
+        name="checkpoint-file",
+    )
+    scm.add(
+        [
+            "dvc.yaml",
+            "dvc.lock",
+            "checkpoint.py",
+            "params.yaml",
+            "metrics.yaml",
+        ]
+    )
+    scm.commit("init")
+
+    new_mock = mocker.spy(dvc.experiments, "new")
+    dvc.reproduce(
+        stage.addressing, experiment=True, checkpoint=True, params=["foo=2"]
+    )
+
+    new_mock.assert_called_once()
+    assert (tmp_dir / "foo").read_text() == "5"
+    assert (
+        tmp_dir / ".dvc" / "experiments" / "metrics.yaml"
+    ).read_text().strip() == "foo: 2"
+
+
+def test_continue_checkpoint(tmp_dir, scm, dvc, mocker):
+    tmp_dir.gen("checkpoint.py", CHECKPOINT_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+    stage = dvc.run(
+        cmd="python checkpoint.py foo 5 params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        outs_persist=["foo"],
+        always_changed=True,
+        name="checkpoint-file",
+    )
+    scm.add(
+        [
+            "dvc.yaml",
+            "dvc.lock",
+            "checkpoint.py",
+            "params.yaml",
+            "metrics.yaml",
+        ]
+    )
+    scm.commit("init")
+
+    results = dvc.reproduce(
+        stage.addressing, experiment=True, checkpoint=True, params=["foo=2"]
+    )
+    exp_rev = first(results)
+
+    dvc.reproduce(
+        stage.addressing,
+        experiment=True,
+        checkpoint=True,
+        checkpoint_continue=exp_rev,
+    )
+
+    assert (tmp_dir / "foo").read_text() == "10"
+    assert (
+        tmp_dir / ".dvc" / "experiments" / "metrics.yaml"
+    ).read_text().strip() == "foo: 2"

tests/func/test_repro.py

@@ -1302,4 +1302,4 @@ 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)
+    m.assert_called_once_with(stage, checkpoint=False)

tests/func/test_repro_multistage.py

@@ -286,7 +286,7 @@ 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)
+    m.assert_called_once_with(stage, checkpoint=False)
 
 
 def test_repro_when_new_deps_is_added_in_dvcfile(tmp_dir, dvc, run_copy):

tests/func/test_stage.py

@@ -9,6 +9,7 @@ from dvc.output.local import LocalOutput
 from dvc.repo import Repo
 from dvc.stage import PipelineStage, Stage
 from dvc.stage.exceptions import StageFileFormatError
+from dvc.stage.run import run_stage
 from dvc.tree.local import LocalTree
 from dvc.utils.serialize import dump_yaml, load_yaml
 from tests.basic_env import TestDvc
@@ -274,3 +275,17 @@ def test_stage_remove_pointer_stage(tmp_dir, dvc, run_copy):
     with dvc.lock:
         stage.remove()
     assert not (tmp_dir / stage.relpath).exists()
+
+
+@pytest.mark.parametrize("checkpoint", [True, False])
+def test_stage_run_checkpoint(tmp_dir, dvc, mocker, checkpoint):
+    stage = Stage(dvc, "stage.dvc", cmd="mycmd arg1 arg2")
+    mocker.patch.object(stage, "save")
+
+    mock_cmd_run = mocker.patch("dvc.stage.run.cmd_run")
+    if checkpoint:
+        callback = mocker.Mock()
+    else:
+        callback = None
+    run_stage(stage, checkpoint_func=callback)
+    mock_cmd_run.assert_called_with(stage, checkpoint=checkpoint)

tests/unit/command/test_experiments.py

@@ -1,5 +1,10 @@
 from dvc.cli import parse_args
-from dvc.command.experiments import CmdExperimentsDiff, CmdExperimentsShow
+from dvc.command.experiments import (
+    CmdExperimentsDiff,
+    CmdExperimentsRun,
+    CmdExperimentsShow,
+)
+from dvc.dvcfile import PIPELINE_FILE
 
 
 def test_experiments_diff(dvc, mocker):
@@ -54,3 +59,34 @@ def test_experiments_show(dvc, mocker):
         all_commits=True,
         sha_only=True,
     )
+
+
+default_run_arguments = {
+    "all_pipelines": False,
+    "downstream": False,
+    "dry": False,
+    "force": False,
+    "run_cache": True,
+    "interactive": False,
+    "no_commit": False,
+    "pipeline": False,
+    "single_item": False,
+    "recursive": False,
+    "force_downstream": False,
+    "params": [],
+    "queue": False,
+    "run_all": False,
+    "jobs": None,
+    "checkpoint": False,
+    "checkpoint_continue": None,
+    "experiment": True,
+}
+
+
+def test_experiments_run(dvc, mocker):
+    cmd = CmdExperimentsRun(parse_args(["exp", "run"]))
+    mocker.patch.object(cmd.repo, "reproduce")
+    cmd.run()
+    cmd.repo.reproduce.assert_called_with(
+        PIPELINE_FILE, **default_run_arguments
+    )