studio: Move post_live_metrics from queue to executor (#8955)

* studio: Move `post_live_metrics` from `queue/workspace` to `executor/base`.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update dvc/repo/experiments/utils.py

Co-authored-by: skshetry <[email protected]>

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: skshetry <[email protected]>

dvc/repo/experiments/executor/base.py

@@ -32,6 +32,7 @@ from dvc.repo.experiments.refs import (
     EXEC_CHECKPOINT,
     ExpRefInfo,
 )
+from dvc.repo.experiments.utils import to_studio_params
 from dvc.stage.serialize import to_lockfile
 from dvc.ui import ui
 from dvc.utils import dict_sha256, env2bool, relpath
@@ -545,6 +546,8 @@ class BaseExecutor(ABC):
         log_errors: bool = True,
         **kwargs,
     ):
+        from dvc_studio_client.post_live_metrics import post_live_metrics
+
         from dvc.repo import Repo
         from dvc.stage.monitor import CheckpointKilledError
 
@@ -561,6 +564,13 @@ class BaseExecutor(ABC):
             os.chdir(dvc.root_dir)
 
         try:
+            post_live_metrics(
+                "start",
+                info.baseline_rev,
+                info.name,
+                "dvc",
+                params=to_studio_params(dvc.params.show()),
+            )
             logger.debug("Running repro in '%s'", os.getcwd())
             yield dvc
             info.status = TaskStatus.SUCCESS
@@ -578,6 +588,14 @@ class BaseExecutor(ABC):
             info.status = TaskStatus.FAILED
             raise
         finally:
+            post_live_metrics(
+                "done",
+                info.baseline_rev,
+                info.name,
+                "dvc",
+                experiment_rev=dvc.experiments.scm.get_ref(EXEC_BRANCH),
+            )
+
             if infofile is not None:
                 info.dump_json(infofile)
             dvc.close()

dvc/repo/experiments/queue/workspace.py

@@ -13,7 +13,7 @@ from dvc.repo.experiments.exceptions import ExpQueueEmptyError
 from dvc.repo.experiments.executor.base import ExecutorInfo, TaskStatus
 from dvc.repo.experiments.executor.local import WorkspaceExecutor
 from dvc.repo.experiments.refs import EXEC_BRANCH, WORKSPACE_STASH
-from dvc.repo.experiments.utils import get_exp_rwlock, to_studio_params
+from dvc.repo.experiments.utils import get_exp_rwlock
 from dvc.rwlock import RWLOCK_FILE, RWLOCK_LOCK, SCHEMA
 from dvc.utils import relpath
 from dvc.utils.fs import remove
@@ -94,8 +94,6 @@ 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)
@@ -103,13 +101,6 @@ class WorkspaceQueue(BaseStashQueue):
         infofile = self.get_infofile_path(exec_name)
         try:
             rev = entry.stash_rev
-            post_live_metrics(
-                "start",
-                executor.info.baseline_rev,
-                executor.info.name,
-                "dvc",
-                params=to_studio_params(self.repo.params.show()),
-            )
             exec_result = executor.reproduce(
                 info=executor.info,
                 rev=rev,
@@ -131,13 +122,6 @@ class WorkspaceQueue(BaseStashQueue):
         except Exception as exc:  # noqa: BLE001
             raise DvcException(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
 

dvc/repo/experiments/utils.py

@@ -327,7 +327,9 @@ def to_studio_params(dvc_params):
         "params.yaml": {"foo": 1}
     }
     """
-    result = {}
+    result: Dict = {}
+    if not dvc_params:
+        return result
     for rev_data in dvc_params.values():
         for file_name, file_data in rev_data["data"].items():
             result[file_name] = file_data["data"]

tests/func/experiments/test_utils.py

@@ -1,4 +1,3 @@
-from dvc_studio_client import env
 from funcy import first
 
 
@@ -28,38 +27,3 @@ 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,
-        "params": {"params.yaml": {"foo": 1}},
-        "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,
-    }

tests/integration/conftest.py

@@ -0,0 +1,2 @@
+# pylint disable=unused-argument
+from tests.unit.repo.experiments.conftest import exp_stage, test_queue  # noqa: F401

tests/integration/test_studio_live_experiments.py

@@ -0,0 +1,41 @@
+import pytest
+from dvc_studio_client import env
+from funcy import first
+
+
+@pytest.mark.parametrize("tmp", [True, False])
+def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp):
+    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"], tmp_dir=tmp)
+    )
+    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,
+        "params": {"params.yaml": {"foo": 1}},
+        "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,
+    }