diff --git a/dvc/repo/experiments/push.py b/dvc/repo/experiments/push.py index 12a7dd32c..677a1c468 100644 --- a/dvc/repo/experiments/push.py +++ b/dvc/repo/experiments/push.py @@ -17,15 +17,18 @@ from .utils import exp_commits, exp_refs, exp_refs_by_baseline, resolve_name logger = logging.getLogger(__name__) -def notify_refs_to_studio(config, git_remote: str, **refs: List[str]): +def notify_refs_to_studio(scm, config, git_remote: str, **refs: List[str]): token = config.get("studio_token") refs = compact(refs) if refs and (token or config["push_exp_to_studio"]) and not env2bool("DVC_TEST"): + from dulwich.porcelain import get_remote_repo + from dvc.utils import studio + _, repo_url = get_remote_repo(scm.dulwich.repo, git_remote) studio_url = config.get("studio_url") studio.notify_refs( - git_remote, + repo_url, default_token=token, studio_url=studio_url, **refs, # type: ignore[arg-type] @@ -88,7 +91,9 @@ def push( # noqa: C901 refs = push_result[SyncStatus.SUCCESS] pushed_refs = [str(r) for r in refs] - notify_refs_to_studio(repo.config["feature"], git_remote, pushed=pushed_refs) + notify_refs_to_studio( + repo.scm, repo.config["feature"], git_remote, pushed=pushed_refs + ) return [ref.name for ref in refs] diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index 27675dc2d..74cb64ae0 100644 --- a/dvc/repo/experiments/remove.py +++ b/dvc/repo/experiments/remove.py @@ -79,7 +79,9 @@ def remove( # noqa: C901, PLR0912 from .push import notify_refs_to_studio removed_refs = [str(r) for r in exp_ref_list] - notify_refs_to_studio(repo.config["feature"], git_remote, removed=removed_refs) + notify_refs_to_studio( + repo.scm, repo.config["feature"], git_remote, removed=removed_refs + ) return removed diff --git a/dvc/utils/studio.py b/dvc/utils/studio.py index 071f1bb0d..ecc1a51b8 100644 --- a/dvc/utils/studio.py +++ b/dvc/utils/studio.py @@ -1,10 +1,9 @@ import logging import os -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional from urllib.parse import urljoin import requests -from dvc_studio_client.post_live_metrics import get_studio_repo_url from funcy import compact from requests.adapters import HTTPAdapter @@ -14,20 +13,9 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) STUDIO_URL = "https://studio.iterative.ai" -STUDIO_REPO_URL = "STUDIO_REPO_URL" STUDIO_TOKEN = "STUDIO_TOKEN" # noqa: S105 -def get_studio_token_and_repo_url( - default_token: Optional[str] = None, - repo_url_finder: Optional[Callable[[], Optional[str]]] = None, -) -> Tuple[Optional[str], Optional[str]]: - token = os.getenv(STUDIO_TOKEN) or default_token - url_finder = repo_url_finder or get_studio_repo_url - repo_url = os.getenv(STUDIO_REPO_URL) or url_finder() - return token, repo_url - - def post( endpoint: str, token: str, @@ -49,36 +37,24 @@ def post( def notify_refs( - git_remote: str, + repo_url: str, + *, default_token: Optional[str] = None, studio_url: Optional[str] = None, - repo_url_finder: Optional[Callable[[], Optional[str]]] = None, **refs: List[str], ) -> None: - # TODO: Should we use git_remote to associate with Studio project - # instead of using `git ls-remote` on fallback? + extra_keys = refs.keys() - {"pushed", "removed"} + assert not extra_keys, f"got extra args: {extra_keys}" + refs = compact(refs) if not refs: return - assert git_remote - token, repo_url = get_studio_token_and_repo_url( - default_token=default_token, - repo_url_finder=repo_url_finder, - ) + token = os.getenv(STUDIO_TOKEN) or default_token if not token: logger.debug("Studio token not found.") return - if not repo_url: - logger.warning( - "Could not detect repository url. " - "Please set %s environment variable " - "to your remote git repository url. ", - STUDIO_REPO_URL, - ) - return - logger.debug( "notifying Studio%s about updated experiments", f" ({studio_url})" if studio_url else "", diff --git a/tests/unit/utils/test_studio.py b/tests/unit/utils/test_studio.py index df9598764..c3767c898 100644 --- a/tests/unit/utils/test_studio.py +++ b/tests/unit/utils/test_studio.py @@ -21,9 +21,8 @@ def test_notify_refs(mocker, status_code): mock_post = mocker.patch("requests.Session.post", return_value=response) notify_refs( - "origin", - "TOKEN", - repo_url_finder=lambda: "git@github.com:iterative/dvc.git", + "git@github.com:iterative/dvc.git", + default_token="TOKEN", pushed=["p1", "p2"], removed=["r1", "r2"], )