diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 3fb27e712..96292e679 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -32,6 +32,7 @@ from dvc.repo.experiments.refs import ( EXEC_CHECKPOINT, ExpRefInfo, ) +from dvc.repo.experiments.utils import to_studio_params from dvc.stage.serialize import to_lockfile from dvc.ui import ui from dvc.utils import dict_sha256, env2bool, relpath @@ -545,6 +546,8 @@ class BaseExecutor(ABC): log_errors: bool = True, **kwargs, ): + from dvc_studio_client.post_live_metrics import post_live_metrics + from dvc.repo import Repo from dvc.stage.monitor import CheckpointKilledError @@ -561,6 +564,13 @@ class BaseExecutor(ABC): os.chdir(dvc.root_dir) try: + post_live_metrics( + "start", + info.baseline_rev, + info.name, + "dvc", + params=to_studio_params(dvc.params.show()), + ) logger.debug("Running repro in '%s'", os.getcwd()) yield dvc info.status = TaskStatus.SUCCESS @@ -578,6 +588,14 @@ class BaseExecutor(ABC): info.status = TaskStatus.FAILED raise finally: + post_live_metrics( + "done", + info.baseline_rev, + info.name, + "dvc", + experiment_rev=dvc.experiments.scm.get_ref(EXEC_BRANCH), + ) + if infofile is not None: info.dump_json(infofile) dvc.close() diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index 4cf7afe57..063f1c54f 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -13,7 +13,7 @@ from dvc.repo.experiments.exceptions import ExpQueueEmptyError from dvc.repo.experiments.executor.base import ExecutorInfo, TaskStatus from dvc.repo.experiments.executor.local import WorkspaceExecutor from dvc.repo.experiments.refs import EXEC_BRANCH, WORKSPACE_STASH -from dvc.repo.experiments.utils import get_exp_rwlock, to_studio_params +from dvc.repo.experiments.utils import get_exp_rwlock from dvc.rwlock import RWLOCK_FILE, RWLOCK_LOCK, SCHEMA from dvc.utils import relpath from dvc.utils.fs import remove @@ -94,8 +94,6 @@ class WorkspaceQueue(BaseStashQueue): def _reproduce_entry( self, entry: QueueEntry, executor: "BaseExecutor" ) -> Dict[str, Dict[str, str]]: - from dvc_studio_client.post_live_metrics import post_live_metrics - from dvc.stage.monitor import CheckpointKilledError results: Dict[str, Dict[str, str]] = defaultdict(dict) @@ -103,13 +101,6 @@ class WorkspaceQueue(BaseStashQueue): infofile = self.get_infofile_path(exec_name) try: rev = entry.stash_rev - post_live_metrics( - "start", - executor.info.baseline_rev, - executor.info.name, - "dvc", - params=to_studio_params(self.repo.params.show()), - ) exec_result = executor.reproduce( info=executor.info, rev=rev, @@ -131,13 +122,6 @@ class WorkspaceQueue(BaseStashQueue): except Exception as exc: # noqa: BLE001 raise DvcException(f"Failed to reproduce experiment '{rev[:7]}'") from exc finally: - post_live_metrics( - "done", - executor.info.baseline_rev, - executor.info.name, - "dvc", - experiment_rev=first(results.get(rev, {})), - ) executor.cleanup(infofile) return results diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index a48aa3d01..e8ed3ba97 100644 --- a/dvc/repo/experiments/utils.py +++ b/dvc/repo/experiments/utils.py @@ -327,7 +327,9 @@ def to_studio_params(dvc_params): "params.yaml": {"foo": 1} } """ - result = {} + result: Dict = {} + if not dvc_params: + return result for rev_data in dvc_params.values(): for file_name, file_data in rev_data["data"].items(): result[file_name] = file_data["data"] diff --git a/tests/func/experiments/test_utils.py b/tests/func/experiments/test_utils.py index 0585bad92..260cc417c 100644 --- a/tests/func/experiments/test_utils.py +++ b/tests/func/experiments/test_utils.py @@ -1,4 +1,3 @@ -from dvc_studio_client import env from funcy import first @@ -28,38 +27,3 @@ def test_generate_random_exp_name(tmp_dir, dvc, scm, exp_stage, mocker): # Can use same name because of different baseline_rev ref = first(dvc.experiments.run(exp_stage.addressing, params=["foo=1"])) assert dvc.experiments.get_exact_name([ref])[ref] == "0-0" - - -def test_post_to_studio(tmp_dir, dvc, scm, exp_stage, mocker, monkeypatch): - 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"])) - name = dvc.experiments.get_exact_name([exp_rev])[exp_rev] - assert mocked_post.call_count == 2 - - start_call, done_call = mocked_post.call_args_list - - 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", - } - - assert done_call.kwargs["json"] == { - "type": "done", - "repo_url": "STUDIO_REPO_URL", - "baseline_sha": baseline_sha, - "name": name, - "client": "dvc", - "experiment_rev": exp_rev, - } diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 000000000..90a77b1c2 --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,2 @@ +# pylint disable=unused-argument +from tests.unit.repo.experiments.conftest import exp_stage, test_queue # noqa: F401 diff --git a/tests/integration/test_studio_live_experiments.py b/tests/integration/test_studio_live_experiments.py new file mode 100644 index 000000000..c43c253f2 --- /dev/null +++ b/tests/integration/test_studio_live_experiments.py @@ -0,0 +1,41 @@ +import pytest +from dvc_studio_client import env +from funcy import first + + +@pytest.mark.parametrize("tmp", [True, False]) +def test_post_to_studio(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) + ) + name = dvc.experiments.get_exact_name([exp_rev])[exp_rev] + assert mocked_post.call_count == 2 + + start_call, done_call = mocked_post.call_args_list + + 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", + } + + assert done_call.kwargs["json"] == { + "type": "done", + "repo_url": "STUDIO_REPO_URL", + "baseline_sha": baseline_sha, + "name": name, + "client": "dvc", + "experiment_rev": exp_rev, + }