Add start and done calls with post_live_metrics.

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
 

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 = [

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,
+    }