From 837c8e8de4243fcc63b19895359daef6d0fce10a Mon Sep 17 00:00:00 2001 From: dberenbaum Date: Thu, 7 Mar 2024 13:33:19 -0500 Subject: [PATCH] exp: use exp git remote for studio repo url --- dvc/repo/experiments/executor/base.py | 5 ++++ dvc/utils/studio.py | 17 +++++++++++++ .../test_studio_live_experiments.py | 10 +++++--- tests/unit/utils/test_studio.py | 24 ++++++++++++++++++- 4 files changed, 52 insertions(+), 4 deletions(-) diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 933353b57..64b670c80 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -25,6 +25,7 @@ from dvc.utils import dict_sha256, env2bool, relpath from dvc.utils.fs import remove from dvc.utils.studio import ( env_to_config, + get_repo_url, get_subrepo_relpath, ) @@ -611,6 +612,8 @@ class BaseExecutor(ABC): # 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 + # override studio repo url if exp git remote set + repo_url = get_repo_url(dvc) try: post_live_metrics( "start", @@ -621,6 +624,7 @@ class BaseExecutor(ABC): dvc_studio_config=dvc_studio_config, message=message, subdir=get_subrepo_relpath(dvc), + studio_repo_url=repo_url, ) logger.debug("Running repro in '%s'", os.getcwd()) yield dvc @@ -646,6 +650,7 @@ class BaseExecutor(ABC): experiment_rev=dvc.experiments.scm.get_ref(EXEC_BRANCH), metrics=_gather_metrics(dvc, on_error="return"), dvc_studio_config=dvc_studio_config, + studio_repo_url=repo_url, ) if infofile is not None: diff --git a/dvc/utils/studio.py b/dvc/utils/studio.py index 470d19549..13065c291 100644 --- a/dvc/utils/studio.py +++ b/dvc/utils/studio.py @@ -1,3 +1,4 @@ +import os from typing import TYPE_CHECKING, Any, Optional from urllib.parse import urljoin @@ -125,3 +126,19 @@ def get_subrepo_relpath(repo: "Repo") -> str: relpath = as_posix(repo.fs.relpath(repo.root_dir, scm_root_dir)) return "" if relpath == "." else relpath + + +def get_repo_url(repo: "Repo") -> str: + from dulwich.porcelain import get_remote_repo + + from dvc.env import DVC_EXP_GIT_REMOTE + + repo_url = os.getenv( + DVC_EXP_GIT_REMOTE, repo.config.get("exp", {}).get("git_remote") + ) + if repo_url: + try: + _, repo_url = get_remote_repo(repo.scm.dulwich.repo, repo_url) + except IndexError: + pass + return repo_url diff --git a/tests/integration/test_studio_live_experiments.py b/tests/integration/test_studio_live_experiments.py index da1136f08..e5851aac7 100644 --- a/tests/integration/test_studio_live_experiments.py +++ b/tests/integration/test_studio_live_experiments.py @@ -3,6 +3,7 @@ from dvc_studio_client import env, post_live_metrics from funcy import first from dvc.env import ( + DVC_EXP_GIT_REMOTE, DVC_STUDIO_OFFLINE, DVC_STUDIO_REPO_URL, DVC_STUDIO_TOKEN, @@ -16,8 +17,9 @@ from dvc.utils.studio import get_subrepo_relpath @pytest.mark.studio @pytest.mark.parametrize("tmp", [True, False]) @pytest.mark.parametrize("offline", [True, False]) +@pytest.mark.parametrize("dvc_exp_git_remote", [None, "DVC_EXP_GIT_REMOTE"]) def test_post_to_studio( - tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp, offline + tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp, offline, dvc_exp_git_remote ): valid_response = mocker.MagicMock() valid_response.status_code = 200 @@ -28,6 +30,8 @@ def test_post_to_studio( monkeypatch.setenv(DVC_STUDIO_TOKEN, "STUDIO_TOKEN") monkeypatch.setenv(DVC_STUDIO_URL, "https://0.0.0.0") monkeypatch.setenv(DVC_STUDIO_OFFLINE, offline) + if dvc_exp_git_remote: + monkeypatch.setenv(DVC_EXP_GIT_REMOTE, dvc_exp_git_remote) baseline_sha = scm.get_rev() exp_rev = first( @@ -52,7 +56,7 @@ def test_post_to_studio( assert start_call.kwargs["json"] == { "type": "start", - "repo_url": "STUDIO_REPO_URL", + "repo_url": dvc_exp_git_remote or "STUDIO_REPO_URL", "baseline_sha": baseline_sha, "name": name, "params": {"params.yaml": {"foo": 1}}, @@ -61,7 +65,7 @@ def test_post_to_studio( assert done_call.kwargs["json"] == { "type": "done", - "repo_url": "STUDIO_REPO_URL", + "repo_url": dvc_exp_git_remote or "STUDIO_REPO_URL", "baseline_sha": baseline_sha, "name": name, "client": "dvc", diff --git a/tests/unit/utils/test_studio.py b/tests/unit/utils/test_studio.py index 68dcfed77..cc9351ce8 100644 --- a/tests/unit/utils/test_studio.py +++ b/tests/unit/utils/test_studio.py @@ -4,12 +4,19 @@ import pytest from requests import Response from dvc.env import ( + DVC_EXP_GIT_REMOTE, 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 +from dvc.utils.studio import ( + STUDIO_URL, + config_to_env, + env_to_config, + get_repo_url, + notify_refs, +) CONFIG = {"offline": True, "repo_url": "repo_url", "token": "token", "url": "url"} @@ -67,3 +74,18 @@ def test_config_to_env(): def test_env_to_config(): assert env_to_config(ENV) == CONFIG + + +@pytest.mark.parametrize( + "exp_git_remote, repo_url", + [ + (None, None), + ("origin", "git@url"), + ("http://url", "http://url"), + ], +) +def test_get_repo_url(dvc, scm, monkeypatch, exp_git_remote, repo_url): + scm.gitpython.repo.create_remote("origin", "git@url") + if exp_git_remote: + monkeypatch.setenv(DVC_EXP_GIT_REMOTE, exp_git_remote) + assert get_repo_url(dvc) == repo_url