From 7868c1301cb2d8c6a45276412be589cb6ffbe530 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Wed, 22 Mar 2023 15:04:26 +0900 Subject: [PATCH] exp: enforce waiting for task completion in CI --- dvc/repo/experiments/__init__.py | 4 +--- dvc/repo/experiments/queue/celery.py | 22 ++++++++++++++++++++++ tests/func/experiments/test_apply.py | 2 +- tests/func/experiments/test_queue.py | 8 +++----- tests/func/experiments/test_show.py | 16 ++++++++++++---- 5 files changed, 39 insertions(+), 13 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index fbc57699f..7057b7121 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -1,7 +1,6 @@ import logging import os import re -import time from typing import TYPE_CHECKING, Dict, Iterable, Optional from funcy import chain, first @@ -205,8 +204,7 @@ class Experiments: ) for entry in entries: # wait for task execution to start - while not self.celery_queue.proc.get(entry.stash_rev): - time.sleep(1) + self.celery_queue.wait_for_start(entry, sleep_interval=1) self.celery_queue.follow(entry) # wait for task collection to complete try: diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index 8d5364429..20750e260 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -295,6 +295,28 @@ class LocalCeleryQueue(BaseStashQueue): # out of the active task list, and needs to be loaded here. return self._get_done_result(entry, timeout) + def wait(self, revs: Collection[str], **kwargs) -> None: + """Block until the specified tasks have completed.""" + revs = [revs] if isinstance(revs, str) else revs + results = self.match_queue_entry_by_name( + revs, self.iter_queued(), self.iter_done(), self.iter_failed() + ) + for entry in results.values(): + if not entry: + continue + self.wait_for_start(entry, **kwargs) + try: + self.get_result(entry) + except FileNotFoundError: + pass + + def wait_for_start(self, entry: QueueEntry, sleep_interval: float = 0.001) -> None: + """Block until the specified task has been started.""" + import time + + while not self.proc.get(entry.stash_rev): + time.sleep(sleep_interval) + def _get_running_task_ids(self) -> Set[str]: running_task_ids: Set[str] = set() active_workers = self.worker_status() diff --git a/tests/func/experiments/test_apply.py b/tests/func/experiments/test_apply.py index fe4e24663..1e5978736 100644 --- a/tests/func/experiments/test_apply.py +++ b/tests/func/experiments/test_apply.py @@ -35,7 +35,7 @@ def test_apply_failed(tmp_dir, scm, dvc, failed_exp_stage, session_queue): dvc.experiments.run( failed_exp_stage.addressing, params=["foo=3"], queue=True, name="foo" ) - dvc.experiments.run(run_all=True) + session_queue.wait(["foo"]) with pytest.raises(InvalidArgumentError): dvc.experiments.apply("bar") diff --git a/tests/func/experiments/test_queue.py b/tests/func/experiments/test_queue.py index aa73e02f7..64fe154db 100644 --- a/tests/func/experiments/test_queue.py +++ b/tests/func/experiments/test_queue.py @@ -22,8 +22,9 @@ def test_celery_logs( test_queue, ): celery_queue = dvc.experiments.celery_queue - dvc.experiments.run(failed_exp_stage.addressing, queue=True) + dvc.experiments.run(failed_exp_stage.addressing, queue=True, name="foo") dvc.experiments.run(run_all=True) + test_queue.wait(["foo"]) done_result = first(celery_queue.iter_done()) @@ -34,10 +35,6 @@ def test_celery_logs( assert "failed to reproduce 'failed-copy-file'" in captured.out -@pytest.mark.xfail( - strict=False, - reason="https://github.com/iterative/dvc/issues/9143", -) def test_queue_remove_done( dvc, exp_stage, @@ -63,6 +60,7 @@ def test_queue_remove_done( name=name_fail, ) dvc.experiments.run(run_all=True) + test_queue.wait(success_tasks + failed_tasks) assert len(celery_queue.failed_stash) == 3 status = to_dict(celery_queue.status()) assert len(status) == 6 diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py index 78be64d09..7c98fa957 100644 --- a/tests/func/experiments/test_show.py +++ b/tests/func/experiments/test_show.py @@ -156,15 +156,18 @@ def test_show_queued(tmp_dir, scm, dvc, exp_stage): @pytest.mark.vscode -def test_show_failed_experiment(tmp_dir, scm, dvc, failed_exp_stage): +def test_show_failed_experiment(tmp_dir, scm, dvc, failed_exp_stage, test_queue): baseline_rev = scm.get_rev() timestamp = datetime.fromtimestamp( scm.gitpython.repo.rev_parse(baseline_rev).committed_date ) - dvc.experiments.run(failed_exp_stage.addressing, params=["foo=2"], queue=True) + dvc.experiments.run( + failed_exp_stage.addressing, params=["foo=2"], queue=True, name="show-failed" + ) exp_rev = dvc.experiments.scm.resolve_rev(f"{CELERY_STASH}@{{0}}") dvc.experiments.run(run_all=True) + test_queue.wait(["show-failed"]) experiments = dvc.experiments.show()[baseline_rev] expected_baseline = { @@ -542,15 +545,20 @@ def test_show_running_celery(tmp_dir, scm, dvc, exp_stage, mocker): assert results["workspace"]["baseline"]["data"]["status"] == "Success" -def test_show_running_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, mocker): +def test_show_running_checkpoint( + tmp_dir, scm, dvc, checkpoint_stage, mocker, test_queue +): from dvc.repo.experiments.executor.local import TempDirExecutor baseline_rev = scm.get_rev() - dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"], queue=True) + dvc.experiments.run( + checkpoint_stage.addressing, params=["foo=2"], queue=True, name="foo" + ) queue = dvc.experiments.celery_queue entries = list(queue.iter_queued()) run_results = dvc.experiments.run(run_all=True) + test_queue.wait(["foo"]) checkpoint_rev = first(run_results) exp_ref = first(exp_refs_by_rev(scm, checkpoint_rev))