diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index c932026c1..2176059ff 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -501,6 +501,7 @@ class BaseExecutor(ABC): infofile, log_errors=log_errors, copy_paths=copy_paths, + message=message, **kwargs, ) as dvc: if auto_push: @@ -631,6 +632,7 @@ class BaseExecutor(ABC): infofile: Optional[str] = None, log_errors: bool = True, copy_paths: Optional[List[str]] = None, + message: Optional[str] = None, **kwargs, ) -> Iterator["Repo"]: from dvc_studio_client.env import STUDIO_REPO_URL, STUDIO_TOKEN @@ -668,6 +670,7 @@ class BaseExecutor(ABC): 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), + message=message, ) logger.debug("Running repro in '%s'", os.getcwd()) yield dvc diff --git a/pyproject.toml b/pyproject.toml index ec6c479df..d61e4bf12 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,<1", - "dvc-studio-client>=0.8.0,<1", + "dvc-studio-client>=0.9.0,<1", "dvc-task>=0.2.1,<1", "flatten_dict<1,>=0.4.1", "flufl.lock>=5", diff --git a/tests/integration/test_studio_live_experiments.py b/tests/integration/test_studio_live_experiments.py index 2ed583b9a..9ad281e19 100644 --- a/tests/integration/test_studio_live_experiments.py +++ b/tests/integration/test_studio_live_experiments.py @@ -47,3 +47,37 @@ def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch, tmp): "experiment_rev": exp_rev, "metrics": {"metrics.yaml": {"data": {"foo": 1}}}, } + + +@pytest.mark.parametrize("tmp", [True, False]) +def test_post_to_studio_custom_message( + 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, message="foo" + ) + ) + name = dvc.experiments.get_exact_name([exp_rev])[exp_rev] + assert mocked_post.call_count == 2 + + start_call = mocked_post.call_args_list[0] + + 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", + "message": "foo", + }