mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-12 07:18:23 +01:00
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:
@@ -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:
|
||||
@@ -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)}")
|
||||
@@ -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 (
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user