Add additional data to Studio live metrics post messages to support live experiments in monorepos

Author: mattseddon

dvc/repo/__init__.py

@@ -13,7 +13,6 @@
 )
 from dvc.ignore import DvcIgnoreFilter
 from dvc.log import logger
-from dvc.utils import as_posix
 from dvc.utils.objects import cached_property
 
 if TYPE_CHECKING:
@@ -351,16 +350,6 @@ def fs(self, fs: "FileSystem"):
         # fs.
         self._reset()
 
-    @property
-    def subrepo_relpath(self) -> str:
-        from dvc.fs import GitFileSystem
-
-        scm_root_dir = "/" if isinstance(self.fs, GitFileSystem) else self.scm.root_dir
-
-        relpath = as_posix(self.fs.relpath(self.root_dir, scm_root_dir))
-
-        return "" if relpath == "." else relpath
-
     @property
     def data_index(self) -> "DataIndex":
         from dvc_data.index import DataIndex

dvc/repo/experiments/executor/base.py

@@ -23,7 +23,11 @@
 from dvc.stage.serialize import to_lockfile
 from dvc.utils import dict_sha256, env2bool, relpath
 from dvc.utils.fs import remove
-from dvc.utils.studio import env_to_config
+from dvc.utils.studio import (
+    env_to_config,
+    get_dvc_experiment_parent_data,
+    get_subrepo_relpath,
+)
 
 if TYPE_CHECKING:
     from queue import Queue
@@ -624,6 +628,10 @@ def _repro_dvc(
                     params=to_studio_params(dvc.params.show()),
                     dvc_studio_config=dvc_studio_config,
                     message=message,
+                    subdir=get_subrepo_relpath(dvc),
+                    dvc_experiment_parent_data=get_dvc_experiment_parent_data(
+                        dvc, info.baseline_rev
+                    ),
                 )
                 logger.debug("Running repro in '%s'", os.getcwd())
                 yield dvc

dvc/scm.py

@@ -1,4 +1,5 @@
 """Manages source control systems (e.g. Git)."""
+
 import os
 from collections.abc import Iterator, Mapping
 from contextlib import contextmanager
@@ -181,7 +182,7 @@ def resolve_rev(scm: Union["Git", "NoSCM"], rev: str) -> str:
         raise RevError(str(exc))  # noqa: B904
 
 
-def _get_n_commits(scm: "Git", revs: list[str], num: int) -> list[str]:
+def get_n_commits(scm: "Git", revs: list[str], num: int) -> list[str]:
     results = []
     for rev in revs:
         if num == 0:
@@ -227,7 +228,7 @@ def iter_revs(
         return {}
 
     revs = revs or []
-    results: list[str] = _get_n_commits(scm, revs, num)
+    results: list[str] = get_n_commits(scm, revs, num)
 
     if all_commits:
         results.extend(scm.list_all_commits())

dvc/utils/studio.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, Union
 from urllib.parse import urljoin
 
 import requests
@@ -12,10 +12,14 @@
     DVC_STUDIO_URL,
 )
 from dvc.log import logger
+from dvc.utils import as_posix
 
 if TYPE_CHECKING:
     from requests import Response
 
+    from dvc.repo import Repo
+
+
 logger = logger.getChild(__name__)
 
 STUDIO_URL = "https://studio.iterative.ai"
@@ -111,3 +115,51 @@ def env_to_config(env: dict[str, Any]) -> dict[str, Any]:
     if DVC_STUDIO_URL in env:
         config["url"] = env[DVC_STUDIO_URL]
     return config
+
+
+def get_subrepo_relpath(repo: "Repo") -> str:
+    from dvc.fs import GitFileSystem
+
+    scm_root_dir = "/" if isinstance(repo.fs, GitFileSystem) else repo.scm.root_dir
+
+    relpath = as_posix(repo.fs.relpath(repo.root_dir, scm_root_dir))
+
+    return "" if relpath == "." else relpath
+
+
+def get_dvc_experiment_parent_data(
+    repo: "Repo", baseline_rev: Union[str, None]
+) -> Union[dict[str, Any], None]:
+    from scmrepo.exceptions import SCMError
+
+    from dvc.scm import NoSCM, get_n_commits
+
+    scm = repo.scm
+
+    try:
+        if (
+            not baseline_rev
+            or not scm
+            or isinstance(scm, NoSCM)
+            or not (commit := scm.resolve_commit(baseline_rev))
+            or not (
+                first_100_parent_shas := get_n_commits(scm, [baseline_rev], 101)[1:]
+            )
+        ):
+            return None
+
+        return {
+            "author": {
+                "name": commit.author_name,
+                "email": commit.author_email,
+            },
+            "date": commit.author_datetime.isoformat(),
+            "message": commit.message,
+            "parent_shas": first_100_parent_shas,
+            "sha": commit.hexsha,
+            "title": commit.message.partition("\n")[0].strip(),
+        }
+
+    except SCMError:
+        # unable to resolve all required baseline information
+        return None

