Repository: iterative/dvc

Issue:
studio: queued experiments never show as done
Not sure whether this is a dvc, dvclive, or studio issue, but when queueing experiments and running them with the `DVC_STUDIO_TOKEN`, they never appear as done:


https://user-images.githubusercontent.com/2308172/233197450-5f7eacd5-537b-4b50-b71d-ba09de34788b.mov

I would expect the spinner in studio to stop after the experiments complete, but it never does. The video only shows until shortly after the end of the experiment, but I've verified across multiple projects that it's reproducible and that the spinner continues to show hours after the experiment completes.


Comments:
So, I found the bug and have a couple of ideas on how to fix it.

But, this got me thinking about whether we really need/want the Studio events in DVC at all.

The original reasons to add `start` and `done` events in DVC instead of handling everything in DVCLive were:

1. Try to send DVC information (metrics, params) beyond the scope of DVCLive.
    a) Because DVCLive was not fully dependent on DVC Python lib back then.
    b) Because Studio was not showing completed exp refs back then.
2. More "correctness" on when the `start` and `done` happened.
3. In preparation for a potential extension to send `stage execution` updates by DVC.

As of today, motivations for `1` no longer hold and `2` and `3` are not feeling such strong reasons to me.

@dberenbaum 

Should we remove these events from DVC and make `live updates` more of a DVCLive-only thing?


No strong opinion from me for now. I think we initially scoped live metrics as DVCLive-only, and I think it's easier to explain that way. On the other hand, we will need to keep reporting the non-DVCLive metrics and other info since that was an almost immediate request AFAIR, and I think we will need some solution for 3 before long.

Patch:
diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py
--- a/dvc/repo/experiments/executor/base.py
+++ b/dvc/repo/experiments/executor/base.py
@@ -627,6 +627,7 @@ def _repro_dvc(  # noqa: C901
         copy_paths: Optional[List[str]] = None,
         **kwargs,
     ) -> Iterator["Repo"]:
+        from dvc_studio_client.env import STUDIO_REPO_URL, STUDIO_TOKEN
         from dvc_studio_client.post_live_metrics import post_live_metrics
 
         from dvc.repo import Repo
@@ -649,12 +650,18 @@ def _repro_dvc(  # noqa: C901
                 os.chdir(dvc.root_dir)
 
             try:
+                args_path = os.path.join(dvc.tmp_dir, cls.PACKED_ARGS_FILE)
+                if os.path.exists(args_path):
+                    _, kwargs = cls.unpack_repro_args(args_path)
+                run_env = kwargs.get("run_env", {})
                 post_live_metrics(
                     "start",
                     info.baseline_rev,
                     info.name,
                     "dvc",
                     params=to_studio_params(dvc.params.show()),
+                    studio_token=run_env.get(STUDIO_TOKEN, None),
+                    studio_repo_url=run_env.get(STUDIO_REPO_URL, None),
                 )
                 logger.debug("Running repro in '%s'", os.getcwd())
                 yield dvc
@@ -680,6 +687,8 @@ def _repro_dvc(  # noqa: C901
                     "dvc",
                     experiment_rev=dvc.experiments.scm.get_ref(EXEC_BRANCH),
                     metrics=get_in(dvc.metrics.show(), ["", "data"]),
+                    studio_token=run_env.get(STUDIO_TOKEN, None),
+                    studio_repo_url=run_env.get(STUDIO_REPO_URL, None),
                 )
 
                 if infofile is not None:
diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py
--- a/dvc/repo/experiments/queue/base.py
+++ b/dvc/repo/experiments/queue/base.py
@@ -17,8 +17,8 @@
     Union,
 )
 
-from dvc_studio_client.env import STUDIO_REPO_URL
-from dvc_studio_client.post_live_metrics import get_studio_repo_url
+from dvc_studio_client.env import STUDIO_REPO_URL, STUDIO_TOKEN
+from dvc_studio_client.post_live_metrics import get_studio_token_and_repo_url
 from funcy import retry
 
 from dvc.dependency import ParamsDependency
@@ -420,12 +420,14 @@ def _stash_exp(  # noqa: PLR0915, C901
                     run_env[DVC_EXP_NAME] = name
                     if resume_rev:
                         run_env[DVCLIVE_RESUME] = "1"
-                    studio_repo_url = get_studio_repo_url()
+
+                    studio_token, studio_repo_url = get_studio_token_and_repo_url()
+                    if studio_token is not None:
+                        run_env[STUDIO_TOKEN] = studio_token
                     if studio_repo_url is not None:
                         run_env[STUDIO_REPO_URL] = studio_repo_url
 
                     self._pack_args(*args, run_env=run_env, **kwargs)
-
                     # save experiment as a stash commit
                     msg = self._stash_msg(
                         stash_head,


Test patch:
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -614,7 +614,8 @@ def test_run_env(tmp_dir, dvc, scm, mocker):
         """
     )
     mocker.patch(
-        "dvc.repo.experiments.queue.base.get_studio_repo_url", return_value="REPO_URL"
+        "dvc.repo.experiments.queue.base.get_studio_token_and_repo_url",
+        return_value=("REPO_TOKEN", "REPO_URL"),
     )
     (tmp_dir / "dump_run_env.py").write_text(dump_run_env)
     baseline = scm.get_rev()
diff --git a/tests/integration/test_studio_live_experiments.py b/tests/integration/test_studio_live_experiments.py
--- a/tests/integration/test_studio_live_experiments.py
+++ b/tests/integration/test_studio_live_experiments.py
@@ -1,5 +1,5 @@
 import pytest
-from dvc_studio_client import env
+from dvc_studio_client import env, post_live_metrics
 from funcy import first
 
 
@@ -7,6 +7,7 @@
 def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp):
     valid_response = mocker.MagicMock()
     valid_response.status_code = 200
+    live_metrics = mocker.spy(post_live_metrics, "post_live_metrics")
     mocked_post = mocker.patch("requests.post", return_value=valid_response)
 
     monkeypatch.setenv(env.STUDIO_ENDPOINT, "https://0.0.0.0")
@@ -18,6 +19,12 @@ def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp):
         dvc.experiments.run(exp_stage.addressing, params=["foo=1"], tmp_dir=tmp)
     )
     name = dvc.experiments.get_exact_name([exp_rev])[exp_rev]
+    assert live_metrics.call_count == 2
+
+    start_call, done_call = live_metrics.call_args_list
+    assert start_call.kwargs["studio_token"] == "STUDIO_TOKEN"
+    assert start_call.kwargs["studio_repo_url"] == "STUDIO_REPO_URL"
+
     assert mocked_post.call_count == 2
 
     start_call, done_call = mocked_post.call_args_list