Repository: iterative/dvc (python)

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.

Pull request:
studio: Ensure `studio_repo_url` works with `exp run --queue`.
Closes #9349

Patch:
diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py
index 5d364a0c6..b9e672a9d 100644
--- a/dvc/repo/experiments/executor/base.py
+++ b/dvc/repo/experiments/executor/base.py
@@ -627,6 +627,7 @@ class BaseExecutor(ABC):
         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 @@ class BaseExecutor(ABC):
                 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 @@ class BaseExecutor(ABC):
                     "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
index 98ff3ace8..321194b2b 100644
--- a/dvc/repo/experiments/queue/base.py
+++ b/dvc/repo/experiments/queue/base.py
@@ -17,8 +17,8 @@ from typing import (
     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 @@ class BaseStashQueue(ABC):
                     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,
diff --git a/pyproject.toml b/pyproject.toml
index e8beab2cf..a86b3a891 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,7 +37,7 @@ dependencies = [
     "dvc-data>=0.47.1,<0.48",
     "dvc-http>=2.29.0",
     "dvc-render>=0.3.1,<0.4.0",
-    "dvc-studio-client>=0.6.1,<1",
+    "dvc-studio-client>=0.8.0,<1",
     "dvc-task>=0.2.0,<1",
     "flatten_dict<1,>=0.4.1",
     "flufl.lock>=5",


Test patch:
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index b4c063c3e..4fe02a3b5 100644
--- 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
index 1f3a91d3b..2ed583b9a 100644
--- 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 @@ from funcy import first
 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