pyproject.toml

@@ -38,7 +38,7 @@ dependencies = [
     "dvc-data>=3.10,<3.11",
     "dvc-http>=2.29.0",
     "dvc-render>=1.0.1,<2",
-    "dvc-studio-client>=0.17.1,<1",
+    "dvc-studio-client@git+https://github.com/iterative/dvc-studio-client.git@refs/pull/144/head",
     "dvc-task>=0.3.0,<1",
     "flatten_dict<1,>=0.4.1",
     # https://github.com/iterative/dvc/issues/9654

tests/integration/test_studio_live_experiments.py

@@ -13,7 +13,7 @@
 @pytest.mark.parametrize("tmp", [True, False])
 @pytest.mark.parametrize("offline", [True, False])
 def test_post_to_studio(
-    tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp, offline
+    M, tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp, offline
 ):
     valid_response = mocker.MagicMock()
     valid_response.status_code = 200
@@ -53,6 +53,17 @@ def test_post_to_studio(
             "name": name,
             "params": {"params.yaml": {"foo": 1}},
             "client": "dvc",
+            "dvc_experiment_parent_data": {
+                "author": {
+                    "email": "dvctester@example.com",
+                    "name": "DVC Tester",
+                },
+                "date": M.any,
+                "message": "init",
+                "parent_shas": M.any,
+                "title": "init",
+                "sha": baseline_sha,
+            },
         }
 
         assert done_call.kwargs["json"] == {
@@ -68,7 +79,7 @@ def test_post_to_studio(
 
 @pytest.mark.parametrize("tmp", [True, False])
 def test_post_to_studio_custom_message(
-    tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp
+    M, tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp
 ):
     valid_response = mocker.MagicMock()
     valid_response.status_code = 200
@@ -97,4 +108,15 @@ def test_post_to_studio_custom_message(
         "params": {"params.yaml": {"foo": 1}},
         "client": "dvc",
         "message": "foo",
+        "dvc_experiment_parent_data": {
+            "author": {
+                "email": "dvctester@example.com",
+                "name": "DVC Tester",
+            },
+            "date": M.any,
+            "message": "init",
+            "parent_shas": M.any,
+            "title": "init",
+            "sha": baseline_sha,
+        },
     }

tests/unit/command/test_studio.py

@@ -1,12 +1,12 @@
-from dvc_studio_client.auth import AuthenticationExpired
+from dvc_studio_client.auth import AuthorizationExpiredError
 
 from dvc.cli import main
 from dvc.utils.studio import STUDIO_URL
 
 
 def test_studio_login_token_check_failed(mocker):
     mocker.patch(
-        "dvc_studio_client.auth.get_access_token", side_effect=AuthenticationExpired
+        "dvc_studio_client.auth.get_access_token", side_effect=AuthorizationExpiredError
     )
 
     assert main(["studio", "login"]) == 1

tests/unit/repo/test_repo.py

@@ -135,51 +135,3 @@ def test_dynamic_cache_initialization(tmp_dir, scm):
     dvc.close()
 
     Repo(str(tmp_dir)).close()
-
-
-def test_monorepo_relpath(tmp_dir, scm):
-    from dvc.repo.destroy import destroy
-
-    tmp_dir.gen({"project_a": {}, "subdir/project_b": {}})
-
-    non_monorepo = Repo.init(tmp_dir)
-    assert non_monorepo.subrepo_relpath == ""
-
-    destroy(non_monorepo)
-
-    monorepo_project_a = Repo.init(tmp_dir / "project_a", subdir=True)
-
-    assert monorepo_project_a.subrepo_relpath == "project_a"
-
-    monorepo_project_b = Repo.init(tmp_dir / "subdir" / "project_b", subdir=True)
-
-    assert monorepo_project_b.subrepo_relpath == "subdir/project_b"
-
-
-def test_virtual_monorepo_relpath(tmp_dir, scm):
-    from dvc.fs.git import GitFileSystem
-    from dvc.repo.destroy import destroy
-
-    tmp_dir.gen({"project_a": {}, "subdir/project_b": {}})
-    scm.commit("initial commit")
-    gfs = GitFileSystem(scm=scm, rev="master")
-
-    non_monorepo = Repo.init(tmp_dir)
-    non_monorepo.fs = gfs
-    non_monorepo.root_dir = "/"
-
-    assert non_monorepo.subrepo_relpath == ""
-
-    destroy(non_monorepo)
-
-    monorepo_project_a = Repo.init(tmp_dir / "project_a", subdir=True)
-    monorepo_project_a.fs = gfs
-    monorepo_project_a.root_dir = "/project_a"
-
-    assert monorepo_project_a.subrepo_relpath == "project_a"
-
-    monorepo_project_b = Repo.init(tmp_dir / "subdir" / "project_b", subdir=True)
-    monorepo_project_b.fs = gfs
-    monorepo_project_b.root_dir = "/subdir/project_b"
-
-    assert monorepo_project_b.subrepo_relpath == "subdir/project_b"

tests/unit/utils/test_studio.py

@@ -1,3 +1,4 @@
+import json
 from urllib.parse import urljoin
 
 import pytest
@@ -9,7 +10,15 @@
     DVC_STUDIO_TOKEN,
     DVC_STUDIO_URL,
 )
-from dvc.utils.studio import STUDIO_URL, config_to_env, env_to_config, notify_refs
+from dvc.repo import Repo
+from dvc.utils.studio import (
+    STUDIO_URL,
+    config_to_env,
+    env_to_config,
+    get_dvc_experiment_parent_data,
+    get_subrepo_relpath,
+    notify_refs,
+)
 
 CONFIG = {"offline": True, "repo_url": "repo_url", "token": "token", "url": "url"}
 
@@ -67,3 +76,97 @@ def test_config_to_env():
 
 def test_env_to_config():
     assert env_to_config(ENV) == CONFIG
+
+
+@pytest.mark.studio
+def test_monorepo_relpath(tmp_dir, scm):
+    from dvc.repo.destroy import destroy
+
+    tmp_dir.gen({"project_a": {}, "subdir/project_b": {}})
+
+    non_monorepo = Repo.init(tmp_dir)
+    assert get_subrepo_relpath(non_monorepo) == ""
+
+    destroy(non_monorepo)
+
+    monorepo_project_a = Repo.init(tmp_dir / "project_a", subdir=True)
+
+    assert get_subrepo_relpath(monorepo_project_a) == "project_a"
+
+    monorepo_project_b = Repo.init(tmp_dir / "subdir" / "project_b", subdir=True)
+
+    assert get_subrepo_relpath(monorepo_project_b) == "subdir/project_b"
+
+
+@pytest.mark.studio
+def test_virtual_monorepo_relpath(tmp_dir, scm):
+    from dvc.fs.git import GitFileSystem
+    from dvc.repo.destroy import destroy
+
+    tmp_dir.gen({"project_a": {}, "subdir/project_b": {}})
+    scm.commit("initial commit")
+    gfs = GitFileSystem(scm=scm, rev="master")
+
+    non_monorepo = Repo.init(tmp_dir)
+    non_monorepo.fs = gfs
+    non_monorepo.root_dir = "/"
+
+    assert get_subrepo_relpath(non_monorepo) == ""
+
+    destroy(non_monorepo)
+
+    monorepo_project_a = Repo.init(tmp_dir / "project_a", subdir=True)
+    monorepo_project_a.fs = gfs
+    monorepo_project_a.root_dir = "/project_a"
+
+    assert get_subrepo_relpath(monorepo_project_a) == "project_a"
+
+    monorepo_project_b = Repo.init(tmp_dir / "subdir" / "project_b", subdir=True)
+    monorepo_project_b.fs = gfs
+    monorepo_project_b.root_dir = "/subdir/project_b"
+
+    assert get_subrepo_relpath(monorepo_project_b) == "subdir/project_b"
+
+
+@pytest.mark.studio
+def test_dvc_experiment_parent_data(M, tmp_dir, scm, dvc):
+    parent_shas = [scm.get_rev()]
+
+    for i in range(5):
+        tmp_dir.scm_gen({"metrics.json": json.dumps({"metric": i})}, commit=f"step {i}")
+        parent_shas.insert(0, scm.get_rev())
+
+    title = "a final commit with a fairly long message"
+    message = f"{title}\nthat is split over two lines"
+
+    tmp_dir.scm_gen({"metrics.json": json.dumps({"metric": 100})}, commit=message)
+
+    head_sha = scm.get_rev()
+
+    assert isinstance(head_sha, str)
+    assert head_sha not in parent_shas
+
+    dvc_experiment_parent_data = get_dvc_experiment_parent_data(dvc, head_sha)
+
+    assert dvc_experiment_parent_data is not None
+    assert isinstance(dvc_experiment_parent_data["date"], str)
+
+    assert dvc_experiment_parent_data == {
+        "author": {
+            "email": "dvctester@example.com",
+            "name": "DVC Tester",
+        },
+        "date": M.any,
+        "message": message,
+        "parent_shas": parent_shas,
+        "title": title,
+        "sha": head_sha,
+    }
+
+
+@pytest.mark.parametrize("func", ["get_rev", "resolve_commit"])
+@pytest.mark.studio
+def test_no_dvc_experiment_parent_data(mocker, scm, dvc, func):
+    mocker.patch.object(scm, func, return_value=None)
+
+    assert get_dvc_experiment_parent_data(dvc, scm.get_rev()) is None