studio: Ensure studio_repo_url works with exp run --queue.

Closes #9349

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:

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,

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",

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()

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