Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

Live metrics config (#9433)

Browse Source 
* exp: live metrics config token and offline

* refactor live metrics config

* use get_studio_config

* revert to using env vars for studio config

* update dvc-studio-client to support live metrics config changes

* upgrade dvc-studio-client

* don't change environ when posting live metrics events
This commit is contained in:
Dave Berenbaum
2023-05-17 09:21:01 +00:00
committed by GitHub
parent 99daf5d113
commit 7f6890ff25
9 changed files with 165 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/config_schema.py
+2
View File
@@ -330,5 +330,7 @@ SCHEMA = {
"studio": {
"token": str,
"url": str,
"repo_url": str,
Optional("offline", default=False): Bool,
},
}
dvc/env.py
+2
View File
@@ -12,6 +12,8 @@ DVC_PAGER = "DVC_PAGER"
DVC_ROOT = "DVC_ROOT"
DVC_SHOW_TRACEBACK = "DVC_SHOW_TRACEBACK"
DVC_SITE_CACHE_DIR = "DVC_SITE_CACHE_DIR"
DVC_STUDIO_OFFLINE = "DVC_STUDIO_OFFLINE"
DVC_STUDIO_REPO_URL = "DVC_STUDIO_REPO_URL"
DVC_STUDIO_TOKEN = "DVC_STUDIO_TOKEN" # noqa: S105 # nosec B105
DVC_STUDIO_URL = "DVC_STUDIO_URL"
DVC_SYSTEM_CONFIG_DIR = "DVC_SYSTEM_CONFIG_DIR"
dvc/repo/experiments/executor/base.py
+13 -11
View File
@@ -43,6 +43,7 @@ from dvc.stage.serialize import to_lockfile
from dvc.ui import ui
from dvc.utils import dict_sha256, env2bool, relpath
from dvc.utils.fs import remove
from dvc.utils.studio import env_to_config
if TYPE_CHECKING:
from queue import Queue
@@ -635,7 +636,6 @@ class BaseExecutor(ABC):
message: Optional[str] = None,
**kwargs,
) -> Iterator["Repo"]:
from dvc_studio_client.env import STUDIO_REPO_URL, STUDIO_TOKEN
from dvc_studio_client.post_live_metrics import post_live_metrics
from dvc.repo import Repo
@@ -657,19 +657,22 @@ class BaseExecutor(ABC):
else:
os.chdir(dvc.root_dir)
args_path = os.path.join(dvc.tmp_dir, cls.PACKED_ARGS_FILE)
if os.path.exists(args_path):
_, kwargs = cls.unpack_repro_args(args_path)
dvc_studio_config = dvc.config.get("studio")
# set missing config options using saved config
# inferring repo url will fail if not set here
run_env_config = env_to_config(kwargs.get("run_env", {}))
dvc_studio_config = {**run_env_config, **dvc_studio_config}
try:
args_path = os.path.join(dvc.tmp_dir, cls.PACKED_ARGS_FILE)
if os.path.exists(args_path):
_, kwargs = cls.unpack_repro_args(args_path)
run_env = kwargs.get("run_env", {})
post_live_metrics(
"start",
info.baseline_rev,
info.name,
info.name, # type: ignore[arg-type]
"dvc",
params=to_studio_params(dvc.params.show()),
studio_token=run_env.get(STUDIO_TOKEN, None),
studio_repo_url=run_env.get(STUDIO_REPO_URL, None),
dvc_studio_config=dvc_studio_config,
message=message,
)
logger.debug("Running repro in '%s'", os.getcwd())
@@ -692,12 +695,11 @@ class BaseExecutor(ABC):
post_live_metrics(
"done",
info.baseline_rev,
info.name,
info.name, # type: ignore[arg-type]
"dvc",
experiment_rev=dvc.experiments.scm.get_ref(EXEC_BRANCH),
metrics=get_in(dvc.metrics.show(), ["", "data"]),
studio_token=run_env.get(STUDIO_TOKEN, None),
studio_repo_url=run_env.get(STUDIO_REPO_URL, None),
dvc_studio_config=dvc_studio_config,
)
if infofile is not None:
dvc/repo/experiments/queue/base.py
+7 -8
View File
@@ -17,8 +17,7 @@ from typing import (
Union,
)
from dvc_studio_client.env import STUDIO_REPO_URL, STUDIO_TOKEN
from dvc_studio_client.post_live_metrics import get_studio_token_and_repo_url
from dvc_studio_client.post_live_metrics import get_studio_config
from funcy import retry
from dvc.dependency import ParamsDependency
@@ -39,6 +38,7 @@ from dvc.repo.experiments.utils import (
)
from dvc.ui import ui
from dvc.utils.objects import cached_property
from dvc.utils.studio import config_to_env
from .utils import get_remote_executor_refs
@@ -421,12 +421,11 @@ class BaseStashQueue(ABC):
if resume_rev:
run_env[DVCLIVE_RESUME] = "1"
studio_token, studio_repo_url = get_studio_token_and_repo_url()
if studio_token is not None:
run_env[STUDIO_TOKEN] = studio_token
if studio_repo_url is not None:
run_env[STUDIO_REPO_URL] = studio_repo_url
# save studio config to read later by dvc and dvclive
studio_config = get_studio_config(
dvc_studio_config=self.repo.config.get("studio")
)
run_env = {**config_to_env(studio_config), **run_env}
self._pack_args(*args, run_env=run_env, **kwargs)
# save experiment as a stash commit
msg = self._stash_msg(
dvc/utils/studio.py
+33
View File
@@ -6,6 +6,13 @@ import requests
from funcy import compact, ignore
from requests.adapters import HTTPAdapter
from dvc.env import (
DVC_STUDIO_OFFLINE,
DVC_STUDIO_REPO_URL,
DVC_STUDIO_TOKEN,
DVC_STUDIO_URL,
)
if TYPE_CHECKING:
from requests import Response
@@ -80,3 +87,29 @@ def notify_refs(
"received response: %s (status=%r)", d, r.status_code
)
return d
def config_to_env(config: Dict[str, Any]) -> Dict[str, Any]:
env = {}
if "offline" in config:
env[DVC_STUDIO_OFFLINE] = config["offline"]
if "repo_url" in config:
env[DVC_STUDIO_REPO_URL] = config["repo_url"]
if "token" in config:
env[DVC_STUDIO_TOKEN] = config["token"]
if "url" in config:
env[DVC_STUDIO_URL] = config["url"]
return env
def env_to_config(env: Dict[str, Any]) -> Dict[str, Any]:
config = {}
if DVC_STUDIO_OFFLINE in env:
config["offline"] = env[DVC_STUDIO_OFFLINE]
if DVC_STUDIO_REPO_URL in env:
config["repo_url"] = env[DVC_STUDIO_REPO_URL]
if DVC_STUDIO_TOKEN in env:
config["token"] = env[DVC_STUDIO_TOKEN]
if DVC_STUDIO_URL in env:
config["url"] = env[DVC_STUDIO_URL]
return config
pyproject.toml
+1 -1
View File
@@ -37,7 +37,7 @@ dependencies = [
"dvc-data>=0.51.0,<0.52",
"dvc-http>=2.29.0",
"dvc-render>=0.3.1,<1",
"dvc-studio-client>=0.9.0,<1",
"dvc-studio-client>=0.9.2,<1",
"dvc-task>=0.2.1,<1",
"flatten_dict<1,>=0.4.1",
"flufl.lock>=5",
tests/func/experiments/test_experiments.py
+39 -11
View File
@@ -9,6 +9,14 @@ from configobj import ConfigObj
from funcy import first
from dvc.dvcfile import PROJECT_FILE
from dvc.env import (
DVC_EXP_BASELINE_REV,
DVC_EXP_NAME,
DVC_STUDIO_OFFLINE,
DVC_STUDIO_REPO_URL,
DVC_STUDIO_TOKEN,
DVC_STUDIO_URL,
)
from dvc.exceptions import DvcException, ReproductionError
from dvc.repo.experiments.exceptions import ExperimentExistsError
from dvc.repo.experiments.queue.base import BaseStashQueue
@@ -606,16 +614,34 @@ def test_run_env(tmp_dir, dvc, scm, mocker):
dump_run_env = dedent(
"""\
import os
from dvc_studio_client.env import STUDIO_REPO_URL
from dvc.env import DVC_EXP_BASELINE_REV, DVC_EXP_NAME
for v in (DVC_EXP_BASELINE_REV, DVC_EXP_NAME, STUDIO_REPO_URL):
from dvc.env import (
DVC_EXP_BASELINE_REV,
DVC_EXP_NAME,
DVC_STUDIO_OFFLINE,
DVC_STUDIO_REPO_URL,
DVC_STUDIO_TOKEN,
DVC_STUDIO_URL
)
for v in (
DVC_EXP_BASELINE_REV,
DVC_EXP_NAME,
DVC_STUDIO_OFFLINE,
DVC_STUDIO_REPO_URL,
DVC_STUDIO_TOKEN,
DVC_STUDIO_URL
):
with open(v, "w") as f:
f.write(os.environ.get(v, ""))
"""
)
mocker.patch(
"dvc.repo.experiments.queue.base.get_studio_token_and_repo_url",
return_value=("REPO_TOKEN", "REPO_URL"),
"dvc.repo.experiments.queue.base.get_studio_config",
return_value={
"token": "TOKEN",
"repo_url": "REPO_URL",
"url": "BASE_URL",
"offline": "false",
},
)
(tmp_dir / "dump_run_env.py").write_text(dump_run_env)
baseline = scm.get_rev()
@@ -624,14 +650,16 @@ def test_run_env(tmp_dir, dvc, scm, mocker):
name="run_env",
)
dvc.experiments.run()
assert (tmp_dir / "DVC_EXP_BASELINE_REV").read_text().strip() == baseline
assert (tmp_dir / "DVC_EXP_NAME").read_text().strip()
assert (tmp_dir / "STUDIO_REPO_URL").read_text().strip() == "REPO_URL"
assert (tmp_dir / DVC_EXP_BASELINE_REV).read_text().strip() == baseline
assert (tmp_dir / DVC_EXP_NAME).read_text().strip()
assert (tmp_dir / DVC_STUDIO_TOKEN).read_text().strip() == "TOKEN"
assert (tmp_dir / DVC_STUDIO_REPO_URL).read_text().strip() == "REPO_URL"
assert (tmp_dir / DVC_STUDIO_URL).read_text().strip() == "BASE_URL"
assert (tmp_dir / DVC_STUDIO_OFFLINE).read_text().strip() == "false"
dvc.experiments.run(name="foo")
assert (tmp_dir / "DVC_EXP_BASELINE_REV").read_text().strip() == baseline
assert (tmp_dir / "DVC_EXP_NAME").read_text().strip() == "foo"
assert (tmp_dir / "STUDIO_REPO_URL").read_text().strip() == "REPO_URL"
assert (tmp_dir / DVC_EXP_BASELINE_REV).read_text().strip() == baseline
assert (tmp_dir / DVC_EXP_NAME).read_text().strip() == "foo"
def test_experiment_unchanged(tmp_dir, scm, dvc, exp_stage):
tests/integration/test_studio_live_experiments.py
+43 -26
View File
@@ -2,51 +2,68 @@ import pytest
from dvc_studio_client import env, post_live_metrics
from funcy import first
from dvc.env import (
DVC_STUDIO_OFFLINE,
DVC_STUDIO_REPO_URL,
DVC_STUDIO_TOKEN,
DVC_STUDIO_URL,
)
@pytest.mark.parametrize("tmp", [True, False])
def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp):
@pytest.mark.parametrize("offline", [True, False])
def test_post_to_studio(
tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp, offline
):
valid_response = mocker.MagicMock()
valid_response.status_code = 200
live_metrics = mocker.spy(post_live_metrics, "post_live_metrics")
mocked_post = mocker.patch("requests.post", return_value=valid_response)
monkeypatch.setenv(env.STUDIO_ENDPOINT, "https://0.0.0.0")
monkeypatch.setenv(env.STUDIO_REPO_URL, "STUDIO_REPO_URL")
monkeypatch.setenv(env.STUDIO_TOKEN, "STUDIO_TOKEN")
monkeypatch.setenv(DVC_STUDIO_REPO_URL, "STUDIO_REPO_URL")
monkeypatch.setenv(DVC_STUDIO_TOKEN, "STUDIO_TOKEN")
monkeypatch.setenv(DVC_STUDIO_URL, "https://0.0.0.0")
monkeypatch.setenv(DVC_STUDIO_OFFLINE, offline)
baseline_sha = scm.get_rev()
exp_rev = first(
dvc.experiments.run(exp_stage.addressing, params=["foo=1"], tmp_dir=tmp)
)
name = dvc.experiments.get_exact_name([exp_rev])[exp_rev]
assert live_metrics.call_count == 2
start_call, done_call = live_metrics.call_args_list
assert start_call.kwargs["studio_token"] == "STUDIO_TOKEN"
assert start_call.kwargs["studio_repo_url"] == "STUDIO_REPO_URL"
assert mocked_post.call_count == 2
if offline:
assert mocked_post.call_count == 0
start_call, done_call = mocked_post.call_args_list
else:
start_call, done_call = live_metrics.call_args_list
assert start_call.kwargs["dvc_studio_config"]["token"] == "STUDIO_TOKEN"
assert start_call.kwargs["dvc_studio_config"]["repo_url"] == "STUDIO_REPO_URL"
assert start_call.kwargs["json"] == {
"type": "start",
"repo_url": "STUDIO_REPO_URL",
"baseline_sha": baseline_sha,
"name": name,
"params": {"params.yaml": {"foo": 1}},
"client": "dvc",
}
assert mocked_post.call_count == 2
assert done_call.kwargs["json"] == {
"type": "done",
"repo_url": "STUDIO_REPO_URL",
"baseline_sha": baseline_sha,
"name": name,
"client": "dvc",
"experiment_rev": exp_rev,
"metrics": {"metrics.yaml": {"data": {"foo": 1}}},
}
start_call, done_call = mocked_post.call_args_list
assert start_call.kwargs["json"] == {
"type": "start",
"repo_url": "STUDIO_REPO_URL",
"baseline_sha": baseline_sha,
"name": name,
"params": {"params.yaml": {"foo": 1}},
"client": "dvc",
}
assert done_call.kwargs["json"] == {
"type": "done",
"repo_url": "STUDIO_REPO_URL",
"baseline_sha": baseline_sha,
"name": name,
"client": "dvc",
"experiment_rev": exp_rev,
"metrics": {"metrics.yaml": {"data": {"foo": 1}}},
}
@pytest.mark.parametrize("tmp", [True, False])
tests/unit/utils/test_studio.py
+25 -1
View File
@@ -3,7 +3,23 @@ from urllib.parse import urljoin
import pytest
from requests import Response
from dvc.utils.studio import STUDIO_URL, notify_refs
from dvc.env import (
DVC_STUDIO_OFFLINE,
DVC_STUDIO_REPO_URL,
DVC_STUDIO_TOKEN,
DVC_STUDIO_URL,
)
from dvc.utils.studio import STUDIO_URL, config_to_env, env_to_config, notify_refs
CONFIG = {"offline": True, "repo_url": "repo_url", "token": "token", "url": "url"}
ENV = {
DVC_STUDIO_OFFLINE: True,
DVC_STUDIO_REPO_URL: "repo_url",
DVC_STUDIO_TOKEN: "token",
DVC_STUDIO_URL: "url",
}
@pytest.mark.parametrize(
@@ -43,3 +59,11 @@ def test_notify_refs(mocker, status_code, side_effect):
timeout=5,
allow_redirects=False,
)
def test_config_to_env():
assert config_to_env(CONFIG) == ENV
def test_env_to_config():
assert env_to_config(ENV) == CONFIG