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")