From b6d8c07d4d453b46d3bee47e2d6d5574654b68c7 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Fri, 21 Apr 2023 10:37:13 +0200 Subject: [PATCH] studio: Ensure `studio_repo_url` works with `exp run --queue`. Closes #9349 --- dvc/repo/experiments/executor/base.py | 9 +++++++++ dvc/repo/experiments/queue/base.py | 10 ++++++---- pyproject.toml | 2 +- tests/func/experiments/test_experiments.py | 3 ++- tests/integration/test_studio_live_experiments.py | 9 ++++++++- 5 files changed, 26 insertions(+), 7 deletions(-) diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 5d364a0c6..b9e672a9d 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -627,6 +627,7 @@ class BaseExecutor(ABC): copy_paths: Optional[List[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 @@ -649,12 +650,18 @@ class BaseExecutor(ABC): os.chdir(dvc.root_dir) 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, "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), ) logger.debug("Running repro in '%s'", os.getcwd()) yield dvc @@ -680,6 +687,8 @@ class BaseExecutor(ABC): "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), ) if infofile is not None: diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 98ff3ace8..321194b2b 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -17,8 +17,8 @@ from typing import ( Union, ) -from dvc_studio_client.env import STUDIO_REPO_URL -from dvc_studio_client.post_live_metrics import get_studio_repo_url +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 funcy import retry from dvc.dependency import ParamsDependency @@ -420,12 +420,14 @@ class BaseStashQueue(ABC): run_env[DVC_EXP_NAME] = name if resume_rev: run_env[DVCLIVE_RESUME] = "1" - studio_repo_url = get_studio_repo_url() + + 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 self._pack_args(*args, run_env=run_env, **kwargs) - # save experiment as a stash commit msg = self._stash_msg( stash_head, diff --git a/pyproject.toml b/pyproject.toml index e8beab2cf..a86b3a891 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ "dvc-data>=0.47.1,<0.48", "dvc-http>=2.29.0", "dvc-render>=0.3.1,<0.4.0", - "dvc-studio-client>=0.6.1,<1", + "dvc-studio-client>=0.8.0,<1", "dvc-task>=0.2.0,<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 b4c063c3e..4fe02a3b5 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -614,7 +614,8 @@ def test_run_env(tmp_dir, dvc, scm, mocker): """ ) mocker.patch( - "dvc.repo.experiments.queue.base.get_studio_repo_url", return_value="REPO_URL" + "dvc.repo.experiments.queue.base.get_studio_token_and_repo_url", + return_value=("REPO_TOKEN", "REPO_URL"), ) (tmp_dir / "dump_run_env.py").write_text(dump_run_env) baseline = scm.get_rev() diff --git a/tests/integration/test_studio_live_experiments.py b/tests/integration/test_studio_live_experiments.py index 1f3a91d3b..2ed583b9a 100644 --- a/tests/integration/test_studio_live_experiments.py +++ b/tests/integration/test_studio_live_experiments.py @@ -1,5 +1,5 @@ import pytest -from dvc_studio_client import env +from dvc_studio_client import env, post_live_metrics from funcy import first @@ -7,6 +7,7 @@ from funcy import first def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp): 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") @@ -18,6 +19,12 @@ def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp): 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 start_call, done_call = mocked_post.call_args_list