From 6a9c0c5ce5988b4caca9d68239eea970d526ed8b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= Date: Sun, 24 Jan 2021 21:06:37 +0100 Subject: [PATCH] live/checkpoints: test continue (#5300) * live: test resuming * initial test, not working * experiments: introduce DVCLIVE_RESUME env variable on resume --- dvc/env.py | 1 + dvc/repo/experiments/__init__.py | 4 +- dvc/stage/__init__.py | 2 +- dvc/stage/run.py | 16 ++-- tests/func/test_live.py | 109 ++++++++++++++++++++++++++++ tests/func/test_repro.py | 4 +- tests/func/test_repro_multistage.py | 4 +- tests/func/test_stage.py | 4 +- 8 files changed, 134 insertions(+), 10 deletions(-) diff --git a/dvc/env.py b/dvc/env.py index ad611f80c..9d2162c98 100644 --- a/dvc/env.py +++ b/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" diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index a86a0b367..15ae734d9 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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( diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index e0854fec0..5480e4eee 100644 --- a/dvc/stage/__init__.py +++ b/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 diff --git a/dvc/stage/run.py b/dvc/stage/run.py index afdb6f31b..f5a059008 100644 --- a/dvc/stage/run.py +++ b/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 diff --git a/tests/func/test_live.py b/tests/func/test_live.py index 1c12817f6..5d3090aac 100644 --- a/tests/func/test_live.py +++ b/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, + ] diff --git a/tests/func/test_repro.py b/tests/func/test_repro.py index 36c5a93a8..d82cd7c33 100644 --- a/tests/func/test_repro.py +++ b/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 + ) diff --git a/tests/func/test_repro_multistage.py b/tests/func/test_repro_multistage.py index cd5d65c4f..306f3329a 100644 --- a/tests/func/test_repro_multistage.py +++ b/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): diff --git a/tests/func/test_stage.py b/tests/func/test_stage.py index 26ebc3ce4..dc19e868d 100644 --- a/tests/func/test_stage.py +++ b/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 + )