From 4a8cb806d29e65feb768c9917814328b4758f5f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= Date: Tue, 2 Mar 2021 16:45:35 +0100 Subject: [PATCH] live integration: get rid of dvc from live (#5466) * live integration: get rid of dvc from live * live: remove from api, add test for html generation during run * stage: run: convert monitors to class context managers * run: make monitors run in single thread * run: move monitor logic out of run --- dvc/{api/__init__.py => api.py} | 4 +- dvc/api/live.py | 24 ----- dvc/repo/experiments/__init__.py | 2 +- dvc/repo/experiments/executor/base.py | 2 +- dvc/repo/live.py | 15 +++ dvc/repo/reproduce.py | 2 +- dvc/stage/monitor.py | 150 ++++++++++++++++++++++++++ dvc/stage/run.py | 115 ++++++-------------- tests/func/api/__init__.py | 0 tests/func/api/test_live.py | 55 ---------- tests/func/experiments/conftest.py | 4 +- tests/func/{api => }/test_api.py | 0 tests/func/test_live.py | 75 +++++++------ 13 files changed, 248 insertions(+), 200 deletions(-) rename dvc/{api/__init__.py => api.py} (97%) delete mode 100644 dvc/api/live.py create mode 100644 dvc/stage/monitor.py delete mode 100644 tests/func/api/__init__.py delete mode 100644 tests/func/api/test_live.py rename tests/func/{api => }/test_api.py (100%) diff --git a/dvc/api/__init__.py b/dvc/api.py similarity index 97% rename from dvc/api/__init__.py rename to dvc/api.py index 8471cf207..c6420b840 100644 --- a/dvc/api/__init__.py +++ b/dvc/api.py @@ -104,14 +104,14 @@ def make_checkpoint(): from time import sleep from dvc.env import DVC_CHECKPOINT, DVC_ROOT - from dvc.stage.run import CHECKPOINT_SIGNAL_FILE + from dvc.stage.monitor import CheckpointTask if os.getenv(DVC_CHECKPOINT) is None: return root_dir = os.getenv(DVC_ROOT, Repo.find_root()) signal_file = os.path.join( - root_dir, Repo.DVC_DIR, "tmp", CHECKPOINT_SIGNAL_FILE + root_dir, Repo.DVC_DIR, "tmp", CheckpointTask.SIGNAL_FILE ) with builtins.open(signal_file, "w") as fobj: diff --git a/dvc/api/live.py b/dvc/api/live.py deleted file mode 100644 index 39fb59c38..000000000 --- a/dvc/api/live.py +++ /dev/null @@ -1,24 +0,0 @@ -import logging -import os -from typing import List - -from dvc.exceptions import NotDvcRepoError -from dvc.repo import Repo -from dvc.utils.html import write - -logger = logging.getLogger(__name__) - - -def summary(path: str, revs: List[str] = None): - try: - root = Repo.find_root() - except NotDvcRepoError: - root = os.getcwd() - - metrics, plots = Repo(root_dir=root, uninitialized=True).live.show( - path, revs - ) - - html_path = path + ".html" - write(html_path, plots, metrics) - logger.info(f"\nfile://{os.path.abspath(html_path)}") diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index c99ed3eac..345dcf807 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -13,7 +13,7 @@ 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 +from dvc.stage.monitor import CheckpointKilledError from dvc.utils import relpath from .base import ( diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 2f96f35ed..ce5cd920e 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -33,7 +33,7 @@ from dvc.repo.experiments.base import ( ) from dvc.scm import SCM from dvc.stage import PipelineStage -from dvc.stage.run import CheckpointKilledError +from dvc.stage.monitor import CheckpointKilledError from dvc.stage.serialize import to_lockfile from dvc.utils import dict_sha256 from dvc.utils.fs import remove diff --git a/dvc/repo/live.py b/dvc/repo/live.py index 42a5e7b03..e3e711703 100644 --- a/dvc/repo/live.py +++ b/dvc/repo/live.py @@ -1,15 +1,30 @@ import contextlib +import logging import os from typing import TYPE_CHECKING, List, Optional from dvc.exceptions import MetricDoesNotExistError, MetricsError +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from dvc.output import BaseOutput from dvc.path_info import PathInfo from dvc.repo import Repo +def create_summary(out): + from dvc.utils.html import write + + assert out.live and out.live["html"] + + metrics, plots = out.repo.live.show(str(out.path_info)) + + html_path = out.path_info.with_suffix(".html") + write(html_path, plots, metrics) + logger.info(f"\nfile://{os.path.abspath(html_path)}") + + def summary_path_info(out: "BaseOutput") -> Optional["PathInfo"]: from dvc.output import BaseOutput diff --git a/dvc/repo/reproduce.py b/dvc/repo/reproduce.py index 9fe61acb4..100143279 100644 --- a/dvc/repo/reproduce.py +++ b/dvc/repo/reproduce.py @@ -186,7 +186,7 @@ def _reproduce_stages( _repro_callback, checkpoint_func, unchanged ) - from dvc.stage.run import CheckpointKilledError + from dvc.stage.monitor import CheckpointKilledError try: ret = _reproduce_stage(stage, **kwargs) diff --git a/dvc/stage/monitor.py b/dvc/stage/monitor.py new file mode 100644 index 000000000..409eb533c --- /dev/null +++ b/dvc/stage/monitor.py @@ -0,0 +1,150 @@ +import functools +import logging +import os +import subprocess +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, List + +from dvc.repo.live import create_summary +from dvc.stage.decorators import relock_repo +from dvc.stage.exceptions import StageCmdFailedError + +if TYPE_CHECKING: + from dvc.output import BaseOutput + from dvc.stage import Stage + + +logger = logging.getLogger(__name__) + + +class CheckpointKilledError(StageCmdFailedError): + pass + + +class LiveKilledError(StageCmdFailedError): + pass + + +@dataclass +class MonitorTask: + stage: "Stage" + execute: Callable + proc: subprocess.Popen + done: threading.Event = threading.Event() + killed: threading.Event = threading.Event() + + @property + def name(self) -> str: + raise NotImplementedError + + @property + def SIGNAL_FILE(self) -> str: + raise NotImplementedError + + @property + def error_cls(self) -> type: + raise NotImplementedError + + @property + def signal_path(self) -> str: + return os.path.join(self.stage.repo.tmp_dir, self.SIGNAL_FILE) + + def after_run(self): + pass + + +class CheckpointTask(MonitorTask): + name = "checkpoint" + SIGNAL_FILE = "DVC_CHECKPOINT" + error_cls = CheckpointKilledError + + def __init__( + self, stage: "Stage", callback_func: Callable, proc: subprocess.Popen + ): + super().__init__( + stage, + functools.partial( + CheckpointTask._run_callback, stage, callback_func + ), + proc, + ) + + @staticmethod + @relock_repo + def _run_callback(stage, callback_func): + stage.save(allow_missing=True) + stage.commit(allow_missing=True) + logger.debug("Running checkpoint callback for stage '%s'", stage) + callback_func() + + +class LiveTask(MonitorTask): + name = "live" + SIGNAL_FILE = "DVC_LIVE" + error_cls = LiveKilledError + + def __init__( + self, stage: "Stage", out: "BaseOutput", proc: subprocess.Popen + ): + super().__init__(stage, functools.partial(create_summary, out), proc) + + def after_run(self): + # make sure summary is prepared for all the data + self.execute() + + +class Monitor: + AWAIT: float = 1.0 + + def __init__(self, tasks: List[MonitorTask]): + self.done = threading.Event() + self.tasks = tasks + self.monitor_thread = threading.Thread( + target=Monitor._loop, args=(self.tasks, self.done,), + ) + + def __enter__(self): + self.monitor_thread.start() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.done.set() + self.monitor_thread.join() + for t in self.tasks: + t.after_run() + + @staticmethod + def kill(proc): + if os.name == "nt": + return Monitor._kill_nt(proc) + proc.terminate() + proc.wait() + + @staticmethod + def _kill_nt(proc): + # windows stages are spawned with shell=True, proc is the shell process + # and not the actual stage process - we have to kill the entire tree + subprocess.call(["taskkill", "/F", "/T", "/PID", str(proc.pid)]) + + @staticmethod + def _loop(tasks: List[MonitorTask], done: threading.Event): + while True: + for task in tasks: + if os.path.exists(task.signal_path): + try: + task.execute() + except Exception: # pylint: disable=broad-except + logger.exception( + "Error running '%s' task, '%s' will be aborted", + task.name, + 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 f5a059008..972a92877 100644 --- a/dvc/stage/run.py +++ b/dvc/stage/run.py @@ -3,23 +3,18 @@ import os import signal import subprocess import threading -from contextlib import contextmanager +from funcy import first + +from dvc.stage.monitor import Monitor from dvc.utils import fix_env -from .decorators import relock_repo, unlocked_repo +from .decorators import unlocked_repo from .exceptions import StageCmdFailedError logger = logging.getLogger(__name__) -CHECKPOINT_SIGNAL_FILE = "DVC_CHECKPOINT" - - -class CheckpointKilledError(StageCmdFailedError): - pass - - def _make_cmd(executable, cmd): if executable is None: return cmd @@ -89,25 +84,45 @@ def _run(stage, executable, cmd, checkpoint_func, **kwargs): exec_cmd = _make_cmd(executable, cmd) old_handler = None - p = None try: p = subprocess.Popen(exec_cmd, **kwargs) if main_thread: old_handler = signal.signal(signal.SIGINT, signal.SIG_IGN) - killed = threading.Event() - with checkpoint_monitor(stage, checkpoint_func, p, killed): + tasks = _get_monitor_tasks(stage, checkpoint_func, p) + + if tasks: + with Monitor(tasks): + p.communicate() + else: p.communicate() + + if p.returncode != 0: + for t in tasks: + if t.killed.is_set(): + raise t.error_cls(cmd, p.returncode) + raise StageCmdFailedError(cmd, p.returncode) finally: if old_handler: signal.signal(signal.SIGINT, old_handler) - retcode = None if not p else p.returncode - if retcode != 0: - if killed.is_set(): - raise CheckpointKilledError(cmd, retcode) - raise StageCmdFailedError(cmd, retcode) + +def _get_monitor_tasks(stage, checkpoint_func, proc): + + result = [] + if checkpoint_func: + from .monitor import CheckpointTask + + result.append(CheckpointTask(stage, checkpoint_func, proc)) + + live = first((o for o in stage.outs if (o.live and o.live["html"]))) + if live: + from .monitor import LiveTask + + result.append(LiveTask(stage, live, proc)) + + return result def cmd_run(stage, dry=False, checkpoint_func=None, run_env=None): @@ -143,69 +158,3 @@ def run_stage( run = cmd_run if dry else unlocked_repo(cmd_run) run(stage, dry=dry, checkpoint_func=checkpoint_func, run_env=run_env) - - -@contextmanager -def checkpoint_monitor(stage, callback_func, proc, killed): - if not callback_func: - yield None - return - - logger.debug( - "Monitoring checkpoint stage '%s' with cmd process '%d'", - stage, - proc.pid, - ) - done = threading.Event() - monitor_thread = threading.Thread( - target=_checkpoint_run, - args=(stage, callback_func, done, proc, killed), - ) - - try: - monitor_thread.start() - yield monitor_thread - finally: - done.set() - monitor_thread.join() - - -def _checkpoint_run(stage, callback_func, done, proc, killed): - """Run callback_func whenever checkpoint signal file is present.""" - signal_path = os.path.join(stage.repo.tmp_dir, CHECKPOINT_SIGNAL_FILE) - while True: - if os.path.exists(signal_path): - try: - _run_callback(stage, callback_func) - except Exception: # pylint: disable=broad-except - logger.exception( - "Error generating checkpoint, %s will be aborted", stage - ) - _kill(proc) - killed.set() - finally: - logger.debug("Remove checkpoint signal file") - os.remove(signal_path) - if done.wait(1): - return - - -def _kill(proc): - if os.name == "nt": - return _kill_nt(proc) - proc.terminate() - proc.wait() - - -def _kill_nt(proc): - # windows stages are spawned with shell=True, proc is the shell process and - # not the actual stage process - we have to kill the entire tree - subprocess.call(["taskkill", "/F", "/T", "/PID", str(proc.pid)]) - - -@relock_repo -def _run_callback(stage, callback_func): - stage.save(allow_missing=True) - stage.commit(allow_missing=True) - logger.debug("Running checkpoint callback for stage '%s'", stage) - callback_func() diff --git a/tests/func/api/__init__.py b/tests/func/api/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/func/api/test_live.py b/tests/func/api/test_live.py deleted file mode 100644 index 3edc69e36..000000000 --- a/tests/func/api/test_live.py +++ /dev/null @@ -1,55 +0,0 @@ -import io -import json -import logging -from typing import Dict, List - -import pytest -from funcy import last - -from dvc.api.live import summary -from tests.utils import dump_sv - - -def _dumps_tsv(metrics: List[Dict]): - stream = io.StringIO() - dump_sv(stream, metrics) - stream.seek(0) - return stream.read() - - -@pytest.fixture -def live_results(tmp_dir): - def make(path="logs"): - datapoints = [{"metric": 0.0, "step": 0}, {"metric": 0.5, "step": 1}] - tmp_dir.gen( - { - (tmp_dir / path).with_suffix(".json"): json.dumps( - last(datapoints) - ), - (tmp_dir / path / "metric.tsv"): _dumps_tsv(datapoints), - } - ) - - yield make - - -def test_live_summary_no_repo(tmp_dir, live_results, caplog): - live_results("logs") - - with caplog.at_level(logging.INFO, logger="dvc"): - summary("logs") - - summary_path = tmp_dir / "logs.html" - assert summary_path.exists() - assert f"file://{str(summary_path)}" in caplog.text - - -def test_live_summary(tmp_dir, dvc, live_results, caplog): - live_results("logs") - - with caplog.at_level(logging.INFO, logger="dvc"): - summary("logs") - - summary_path = tmp_dir / "logs.html" - assert summary_path.exists() - assert f"file://{str(summary_path)}" in caplog.text diff --git a/tests/func/experiments/conftest.py b/tests/func/experiments/conftest.py index 5d35e33c3..455aa687c 100644 --- a/tests/func/experiments/conftest.py +++ b/tests/func/experiments/conftest.py @@ -67,7 +67,9 @@ def exp_stage(tmp_dir, scm, dvc): @pytest.fixture -def checkpoint_stage(tmp_dir, scm, dvc): +def checkpoint_stage(tmp_dir, scm, dvc, mocker): + mocker.patch("dvc.stage.run.Monitor.AWAIT", 0.01) + tmp_dir.gen("checkpoint.py", CHECKPOINT_SCRIPT) tmp_dir.gen("params.yaml", "foo: 1") stage = dvc.run( diff --git a/tests/func/api/test_api.py b/tests/func/test_api.py similarity index 100% rename from tests/func/api/test_api.py rename to tests/func/test_api.py diff --git a/tests/func/test_live.py b/tests/func/test_live.py index b6bb3b727..54134c4aa 100644 --- a/tests/func/test_live.py +++ b/tests/func/test_live.py @@ -8,7 +8,7 @@ from funcy import first from dvc import stage as stage_module from dvc.exceptions import MetricsError -LIVE_SCRITP = dedent( +LIVE_SCRIPT = dedent( """ import dvclive import sys @@ -53,15 +53,23 @@ LIVE_CHECKPOINT_SCRIPT = dedent( @pytest.fixture -def live_stage(tmp_dir, scm, dvc): +def live_stage(tmp_dir, scm, dvc, mocker): try: import dvclive # noqa, pylint:disable=unused-import except ImportError: pytest.skip("no dvclive") - def make(summary=True, html=True, live=None, live_no_cache=None): + mocker.patch("dvc.stage.run.Monitor.AWAIT", 0.01) + + def make( + summary=True, + html=True, + live=None, + live_no_cache=None, + code=LIVE_SCRIPT, + ): assert bool(live) != bool(live_no_cache) - tmp_dir.gen("train.py", LIVE_SCRITP) + tmp_dir.gen("train.py", code) tmp_dir.gen("params.yaml", "foo: 1") stage = dvc.run( cmd="python train.py", @@ -81,19 +89,11 @@ def live_stage(tmp_dir, scm, dvc): yield make -@pytest.mark.parametrize("report", (True, False)) +@pytest.mark.parametrize("html", (True, False)) @pytest.mark.parametrize("summary", (True, False)) -def test_export_config_tmp(tmp_dir, dvc, mocker, summary, report): +def test_export_config(tmp_dir, dvc, mocker, live_stage, summary, html): run_spy = mocker.spy(stage_module.run, "_run") - tmp_dir.gen("src", "dependency") - dvc.run( - cmd="mkdir logs && touch logs.json", - deps=["src"], - name="run_logger", - live="logs", - live_no_summary=not summary, - live_no_html=not report, - ) + live_stage(summary=summary, html=html, live="logs") assert run_spy.call_count == 1 _, kwargs = run_spy.call_args @@ -105,22 +105,7 @@ def test_export_config_tmp(tmp_dir, dvc, mocker, summary, report): assert kwargs["env"]["DVCLIVE_SUMMARY"] == str(int(summary)) assert "DVCLIVE_HTML" in kwargs["env"] - assert kwargs["env"]["DVCLIVE_HTML"] == str(int(report)) - - -@pytest.mark.parametrize("summary", (True, False)) -def test_export_config(tmp_dir, dvc, mocker, summary, live_stage): - run_spy = mocker.spy(stage_module.run, "_run") - live_stage(summary=summary, live="logs") - - assert run_spy.call_count == 1 - _, kwargs = run_spy.call_args - - assert "DVCLIVE_PATH" in kwargs["env"] - assert kwargs["env"]["DVCLIVE_PATH"] == "logs" - - assert "DVCLIVE_SUMMARY" in kwargs["env"] - assert kwargs["env"]["DVCLIVE_SUMMARY"] == str(int(summary)) + assert kwargs["env"]["DVCLIVE_HTML"] == str(int(html)) def test_live_provides_metrics(tmp_dir, dvc, live_stage): @@ -171,12 +156,14 @@ def test_live_html(tmp_dir, dvc, live_stage, html): @pytest.fixture -def live_checkpoint_stage(tmp_dir, scm, dvc): +def live_checkpoint_stage(tmp_dir, scm, dvc, mocker): try: import dvclive # noqa, pylint:disable=unused-import except ImportError: pytest.skip("no dvclive") + mocker.patch("dvc.stage.run.Monitor.AWAIT", 0.01) + def make(live=None, live_no_cache=None): assert bool(live) != bool(live_no_cache) @@ -247,3 +234,27 @@ def test_live_checkpoints_resume( 4, 2, ] + + +def test_dvc_generates_html_during_run(tmp_dir, dvc, mocker, live_stage): + show_spy = mocker.spy(dvc.live, "show") + + # make sure script takes more time to execute than one monitor sleep cycle + monitor_await_time = 0.01 + mocker.patch("dvc.stage.run.Monitor.AWAIT", monitor_await_time) + + script = dedent( + """ + import dvclive + import sys + import time + dvclive.log("loss", 1/2) + dvclive.log("accuracy", 1/2) + dvclive.next_step() + time.sleep({})""".format( + str(monitor_await_time * 2) + ) + ) + live_stage(summary=True, live="logs", code=script) + + assert show_spy.call_count == 2