From e73cc9b40ac6a2b21e5adfb0890cfe15b35dbcb5 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Tue, 14 Mar 2023 10:32:07 +0100 Subject: [PATCH] exp: Include `STUDIO_REPO_URL` in run_env. Closes #8951. --- dvc/repo/experiments/queue/base.py | 6 ++++++ tests/func/experiments/test_experiments.py | 14 ++++++++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 82d8061818..72913d59df 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -17,6 +17,8 @@ Union, ) +from dvc_studio_client.env import STUDIO_REPO_URL +from dvc_studio_client.post_live_metrics import get_studio_repo_url from funcy import retry from dvc.dependency import ParamsDependency @@ -415,6 +417,10 @@ def _stash_exp( # noqa: PLR0915, C901 run_env[DVC_EXP_NAME] = name if resume_rev: run_env[DVCLIVE_RESUME] = "1" + studio_repo_url = get_studio_repo_url() + 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 diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 4a1bbe7fa5..f49336c8a3 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -597,17 +597,21 @@ def test_experiments_workspace_not_log_exception(caplog, dvc, scm): assert not caplog.text +@pytest.mark.vscode 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 - with open("DVC_EXP_BASELINE_REV", "w") as f: - f.write(os.environ.get(DVC_EXP_BASELINE_REV, "")) - with open("DVC_EXP_NAME", "w") as f: - f.write(os.environ.get(DVC_EXP_NAME, "")) + for v in (DVC_EXP_BASELINE_REV, DVC_EXP_NAME, STUDIO_REPO_URL): + with open(v, "w") as f: + f.write(os.environ.get(v, "")) """ ) + mocker.patch( + "dvc.repo.experiments.queue.base.get_studio_repo_url", return_value="REPO_URL" + ) (tmp_dir / "dump_run_env.py").write_text(dump_run_env) baseline = scm.get_rev() dvc.stage.add( @@ -617,10 +621,12 @@ def test_run_env(tmp_dir, dvc, scm, mocker): 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" 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" def test_experiment_unchanged(tmp_dir, scm, dvc, exp_stage):