diff --git a/dvc/config_schema.py b/dvc/config_schema.py index 4da2d662f..2e36e9019 100644 --- a/dvc/config_schema.py +++ b/dvc/config_schema.py @@ -330,5 +330,7 @@ SCHEMA = { "studio": { "token": str, "url": str, + "repo_url": str, + Optional("offline", default=False): Bool, }, } diff --git a/dvc/env.py b/dvc/env.py index 9646278d3..081ec9d79 100644 --- a/dvc/env.py +++ b/dvc/env.py @@ -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" diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 2176059ff..5cbe248b2 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -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: diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 739852b88..0df531d69 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -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( diff --git a/dvc/utils/studio.py b/dvc/utils/studio.py index a406fcef6..da6a9d994 100644 --- a/dvc/utils/studio.py +++ b/dvc/utils/studio.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 9367c89f8..3cf93cd9f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 62a46f8ee..2d5ab0488 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -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): diff --git a/tests/integration/test_studio_live_experiments.py b/tests/integration/test_studio_live_experiments.py index 9ad281e19..c56079f52 100644 --- a/tests/integration/test_studio_live_experiments.py +++ b/tests/integration/test_studio_live_experiments.py @@ -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]) diff --git a/tests/unit/utils/test_studio.py b/tests/unit/utils/test_studio.py index a10c88bb0..68dcfed77 100644 --- a/tests/unit/utils/test_studio.py +++ b/tests/unit/utils/test_studio.py @@ -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