diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index 4331c4562..7de537e9e 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -96,12 +96,17 @@ class WorkspaceQueue(BaseStashQueue): def _reproduce_entry( self, entry: QueueEntry, executor: BaseExecutor ) -> Dict[str, Dict[str, str]]: + from dvc_studio_client.post_live_metrics import post_live_metrics + from dvc.stage.monitor import CheckpointKilledError results: Dict[str, Dict[str, str]] = defaultdict(dict) exec_name = self._EXEC_NAME or entry.stash_rev infofile = self.get_infofile_path(exec_name) try: + post_live_metrics( + "start", executor.info.baseline_rev, executor.info.name, "dvc" + ) rev = entry.stash_rev exec_result = executor.reproduce( info=executor.info, @@ -130,6 +135,13 @@ class WorkspaceQueue(BaseStashQueue): f"Failed to reproduce experiment '{rev[:7]}'" ) from exc finally: + post_live_metrics( + "done", + executor.info.baseline_rev, + executor.info.name, + "dvc", + experiment_rev=first(results.get(rev, {})), + ) executor.cleanup(infofile) return results diff --git a/pyproject.toml b/pyproject.toml index 48f7f175f..c5504b531 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,7 @@ dependencies = [ "dvc-http==2.27.2", "hydra-core>=1.1.0", "iterative-telemetry==0.0.6", + "dvc-studio-client>=0.1.1" ] dynamic = ["version"] classifiers = [ diff --git a/tests/func/experiments/test_utils.py b/tests/func/experiments/test_utils.py index 260cc417c..76dc9df6a 100644 --- a/tests/func/experiments/test_utils.py +++ b/tests/func/experiments/test_utils.py @@ -1,3 +1,4 @@ +from dvc_studio_client import env from funcy import first @@ -27,3 +28,39 @@ def test_generate_random_exp_name(tmp_dir, dvc, scm, exp_stage, mocker): # Can use same name because of different baseline_rev ref = first(dvc.experiments.run(exp_stage.addressing, params=["foo=1"])) assert dvc.experiments.get_exact_name([ref])[ref] == "0-0" + + +def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch): + valid_response = mocker.MagicMock() + valid_response.status_code = 200 + mocked_post = mocker.patch("requests.post", return_value=valid_response) + + monkeypatch.setenv(env.STUDIO_ENDPOINT, "https://0.0.0.0") + monkeypatch.setenv(env.STUDIO_REPO_URL, "STUDIO_REPO_URL") + monkeypatch.setenv(env.STUDIO_TOKEN, "STUDIO_TOKEN") + + baseline_sha = scm.get_rev() + exp_rev = first( + dvc.experiments.run(exp_stage.addressing, params=["foo=1"]) + ) + name = dvc.experiments.get_exact_name([exp_rev])[exp_rev] + assert mocked_post.call_count == 2 + + start_call, done_call = mocked_post.call_args_list + + assert start_call.kwargs["json"] == { + "type": "start", + "repo_url": "STUDIO_REPO_URL", + "baseline_sha": baseline_sha, + "name": name, + "client": "dvc", + } + + assert done_call.kwargs["json"] == { + "type": "done", + "repo_url": "STUDIO_REPO_URL", + "baseline_sha": baseline_sha, + "name": name, + "client": "dvc", + "experiment_rev": exp_rev, + }