Make temp dir executed checkpoint experiment return result to workspace.

fix: #8612
For checkpoint experiment in some case users might want to give it a early
stopping to reduce variance.  But currently if we interrupt/kill the
experiment it will be marked as failed, and all of the completed
checkpoints will be removed as we cleanup all the running directly after
the process failed.

1. We raise CheckpointKilledError intead of StageCmdFailed error if at
   least one checkpoint had been commited.
2. Temp_dir executor will continue collecting data if the Checkpoint
   stage was interrupted.
3. Raise warnings if a checkpoint stage was incomplete and the other
   stages were not forwarded.
4. Add a new functional test for this

dvc/repo/experiments/queue/tempdir.py

@@ -1,9 +1,12 @@
 import logging
 import os
+from collections import defaultdict
 from typing import TYPE_CHECKING, Dict, Generator, Optional
 
 from funcy import cached_property, first
 
+from dvc.exceptions import DvcException
+
 from ..exceptions import ExpQueueEmptyError
 from ..executor.base import (
     BaseExecutor,
@@ -90,6 +93,51 @@ class TempDirQueue(WorkspaceQueue):
                     None,
                 )
 
+    def _reproduce_entry(
+        self, entry: QueueEntry, executor: BaseExecutor
+    ) -> Dict[str, Dict[str, str]]:
+        from dvc.stage.monitor import CheckpointKilledError
+
+        results: Dict[str, Dict[str, str]] = defaultdict(dict)
+        exec_name = self._EXEC_NAME or entry.stash_rev
+        infofile = self.get_infofile_path(exec_name)
+        try:
+            rev = entry.stash_rev
+            exec_result = executor.reproduce(
+                info=executor.info,
+                rev=rev,
+                infofile=infofile,
+                log_level=logger.getEffectiveLevel(),
+                log_errors=True,
+            )
+            if not exec_result.exp_hash:
+                raise DvcException(
+                    f"Failed to reproduce experiment '{rev[:7]}'"
+                )
+            if exec_result.ref_info:
+                results[rev].update(
+                    self.collect_executor(
+                        self.repo.experiments, executor, exec_result
+                    )
+                )
+        except CheckpointKilledError:
+            results[rev].update(
+                self.collect_executor(
+                    self.repo.experiments, executor, exec_result
+                )
+            )
+
+            return results
+        except DvcException:
+            raise
+        except Exception as exc:
+            raise DvcException(
+                f"Failed to reproduce experiment '{rev[:7]}'"
+            ) from exc
+        finally:
+            executor.cleanup(infofile)
+        return results
+
     @staticmethod
     def collect_executor(
         exp: "Experiments",

dvc/repo/reproduce.py

@@ -1,13 +1,14 @@
 import logging
-import typing
 from functools import partial
+from typing import TYPE_CHECKING, Iterator
 
 from dvc.exceptions import DvcException, ReproductionError
 from dvc.repo.scm_context import scm_context
+from dvc.stage.exceptions import CheckpointKilledError
 
 from . import locked
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from dvc.stage import Stage
 
     from . import Repo
@@ -58,7 +59,7 @@ def _dump_stage(stage):
     dvcfile.dump(stage, update_pipeline=False)
 
 
-def _get_stage_files(stage: "Stage") -> typing.Iterator[str]:
+def _get_stage_files(stage: "Stage") -> Iterator[str]:
     yield stage.dvcfile.relpath
     for dep in stage.deps:
         if (
@@ -183,7 +184,8 @@ def _reproduce_stages(
     # `ret` is used to add a cosmetic newline.
     ret = []
     checkpoint_func = kwargs.pop("checkpoint_func", None)
-    for stage in steps:
+
+    for i, stage in enumerate(steps):
         if ret:
             logger.info("")
 
@@ -192,8 +194,6 @@ def _reproduce_stages(
                 _repro_callback, checkpoint_func, unchanged
             )
 
-        from dvc.stage.monitor import CheckpointKilledError
-
         try:
             ret = _reproduce_stage(stage, **kwargs)
 
@@ -210,7 +210,18 @@ def _reproduce_stages(
             if ret:
                 result.extend(ret)
         except CheckpointKilledError:
-            raise
+            result.append(stage)
+            logger.warning(
+                "Checkpoint stage '%s' was interrupted remaining stages in "
+                "pipeline will not be reproduced.",
+                stage.addressing,
+            )
+            logger.warning(
+                "skipped stages '%s'",
+                ", ".join(s.addressing for s in steps[i + 1 :]),
+            )
+
+            break
         except Exception as exc:
             raise ReproductionError(stage.addressing) from exc
 

dvc/stage/exceptions.py

@@ -9,6 +9,14 @@ class StageCmdFailedError(DvcException):
         super().__init__(msg)
 
 
+class CheckpointKilledError(DvcException):
+    def __init__(self, cmd, status=None):
+        msg = f"failed to finish: {cmd}"
+        if status is not None:
+            msg += f", exited with {status}"
+        super().__init__(msg)
+
+
 class StageFileDoesNotExistError(DvcException):
     DVC_IGNORED = "is dvc-ignored"
     DOES_NOT_EXIST = "does not exist"

dvc/stage/monitor.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable, List
 
 from dvc.stage.decorators import relock_repo
-from dvc.stage.exceptions import StageCmdFailedError
+from dvc.stage.exceptions import CheckpointKilledError
 
 if TYPE_CHECKING:
     from dvc.stage import Stage
@@ -16,17 +16,13 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class CheckpointKilledError(StageCmdFailedError):
-    pass
-
-
 @dataclass
 class MonitorTask:
     stage: "Stage"
     execute: Callable
     proc: subprocess.Popen
     done: threading.Event = threading.Event()
-    killed: threading.Event = threading.Event()
+    updated: threading.Event = threading.Event()
 
     @property
     def name(self) -> str:
@@ -77,7 +73,7 @@ class CheckpointTask(MonitorTask):
 class Monitor:
     AWAIT: float = 1.0
 
-    def __init__(self, tasks: List[MonitorTask]):
+    def __init__(self, tasks: List[CheckpointTask]):
         self.done = threading.Event()
         self.tasks = tasks
         self.monitor_thread = threading.Thread(
@@ -113,6 +109,7 @@ class Monitor:
                 if os.path.exists(task.signal_path):
                     try:
                         task.execute()
+                        task.updated.set()
                     except Exception:  # pylint: disable=broad-except
                         logger.exception(
                             "Error running '%s' task, '%s' will be aborted",
@@ -120,11 +117,11 @@ class Monitor:
                             task.stage,
                         )
                         Monitor.kill(task.proc)
-                        task.killed.set()
                     finally:
                         logger.debug(
                             "Removing signal file for '%s' task", task.name
                         )
                         os.remove(task.signal_path)
+
             if done.wait(Monitor.AWAIT):
                 return

dvc/stage/run.py

@@ -3,13 +3,17 @@ import os
 import signal
 import subprocess
 import threading
+from typing import TYPE_CHECKING, List
 
-from dvc.stage.monitor import Monitor
+from dvc.stage.monitor import CheckpointTask, Monitor
 from dvc.utils import fix_env
 
 from .decorators import unlocked_repo
 from .exceptions import StageCmdFailedError
 
+if TYPE_CHECKING:
+    from dvc.stage import Stage
+
 logger = logging.getLogger(__name__)
 
 
@@ -74,10 +78,11 @@ def get_executable():
     return (os.getenv("SHELL") or "/bin/sh") if os.name != "nt" else None
 
 
-def _run(stage, executable, cmd, checkpoint_func, **kwargs):
+def _run(stage: "Stage", executable, cmd, checkpoint_func, **kwargs):
+    # pylint: disable=protected-access
     main_thread = isinstance(
         threading.current_thread(),
-        threading._MainThread,  # pylint: disable=protected-access
+        threading._MainThread,  # type: ignore[attr-defined]
     )
 
     exec_cmd = _make_cmd(executable, cmd)
@@ -98,7 +103,7 @@ def _run(stage, executable, cmd, checkpoint_func, **kwargs):
 
         if p.returncode != 0:
             for t in tasks:
-                if t.killed.is_set():
+                if t.updated.is_set():
                     raise t.error_cls(cmd, p.returncode)
             raise StageCmdFailedError(cmd, p.returncode)
     finally:
@@ -106,12 +111,10 @@ def _run(stage, executable, cmd, checkpoint_func, **kwargs):
             signal.signal(signal.SIGINT, old_handler)
 
 
-def _get_monitor_tasks(stage, checkpoint_func, proc):
+def _get_monitor_tasks(stage, checkpoint_func, proc) -> List[CheckpointTask]:
 
     result = []
     if checkpoint_func:
-        from .monitor import CheckpointTask
-
         result.append(CheckpointTask(stage, checkpoint_func, proc))
 
     return result

tests/func/experiments/conftest.py

@@ -3,6 +3,7 @@ import pytest
 from tests.unit.repo.experiments.conftest import (  # noqa, pylint disable=unused-argument
     checkpoint_stage,
     exp_stage,
+    failed_checkpoint_stage,
     failed_exp_stage,
     test_queue,
 )

tests/func/experiments/test_checkpoints.py

@@ -275,3 +275,22 @@ def test_auto_push_self_remote(
             "automatically be pushed to the default DVC remote (if any) "
             "on each experiment commit." in caplog.text
         )
+
+
+def test_tmp_dir_failed_checkpoint(
+    tmp_dir, scm, dvc, failed_checkpoint_stage, caplog
+):
+    baseline = scm.get_rev()
+
+    results = dvc.experiments.run(
+        failed_checkpoint_stage.addressing, params=["foo=2"], tmp_dir=True
+    )
+    exp = first(results)
+
+    result = dvc.experiments.show()[baseline]
+    assert len(result) == 1 + 2
+    assert exp in result
+    assert (
+        "Checkpoint stage 'failed-checkpoint-file' was interrupted remaining "
+        "stages in pipeline will not be reproduced." in caplog.text
+    )

tests/unit/repo/experiments/conftest.py

@@ -30,15 +30,23 @@ CHECKPOINT_SCRIPT_FORMAT = dedent(
     shutil.copyfile({}, {})
 
     if os.getenv("DVC_CHECKPOINT"):
-        for _ in range(checkpoint_iterations):
+        for index 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]"
+    "sys.argv[1]", "sys.argv[2]", "sys.argv[3]", "sys.argv[4]", ""
+)
+FAILED_CHECKPOINT_SCRIPT = CHECKPOINT_SCRIPT_FORMAT.format(
+    "sys.argv[1]",
+    "sys.argv[2]",
+    "sys.argv[3]",
+    "sys.argv[4]",
+    "if index == (checkpoint_iterations - 2): raise Exception",
 )
 
 
@@ -90,6 +98,30 @@ def checkpoint_stage(tmp_dir, scm, dvc, mocker):
     return stage
 
 
+@pytest.fixture
+def failed_checkpoint_stage(tmp_dir, scm, dvc, mocker):
+    mocker.patch("dvc.stage.run.Monitor.AWAIT", 0.01)
+
+    tmp_dir.gen("checkpoint.py", FAILED_CHECKPOINT_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+    stage = dvc.run(
+        cmd=(
+            f"python checkpoint.py foo {DEFAULT_ITERATIONS+2} "
+            "params.yaml metrics.yaml"
+        ),
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        checkpoints=["foo"],
+        deps=["checkpoint.py"],
+        no_exec=True,
+        name="failed-checkpoint-file",
+    )
+    scm.add(["dvc.yaml", "checkpoint.py", "params.yaml", ".gitignore"])
+    scm.commit("init")
+    stage.iterations = DEFAULT_ITERATIONS
+    return stage
+
+
 @pytest.fixture
 def failed_exp_stage(tmp_dir, scm, dvc, copy_script):
     tmp_dir.gen("params.yaml", "foo: 1")