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
This commit is contained in:
Paweł Redzyński
2021-03-03 00:45:35 +09:00
committed by GitHub
parent e385e59a09
commit 4a8cb806d2
13 changed files with 248 additions and 200 deletions
+2 -2
View File
@@ -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:
-24
View File
@@ -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)}")
+1 -1
View File
@@ -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 (
+1 -1
View File
@@ -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
+15
View File
@@ -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
+1 -1
View File
@@ -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)
+150
View File
@@ -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
+32 -83
View File
@@ -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()
View File
-55
View File
@@ -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
+3 -1
View File
@@ -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(
+43 -32
View File
@@ -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