diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index e5a2bddee..873f485f3 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -246,10 +246,14 @@ class LocalCeleryQueue(BaseStashQueue): infofile = self.get_infofile_path(rev) return ExecutorInfo.load_json(infofile) - try: - executor_info = _load_info(entry.stash_rev) + def _load_collected(rev: str) -> Optional[ExecutorResult]: + executor_info = _load_info(rev) if executor_info.collected: return executor_info.result + raise FileNotFoundError + + try: + return _load_collected(entry.stash_rev) except FileNotFoundError: # Infofile will not be created until execution begins pass @@ -269,7 +273,15 @@ class LocalCeleryQueue(BaseStashQueue): ) from exc executor_info = _load_info(entry.stash_rev) return executor_info.result - raise DvcException("Invalid experiment.") + + # NOTE: It's possible for an exp to complete while iterating through + # other queued and active tasks, in which case the exp will get moved + # out of the active task list, and needs to be loaded here. + try: + return _load_collected(entry.stash_rev) + except FileNotFoundError: + pass + raise DvcException("Invalid experiment '{entry.stash_rev[:7]}'.") def kill(self, revs: Collection[str]) -> None: to_kill: Set[QueueEntry] = set() diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index d74ef89bb..3aec05376 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -199,8 +199,7 @@ def test_apply(tmp_dir, scm, dvc, exp_stage): assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 3" -def test_apply_queued(tmp_dir, scm, dvc, exp_stage, test_queue): - # from dvc.exceptions import InvalidArgumentError +def test_apply_queued(tmp_dir, scm, dvc, exp_stage): from dvc.repo.experiments.exceptions import ApplyConflictError metrics_original = (tmp_dir / "metrics.yaml").read_text().strip() diff --git a/tests/unit/repo/experiments/queue/test_local.py b/tests/unit/repo/experiments/queue/test_local.py index 0b91d31c7..e24d3a92a 100644 --- a/tests/unit/repo/experiments/queue/test_local.py +++ b/tests/unit/repo/experiments/queue/test_local.py @@ -2,6 +2,7 @@ import time import pytest from celery import shared_task +from flaky.flaky_decorator import flaky from dvc.repo.experiments.exceptions import UnresolvedExpNamesError @@ -48,6 +49,8 @@ def test_shutdown_with_kill(test_queue, mocker): shutdown_spy.assert_called_once() +# pytest-celery worker thread may finish the task before we check for PENDING +@flaky(max_runs=3, min_passes=1) def test_post_run_after_kill(test_queue): from celery import chain