-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
exp share refactor - ping on exp:remove, do not fail on error (#9248)
studio experiments share refactor 1) Pings Studio on exp remove 2) Does not fail on any error during request.post 3) UI improvement on exp remove
- Loading branch information
Showing
6 changed files
with
171 additions
and
100 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,93 @@ | ||
import logging | ||
import os | ||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple | ||
from urllib.parse import urljoin | ||
|
||
from dvc_studio_client.post_live_metrics import get_studio_repo_url | ||
from funcy import compact | ||
from requests import RequestException, Session | ||
from requests.adapters import HTTPAdapter | ||
|
||
if TYPE_CHECKING: | ||
from requests import Response | ||
|
||
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, | ||
data: Dict[str, Any], | ||
url: Optional[str] = STUDIO_URL, | ||
max_retries: int = 3, | ||
timeout: int = 5, | ||
) -> "Response": | ||
endpoint = urljoin(url or STUDIO_URL, endpoint) | ||
session = Session() | ||
session.mount(endpoint, HTTPAdapter(max_retries=max_retries)) | ||
|
||
logger.trace("Sending %s to %s", data, endpoint) # type: ignore[attr-defined] | ||
|
||
headers = {"Authorization": f"token {token}"} | ||
resp = session.post(endpoint, json=data, headers=headers, timeout=timeout) | ||
resp.raise_for_status() | ||
return resp | ||
|
||
|
||
def notify_refs( | ||
git_remote: 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? | ||
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, | ||
) | ||
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 "", | ||
) | ||
data = {"repo_url": repo_url, "client": "dvc", "refs": refs} | ||
|
||
try: | ||
post("/webhook/dvc", token=token, data=data, url=studio_url) | ||
except RequestException: | ||
# TODO: handle expected failures and show appropriate message | ||
# TODO: handle unexpected failures and show appropriate message | ||
logger.debug("failed to notify Studio", exc_info=True) |
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,44 @@ | ||
from urllib.parse import urljoin | ||
|
||
import pytest | ||
from requests import Response | ||
|
||
from dvc.utils.studio import STUDIO_URL, notify_refs | ||
|
||
|
||
@pytest.mark.parametrize( | ||
"status_code", | ||
[ | ||
200, # success | ||
401, # should not fail on client errors | ||
500, # should not fail even on server errors | ||
], | ||
) | ||
def test_notify_refs(mocker, status_code): | ||
response = Response() | ||
response.status_code = status_code | ||
|
||
mock_post = mocker.patch("requests.Session.post", return_value=response) | ||
|
||
notify_refs( | ||
"origin", | ||
"TOKEN", | ||
repo_url_finder=lambda: "[email protected]:iterative/dvc.git", | ||
pushed=["p1", "p2"], | ||
removed=["r1", "r2"], | ||
) | ||
|
||
assert mock_post.called | ||
assert mock_post.call_args == mocker.call( | ||
urljoin(STUDIO_URL, "/webhook/dvc"), | ||
json={ | ||
"repo_url": "[email protected]:iterative/dvc.git", | ||
"client": "dvc", | ||
"refs": { | ||
"pushed": ["p1", "p2"], | ||
"removed": ["r1", "r2"], | ||
}, | ||
}, | ||
headers={"Authorization": "token TOKEN"}, | ||
timeout=5, | ||
) |