From 08ae68caf46ad1c2343b43e2994e7dbbc2ef9a4c Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Fri, 16 Dec 2022 10:23:46 +0800 Subject: [PATCH] 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 | 48 ++++++++++++++++++++++ dvc/repo/reproduce.py | 25 +++++++---- dvc/stage/exceptions.py | 8 ++++ dvc/stage/monitor.py | 13 +++--- dvc/stage/run.py | 17 ++++---- tests/func/experiments/conftest.py | 1 + tests/func/experiments/test_checkpoints.py | 19 +++++++++ tests/unit/repo/experiments/conftest.py | 36 +++++++++++++++- 8 files changed, 143 insertions(+), 24 deletions(-) diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index ae8b33048..2261216fb 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/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", diff --git a/dvc/repo/reproduce.py b/dvc/repo/reproduce.py index 1c8bd35df..59e4a70bb 100644 --- a/dvc/repo/reproduce.py +++ b/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 diff --git a/dvc/stage/exceptions.py b/dvc/stage/exceptions.py index d89142ab6..4ed800371 100644 --- a/dvc/stage/exceptions.py +++ b/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" diff --git a/dvc/stage/monitor.py b/dvc/stage/monitor.py index 7d2675f91..b0869ba2b 100644 --- a/dvc/stage/monitor.py +++ b/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 diff --git a/dvc/stage/run.py b/dvc/stage/run.py index ca76dcb2f..d4478cfac 100644 --- a/dvc/stage/run.py +++ b/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 diff --git a/tests/func/experiments/conftest.py b/tests/func/experiments/conftest.py index 8f0228b7f..5ece8fd2d 100644 --- a/tests/func/experiments/conftest.py +++ b/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, ) diff --git a/tests/func/experiments/test_checkpoints.py b/tests/func/experiments/test_checkpoints.py index 50675b21f..77922d996 100644 --- a/tests/func/experiments/test_checkpoints.py +++ b/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 + ) diff --git a/tests/unit/repo/experiments/conftest.py b/tests/unit/repo/experiments/conftest.py index fe4055374..94679e83e 100644 --- a/tests/unit/repo/experiments/conftest.py +++ b/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")